Non-linear Curve Fitting

If you are running this notebook locally, make sure you’ve followed steps here to set up the environment. (This environment.yml file specifies a list of packages required to run the notebooks)

In this notebook, we use cofi to run a non-linear curve fitting problem:

\[f(x)=\exp(a*x)+b\]

Import modules

# -------------------------------------------------------- #
#                                                          #
#     Uncomment below to set up environment on "colab"     #
#                                                          #
# -------------------------------------------------------- #

# !pip install -U cofi

import numpy as np
import matplotlib.pyplot as plt
import arviz as az

from cofi import BaseProblem, InversionOptions, Inversion

np.random.seed(42)

Define the problem

def my_forward(m, x):
    return np.exp(m[0] * x) + m[1]

def my_jacobian(m, x):
    G=np.zeros([len(x),2])
    G[:,0]=x*np.exp(m[0]*x) # derivative with respect to m[0]
    G[:,1]=np.ones(len(x))  # derivtavie with respect to m[1]
    return G

def my_residuals(m, x, y):
    yhat = my_forward(m,x)
    return yhat-y

# Choose the "true" parameters.
a_true = 5.0
b_true = 4.0
f_true = 0.1

m_true = [a_true,b_true]
mf_true= [a_true,b_true,f_true]

# Generate some synthetic data from the model.
N = 50
x = np.sort(1 * np.random.rand(N))
yerr = 0.1 + 0.5 * np.random.rand(N)
y = my_forward(m_true,x)
y += np.abs(f_true * y) * np.random.randn(N)
y += yerr * np.random.randn(N)
plt.errorbar(x, y, yerr=yerr, fmt=".k", capsize=0)
x0 = np.linspace(0, 1, 500)
plt.plot(x0, my_forward(m_true,x0), "k", alpha=0.3, lw=3)
plt.xlim(0, 1)
plt.xlabel("x")
plt.ylabel("y");

Text(38.222222222222214, 0.5, 'y')

# define the problem in cofi
inv_problem = BaseProblem()
inv_problem.name = "Curve Fitting"
inv_problem.set_data(y)
inv_problem.set_forward(my_forward, args=[x])
inv_problem.set_jacobian(my_jacobian, args=[x])
inv_problem.set_residual(my_residuals, args=[x,y])
inv_problem.set_initial_model([3,3])

Example 1. least squares optimizer (levenber marquardt)

inv_options = InversionOptions()
inv_options.set_tool("scipy.optimize.least_squares")
inv_options.set_params(method="lm", max_nfev=10)

######## Run it
inv = Inversion(inv_problem, inv_options)
inv_result = inv.run()

######## Check result
print(f"The inversion result from `scipy.optimize.minimize`: {inv_result.model}\n")
inv_result.summary()

The inversion result from `scipy.optimize.minimize`: [5.06442618 3.54842171]

============================
Summary for inversion result
============================
SUCCESS
----------------------------
cost: 751.5703778228752
fun: [ 8.46834890e-02 -1.77231040e-02 -5.52853301e-01  8.89806495e-01
  2.91152911e-01 -6.80792325e-01 -1.14702072e+00 -2.15801098e-01
  1.82952932e-01 -5.26482038e-01 -7.76017787e-01 -5.59530390e-01
 -4.95847940e-01 -4.13394800e-01 -5.36314278e-01 -1.56467760e+00
  4.20608340e-01 -1.91245192e-01 -7.95757156e-02  4.30437719e-01
 -1.36307950e-02 -3.20414165e-01 -3.61292261e-01 -1.97016385e-01
  1.47256652e+00  1.95462591e-01  6.42560473e-01  1.17710109e+00
  1.82720274e-01 -5.85651739e-01 -4.32433162e+00 -4.33451436e-01
  1.59206971e-02  4.24747091e-01  5.23801008e+00  2.40244375e-01
 -2.85673023e-01 -6.65912029e+00  1.06971709e+00 -1.41328840e-01
  1.44236334e+00  7.70525926e+00 -4.25388811e+00 -1.75601283e+00
 -1.98652705e+00  1.44619318e+01 -9.86284706e+00  2.35903629e+01
 -2.98371239e-02 -2.11903104e+01]
jac: [[2.28462443e-02 1.00000000e+00]
 [4.09307227e-02 1.00000000e+00]
 [5.87699128e-02 1.00000000e+00]
 [7.79481395e-02 1.00000000e+00]
 [9.04348484e-02 1.00000000e+00]
 [1.60175374e-01 1.00000000e+00]
 [2.26419173e-01 1.00000000e+00]
 [2.82725638e-01 1.00000000e+00]
 [3.43725582e-01 1.00000000e+00]
 [3.43820726e-01 1.00000000e+00]
 [4.04431999e-01 1.00000000e+00]
 [4.56634655e-01 1.00000000e+00]
 [4.64300862e-01 1.00000000e+00]
 [4.71420526e-01 1.00000000e+00]
 [5.48900988e-01 1.00000000e+00]
 [6.22385905e-01 1.00000000e+00]
 [9.59632427e-01 1.00000000e+00]
 [1.27285650e+00 1.00000000e+00]
 [1.28279178e+00 1.00000000e+00]
 [1.42031878e+00 1.00000000e+00]
 [1.42473141e+00 1.00000000e+00]
 [1.51128333e+00 1.00000000e+00]
 [2.34264052e+00 1.00000000e+00]
 [2.49621211e+00 1.00000000e+00]
 [3.85009064e+00 1.00000000e+00]
 [4.08975791e+00 1.00000000e+00]
 [4.59341502e+00 1.00000000e+00]
 [6.07964753e+00 1.00000000e+00]
 [6.95336936e+00 1.00000000e+00]
 [7.24310829e+00 1.00000000e+00]
 [7.48401301e+00 1.00000000e+00]
 [8.71405798e+00 1.00000000e+00]
 [1.19018190e+01 1.00000000e+00]
 [1.24136629e+01 1.00000000e+00]
 [1.26206405e+01 1.00000000e+00]
 [1.31778419e+01 1.00000000e+00]
 [1.35640163e+01 1.00000000e+00]
 [1.89839544e+01 1.00000000e+00]
 [2.18847690e+01 1.00000000e+00]
 [2.55534573e+01 1.00000000e+00]
 [2.98190140e+01 1.00000000e+00]
 [4.18720384e+01 1.00000000e+00]
 [4.84904621e+01 1.00000000e+00]
 [5.63991081e+01 1.00000000e+00]
 [6.96176523e+01 1.00000000e+00]
 [9.09334829e+01 1.00000000e+00]
 [1.15942416e+02 1.00000000e+00]
 [1.17246757e+02 1.00000000e+00]
 [1.28432017e+02 1.00000000e+00]
 [1.31826241e+02 1.00000000e+00]]
grad: [1.72526347e-04 1.33278988e-09]
optimality: 0.0001725263465452648
active_mask: [0 0]
nfev: 8
njev: 6
status: 2
message: `ftol` termination condition is satisfied.
model: [5.06442618 3.54842171]

Example 2. emcee

sigma = 10                                     # common noise standard deviation
Cdinv = np.eye(len(y))/(sigma**2)      # inverse data covariance matrix

def my_log_likelihood(m,x,y,Cdinv):
    yhat = my_forward(m,x)
    residual = y-yhat
    return -0.5 * residual @ (Cdinv @ residual).T

m_min = [0,0]             # lower bound for uniform prior
m_max = [10,10]          # upper bound for uniform prior

def my_log_prior(m,m_min,m_max):    # uniform distribution
    for i in range(len(m)):
        if m[i] < m_min[i] or m[i] > m_max[i]: return -np.inf
    return 0.0 # model lies within bounds -> return log(1)

nwalkers = 12
ndim = 2
nsteps = 500
walkers_start = np.array([5.,4.]) + 1e-1 * np.random.randn(nwalkers, ndim)

inv_problem.set_log_prior(my_log_prior,args=[m_min,m_max])
inv_problem.set_log_likelihood(my_log_likelihood,args=[x,y,Cdinv])
inv_problem.set_model_shape(ndim)

inv_options = InversionOptions()
inv_options.set_tool("emcee")
inv_options.set_params(nwalkers=nwalkers, nsteps=nsteps, initial_state=walkers_start)

######## Run it
inv = Inversion(inv_problem, inv_options)
inv_result = inv.run()

######## Check result
print(f"The inversion result from `emcee`:")
inv_result.summary()

The inversion result from `emcee`:
============================
Summary for inversion result
============================
SUCCESS
----------------------------
sampler: <emcee.ensemble.EnsembleSampler object>
blob_names: ['log_likelihood', 'log_prior']

sampler = inv_result.sampler
az_idata = inv_result.to_arviz()

labels = ["m0", "m1"]
#az.plot_trace_dist(az_idata, visuals={"xlabel_trace": False});
az.plot_trace_dist(az_idata, visuals={"xlabel_trace": False, "trace": {"lw": 0.5}, "dist": {"lw": 0.5}});

<arviz_plots.plot_collection.PlotCollection object at 0x7efdbd251370>

pm = az.plot_pair(
    az_idata.sel(draw=slice(300, None)),
    marginal=True,
    triangle="lower",
    visuals={"scatter": {"s": 3}},
)

# Add reference dots for true model values
ref_values = list(m_true)
n = len(ref_values)
for i in range(n):
    for j in range(n):
        try:
            ax = pm.iget_target(i, j)
        except (ValueError, IndexError):
            continue
        if i == j:
            ax.axvline(ref_values[i], color="red", linestyle="--", lw=1, alpha=0.5)
        elif i > j:
            ax.plot(
                ref_values[j], ref_values[i], "o",
                color="red", markeredgecolor="black", markeredgewidth=1.5, ms=5, zorder=5,
            )

flat_samples = sampler.get_chain(discard=300, thin=30, flat=True)
inds = np.random.randint(len(flat_samples), size=100) # get a random selection from posterior ensemble
_x_plot = np.linspace(0,1.0)
_y_plot =  my_forward(m_true,_x_plot)
plt.figure(figsize=(12,8))
sample = flat_samples[0]
_y_synth =  my_forward(sample,_x_plot)
plt.plot(_x_plot, _y_synth, color="seagreen", label="Posterior samples",alpha=0.1)
for ind in inds:
    sample = flat_samples[ind]
    _y_synth =  my_forward(sample,_x_plot)
    plt.plot(_x_plot, _y_synth, color="seagreen", alpha=0.1)
plt.plot(_x_plot, _y_plot, color="darkorange", label="true model")
plt.scatter(x, y, color="lightcoral", label="observed data")
plt.xlabel("X")
plt.ylabel("Y")
plt.legend();

<matplotlib.legend.Legend object at 0x7efd94226490>

---

Watermark

watermark_list = ["cofi", "numpy", "scipy", "matplotlib", "emcee", "arviz"]
for pkg in watermark_list:
    pkg_var = __import__(pkg)
    print(pkg, getattr(pkg_var, "__version__"))

cofi 0.2.11+71.gb28b5b0
numpy 2.2.6
scipy 1.17.1
matplotlib 3.10.9
emcee 3.1.6
arviz 1.1.0

sphinx_gallery_thumbnail_number = -1