pyedwin

Python Dynamic Maximum Entropy — a PyTorch implementation of Edwin (entropy-driven compression with interpretable nonlinear model discovery), named in honor of E. T. Jaynes.

Edwin simultaneously performs dimensionality reduction via the Dynamic Maximum Entropy (DME) principle and discovers sparse symbolic models governing latent dynamics and feature structure, using the SLIC sparse regression algorithm.

Overview

Edwin decomposes a non-negative data matrix X (shape n_features × n_timepoints) into interpretable latent factors:

q = softmax_col( -Y @ z.T )

minimising the column-normalised KL divergence D_KL(X ‖ q). Here z ∈ ℝ^(T×K) are time-domain latents and Y ∈ ℝ^(N×K) are feature-domain latents, with K ≪ min(T, N).

Four model classes are provided:

Class	Description
Edwin	Core DME compression: learn z and Y
EdwinWithDynamics	Adds sparse dynamics DW @ z ≈ W @ Θ(z) @ ξz
EdwinWithFeatureModel	Also constrains Y ≈ θy @ ξy via a feature library
EdwinWithStableDynamics	Enforces Lyapunov stability: ż = -M∇V(z)

Sparse models (ξz, ξy) are discovered automatically via EnAdSR — an ensemble sparse regression method with SLIC / AICc / BIC model selection.

Backward-compatible aliases TMI, TMIWithDynamics, TMIWithFeatureModel are exported for existing code.

Installation

pip install -e ".[dev]"

Requires Python ≥ 3.10, PyTorch ≥ 2.0, NumPy ≥ 1.24. Optional: SciPy (for scipy-based ODE integration in forecast), scikit-learn (for KDE in traj_to_dist).

Quick start

Core Edwin (2-D data)

import numpy as np
from pyedwin import Edwin

X = np.random.exponential(1.0, (50, 100))   # (n_features, n_timepoints)
model = Edwin(K=5, normalize=True, lr=1e-3, max_iters=10_000)
model.fit(X)

z  = model.z_      # (n_timepoints, K) — temporal latents
Y  = model.Y_      # (n_features, K)  — feature-domain latents
q  = model.q_      # (n_features, n_timepoints) — reconstruction
kld = model.score(X)  # mean per-column KL divergence (nats)

Edwin with sparse dynamics

from pyedwin import EdwinWithDynamics
from pyedwin.galerkin import build_weight_matrix, build_poly_degree_matrix
import numpy as np

t = np.linspace(0.2, 3.0, 100)
W  = build_weight_matrix(t, window=11, d=0)
DW = build_weight_matrix(t, window=11, d=1)
N  = build_poly_degree_matrix(K=2, order=3, include_constant=False)

model = EdwinWithDynamics(
    K=2, W=W, DW=DW, N=N,
    lambda_z=10.0,
    sparse_iter=5_000,
    lr=1e-3, max_iters=50_000,
)
model.fit(X)

xi_z = model.xi_z_   # (L, K) sparse dynamics coefficients

Edwin with dynamics and feature model

from pyedwin import EdwinWithFeatureModel

x = np.linspace(-6, 6, 50)
theta_y = np.column_stack([x, x**2, x**3])   # (n_features, P)

model = EdwinWithFeatureModel(
    K=2, W=W, DW=DW, N=N, theta_y=theta_y,
    lambda_z=10.0, lambda_y=10.0,
    lr=1e-3, max_iters=50_000,
)
model.fit(X)

xi_y = model.xi_y_   # (P, K) sparse feature coefficients

Equation pretty-printing

from pyedwin import print_equations, equations_str

# Unicode (default) — readable in terminals and Jupyter
print_equations(model)
# ─────────────────────────────────────────────
# Discovered equations (EdwinWithDynamics)
# ─────────────────────────────────────────────
#   dz₀/dt = -0.4982 z₀²

# LaTeX strings (no $ delimiters)
for eq in equations_str(model, fmt="latex"):
    print(eq)
# \dot{z}_{0} = -0.4982 z_{0}^{2}

# Plain ASCII
equations_str(model, fmt="plain")
# ['dz0/dt = -0.4982 z0**2']

# Attach human-readable names for the feature library columns
model.theta_y_names_ = ["x", "x²", "x³"]
print_equations(model)   # Y₀ = 0.51 x²

Trajectory forecasting

from pyedwin import forecast, adapt_and_forecast
import numpy as np

t_fc = np.linspace(1.0, 5.0, 100)
z0 = model.z_[0].cpu()

# Integrate dz/dt = Θ(z) @ ξz forward from z0
z_forecast = forecast(model, z0, t_fc, method="rk4")          # (100, K)
z_fc, q_fc = forecast(model, z0, t_fc, return_data=True)      # also get data recon

# Adapt ξz from a short observed prefix, then forecast
X_prefix = X[:, :30]
z_adapt, q_adapt = adapt_and_forecast(
    model, X_prefix, t_fc,
    prefix_iters=300, lr=5e-3, return_data=True,
)

Trajectory → distribution

from pyedwin import traj_to_dist, traj_to_dist_grid
import numpy as np

z_np = model.z_.cpu().numpy().squeeze()   # (T,)
x_grid = np.linspace(-6, 6, 60)

# KDE: produces a normalised distribution at each time point
dist = traj_to_dist(z_np, x_grid, bandwidth="scott")  # (60, T)

# Histogram: faster, no extra dependencies
dist_hist, bin_centres = traj_to_dist_grid(z_np, bins=40)  # (40, T)

Lyapunov-stable dynamics

from pyedwin import EdwinWithStableDynamics

model_stable = EdwinWithStableDynamics(
    K=2, W=W, DW=DW, N=N,
    lambda_z=10.0, lambda_V0=1.0,
    lr=1e-3, max_iters=50_000,
)
model_stable.fit(X)

# V(z) = (Θ(z)·ξV)² + eps_V‖z‖²  is positive definite and radially unbounded;
# M ⪰ 0 gives V̇ = -∇Vᵀ M ∇V ≤ 0, so trajectories converge to a fixed point.
print_equations(model_stable, lyapunov=True)
V_coeffs = model_stable.xi_V_   # (L, 1)
M = model_stable.M_             # (K, K) PSD matrix
xi_z = model_stable.xi_z_       # (L, K) projection of -M∇V onto the library

See The task axis for per-task V⁽ʲ⁾ and M⁽ʲ⁾.

3-D data (task axis)

The third axis indexes tasks — the same features observed under a different stimulus or condition, each slice already averaged over trials. Y is shared across tasks; z gets one trajectory per task.

X3 = np.random.exponential(1.0, (50, 100, 20))   # 20 tasks
model = Edwin(K=3).fit(X3)
# model.z_.shape == (n_timepoints, K, n_tasks)
# model.Y_.shape == (n_features, K)

Whether the dynamics are shared across tasks is controlled by share_dynamics=; see The task axis.

Hyperparameter search

from pyedwin.models import scan_hyperparams, score_models

results = scan_hyperparams(
    X, EdwinWithFeatureModel,
    param_grid={"lambda_z": [1., 10., 100.], "lambda_y": [1., 10.]},
    K=3, W=W, DW=DW, N=N, theta_y=theta_y,
    max_iters=20_000, verbose=False,
)
scores = score_models(results, N=N, W=W, DW=DW, theta_y=theta_y)
best = results[scores.index(min(scores))]

GPU acceleration

All model classes accept GPU tensors natively:

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
X_gpu = torch.tensor(X, dtype=torch.float64, device=device)
model = EdwinWithDynamics(K=2, W=W, DW=DW, N=N, lambda_z=10.0)
model.fit(X_gpu)
# model.z_, model.Y_, model.xi_z_ are all on the same device

Standalone sparse regression

from pyedwin import AdSR, EnAdSR
import torch

theta = torch.randn(200, 10)
y     = torch.randn(200, 3)

reg  = AdSR(ic="slic", max_iter=10).fit(theta, y)
ereg = EnAdSR(ic="slic", num_batches=20, tol=0.7).fit(theta, y)
print(ereg.coef_)            # (10, 3) sparse
print(ereg.inclusion_probs_) # (10, 3) fraction of ensemble runs each term was active

Observation models (link=)

Edwin factorises data through a bilinear natural parameter eta = Y z (+ E). The link decides what distribution eta parameterises. Everything else -- dynamics discovery, feature models, Lyapunov stability, forecasting -- acts on z and Y and is link-agnostic.

	link="softmax" (default)	link="sigmoid"
Reconstruction	q = softmax_col(-eta)	pi = sigmoid(-eta)
Loss	`D_KL(X
Data	columns sum to 1	each entry in [0, 1]
normalize	True	False (forced)

link="sigmoid" is the SiGMoiD observation model (Xu & Dixit, PLoS Comput Biol 2021, Eq. 1) with beta <-> z and E <-> Y.

# Static Bernoulli factorisation (no dynamics) -- time axis is a sample axis
model = Edwin(K=6, link="sigmoid", use_E=True).fit(X)
model.z_   # (n_samples, K)    -- SiGMoiD's beta
model.Y_   # (n_features, K)   -- SiGMoiD's energies
model.q_   # sigmoid(-Yz)      -- reconstructed P(feature = 1)

# Trajectory Bernoulli -- binary observations evolving in time
model = EdwinWithDynamics(
    K=2, link="sigmoid", W=W, DW=DW, N=N, lambda_z=10.0,
).fit(X)
model.print_equations()

score() returns mean per-column KL for softmax and mean per-element BCE for sigmoid. Scores are not comparable across links.

Frequency smoothing

When trial-averaged frequencies are exactly 0 or 1, the maximum-likelihood estimate drives |eta| -> inf -- perfect separation, as in logistic regression. Clamping the reconstruction does not fix it. Shrink the data instead:

model = Edwin(K=6, link="sigmoid", smoothing=0.0122).fit(X)   # Jeffreys, n=40 trials

X -> X(1-2s) + s with s = alpha / (n + 2 alpha) reproduces a Beta(alpha, alpha) posterior mean exactly. Jeffreys is alpha = 1/2; Laplace is alpha = 1. At 5 trials per task and a 10% firing rate, ~60% of entries saturate, and smoothing drops max|eta| from ~650 to ~3.

The task axis (share_dynamics=)

For 3-D input X of shape (n_features, n_timepoints, n_tasks), the third axis is a task axis: the same features observed under a different stimulus or condition, with each slice already averaged over trials. Y is always shared across tasks; z always has one trajectory per task. What varies is the flow field.

# (a) Shared flow field  [default] -- the stimulus sets the initial condition
m = EdwinWithDynamics(K=2, W=W, DW=DW, N=N, lambda_z=10.0).fit(X)
m.xi_z_          # (L, K)          -- one flow field

# (d) Independent flow fields -- each stimulus induces its own dynamics
m = EdwinWithDynamics(K=2, W=W, DW=DW, N=N, lambda_z=10.0,
                      share_dynamics=False).fit(X)
m.xi_z_          # (L, K, n_tasks) -- one flow field per task, own sparsity pattern
m.print_equations(task=0)

Under share_dynamics=False each task selects its own support, so the discovered terms may differ across tasks. Forecasting then requires a task index:

z_fc, q_fc = forecast(m, z0=m.z_[0, :, 1], t=t, task=1, return_data=True)

share_dynamics is ignored for 2-D data. Equal trial counts per task are assumed; with unequal counts the per-task loss weights should be the trial counts (the Binomial log-likelihood equals n_j x BCE), which is not yet implemented.

Per-task Lyapunov stability

EdwinWithStableDynamics parameterises the flow as dz/dt = -M grad V(z) with

V(z) = (Theta(z) . xi_V)^2 + eps_V ||z||^2       M = tril(L) tril(L)^T

M >= 0 gives dV/dt = -grad V^T M grad V <= 0; the squared form plus the radial eps_V ||z||^2 term makes V positive definite and radially unbounded, which is the other half of Lyapunov's theorem. Both are required: dV/dt <= 0 alone can be satisfied by a concave V whose trajectories diverge.

Under share_dynamics=False, each task gets its own V^(j) and M^(j):

m = EdwinWithStableDynamics(K=2, W=W, DW=DW, N=N, lambda_z=5.0,
                            share_dynamics=False).fit(X)
m.xi_V_   # (L, 1, n_tasks)
m.M_      # (K, K, n_tasks)
m.xi_z_   # (L, K, n_tasks) -- projection of -M grad V onto the library

Examples and their exact ground truth

examples/systems.py holds the data generators and the analytically derived latent law for each system in the manuscript. Both the examples and pyedwin/tests/test_examples.py import from it, so they cannot drift apart.

Example	K	Feature mode Y	Latent law	Support
gaussian_diffusion.py	1	x²	dZ/dt = -4D Z², Z = 1/(4Dt)	{Z²}
brownian_2d.py	1	x² + y²	dZ/dt = -4D Z²	{Z²}
self_assembly.py	1	k	dZ/dt = -4c sinh²(Z/2) ≈ -c Z², c = K·N₀/2	{Z²}
ou_process.py	2	x, x²	dZ₁/dt = βZ₁ - 2βμZ₂ - 2σ²Z₁Z₂ dZ₂/dt = 2βZ₂ - 2σ²Z₂²	{Z₁,Z₂,Z₁Z₂}, {Z₂,Z₂²}

Every dynamics library uses include_constant=False. A Z⁰ row lets the sparse regression satisfy the Galerkin residual trivially and destroys sparsity.

OU feature library and pruning

The OU example uses the manuscript's exact settings: an order-5 feature library [x, x², x³, x⁴, x⁵], a second-order dynamics library with no constant term (N = [1 0; 0 1; 2 0; 1 1; 0 2]), and c_y = 10, c_z = 1e-6, window = 11. SLIC's scoring is automatic, but the penalty scales as c · cond(θ), so c_y = 10 is what prunes the (ill-conditioned) order-5 library cleanly to the two modes the paper reports, Y₁ ∝ x and Y₂ ∝ x². The default c_y = 1e-4 is too permissive and retains x³, x⁴.

The OU dynamics are nonlinear

Both OU equations carry a quadratic term: a bilinear Z₁Z₂ coupling in the mean equation and a logistic Z₂² in the variance equation. This follows exactly from the moment ODEs (verified to 1e-14) and matches the supports of the manuscript's Eqs. (14)–(15).

Gauge

q ∝ exp(-Y Zᵀ) is invariant under (Y, Z) → (Y B, Z B⁻ᵀ).

• K = 1 — B is a scalar a. The support of xi_z is gauge-invariant; a Z^n coefficient scales as a^{n-1}. Divide by a = xi_y[dominant] to compare against theory.
• K > 1 — B is a full GL(K) matrix, so neither the support nor the coefficients of xi_z mean anything on their own. systems.canonicalize(model) fixes the gauge by demanding Y = theta_y (requires a feature library with exactly K columns). Only then are the latents comparable to theory. This is why the manuscript's OU coefficients differ numerically from the ones above while the supports agree.

2-D Brownian: isotropy and the bootstrap tie

The 2-D grid is flattened, so the feature library carries x² and y² as separate terms and isotropy has to be discovered (equal coefficients), not imposed. Two things bite here:

1. The bootstrap tie. For an isotropic Gaussian x² and y² are exactly interchangeable, so EnAdSR's inclusion vote splits between them: each lands near probability 0.5, the default threshold of 0.7 drops one at random, and isotropy is destroyed on roughly half of seeds (KLD jumps from ~1e-2 to ~2e-1). The example uses inclusion_tol=0.5, which makes recovery deterministic.
2. 2-D histogram undersampling. brownian_2d_kde bins counts, so counts per bin scale as n_traj / n_grid². A 31×31 grid with 6000 trajectories gives ~6 counts per bin and a useless estimate (KL to truth ~0.25). The defaults are now 15×15 with 40k trajectories (~180 per bin, KL ~0.003); the paper's 41×41 would need ~300k.

Timepoint-count robustness

Support recovery in the joint fit is seed-dependent when n_time is small and becomes deterministic once there are enough timepoints. A sweep over analytic data at 80k iterations found:

n_time	diffusion recovery	coefficient error
30	33% (1 of 3 seeds)	8.3%
50	100%	7.1%
80	100%	4.7%
120	100%	3.0%

At n_time=30, seed 0 recovered dZ/dt ∝ Z¹ and seed 2 pruned to an empty model; at n_time=80 all seeds gave a bit-identical coefficient. Self-assembly and OU show the same threshold. An oracle test — feeding the sparse regression the exact analytic latents — recovers the correct support in every case, so the fragility is in the joint optimisation producing an imperfect Z(t) at low resolution, not in the sparse regression.

The examples ship with n_time=80 for this reason.

Note: build_weight_matrix assumes a uniformly spaced time grid (it does not validate this). Non-uniform grids — e.g. np.logspace — silently produce a wrong Galerkin derivative. All examples use np.linspace.

Running the verification suite

pytest pyedwin/tests/test_examples.py -v            # everything (~6 min)
pytest pyedwin/tests/test_examples.py -m "not slow" # algebra + integrity only (~4 s)

The slow tests fit each system and assert that Edwin recovers the correct support, the correct feature mode, and a gauge-corrected coefficient within tolerance. Tolerances are loose because the Galerkin estimator smooths, and because Z_var in the OU system varies only from ~4.4 to ~5.5, making Z_var and Z_var² strongly collinear.

Testing

# CPU tests (always run)
pytest tests/test_gpu.py -v

# GPU tests (automatically skipped if CUDA unavailable)
pytest tests/test_gpu.py -v -k "gpu"

# Force CPU-only mode (tests device-propagation logic without CUDA)
PYDME_TEST_DEVICE=cpu pytest tests/test_gpu.py -v

Design notes

• Shared training loop: all Edwin variants share a single _fit_loop in the base class; subclasses customise behaviour via _init_state, _extra_loss, and _post_step hooks.
• Galerkin projection: rather than differentiating z directly, build_weight_matrix computes integration-by-parts weight matrices that avoid numerical differentiation, improving noise robustness.
• Sparse regression via EnAdSR: ξz and ξy are updated by ensemble sparse regression (not autograd), providing automatic model-order selection via SLIC.
• scikit-learn API: hyperparameters in __init__, fit() returns self, fitted attributes end in _.
• GPU-native: all tensor operations respect the device of the input data; no explicit .to(device) calls are needed inside model code.