条件曲线与 scipy 拟合？

Question

假设我想对关闭灯时记录的数据拟合一条直线。现在我不小心把灯打开了，我的数据从数据点 101 开始有一个恒定的偏移量。

我怎样才能适应这个？我尝试为 x 加入一个条件，但出现错误

ValueError：具有多个元素的数组的真值不明确。使用 a.any() 或 a.all()

请记住取消注释其余代码（以遇到错误）。

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()

预期输出：

Answer 1

我认为标准技巧是将布尔条件转换为整数因子：

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