casuality-2.py

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

def func(x, a, b, c):
    x[x < b] = b
    return a * np.sqrt(x - b) + c

sample = 500
X = np.random.normal(5, 1, sample)
X = np.sort(X)
N = np.random.uniform(-1, 1, sample)
Y = X * X + N

p = np.polyfit(X, Y, 2)

plt.plot(X, Y, 'o')
plt.plot(X, np.polyval(p, X))
plt.show()

plt.plot(X, Y - np.polyval(p, X), 'o')
plt.plot([np.min(X), np.max(X)], [0, 0])
plt.show()

popt, pcov = curve_fit(func, Y, X)
x_pred = func(Y, *popt)
plt.plot(X, Y, 'o')
plt.plot(x_pred, Y)
plt.show()

plt.plot(X - x_pred, Y, 'o')
plt.plot([0, 0], [np.min(Y), np.max(Y)])
plt.show()