Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Solve a nonlinear pendulum for a trajectory on the separatrix.
# -*- 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"])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment