How do I correctly write the syntax for performing and plotting a for loop operation?

Question

I am trying to create a for loop which uses a defined function (B_lambda) and takes in values of wavelength and temperature to produce values of intensity. i.e. I want the loop to take the function B_lambda and to run through every value within my listed wavelength range for each temperature in the temperature list. Then I want to plot the results. I am not very good with the syntax and have tried many ways but nothing is producing what I need and I am mostly getting errors. I have no idea how to use a for loop to plot and all online sources that I have checked out have not helped me with using a defined function in a for loop. I will put my latest code that seems to have the least errors down below with the error message:

import matplotlib.pylab as plt
import numpy as np
from astropy import units as u
import scipy.constants
%matplotlib inline

#Importing constants to use.
h = scipy.constants.h
c = scipy.constants.c
k = scipy.constants.k

wavelengths= np.arange(1000,30000)*1.e-10
temperature=[3000,4000,5000,6000]

for lam in wavelengths:
    for T in temperature:
        B_lambda = ((2*h*c**2)/(lam**5))*((1)/(np.exp((h*c)/(lam*k*T))-1))
        plt.figure()
        plt.plot(wavelengths,B_lambda)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-6-73b866241c49> in <module>
     17         B_lambda = ((2*h*c**2)/(lam**5))*((1)/(np.exp((h*c)/(lam*k*T))-1))
     18         plt.figure()
---> 19         plt.plot(wavelengths,B_lambda)
     20 
     21 

/usr/local/lib/python3.6/dist-packages/matplotlib/pyplot.py in plot(scalex, scaley, data, *args, **kwargs)
   2787     return gca().plot(
   2788         *args, scalex=scalex, scaley=scaley, **({"data": data} if data
-> 2789         is not None else {}), **kwargs)
   2790 
   2791 

/usr/local/lib/python3.6/dist-packages/matplotlib/axes/_axes.py in plot(self, scalex, scaley, data, *args, **kwargs)
   1663         """
   1664         kwargs = cbook.normalize_kwargs(kwargs, mlines.Line2D._alias_map)
-> 1665         lines = [*self._get_lines(*args, data=data, **kwargs)]
   1666         for line in lines:
   1667             self.add_line(line)

/usr/local/lib/python3.6/dist-packages/matplotlib/axes/_base.py in __call__(self, *args, **kwargs)
    223                 this += args[0],
    224                 args = args[1:]
--> 225             yield from self._plot_args(this, kwargs)
    226 
    227     def get_next_color(self):

/usr/local/lib/python3.6/dist-packages/matplotlib/axes/_base.py in _plot_args(self, tup, kwargs)
    389             x, y = index_of(tup[-1])
    390 
--> 391         x, y = self._xy_from_xy(x, y)
    392 
    393         if self.command == 'plot':

/usr/local/lib/python3.6/dist-packages/matplotlib/axes/_base.py in _xy_from_xy(self, x, y)
    268         if x.shape[0] != y.shape[0]:
    269             raise ValueError("x and y must have same first dimension, but "
--> 270                              "have shapes {} and {}".format(x.shape, y.shape))
    271         if x.ndim > 2 or y.ndim > 2:
    272             raise ValueError("x and y can be no greater than 2-D, but have "

ValueError: x and y must have same first dimension, but have shapes (29000,) and (1,)```

warped · Accepted Answer · 2020-03-08T15:43:29.620

First thing to note (and this is minor) is that astropy is not required to run your code. So, you can simplify the import statements.

import matplotlib.pylab as plt
import numpy as np
import scipy.constants
%matplotlib inline

#Importing constants to use.
h = scipy.constants.h
c = scipy.constants.c
k = scipy.constants.k


wavelengths= np.arange(1000,30000,100)*1.e-10 # here, I chose steps of 100, because plotting 29000 datapoints takes a while
temperature=[3000,4000,5000,6000]

Secondly, to tidy up the loop a bit, you can write a helper function, that youn call from within you loop:

def f(lam, T):
    return ((2*h*c**2)/(lam**5))*((1)/(np.exp((h*c)/(lam*k*T))-1))

now you can collect the output of your function, together with the input parameters, e.g. in a list of tuples:

outputs = []

for lam in wavelengths:
    for T in temperature:
        outputs.append((lam, T, f(lam, T)))

Since you vary both wavelength and temperature, a 3d plot makes sense:

from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure(figsize=(10,10))
ax = fig.add_subplot(111, projection='3d')

ax.plot(*zip(*outputs))

An alternative would be to display the data as an image, using colour to indicate the function output.

I am also including an alternative method to generate the data in this one. Since the function f can take arrays as input, you can feed one temperature at a time, and with it, all the wavelengths simultaneously.

# initialise output as array with proper shape
outputs = np.zeros((len(wavelengths), len(temperature)))

for i, T in enumerate(temperature):
    outputs[:,i] = f(wavelengths, T)

The output now is a large matrix, which you can visualise as an image:

fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(outputs, aspect=10e8, interpolation='none', 
          extent=[
              np.min(temperature),
              np.max(temperature),
              np.max(wavelengths),
              np.min(wavelengths)]
         )

How do I correctly write the syntax for performing and plotting a for loop operation?

1 Answers1