Non-linear least square for piecewise linear function

nnls-piecewise.py

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

x = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10 ,11, 12, 13, 14, 15], dtype=float)
y = np.array([5, 7, 9, 11, 13, 15, 28.92, 42.81, 56.7, 70.59, 84.47, 90, 80, 70, 60])

def piecewise_linear(x, y0, y1, b0, b1, b2):
    x0 = 6
    x1 = 12
    return np.piecewise(x, 
                        [x < x0, 
                         (x >= x0) & (x < x1), 
                         x >= x1], 
                        [lambda x: b0*x + y0, 
                         lambda x: b1*x + y1-b1*x1,
                         lambda x: b2*x + y1-b2*x1])
  
p, e = optimize.curve_fit(piecewise_linear, x, y)
  
new_x = np.linspace(0, 15, 100)
plt.plot(x, y, "o")
plt.plot(new_x, piecewise_linear(new_x, *p))