PySINDy-RSF: Data-Driven Equation Discovery for a Chaotic Rate-and-State Spring-Slider

Code and results supporting the manuscript "Data-Driven Equation Discovery for a Chaotic Rate-and-State Spring-Slider" (Gualandi, Petrillo, Gao, Brunton, Kutz, Marone, Mengaldo, Dal Zilio — submitted to Journal of Geophysical Research: Machine Learning and Computation).

This repository benchmarks the Sparse Identification of Nonlinear Dynamics (SINDy) framework, and in particular the TrappingSR3 optimizer, on a two-state rate-and-state friction spring-slider that reproduces the chaotic stick-slip dynamics of Becker (2000). The code quantifies how the ability to recover the correct governing equations, calibrate their coefficients, and forecast the system's future evolution depends on the observation sampling interval, the level of measurement noise, and the fraction of a trajectory used for training.

The model

The spring-slider obeys the non-dimensional, autonomous system (see Section 2 of the manuscript for the full derivation from the dimensional Dieterich-Ruina friction law):

dx/dt = e^x [ (beta1 - 1) x + y - z ] + dy/dt - dz/dt
dy/dt = kappa (1 - e^x)
dz/dt = -rho * e^x * (beta2 * x + z)

implemented in src/odes.py as odes.becker. With (beta1, beta2, rho) = (1, 0.84, 0.048) (Gu et al., 1984) and normalized stiffness kappa ≈ 0.06844 (Becker, 2000), the system sits in a deterministic chaotic regime with leading Lyapunov exponent lambda_max ≈ 0.0111 (1/lambda_max ≈ 90.4 non-dimensional time units).

Two distinct time symbols are used throughout this repository and the manuscript: dimensional time $t$ (seconds, as in the raw Dieterich-Ruina friction law) and non-dimensional time $T := (V_*/L_1),t$, the rescaled time used on every figure axis and in every sweep script. Figures are always labeled "Non-dimensional time $T$" for this reason.

Figures

Figure 1. Schematic of the spring-slider system: a rigid block driven at constant velocity through a spring, sliding on a rate-and-state frictional interface.

Figure 2. Overview of the SINDy / SR3 workflow: noisy observations of the system state are used to build a candidate function library, from which a sparse regression step (TrappingSR3) recovers the governing equations.

Figure 3. Reference trajectory $x(T)$, $y(T)$, $z(T)$ of the noise-free Becker spring-slider, showing the chaotic stick-slip limit cycle used throughout this study.

Figure 4. Model recovery skill as a function of the observation sampling interval and measurement noise level: F1 score on the recovered library terms, coefficient error $E_{\text{coef}}$, forecast NRMSE, and predictability horizon $T_{\text{pred}}$.

Figure 5. The same four metrics as Figure 4, shown as a function of the fraction of the trajectory used for training, at fixed sampling interval and across several noise levels.

Figure 6. Example rollout and recovered coefficients for a short training window, comparing the SINDy-identified model against the reference trajectory.

Figure 7. Same as Figure 6, for a longer training window, showing the corresponding improvement in coefficient recovery and forecast accuracy.

Figure A1. Test-set $R^2$ as a function of the sparsity threshold for four optimizers (STLSQ, SR3, SSR, TrappingSR3) under 1% Gaussian observation noise.

Figure A2. Distribution of test-set $R^2$ across the threshold/alpha grid, by optimizer, under 1% Gaussian observation noise.

Repository structure

PySINDy_RSF_clean/
├── src/
│   ├── odes.py                    # Becker/Becker-QD/Lorenz right-hand sides + solve_ivp wrapper
│   ├── optimizers_trapping.py     # TrappingSR3 optimizer wrapper, compatible with pysindy 1.7.x and >= 2.0
│   ├── pysindy_compat.py          # Version-robust SINDy() construction / .fit() call
│   └── plot_style.py              # Shared matplotlib style used by every figure script
├── scripts/
│   ├── 01_run_spring_slider.py            # Runs the Becker (2000) spring-slider model; produces Figure 3
│   ├── 02_sweep_sampling_noise.py         # Figure 4 data: metrics vs sampling interval & noise
│   ├── 03_plot_figure4_heatmaps.py        # Figure 4 plot
│   ├── 04_sweep_training_fraction.py      # Figure 5 data: metrics vs training fraction
│   ├── 05_plot_figure5_lines.py           # Figure 5 plot
│   ├── 06_appendix_optimizer_threshold_scan.py  # Appendix A data: optimizer/threshold/alpha sweep, 1% noise
│   ├── 07_plot_figure6_rollout.py         # Figures 6 & 7: rollout + recovered coefficients
│   └── 08_plot_appendixA1_A2.py           # Appendix A plots
├── tests/
│   └── test_lorenz_sanity_check.py        # Sanity check: recovers the Lorenz system with SINDy
├── spring-slider_sketch/                  # Figure 1 source assets
├── SINDy_sketch/                          # Figure 2 source assets
├── figures/                                # Final manuscript figures (figure1-7.png, figureA1-2.png)
├── results/                                # Numerical outputs and vector (SVG) figures
├── requirements.txt
└── README.md

Installation

python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt

Tested with Python 3.10–3.11. No GPU is required. Each sweep script runs joblib-parallelized over CPU cores (N_JOBS = 8 by default — edit this constant to match your machine).

Reproducing the figures

All scripts are run from inside scripts/:

cd scripts

# Figures 1 and 2 are static assets (see spring-slider_sketch/ and SINDy_sketch/).

# Figure 3 — runs the Becker (2000) spring-slider and plots the reference trajectory
python 01_run_spring_slider.py

# Figure 4 — sampling interval x noise grid
python 02_sweep_sampling_noise.py
python 03_plot_figure4_heatmaps.py

# Figure 5 — training-fraction sweep
python 04_sweep_training_fraction.py
python 05_plot_figure5_lines.py

# Figures 6 & 7 — example rollouts and recovered coefficients (one run writes both)
python 07_plot_figure6_rollout.py

# Appendix A — optimizer/threshold/alpha sweep under 1% noise, then the plots
python 06_appendix_optimizer_threshold_scan.py
python 08_plot_appendixA1_A2.py

Every figure is written in two formats: a .png (raster) and a .svg (vector, editable in Inkscape, Illustrator, or Affinity Designer, and directly importable into PowerPoint/Keynote). Raster figures are rendered at 390 dpi.

02_sweep_sampling_noise.py and 04_sweep_training_fraction.py are the two most expensive scripts (tens of minutes on a laptop-class 8-core machine), since each grid cell requires fitting and simulating a model.

Appendix A environment

06_appendix_optimizer_threshold_scan.py uses PySINDy's built-in TrappingSR3 optimizer, whose constructor API was substantially reworked between PySINDy 1.7.5 and PySINDy ≥ 2.0, and which requires the optional cvxpy dependency. This script should therefore be run in a separate, pinned environment:

python3 -m venv legacy_venv
legacy_venv/bin/pip install "pysindy==1.7.5" "numpy<1.24" cvxpy pandas \
    scikit-learn scipy tqdm
legacy_venv/bin/python 06_appendix_optimizer_threshold_scan.py

The script also accepts optional <OPTIMIZER> [THRESHOLD] arguments to restrict a run to a subset of the full sweep, useful on machines with a short per-process time budget; repeated runs update the results file in place rather than overwriting it.

Compatibility notes

This code targets PySINDy ≥ 2.0 by default (pysindy.optimizers.SR3, pysindy.SmoothedFiniteDifference), with two exceptions:

• src/optimizers_trapping.py and src/pysindy_compat.py translate a small number of constructor and .fit() keyword-argument changes between PySINDy 1.7.x and ≥ 2.0 automatically (SR3 kwarg renames, SINDy() constructor/.fit() signature changes), so the sweep and plotting scripts run unmodified against either PySINDy version.
• 06_appendix_optimizer_threshold_scan.py requires the legacy PySINDy 1.7.5 environment described above, because PySINDy's built-in TrappingSR3 optimizer is not backward-compatible across that version boundary (see Appendix A environment).

SSR is run with kappa=1e-3; without it, some fits produce a rollout that is numerically unstable under scipy's LSODA integrator. SR3's L0 thresholding can trigger a benign ConvergenceWarning after its fixed inner_max_iter=200, which does not affect the reported results and is suppressed in the sweep scripts.

Model parameters

Quantity	Value
beta1, beta2, rho	1.0, 0.84, 0.048
kappa	0.068435
Integrator	implicit Radau IIA (scipy.integrate.solve_ivp(method='Radau')), atol=rtol=1e-9
Fine generation time step dtgen	0.005 (non-dim. time T)
Integration window	T = 0 to 1000, burn-in discarded up to T = 315
Analysis trajectory	T = 315 to 1000, containing 17 stick-slip cycles (mean recurrence ≈ 40.3 non-dim. time units)
Leading Lyapunov exponent	lambda_max ≈ 0.01106 (1/lambda_max ≈ 90.4 non-dim. time units)
Figure 4 grid	sampling interval {0.01,...,0.10} (step 0.01) x noise {0%,...,9%} (step 1%); 5 repeats; fixed 50/50 train/test split
Figure 5 grid	sampling interval 0.01 fixed; training fraction {5%,10%,20%,30%,50%}; noise {0.5%,1%,1.5%,2%}; 10 repeats
TrappingSR3 hyperparameters	threshold=0.01, eta=0.01, gamma=0, nu=0.1, outer_max_iter=60, inner_max_iter=200, tol=1e-6
Differentiation	pysindy.SmoothedFiniteDifference, Savitzky-Golay window length 21
Function library	1, x, y, z, e^(x-x̄), x e^(x-x̄), y e^(x-x̄), z e^(x-x̄)

Citation

If you use this code, please cite the manuscript (full citation to be added once published) and, if applicable, the underlying PySINDy package:

Kaptanoglu, A. A., de Silva, B. M., Fasel, U., et al. (2022). PySINDy: A comprehensive Python package for robust sparse system identification. Journal of Open Source Software, 7(69), 3994.