scipy optimize.curve_fit cannot fit a function whose return value depends on a conditional

Question

I want to fit a function, defined as follows, to a time series data:

def func(t, a0, a1, a2, T, tau1, tau2):
    if t < T:
        return a0 + a1 * np.exp(-t/tau1) + a2 * np.exp(-t/tau2)
    else:
        return a0 + a1 * np.exp(-T/tau1) * (1 - t/tau1 + T/tau1) + a2 * np.exp(-T/tau2) * (1 - t/tau2 + T/tau2) 

Here, t represents the time at which a measurement is made, and the rest of the arguments are the parameters of the function. The problem is that when I feed it into curve_fit, Python complains about the ambiguity in the t < T comparison. I believe this happens because t becomes a list of data points when func is called inside curve_fit, whereas T is a number (not a list):

popt, pcov = curve_fit(func, t1, d1)

where t1 is a list of times and d1 is a list of the data values measured at the corresponding times. I have attempted a number of ways to get around this problem, but to no avail. Any suggestion? Thanks a lot!

Answer 1

That's right, t < T is a boolean array. NumPy refuses to assign a truth value to boolean arrays because there are many possible choices -- should it be True if all elements are True, or if any element is True?

But that's okay. In this case, NumPy provides a nice function to replace the if ... else ... blocks, namely, np.where:

def func(t, a0, a1, a2, T, tau1, tau2):
    return np.where(
        t < T,
        a0 + a1 * np.exp(-t/tau1) + a2 * np.exp(-t/tau2),
        a0 + a1 * np.exp(-T/tau1) * (1 - t/tau1 + T/tau1) + a2 * np.exp(-T/tau2) * (1 - t/tau2 + T/tau2) )