Solve a nonlinear pendulum for a trajectory on the separatrix.

pendulum.py

# -*- coding: utf-8 -*-
"""
Solve a nonlinear pendulum for a trajectory on the separatrix.

@author: Nicolás Guarín-Zapata
@date: April 2020
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from scipy.optimize import curve_fit

#%% Setup

repo = "https://raw.githubusercontent.com/nicoguaro/matplotlib_styles/master"
style = repo + "/styles/clean.mplstyle"
plt.style.use(style)

#%% Functions
def fun(t, x):
    x0, x1 = x
    return x1, -np.sin(x0)


def fit_fun(x, a, b, c):
    return a + b*np.exp(-c*x)


def fit_jac(x, a, b, c):
    jac = np.vstack((np.ones_like(x), np.exp(-c*x), -b*x*np.exp(-c*x)))
    return jac.T


#%% ODE
npts = 1000
t_span = (0, 10)
x_ini = (0, 2.0)
t_eval = np.linspace(*t_span, npts)
sol = solve_ivp(fun, t_span, x_ini, t_eval=t_eval, method='Radau',
                rtol=1e-6, atol=1e-12)


#%% Fitting
p0 = 3, -3, 1
popt, pcov = curve_fit(fit_fun, sol.t, sol.y[0, :], jac=fit_jac)
a, b, c = popt


#%% Visualization
plt.plot(sol.t, sol.y[0, :])
plt.plot(sol.t, a + b*np.exp(-c*sol.t), lw=2, linestyle="dashed")
plt.xlabel("Dimensionless time")
plt.ylabel("Angle")
plt.legend(["ODE solution", "Exponential fit"])