PolyChristoffel Network: Interpretable Dynamics via Learnable Curvatures

This project implements an interpretable manifold dynamics model where the learnable object is a torsion-free connection (Christoffel symbols) parameterized as low-degree polynomials of the state. Trajectories are generated by integrating the geodesic equation, and additional losses encourage the learned connection to be consistent with a (pseudo-)Riemannian metric via metric-compatibility reconstruction and path-independence regularization.

The implementation is PyTorch + CUDA, batched/parallel, with an optional jointly-trained autoencoder whose latent can be constrained to lie on a chosen manifold (sphere / hyperboloid / none).

---

Quick Start

1) Install

pip install -r requirements.txt

Install PyTorch with CUDA using the official command for your platform, then install the remaining dependencies via requirements.txt.

2) Train (direct state; no autoencoder)

Default demo uses a synthetic dataset: inertial motion in Cartesian coordinates expressed in polar coordinates (r, theta).

python train.py --use_ae 0 --manifold none --p 2 --q 0 --epochs 30

3) Inspect learned Christoffels

python inspect_christoffel.py --ckpt checkpoints/latest.pt --use_ae 0 --manifold none --p 2 --q 0

Save plots if desired:

python inspect_christoffel.py --ckpt checkpoints/latest.pt --save_fig figs/gamma.png

---

Repository Structure

Training / orchestration

• train.py Main training loop. Loads data, encodes to latent (if AE enabled), computes initial (x0, v0), rolls out geodesic dynamics under learned Christoffels, computes losses (trajectory + metric/geometry regularizers), backpropagates, and saves checkpoints.

• config.py Central hyperparameters and switches: latent dimension d, polynomial degree, training schedule, loss weights, and pseudo-Riemannian signature (p, q) for metric validity penalties.

• utils.py Small utilities: seeding, device selection, and angle unwrapping (needed for polar-coordinate trajectories to avoid discontinuities at theta).

Data

• data.py Synthetic dataset for sanity-check experiments. Default: inertial Cartesian motion represented in polar coordinates (r, theta); the induced Christoffels are known in closed form.

Geometry + dynamics core

• poly_features.py Deterministic polynomial feature map phi(x) (monomial basis up to degree 2 by default). This is the interpretability anchor: Christoffels are explicit polynomials in state.

• christoffel.py The Christoffel network. Computes Gamma^i_{jk}(x) as a linear combination of features phi(x) with learned coefficients. Enforces torsion-free symmetry (Gamma^i_{jk} = Gamma^i_{kj}) by construction.

• integrators.py Differentiable geodesic rollout. Integrates the second-order geodesic equation (converted to a first-order system) via a ResNet-style discrete update. Optionally applies manifold projection/retraction each step.

• metric_recon.py Metric reconstruction by integrating the metric-compatibility equation along a path, yielding a reconstructed g(x) given a base metric g(x0) and learned Gamma(x). Supports multi-path reconstruction to quantify path dependence (loop loss).

• losses.py Loss functions: trajectory loss + geometric regularizers: loop loss, metric symmetry, nondegeneracy (log|det g| barrier), fixed signature penalty (pseudo-Riemannian), coefficient weight decay.

Optional representation learning

• autoencoder.py Optional MLP autoencoder trained jointly with the manifold dynamics. Latents are projected onto a chosen manifold via manifold.py.

• manifold.py Manifold constraint operators:

  • project(z) retraction/projection to the manifold
  • tangent_project(z, v) projects velocities onto the tangent space (optional stabilization)

Diagnostics

• inspect_christoffel.py Evaluates learned Christoffel components on a grid and compares to ground truth for the polar demo:

  • Gamma^r_{theta,theta} = -r
  • Gamma^theta_{r,theta} = Gamma^theta_{theta,r} = 1/r

• inspect_metric.py Reconstructs the metric on a grid and compares to the polar target metric diag(1, r^2) with optional loop error maps.

• inspect_metric_ae.py Pulls back the reconstructed latent metric via the encoder Jacobian for AE experiments.

---

Modeled Mathematics

1) State, trajectories, and coordinates

Let x(t) in R^d be the state in a coordinate chart. Two modes:

• Direct-state mode: x(t) are observed coordinates.
• Autoencoder mode: you observe y(t) and set x(t) = E(y(t)) (latent coordinates).

Goal: model time evolution as geodesic flow governed by a learned connection.

---

2) Christoffel Network: learning a torsion-free connection

Polynomial basis

Define monomial features phi(x) up to degree p:

phi(x) = [1, x_1, ..., x_d, x_1^2, x_1 x_2, ..., x_d^2] in R^P

Christoffel parameterization

Learn Christoffels as explicit polynomials:

Gamma^i_{jk}(x) = sum_{m=1}^P C^i_{jk,m} phi_m(x)

with trainable coefficients C^i_{jk,m}.

Torsion-free constraint

Enforce symmetry in the lower indices:

Gamma^i_{jk}(x) = Gamma^i_{kj}(x)

This makes the connection a plausible Levi-Civita candidate and reduces degrees of freedom.

---

3) Dynamics: geodesic equation as an equation of motion

Geodesics satisfy:

d2 x^i / dt^2 + Gamma^i_{jk}(x) dx^j/dt dx^k/dt = 0

Let v = dx/dt. Then the first-order system is:

dx/dt = v
dv^i/dt = -Gamma^i_{jk}(x) v^j v^k

Define acceleration:

a^i(x, v) = -Gamma^i_{jk}(x) v^j v^k

Interpretability: acceleration is quadratic in velocity and structured entirely by Gamma(x).

---

4) Integrator: ResNet-style discrete geodesic rollout

We discretize the geodesic ODE using a simple differentiable update:

v_{t+1} = v_t + dt * a(x_t, v_t)
x_{t+1} = x_t + dt * v_{t+1}

This is equivalent to stacking residual blocks that implement the geodesic flow.

---

5) Manifold projection / latent validity constraints (optional)

If latents must lie on a manifold M subset R^d, apply a projection/retraction:

x <- Pi_M(x)

Supported:

• none: Pi(x) = x
• sphere: Pi(x) = x / ||x|| enforcing ||x|| = 1
• hyperboloid (Lorentz model): enforce <x, x>_L = -1 with x_0 > 0

Optionally project velocity to the tangent space:

• sphere: v <- v - <v, x> x
• hyperboloid: v <- v + <v, x>_L x (since <x, x>_L = -1)

This keeps discrete rollouts on the constraint manifold.

---

6) Metric reconstruction and pseudo-Riemannian consistency

A torsion-free connection Gamma is not necessarily Levi-Civita for any metric. To encourage (pseudo-)Riemannian structure, reconstruct a metric g(x) compatible with Gamma(x).

Metric compatibility equation

partial_i g_{jk} = Gamma^l_{ij} g_{lk} + Gamma^l_{ik} g_{jl}

This is a first-order linear PDE. We convert it to a path ODE by integrating along a curve x(s) from basepoint x_0 to x:

d/ds g_{jk}(s) =
  (Gamma^l_{ij}(x(s)) g_{lk}(s) + Gamma^l_{ik}(x(s)) g_{jl}(s)) * dx^i/ds

Given:

• basepoint (x_0)
• base metric g(x_0) = g_0 (parameterized to have fixed signature)
• a chosen path

we integrate to get a reconstructed g(x).

Loop / path-independence loss

If Gamma is not metric-realizable, reconstructed g(x) depends on path. We penalize this by using two paths:

• A: (x_0 -> x)
• B: (x_0 -> x_m -> x)

L_loop = || g^(A)(x) - g^(B)(x) ||_F^2

This provides extra training signal pushing Gamma toward metrizable behavior.

---

7) Losses / objective

Primary loss

• Direct-state mode:

L_traj = (1 / (B T)) sum_{b,t} || x_hat_{b,t} - x_{b,t} ||^2

• AE mode (reconstruction in observation space):

L_recon = (1 / (B T)) sum_{b,t} || D(z_hat_{b,t}) - y_{b,t} ||^2

Metric validity penalties

Given reconstructed g(x):

• symmetry: L_sym = || g - g^T ||^2
• nondegeneracy barrier: L_det = softplus(alpha - log|det g|)
• fixed signature: penalize eigenvalues that violate desired (p, q)

Simplicity penalty

Coefficient weight decay:

L_wd = || C ||_2^2

Full objective

L =
  L_traj/recon
  + lambda_loop L_loop
  + lambda_sym L_sym
  + lambda_det L_det
  + lambda_sig L_sig
  + lambda_wd L_wd

---

How the pieces fit together (forward pass)

For each batch:

1. Encode (optional): x_{0:T} = E(y_{0:T})
2. Initialize: x_0 = x(0), v_0 approx (x_1 - x_0) / dt
3. Rollout: integrate geodesic dynamics under Gamma(x) to get x_hat_{0:T}
4. Decode (optional): y_hat_{0:T} = D(x_hat_{0:T})
5. Reconstruct metric g(x) along paths and compute loop/validity losses
6. Backprop to update Gamma coefficients (and AE / base metric parameters if enabled)

---

Notes / recommended workflow

1. Start with the polar-coordinate synthetic benchmark (no AE, SPD signature).
2. Confirm inspect_christoffel.py shows Gamma^r_{theta theta} approx -r and Gamma^theta_{r theta} approx 1/r.
3. Enable AE and then move to more complex observation models (e.g., CNN encoders).

---

WIP

• Swap integrator for RK4 or differentiable ODE solvers (torchdiffeq) after pipeline validation.
• Co-train an explicit metric network g(x) and enforce Gamma_poly approx Gamma(g) (strong Levi-Civita guarantee).
• Apply to domains with strong existing autoencoders (e.g., CNN AEs) to learn curved latent dynamics in concept/feature latent spaces.