Matplotlib and Python loop over two dimensions

Question

I have 6 cubic polynomial curves I'd like to plot. And I have the following code::

import numpy as np
import matplotlib.pyplot as plt 

a0 = np.array([ 0.04438, -0.01808, -0.01836,  0.01170, -0.01062, -0.01062])
a1 = np.array([-2.26095, -0.13595, -0.03577, -0.00400, -0.03577, -0.00400])
a2 = np.array([-0.13387,  0.01941,  0.02612,  0.00066,  0.02612,  0.00066])
a3 = np.array([ 0.00066, -0.00183, -0.00558, -0.00558,  0.00890,  0.00890])

x_max = 2.80
x_min = 0.30
diff  = 0.01

y = int(((x_max - x_min)/diff))
m_diff = np.zeros((y,6))

xmin = int(x_min * 100)
xmax = int(x_max * 100)
di   = int(diff  * 100)

for xx in range(xmin, xmax, di):
   x = xx*diff
   m_diff[xx] = a0 + (a1*x) + (a2*x*x) + (a3*x*x*x)

Why do I get an "IndexError: index 250 is out of bounds for axis 0 with size 250" error? What should this be? Ultimately I'd like to just do:

plt(xx, m_diff[:,0])
plt(xx, m_diff[:,1])
plt(xx, m_diff[:,2])
plt(xx, m_diff[:,3])
plt(xx, m_diff[:,4])
plt(xx, m_diff[:,5])

Thanks!!

Answer 1

Here is the vectored approach, which prevents mistakes in the array shapes :

x=np.arange(x_min,x_max,diff)
xs=np.power.outer(x,range(4))
coeffs=np.vstack((a0,a1,a2,a3))
curves=np.dot(xs,coeffs)
plt.plot(x,curves)

You just have to learn to play with dimensions and all is simple ;):

for array in x,xs,coeffs,curves : print (array.shape) 
# (250,)
# (250, 4)
# (4, 6)
# (250, 6)

Results :

Answer 2

Numpy has a function numpy.polyval to evaluate polynomials. You can use it in this case to plot the polynomials

import numpy as np
import matplotlib.pyplot as plt 

a0 = np.array([ 0.04438, -0.01808, -0.01836,  0.01170, -0.01062, -0.01062])
a1 = np.array([-2.26095, -0.13595, -0.03577, -0.00400, -0.03577, -0.00400])
a2 = np.array([-0.13387,  0.01941,  0.02612,  0.00066,  0.02612,  0.00066])
a3 = np.array([ 0.00066, -0.00183, -0.00558, -0.00558,  0.00890,  0.00890])

x_max = 2.80
x_min = 0.30
diff  = 0.01

coeff = np.c_[a3,a2,a1,a0]

x = np.arange(x_min,x_max,diff)
y = np.array([np.polyval(coeff[i],x) for i in range(coeff.shape[0])])

plt.plot(x,y.T)

plt.show()