Polynomial curve-fitting over a large 3D data set

Question

I have a list of 4 images, called listfile.list, which looks like this:

image1
image2
image3
image4

Each image has 10 frames containing a 2000 x 2000 array of pixels, so the size of each image is [10,2000,2000]. The pixel value for each frame increases from 0 to 10, so for example for one pixel in image1:

image1[:,150,150] = [435.8, 927.3, 1410. , 1895.1, 2374.6, 2847.1, 
                     3340.5, 3804.8, 4291.6, 4756.1]

The other pixels show similar values and a roughly linear increase across the frames. I need to fit several functions to pixel values across the frames for each pixel in each image and then average over the images. The data sets are 3D and very large, so I don't know how to paste them here. My code does exactly what I want it to, the issue is that it can take days (> 3) to run my full script.

One function is frame_fit to return rates and intercepts. There are several other functions. My code is structured as follows:

import itertools
import numpy as np
from scipy.optimize import curve_fit

def frame_fit(xdata, ydata, poly_order):
    '''Function to fit the frames and determine rate.'''

    # Define polynomial function. Here, 'b' will be ideal rate and
    # 'c', 'd', 'e', etc. describe magnitudes of deviation from
    # ideal rate.
    if poly_order == 5:
        def func(t, a, b, c, d, e, f):
            return a+ b*t+ c*(b*t)**2+ d*(b*t)**3+ e*(b*t)**4+ f*(b*t)**5

    elif poly_order == 4:
        def func(t, a, b, c, d, e):
            return a + b*t + c*(b*t)**2 + d*(b*t)**3 + e*(b*t)**4

    # Initial values for curve-fitting.
    initvals = np.array([100, 4.e+01, 7.e-03, -6.e-06, 3.e-08, -8.e-11])

    # Provide uncertainty estimate
    unc = np.sqrt(64 + ydata)

    beta, pcov = curve_fit(func, xdata, ydata,
                           sigma=unc, absolute_sigma=False,
                           p0=initvals[:poly_order+1],
                           maxfev=20000)

    # beta[1] is rate, beta[0] is intercept

    return beta[1], beta[0]


all_rates = np.zeros((number_of_exposures, 2000, 2000),dtype=np.float64)
all_intercepts = np.zeros((number_of_exposures, 2000, 2000),dtype=np.float64)
all_results = np.zeros((2, 2000, 2000),dtype=np.float64)
pix_x_min = 0
pix_y_min = 0
pix_x_max = 2000    # max number of pixels in one direction
pix_y_max = 2000
xdata = np.arange(0,number_of_frames)    # so xdata [0,1,2,3,4,5,6,7,8,9]


# Here is where I need major speed improvements

for exposure in list_of_exposures:

    for i,j in itertools.product(np.arange(pix_x_min,pix_x_max),
                                     np.arange(pix_y_min,pix_y_max)): 

        ydata = exposure[:,i,j]

        rate, intercept = frame_fit(xdata,ydata,5)

        # Plus other similar functions ... 
        # results2, residuals2 = function2(rate, intercept)
        # results3, residuals3 = function3(results2, specifications)

        all_rates[exposure,i,j] = rate
        all_intercepts[exposure,i,j] = intercept

avg_rates = np.average(all_rates, axis=0)
avg_intercepts np.average(all_intercepts, axis=0)

all_results[0,:,:] = avg_rates
all_results[1,:,:] = avg_intercepts          

# all_results is saved to a file after everything is finished. 

There are usually at least 4 images with 2000x2000 arrays of pixels in my input list. Besides the frame_fit function, there are several others that I have to run for each pixel. I am a relatively new Python programmer so I often don't know about all available tools or best practices to improve speed. I know a nested for-loop as I've written it is going to be slow. I have considered cython or multiprocessing, but those are complicated to tackle for a new programmer. Is there a better way to perform a function on each pixel and then average over the results? I am hoping to stick with standard python 3.5 modules (rather than installing extra packages).

Answer 1

In terms of possible speed optimisations there's not much to work with, but one thing does stand out:

        def func(t, a, b, c, d, e, f):
            return a+ b*t+ c*(b*t)**2+ d*(b*t)**3+ e*(b*t)**4+ f*(b*t)**5

(and similarly for the four-element case). It's fairly standard knowledge that polynomials evaluation can be optimised for speed using Horner's method. I'll also pull out that multiplication by b:

        def func(t, a, b, c, d, e, f):
            bt = b*t
            return a + bt*(1 + bt*(c + bt*(d + bt*(e + bt*f))))