Conditional curve fit with scipy?

问题

Let's say I want to fit a straight line to my data recorded with lights off. Now I accidentally left the lights on, and my data has a constant offset from datapoint 101 and onwards.

How can I fit this? I've tried to incorporate a condition for x, but I get the error

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Remember to uncomment the remainder of the code (to encounter the error).

import numpy as np
from scipy import optimize
import matplotlib.pyplot as plt

d1 = np.random.normal(0,0.1, 100)
d2 = np.random.normal(3,0.1, 100)

x = np.arange(0,200)
y = np.concatenate((d1,d2))

plt.plot(x, y)

# def line(x, a, b, offset):
#     if x < 101:
#         y = a * x + b
#     else:
#         y = (a * x + b) + offset
#     return y
# 
# popt, pcov = optimize.curve_fit(line, xdata = x, ydata = y)
# 
# plt.plot(x, line(x, *popt), color = "firebrick")
plt.show()

Expected output:

回答1:

I think the standard trick would be to convert boolean condition to an integer factor:

def line(x, a, b, offset):
    return (a * x + b) + offset * (x>100)

回答2:

The reason you are getting that error is that optimize is calling your line function by passing it an array of values, not just a single value. To fix this issue, your line function must me able to handle an array of values. Fortunately, numpy has a function to help you.

def line(x, a, b, offset):
    return np.piecewise(x, 
                        [x < 101, x >= 101],
                        [lambda x: a * x + b, lambda x: a * x + b + offset])

I should note that it still does not converge but that is a different issue.

来源：https://stackoverflow.com/questions/47892502/conditional-curve-fit-with-scipy