import matplotlib.pyplot as plt
import numpy as np
import matplotlib.gridspec as gridspec
import scipy.special
from scipy.optimize import curve_fit
import matplotlib.patches as patches
import seaborn as sns
set(style="ticks", font_scale=1.5) sns.
Least Squares
How to fit a nonlinear function to data
Introduction
This code produces the figure above. It’s main tool is the curve_fit
method, that allows us to fit any function to data, and get optimal parameter values.
The code
Define some functions
= np.arange(0, 12, 0.4)
x
def func(x, par0, par1, par2):
return par0 + np.cos(par1 * x + par2)
def add_rec(ax, c, v, col):
ax.add_patch(
patches.Rectangle(# (x,y)
c, abs(v), # width
np.# height
v, =0.4,
alpha=col
color
) )
Now let’s plot some stuff
%matplotlib notebook
= plt.subplots(2, 1, figsize=(8,6), sharex=True, sharey=True,
fig, [ax1, ax2] ={'aspect':'equal'}
subplot_kw
)=0.04, right=0.98, top=0.93, bottom=0.15,
fig.subplots_adjust(left=0.05, wspace=0.02)
hspace= {#'backend': 'ps',
params 'font.family': 'serif',
'font.serif': ['Computer Modern Roman'],
'ps.usedistiller': 'xpdf',
'text.usetex': True,
# include here any neede package for latex
'text.latex.preamble': [r'\usepackage{amsmath}'],
}# the parameter values
= (1.3, 2, 1)
par # generating data with noise
= func(x, *par) + (np.random.random(len(x)) - 0.5)
y ='o', ls='None', markerfacecolor="blue", markeredgecolor="black")
ax1.plot(x, y, marker='o', ls='None', markerfacecolor="red", markeredgecolor="black")
ax2.plot(x, y, marker
# best fit
= curve_fit(func, x, y, p0=(1.5, 1.5, 2.5)) # p0 = initial guess
popt, pcov = popt
p0, p1, p2 # The total sum of squares (proportional to the variance of the data)
= ((y - y.mean()) ** 2).sum()
SStot # The sum of squares of residuals
= ((y - func(x, p0, p1, p2)) ** 2).sum()
SSres = 1 - SSres / SStot
Rsquared # plot best fit
= np.linspace(x.min(), x.max(), 1001)
h = ax1.plot(h, func(h, p0, p1, p2), color='black', linewidth=2)
fit, "Best fit"], loc="upper right",
ax1.legend([fit], [=False, handlelength=4)
frameon# plot mean
= ax2.plot(h, h * 0 + np.mean(y), ls='--', color='black', linewidth=2)
mean, "Mean"], loc="upper right", frameon=False, handlelength=4)
ax2.legend([mean], [
# plot blue and red squares
for ind in np.arange(len(x)):
= x[ind]
x0 = y[ind]
y0 # print(x0,y0)
= y0 - func(x0, p0, p1, p2)
v1 = y0 - y.mean()
v2 -v1, "blue")
add_rec(ax1, (x0, y0), -v2, "red")
add_rec(ax2, (x0, y0),
0.5, 2.9, r"Total sum of squares: {:.1f}".format(SStot))
ax2.text(0.5, -0.8, r"Sum of squares of residuals: {:.1f}".format(SSres))
ax1.text(# ax2.set_xlabel(
# r"R-squared = $1 - \frac{\text{blue area}}{\text{red area}}$ = " +
# "{:.2f}".format(Rsquared))
ax2.set_xlabel(r"R-squared = $1 - \frac{\mathrm{blue\ area}}{\mathrm{red\ area}}$"
)
ax1.set_xlabel(r"Data: $f(x) = p_0 + \cos(p_1 x + p_2)+ $ noise ",
=12
labelpad
)"top")
ax1.xaxis.set_label_position(
set(xlim=[x.min(), x.max()],
ax1.=[-1, 3.5],
ylim=[],
xticklabels=np.arange(-1, 4)
yticks
)set(
ax2.=np.arange(0,13,2)
xticklabels
)
"least-squares.png", dpi=300) fig.savefig(
/Users/yairmau/opt/anaconda3/lib/python3.7/site-packages/ipykernel_launcher.py:70: UserWarning: FixedFormatter should only be used together with FixedLocator