Characterizing Nonlinear Dynamics via Smooth Prototype Equivalences

Roy Friedman, Noa Moriel, Matthew Ricci, Guy Pelc, Yair Weiss, Mor Nitzan

Abstract: Characterizing the long term behavior of dynamical systems given limited measurements is a common challenge throughout the physical and biological sciences. This is a challenging task due to the sparsity and noise inherent to empirical observations, as well as the variability of possible longterm dynamics. We address this by introducing smooth prototype equivalences (SPE), a framework for matching sparse observations to prototypical behaviors using invertible neural networks which model smooth phase space deformations. SPE can localize the invariant sets describing long-term behavior of the observed dynamics through the learned mapping from prototype space to data space. Furthermore, SPE can classify dynamical regimes by comparing the data residual of the deformed measurements to prototype dynamics. Our method outperforms existing techniques in the classification of oscillatory systems and can efficiently identify invariant structures like limit cycles and fixed points in an equation-free manner, even when only a small, noisy subset of the phase space is observed. SPE further reveals driving genes in synthetic oscillators such as the repressilator regulatory circuit, and traces cyclic biological processes like the cell cycle trajectory directly from experimental high-dimensional single-cell gene expression data.

---

Installation

This package requires Python 3.9 or later. Clone the repository and install with pip:

git clone https://github.com/nitzanlab/prototype-equivalences.git
cd prototype-equivalences
pip install .

The main dependencies (torch, numpy, scipy, matplotlib, tqdm, POT, torchdiffeq, pysindy) are installed automatically.

Basic Usage

SPE characterizes observations from a vector field, i.e. the set $D=\{(x_i,\dot{x}i)\}{i=1}^N$ comprised of positions ($x_i$) and velocities ($\dot{x}_i$). To do so, the data is matched to a prototype, a simple vector field governed by a known equation $\dot{y}=g(y)$ which fully describes the behaviors expected in the data. The matching is carried out by a diffeomorphism $H$ (an invertible and differentiable function) parameterized with a normalizing flow, fitted so that the observed field is smoothly (orbitally) equivalent to the prototype:

$$\partial_x H(x)\dot{x}\ \propto\ g(H(x))$$

Once $H$ is fitted, everything known analytically about the prototype (its fixed points, limit cycles, or future trajectories) can be pushed back to data space through $H^{-1}$.

Every use of the package follows the same three lines:

from spe.SPE import SPEModel, fit_prototype

model = SPEModel(dim=x.shape[-1], prototype=proto, n_layers=3, K=4)
model, loss, score = fit_prototype(model, x, xdot, its=1000, lr=1e-3)
invariant = model.get_invariant(N=100)   # the prototype's invariant set, in data coordinates

where x and xdot are torch tensors of shape [N, dim], and proto is one of the prototypes described below. fit_prototype returns the fitted model, the final loss and the score, which is the equivalence error (used to compare how well different prototypes explain the same data).

Package structure

The package lives in src/spe/ and has four submodules:

module	contents
spe.SPE	SPEModel (diffeomorphism + prototype), fit_prototype (the fitting procedure), and fit_all_prototypes (fit a list of prototypes to the same data, for classification)
spe.dynamics	prototypes.py — the Prototype base class and all implemented prototypes; systems.py — simulated systems (Van der Pol, repressilator, Selkov, …); utils.py — trajectory simulation and distance/comparison metrics
spe.models	NFDiffeo.py — the Diffeo normalizing flow used as $H$; baselines.py and SINDY_AE.py — neural ODE, SINDy and SINDy-autoencoder baselines
spe.pp	preprocessing for single-cell data: RNA_velocity / prepare_RNA_velocity_embedding and ot_velocities (optimal-transport velocities from time-resolved snapshots)

The SPEModel wraps the two pieces: model.H is the diffeomorphism (with forward, reverse, jvp_forward and jvp_reverse) and model.proto is the prototype. Any keyword argument not recognized by SPEModel is passed to Diffeo, the most useful being n_layers (number of coupling layers), K (hidden units), and n_householder (number of Householder rotations).

Runnable examples are in notebooks/; the three below are the ones worth reading first.

Fitting a cycle

The simplest possible use of SPE is to recover a limit cycle from sampled vectors. The prototype here is the simple oscillator (SO), defined in polar coordinates as:

• $\dot{r}=r(a-r^2)$
• $\dot{\theta}=\omega$

so that $a>0$ gives a limit cycle of radius $\sqrt{a}$ around the origin ($\omega>0$ is counter-clockwise, and $a<0$ gives a single node attractor at the origin). Because it is the default prototype, it can be specified with a plain dict of its parameters:

import torch
from spe.dynamics.systems import VanDerPol
from spe.SPE import SPEModel, fit_prototype

# sample positions and velocities from a Van der Pol system with a limit cycle
system = VanDerPol(**VanDerPol.random_cycle_params())
x = system.rand_on_traj(system.random_x(200), T=3)
xdot = system(x)

# fit a simple-oscillator prototype to the observations
model = SPEModel(dim=2, prototype={'a': .25, 'omega': -.5}, n_layers=3, K=4)
model, loss, score = fit_prototype(model, x, xdot, its=1500, lr=5e-3,
                                   det_reg=1e-3, weight_decay=1e-3, verbose=True)

# the estimated limit cycle, in the coordinates of the data
cycle = model.get_invariant(N=100).detach().numpy()

See notebooks/demo.ipynb for a more in depth example, which shows the fit to both the 2D Van der Pol system and a 6D repressilator and then compares the recovered cycle against the ground-truth trajectory gene by gene.

Fitting a single node

To find sinks in the dynamical system, use the Attractor prototype, which models the dynamics as a pure decay towards the origin, $\dot{y}=-\lambda y$. Mapping the origin back through $H^{-1}$ gives the location of the attractor in data space:

import torch
from spe.SPE import SPEModel, fit_prototype
from spe.dynamics.prototypes import Attractor

x = torch.from_numpy(positions).float()           # positions, [N, dim]
xdot = torch.from_numpy(velocities).float()       # velocities, [N, dim]

model = SPEModel(dim=x.shape[-1], prototype=Attractor(), n_layers=2, K=4)
model, _, _ = fit_prototype(model, x, xdot, its=1000, lr=1e-3, weight_decay=1e-3,
                            det_reg=1e-3, freeze_frac=50)

# the fixed point of the prototype (the origin) pulled back into data coordinates
with torch.no_grad():
    attractor = model.H.reverse(torch.zeros(1, x.shape[-1]))

Note that freeze_frac can be given either as a fraction of the iterations or, as here, as a number of iterations; it keeps the nonlinear part of $H$ fixed at the start of fitting so the linear part settles first.

Other prototypes

Any prototype can be swapped in for the two above; the fitting code is identical, only the invariant set and what it means change. The implemented prototypes, all in spe/dynamics/prototypes.py, are:

prototype	behavior
SOPrototype	simple oscillator: limit cycle ($a>0$) or node ($a<0$)
Attractor	a single stable fixed point
HeteroclinicFlip	one source whose unstable manifold connects to two sinks — a bifurcation
DualCusp	cusp-type bistability

As an example, notebooks/mouse_endoderm.ipynb fits HeteroclinicFlip to the ICM $\to$ (EPI, PrE) decision in early mouse endoderm, using optimal-transport velocities:

import torch
from spe.pp import ot_velocities
from spe.SPE import SPEModel, fit_prototype
from spe.dynamics.prototypes import HeteroclinicFlip

velocities = ot_velocities(adata.obsm['X_pca'], timepoints, extend_last=False, k=30)

x = torch.from_numpy(adata.obsm['X_pca'].copy()).float()
xdot = torch.from_numpy(velocities).float()

proto = HeteroclinicFlip()
model = SPEModel(dim=x.shape[-1], prototype=proto, n_layers=2, K=4, n_householder=6)
model, _, _ = fit_prototype(model, x, xdot, its=1000, lr=5e-4, verbose=True,
                            det_reg=1e-3, weight_decay=1e-3, poten_reg=1e-3, freeze_frac=300)

# the invariant set: one source (the progenitor) and two sinks (the two fates)
inv = model.proto.get_invariant(N=3, dim=x.shape[-1])

# embed the data in prototype ("SPE") coordinates, for plotting
lat, latdot, _ = model.H.jvp_forward(x, xdot)

# trajectories of the prototype, started on either side of the source, describe
# the two differentiation branches; H.reverse maps them back to data coordinates
inits = torch.zeros(2, x.shape[-1])
inits[:, 1], inits[0, 0], inits[1, 0] = 1, -.1, .1
branches = model.proto.trajectories(inits, T=2.)
pca_branches = model.H.reverse(branches.reshape(-1, x.shape[-1]))

Since H is invertible, points sampled anywhere in prototype space can be mapped back through $H^{-1}$, and then through the PCA loadings, into predicted gene expression along each branch. poten_reg is available for gradient-like prototypes such as this one, and penalizes the fitted field by the prototype's potential.

Defining your own prototype

A prototype is an nn.Module subclassing Prototype, and needs at minimum a forward that returns $\dot{y}=g(y)$. Implementing get_invariant (points on the invariant set) and project_onto_invariant (nearest point of the invariant set) makes SPEModel.get_invariant work, and potential enables poten_reg. By convention the dynamics of interest live in the first two coordinates and the remaining dim-2 decay exponentially, which is what lets the same prototype be used at any dimension.

Passing optimize=True to the built-in prototypes turns their parameters into nn.Parameters, so they are fitted jointly with $H$ instead of being held fixed.

Classifying dynamics

Because the score returned by fit_prototype measures how well a prototype explains the data, fitting several prototypes and comparing scores classifies the observed dynamics. fit_all_prototypes handles this natively:

from spe.SPE import fit_all_prototypes

results = fit_all_prototypes(
    x, xdot,
    prototypes=[{'a': -.25, 'omega': .5}, {'a': .25, 'omega': .5}],   # node vs. cycle
    diffeo_args=dict(n_layers=3, K=4),
    fitting_args=dict(its=1000, lr=5e-3, det_reg=1e-3),
)
best = results['prototypes'][int(np.argmin(results['scores']))]

It returns a dictionary with the prototypes, their losses and scores, and the fitted models under Hs.

Contact

Please be in touch if you have any questions! For contact, you can email Roy Friedman at roy.friedman@mail.huji.ac.il .