Iterating an Array in Python using the brentq Func

Question

I am having trouble iterating every element of an array using the brentq function. q in the defined function below is a FITS file array, and we are using every element in this array as inputs to run through the brentqfunction in order to solve for T.

Essentially, my problem lies in not particularly knowing where or how to implement the appropriate for loop to iterate the function over every element of q.

Any suggestions on how to go about this problem?

def f(T,q,coeff1,coeff2,coeff3):
    return q*const3 - ((exp(const2/T)-1)/(exp(const/T)-1))

a = brentq(f, 10, 435.1, args=(q,4351.041,4262.570,0.206))
print a

newhdu = fits.PrimaryHDU(a)
newhdulist = fits.HDUList([newhdu])
newhdulist.writeto('Temp21DCOT.fits')

Further explanation: The basis of what I'm trying to do is to use brentq to solve for temperature values using the intensity values of our initial array (our FITS file).

The equation is derived from a ratio of two wavelengths of Plank's Equation, so q = B_1/B_2 if we want to be true to physics, where every element in q are intensity values. brentq would solve this analytically insolvable equation for T (temperature) for every element in q, and make a new temperature array of the same size as q. In other words, I am trying to solve for the temperature of every pixel in a FITS file using Plank's Equation.

Note: I re-posted this to clarify the problem in a more efficient manner.

Answer 1

Are you having problems with iteration, or with efficiency?

This iteration works for me:

In [485]: from scipy import optimize

In [486]: def f(T,q,coeff1,coeff2,coeff3):
        return q*coeff3 - ((np.exp(coeff2/T)-1)/(np.exp(coeff1/T)-1))
        # corrected the coeff use

In [487]: q=np.linspace(1,3,10)
# q range chosen to avoid the different signs ValueError

In [488]: A=[optimize.brentq(f, 10, 435.1, args=(i,4351.041,4262.570,0.206),full_output=True) for i in q]

In [489]: A
Out[489]: 
[(55.99858839149886, <scipy.optimize.zeros.RootResults at 0xa861284c>),
 (64.14621536172528, <scipy.optimize.zeros.RootResults at 0xa861286c>),
 (72.98658083834341, <scipy.optimize.zeros.RootResults at 0xa861288c>),
 (82.75638321495505, <scipy.optimize.zeros.RootResults at 0xa86128ac>),
 (93.73016750496367, <scipy.optimize.zeros.RootResults at 0xa86128cc>),
 (106.25045004489637, <scipy.optimize.zeros.RootResults at 0xa86128ec>),
 (120.76612665626851, <scipy.optimize.zeros.RootResults at 0xa861290c>),
 (137.88917389176325, <scipy.optimize.zeros.RootResults at 0xa861292c>),
 (158.4854607193551, <scipy.optimize.zeros.RootResults at 0xa861294c>),
 (183.82941862839408, <scipy.optimize.zeros.RootResults at 0xa861296c>)]

In [490]: [a[1].iterations for a in A]
Out[490]: [8, 9, 10, 10, 10, 10, 10, 9, 8, 10]

In the brentq docs f returns one value, for one set of args. There are some solvers, such as the ode ones, that let you define a function that takes a vector variable, and returns a matching vector derivative. It doesn't look like this root finder allows that. So you are stuck with iterating over the args values, and solving for each case. I wrote the iteration as a list comprehension. Other iteration formats are possible (for loop etc). We might even be able the wrap this brentq call in a function that could be passed through np.vectorize. But that is still going to be a iteration with minor time savings.

---

There are various ways of handling a multidimensional array. One simple one is to flatten the input, do the 1d iteration, and then reshape the result. For example:

In [517]: q1=q.reshape(2,5)

In [518]: q1
Out[518]: 
array([[ 1.        ,  1.22222222,  1.44444444,  1.66666667,  1.88888889],
       [ 2.11111111,  2.33333333,  2.55555556,  2.77777778,  3.        ]])

In [519]: np.array([optimize.brentq(f, 10, 435.1, args=(i,4351.041,4262.570,0.206)) for i in q1.flat]).reshape(q1.shape)
Out[519]: 
array([[  55.99858839,   64.14621536,   72.98658084,   82.75638321,
          93.7301675 ],
       [ 106.25045004,  120.76612666,  137.88917389,  158.48546072,
         183.82941863]])

I left off the full_output flag, since that adds complications.