Dynamical System Discovery with SINDy

Run in Colab….

Introduction to SINDy

• Given data \(\hat{x}\) we want to find the analytical expression of the dynamical system \(\dot{x}(t) = f(x(t))\) that best explains the data.

• For this we have mesurements of all dimensions \(x_1,...,x_n\) for time points \(t_1,...,t_m\)

• These can be put into a martrix $\( X = \begin{bmatrix} x_1(t_1) & x_2(t_1) & \cdots & x_n(t_1) \\ x_1(t_2) & x_2(t_2) & \cdots & x_n(t_2) \\ \vdots & \vdots & & \vdots \\ x_1(t_m) & x_2(t_m) & \cdots & x_n(t_m) \end{bmatrix}. \)$

• Using difference quotients \(\dot{x}_k(t_i) \approx \frac{x_k(t_{i+1}) - x_k(t_i)}{t_{i+1} - t_i}\) we can also estimate the left-hand-side of our ode and agragate all these terms into the matrix $\( \dot{X} = \begin{bmatrix} \dot{x}_1(t_1) & \dot{x}_2(t_1) & \cdots & \dot{x}_n(t_1) \\ \dot{x}_1(t_2) & \dot{x}_2(t_2) & \cdots & \dot{x}_n(t_2) \\ \vdots & \vdots & & \vdots \\ \dot{x}_1(t_m) & \dot{x}_2(t_m) & \cdots & \dot{x}_n(t_m) \end{bmatrix}. \)$

• We can also set up a library of possible terms, e.g. monomials of order

\[\begin{split} \Theta(X) = \begin{bmatrix} \mid & \mid & & \mid \\ \theta_1(X) & \theta_2(X) & \cdots & \theta_\ell(X) \\ \mid & \mid & & \mid \end{bmatrix}. \end{split}\]

For example, if \(\theta_1(x), \theta_2(x), \dots, \theta_\ell(x)\) are monomials (\(\theta_i(x) = x^{i-1}\)), then

\[\begin{split} \theta_3(x) = \begin{bmatrix} \mid & \mid & & \mid & \mid & & \mid\\ x_1(t)^2 & x_1(t)x_2(t) & \cdots & x_2(t)^2 & x_2(t)x_3(t) & \cdots & x_n(t)^2\\ \mid & \mid & & \mid & \mid & & \mid \end{bmatrix}, \end{split}\]

where vector products and powers are understood to be element-wise.

\[\begin{split} \Xi = \begin{bmatrix} \mid & \mid & & \mid\\ \xi_1 & \xi_2 & \cdots & \xi_n\\ \mid & \mid & & \mid \end{bmatrix}. \end{split}\]

Finally, the approximation problem underlying SINDy is

\[ \dot{X} \approx \Theta(X)\Xi. \]

Install and Load Packages

Run the cells below to install and run everything needed for the rest of the Tutorial.

# Install necessary packages if running in Colab
!pip install pysindy matplotlib pandas numpy scipy scikit-learn

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

from sklearn.linear_model import LinearRegression, Lasso
import pysindy as ps

Lotka-Volterra Model

The Lotka-Volterra model, also called preditor-prey model describes the population size of a preditor population \(y\) and a prey population \(x\) as a mean-field model. The following set of differential equations describes the dynamics of the population sizes:

\begin{align} \dot{x} &= \alpha x - \beta x y,\ \dot{y} &= -\gamma y + \delta x y. \end{align}

Here the parameters have the following meaning:

• \(\alpha\) is the per-capita growth rate of prey animals, i.e. how many new prey are born per capita and unit of time,

• \(\beta\) is the effect of the presence of preditors on prey death rate,

• \(\gamma\) is the per-capita death rate of the preditor population,

• \(\delta\) is the effect of the presence of prey on the preditor growth rate.

The cell below defines the Lotka-Volterra equations.

def lotka_volterra(t, z, alpha, beta, gamma, delta):
    """
    Defines the Lotka-Volterra system of differential equations.

    Args:
        t: Time (scalar).
        z: A list or array containing the state variables [x, y].
        alpha: Rate of prey reproduction.
        beta: Rate of prey-predator interaction leading to prey death.
        gamma: Rate of predator death.
        delta: Rate of prey-predator interaction leading to predator reproduction.

    Returns:
        A NumPy array containing the derivatives [dx/dt, dy/dt].
    """
    x, y = z
    dxdt = alpha * x - beta * x * y
    dydt = -gamma * y + delta * x * y
    return np.array([dxdt, dydt])

We can use this function to generate mock data. Feel free to play around with the parameters.

tspan = (0, 50)
z0 = [10, 5]
alpha = 1.1
beta = 0.4
gamma = 0.2
delta = 0.1
t_eval = np.linspace(tspan[0], tspan[1], 500)

sol = solve_ivp(
    fun=lambda t, z: lotka_volterra(t, z, alpha, beta, gamma, delta),
    t_span=tspan,
    y0=z0,
    t_eval=t_eval
)

time = sol.t
data = sol.y.T

Using matplotlib we’ll visualize the data.

plt.figure(figsize=(10, 6))
plt.plot(time, data[:, 0], label='Prey')
plt.plot(time, data[:, 1], label='Predator')
plt.xlabel('Time')
plt.ylabel('Population')
plt.title('Lotka-Volterra Model Simulation')
plt.legend()
plt.show()

A basic SINDy

Let us implement a basic SINDy ourselves. We will follow the following steps.

1. Get data (which we already did by simulating).

2. Compute estimated derivatives using the difference quotient \(\dot{x}(t_i) \approx \frac{x(t_{i+1}) - x(t_i)}{t_{i+1} - t_i}\). Make sure it has the same dimensions as our data.

Xdot_sim = ...  # TODO: calculate the difference quotient, account for array size

Reveal

Xdot_sim = (data[1:,:] - data[:-1,:])/(time[1:] - time[:-1])[:, np.newaxis]
Xdot_sim = np.concatenate((Xdot_sim,Xdot_sim[-2:-1,:]), axis=0)

X, Y = data[:, 0], data[:, 1]
Theta = np.column_stack([
      # TODO: calclulate monomes up to second order
      # constant
      # x
      # y
      # x^2
      # xy
      # y^2
])

feature_names = ['1', 'x', 'y', 'x^2', 'xy', 'y^2']

Reveal

Theta = np.column_stack([
    np.ones_like(X),     # constant
    X,                   # x
    Y,                   # y
    X**2,                # x^2
    X*Y,                 # xy
    Y**2                 # y^2
])

4. Least squares regression for each variable

coefs = []
for i in range(2):  # prey and predator
    lr = LinearRegression(fit_intercept=False)
    lr.fit(
        # TODO: fit basis to derivative
    )

coefs = np.array(coefs)

def print_sindy(coefs, name="model"):
    print("-----------------------")
    print(name+":")
    for i, var in enumerate(['Prey', 'Predator']):
        terms = []
        for c, name in zip(coefs[i], feature_names):
            if c != 0:
                terms.append(f"{c:+.2f}*{name}")
        eq = " ".join(terms)
        print(f"d{var}/dt = {eq}")
    print("-----------------------")

print_sindy(coefs, "Barebones SINDy without sparsity, aka INDy")

# 5. Apply thresholding for sparsity
threshold = ...  # Tune this threshold
coefs[np.abs(coefs) < threshold] = 0

print_sindy(coefs, "Barebones SINDy")

Reveal

    lr.fit(Theta, Xdot_sim[:, i])
    coefs.append(lr.coef_)

threshold = 0.09

To see if this was successful, we will simulate the inferred model and plot this model against the data.

# Define the discovered system as a function
def discovered_system(t, z, coefs, feature_names):
    x, y = z
    # TODO: Recreate the Theta matrix for the current state (x, y)

# Simulate the discovered system
# Use the same time points and initial conditions as the original simulation
sol_discovered = solve_ivp(
    # TODO: integrate discoverd system with the learned coefs
)

sim_barebones = sol_discovered.y.T

Reveal

def discovered_system(t, z, coefs, feature_names):
    x, y = z
    Theta_current = np.array([1, x, y, x**2, x*y, y**2])
    dxdt = np.dot(Theta_current, coefs[0])
    dydt = np.dot(Theta_current, coefs[1])
    return np.array([dxdt, dydt])


sol_discovered = solve_ivp(
    fun=lambda t, z: discovered_system(t, z, coefs, feature_names),
    t_span=tspan,
    y0=z0,
    t_eval=t_eval
)

plt.figure(figsize=(10, 6))
plt.plot(time, data[:, 0], label='Prey (Original Data)')
plt.plot(time, data[:, 1], label='Predator (Original Data)')
plt.plot(time, sim_barebones[:, 0], "c--", label="Prey (Barebones SINDy)", linewidth=3)
plt.plot(time, sim_barebones[:, 1], "m--", label="Predator (Barebones SINDy)", linewidth=3)
plt.ylim([0, 14])
plt.xlabel('Time')
plt.ylabel('Population')
plt.title('Barebones SINDy Simulation vs. Original Data')
plt.legend()
plt.show()

PySINDy

We saw that our barebones SINDy implementation did infer a model very similar to the data, but not quite. So now we will try again with PySINDy. We will follow the same steps but using more sophisticated tools. PySINDy has different implementations of estimating the derivates for our data. We will chose FiniteDifference but go to the second order this time:

differentiation_method = ...

Reveal

differentiation_method = ps.FiniteDifference(order=2)

PySINDy also comes with predefined feature libraries. Many are possible but here we will stay with PolynomialLibrary to order 3.

feature_library = ...

Reveal

feature_library = ps.PolynomialLibrary(degree=3)

These are both higher ordered versions of what we did above. But the real advantage are the different optimizers that are implemented in PySINDy. For now we will use the Sequentially thresholded least squares algorithm STLSQ with the threshold = 0.01.

optimizer = ...

Reveal

optimizer = ps.STLSQ(threshold=0.01)

Now, we put everything together into a SINDy model.

model = ps.SINDy(
    # TODO: give the model differentiation method, feature library, and and optimizer
)

Reveal

model = ps.SINDy(
    differentiation_method=differentiation_method,
    feature_library=feature_library,
    optimizer=optimizer,
)

When fitting the model to data we can also name the variables of our ODE for nicer outputs

model.fit(
    # TODO: fit the model to the data and give the variable names
)

Reveal

model.fit(
    data,
    t=time,
    feature_names=["x", "y"],
)

We can now print the set of equations that PySINDy found for us.

model.print()

This closer to the ground truth than our barebones version. We can visualize by first simulating and then plotting with matplotlib again.

sim = ...

Reveal

sim = model.simulate(z0,t=time)

plt.figure(figsize=(10, 6))
plt.scatter(time, data[:, 0], label='Prey')
plt.scatter(time, data[:, 1], label='Predator')
plt.plot(time, sim[:, 1], "m--", label="SINDy model, preditor", linewidth=3)
plt.plot(time, sim[:, 0], "c--", label="SINDy model, pray", linewidth=3)
plt.xlabel("x")
plt.ylabel("y")
plt.legend()
plt.show()

FitzHugh–Nagumo

The FitzHugh-Nagumo model describes the prototype of an excitable system e.g. a neuron.

\begin{align} \dot{v} &= v-\frac{v^3}{3}-w+I_\text{ext},\ \tau \dot{w} &= v + a - bw. \end{align}

Where

• v is the fast activator variable, e.g. the membrane potential

• w is the slow recovery variable, e.g. the sodium channel reactivation

• a controlls the excitability threshold

• b controlls the coupling of the recovery back onto itself

• \(I_\text{ext}\) is an external stimulus current

• \(\tau\) seperates time scales of \(v\) and \(w\).

def fitzhugh_nagumo(t, z, a, b, tau, I_ext):
    """
    Defines the FitzHugh-Nagumo system of differential equations.

    Args:
        t: Time (scalar).
        z: A list or array containing the state variables [v, w].
        a: Parameter 'a'.
        b: Parameter 'b'.
        tau: Parameter 'tau'.
        I_ext: External current term.

    Returns:
        A NumPy array containing the derivatives [dv/dt, dw/dt].
    """
    v, w = z
    dvdt = v - (v**3)/3 - w + I_ext
    dwdt = (v + a - b*w) / tau
    return np.array([dvdt, dwdt])

With this set of differential equations we can generate data.

# Standard parameter values for FitzHugh-Nagumo
a = 0.7
b = 0.8
tau = 12.5
I_ext = 0.4

# Initial conditions and time span for simulation
z0_fn = [-1.0, 1.0]
tspan_fn = (0, 100)
t_eval_fn = np.linspace(tspan_fn[0], tspan_fn[1], 1000)

# Simulate the FitzHugh-Nagumo system
sol_fn = solve_ivp(
    # TODO: integrate the FitzHugh-Nagumo model and evaluate at t_eval_fn
)

Reveal

sol_fn = solve_ivp(
    fun=lambda t, z: fitzhugh_nagumo(t, z, a, b, tau, I_ext),
    t_span=tspan_fn,
    y0=z0_fn,
    t_eval=t_eval_fn
)

Real-world data has noise. So we will add a small amount of gaussian noise to this data.

time_fn = sol_fn.t
data_fn = ... # TODO: Bring the data in form and add a small amount noise

Reveal

data_fn = sol_fn.y.T + np.random.normal(0, 0.01, sol_fn.y.shape).T

# You can display or plot the data here if you like
plt.figure(figsize=(10, 6))
plt.plot(time_fn, data_fn[:, 0], label='v (Membrane Potential)')
plt.plot(time_fn, data_fn[:, 1], label='w (Recovery Variable)')
plt.xlabel('Time')
plt.ylabel('State Variable')
plt.title('FitzHugh-Nagumo Model Simulation')
plt.legend()
plt.grid(True)
plt.show()

SINDy

We will set up PySINDy the same as before.

feature_library_fn = ...

Reveal

feature_library_fn = ps.PolynomialLibrary(degree=3)

optimizer_fn = ...

Reveal

optimizer_fn = ps.STLSQ(threshold=0.04)

model_fn = ... # TODO: set up the model. We will reuse differentiation_method from before

Reveal

model_fn = ps.SINDy(
    differentiation_method=differentiation_method,
    feature_library=feature_library_fn,
    optimizer=optimizer_fn,
)

model_fn.fit(
    # TODO: fit the model
)

Reveal

model_fn.fit(
    data_fn,
    t=time_fn,
    feature_names=["v", "w"],
)

model_fn.print()

sim_fn = model_fn.simulate(z0_fn,t=time_fn)

Reveal

Xdot_sim = (data[1:,:] - data[:-1,:])/(time[1:] - time[:-1])[:, np.newaxis]
Xdot_sim = np.concatenate((Xdot_sim,Xdot_sim[-2:-1,:]), axis=0)

plt.figure(figsize=(10, 6))
plt.plot(time_fn, data_fn[:, 0], label='v (Membrane Potential)')
plt.plot(time_fn, data_fn[:, 1], label='w (Recovery Variable)')
plt.plot(time_fn, sim_fn[:, 1], "m--", label="SINDy model, v", linewidth=3)
plt.plot(time_fn, sim_fn[:, 0], "c--", label="SINDy model, w", linewidth=3)
plt.xlabel('Time')
plt.ylabel('State Variable')
plt.title('FitzHugh-Nagumo Model Simulation')
plt.legend()
plt.grid(True)
plt.show()

Analyzing the results

We will compare the derivatives, that we estimated from the data, with the ones SINDy found. We can save the estimated derivates by calling our differentiation_method on the data.

X_dot_computed = ...

Reveal

X_dot_computed = differentiation_method(data_fn, t = time_fn)

Calling the predict method of our fitted model returns the right-hand side of the learned model.

X_dot_predicted = ...

Reveal

X_dot_predicted = model_fn.predict(data_fn)

fig, axs = plt.subplots(data_fn.shape[1], 1, sharex=True, figsize=(7, 9))
for i in range(data_fn.shape[1]):
    axs[i].plot(time_fn, X_dot_computed[:, i], "k", label="numerical derivative")
    axs[i].plot(time_fn, X_dot_predicted[:, i], "r--", label="model prediction")
    axs[i].legend()
axs[0].set(xlabel="t", ylabel=r"$\dot v$")
axs[1].set(xlabel="t", ylabel=r"$\dot w$")
fig.show()

Predictions

We would also like the model to generalize to a different set of initial conditions. Let’s start some place different and simulate both the ground truth and the learned model.

z0_pred = ...
sol_pred = solve_ivp(
    # TODO: integrate the FitzHugh-Nagamu model with new initial conditions
    )
pred_fn = ... # TODO: simulate the identified system

Reveal

z0_pred = np.random.randn(2)
sol_pred = solve_ivp(
    fun=lambda t, z: fitzhugh_nagumo(t, z, a, b, tau, I_ext),
    t_span=tspan_fn,
    y0=z0_pred,
    t_eval=time_fn)
pred_fn = model_fn.simulate(z0_pred,t=time_fn)

plt.figure(figsize=(10, 6))
plt.plot(time_fn, sol_pred.y[0,:], label='v (Membrane Potential)')
plt.plot(time_fn, sol_pred.y[1,:], label='w (Recovery Variable)')
plt.plot(time_fn, pred_fn[:, 1], "m--", label="SINDy model, v", linewidth=3)
plt.plot(time_fn, pred_fn[:, 0], "c--", label="SINDy model, w", linewidth=3)
plt.xlabel('Time')
plt.ylabel('State Variable')
plt.title('FitzHugh-Nagumo Model Simulation')
plt.legend()
plt.grid(True)
plt.show()