曲線擬合分段函數？

Question

正如標題所示，我試圖用分段函數進行曲線擬合。 這是一個例子：

import numpy as np
import scipy.optimize as sio

x = np.array([0, 1, 2, 3, 4, 5])
y = np.array([2, 3, 4, 11, 18, 27])

def f(x, a):
    if x <= 2:
        return x + a
    elif x >= 3:
        return x**2 + a

sio.curve_fit(f, x, y)

現在，這將返回以下錯誤：

ValueError: The truth value of an array with more than one element is ambiguous

通過添加打印函數，我發現這是因為傳遞給函數'f'的'x'值是numpy數組[0, 1, 2, 3, 4, 5]

我想知道這個問題是否有解決方法。

謝謝！

Answer 1

整體問題是你為標量寫了f ，但x和a是numpy數組。 這是if行的情況，也是if語句中的return 。

與您的代碼最相似的解決方案是：

import numpy as np
import scipy.optimize as sio

x = np.array([0, 1, 2, 3, 4, 5])
y = np.array([2, 3, 4, 11, 18, 27])

def f(x, a):
    y = np.zeros(len(x))
    for i, xi in enumerate(x): 
        if xi <= 2:
            y[i] =  xi + a
        elif xi >= 3:
            y[i] =  xi**2 + a
    return y

sio.curve_fit(f, x, y)

f更矢量化定義是：

def f(x, a):
    y = np.zeros(len(x))
    y[:3] += a
    y[3:] = x[3:]**2 + a
    return y

Answer 2

這是你可以使用np.boolean數組的問題。 如果您有大量數據，它比循環數組更有效，並允許您使用任意表達式作為if / else條件。

import numpy as np
import scipy.optimize as sio

x = np.array([0, 1, 2, 3, 4, 5])
y = np.array([2, 3, 4, 11, 18, 27])

def f(x, a):
    y = np.zeros(len(x))
    y += (x + a) * (x <= 2)
    y += (x ** 2 + a) * (x >= 3)
    return y

sio.curve_fit(f, x, y)

其中x <= 2 =

[True, True, True, False, False, False]

可以在numpy表達式中使用：

[1, 1, 1, 0, 0, 0]