g-potential equilibrium map

examples/duffing/pkg_learned_splitting_control.ipynb

@@ -4,27 +4,7 @@
    "cell_type": "markdown",
    "id": "aeae1229",
    "metadata": {},
-   "source": [
-    "# Learned splitting control, via the `ftnode.control` package\n",
-    "\n",
-    "Thin driver reproducing `duffing_learned_splitting_control.ipynb`, which stays\n",
-    "untouched as the full write-up (it carries ~350 lines of markdown deriving the\n",
-    "method, plus the k_trunc ablation, the reachability study and the certificate\n",
-    "analysis that are not repeated here).\n",
-    "\n",
-    "**The idea.** Identification fixes the plant as a latent field `F_theta`, but the\n",
-    "splitting is *not unique*: for any invertible `f`,\n",
-    "\n",
-    "    g_psi = z - f_psi^{-1} F_theta   ==>   f_psi (z - g_psi) = F_theta   exactly.\n",
-    "\n",
-    "So `psi` is a gauge choice. The control stage freezes `F_theta` and trains only\n",
-    "that gauge, which means the represented plant provably cannot move while the\n",
-    "control cost landscape is reshaped. `SplitOperator` is the only trainable object\n",
-    "in this notebook.\n",
-    "\n",
-    "The control law is gradient flow on `J(u) = 1/2 ||g_psi(z,u) - z*||^2`: steer the\n",
-    "input so the equilibrium the plant is currently heading toward lands on target."
-   ]
+   "source": "# Learned splitting control, via the `ftnode.control` package\n\nDriver for learned-splitting control built on the `ftnode.control` package.\n\n**The idea.** Identification fixes the plant as a latent field `F_theta`, but the\nsplitting is *not unique*: for any invertible `f`,\n\n    g_psi = z - f_psi^{-1} F_theta   ==>   f_psi (z - g_psi) = F_theta   exactly.\n\nSo `psi` is a gauge choice. The control stage freezes `F_theta` and trains only\nthat gauge, which means the represented plant provably cannot move while the\ncontrol cost landscape is reshaped. `SplitOperator` is the only trainable object\nin this notebook.\n\nThe control law is gradient flow on `J(u) = 1/2 ||g_psi(z,u) - z*||^2`: steer the\ninput so the equilibrium the plant is currently heading toward lands on target."
   },
   {
    "cell_type": "code",
@@ -70,16 +50,7 @@
    "cell_type": "markdown",
    "id": "98f4ea91",
    "metadata": {},
-   "source": [
-    "## Configuration\n",
-    "\n",
-    "`u_bound = 0.5` is wider than the identification excitation range (`0.25`) on\n",
-    "purpose. The frozen notebook's reachability study finds a saturated LQR at the\n",
-    "saddle stabilizes 19 of 39 initial conditions at `|u| <= 0.25`, but all 39 at\n",
-    "`|u| <= 0.5` -- 0.25 is simply not enough authority for this task. The package\n",
-    "makes `u_lo`/`u_hi` required arguments so that value can never be inherited by\n",
-    "accident."
-   ]
+   "source": "## Configuration\n\n`u_bound = 0.5` is wider than the identification excitation range (`0.25`) on\npurpose. The reachability study finds a saturated LQR at the saddle stabilizes 19\nof 39 initial conditions at `|u| <= 0.25`, but all 39 at `|u| <= 0.5` -- 0.25 is\nsimply not enough authority for this task. The package makes `u_lo`/`u_hi`\nrequired arguments so that value can never be inherited by accident."
   },
   {
    "cell_type": "code",
@@ -126,14 +97,7 @@
    "cell_type": "markdown",
    "id": "1e45ae0c",
    "metadata": {},
-   "source": [
-    "## The frozen plant\n",
-    "\n",
-    "Loaded from the checkpoint committed alongside the frozen notebook. It is a bare\n",
-    "state dict carrying no architecture metadata, so the class has to be rebuilt from\n",
-    "`model_cfg` and `budget`; `from_checkpoint` loads with `strict=True` so any drift\n",
-    "fails loudly instead of leaving layers at their init."
-   ]
+   "source": "## The frozen plant\n\nLoaded from the checkpoint committed under `examples/duffing/`. It is a bare\nstate dict carrying no architecture metadata, so the class has to be rebuilt from\n`model_cfg` and `budget`; `from_checkpoint` loads with `strict=True` so any drift\nfails loudly instead of leaving layers at their init."
   },
   {
    "cell_type": "code",
@@ -607,34 +571,7 @@
    "cell_type": "markdown",
    "id": "51a55e49",
    "metadata": {},
-   "source": [
-    "## Takeaway\n",
-    "\n",
-    "**The gauge argument holds numerically.** `psi` is trained against the frozen\n",
-    "model only, and `f_psi (z - g_psi) - F_theta` stays at ~2e-7 -- float32 noise --\n",
-    "so the represented plant is provably untouched by control training. The kappa\n",
-    "bound on `f_psi` holds before and after training, because the parameterization\n",
-    "cannot express a violating operator.\n",
-    "\n",
-    "**Training the splitting does real work.** Trained `psi` regulates far better\n",
-    "than untrained `psi`, which is the same operator class with no control design in\n",
-    "it. That gap is attributable entirely to the gauge choice.\n",
-    "\n",
-    "**It does not beat LQR from every initial condition.** With `z0_spread = 0.35`\n",
-    "the initial conditions are drawn tightly around `z*`, so `psi` is trained for\n",
-    "*regulation*, not global capture -- and the numbers show exactly that shape:\n",
-    "near-target initial conditions land around 0.05, the far ones (|q0| = 0.9) do\n",
-    "markedly worse, and the static latent LQR is better on average over this mixed\n",
-    "set. Read the table as a statement about the training distribution, not as a\n",
-    "verdict on the method.\n",
-    "\n",
-    "To probe that directly, set `z0_spread=None` (uniform over the latent box) and\n",
-    "retrain, or lengthen `T_LIST`. The frozen\n",
-    "`duffing_learned_splitting_control.ipynb` runs a longer horizon over a different\n",
-    "initial-condition set and reports a different ordering; it also carries the\n",
-    "derivation, the `k_trunc` ablation, the `lam_g` sweep, the reachability study and\n",
-    "the (vacuous) forward-invariance certificate."
-   ]
+   "source": "## Takeaway\n\n**The gauge argument holds numerically.** `psi` is trained against the frozen\nmodel only, and `f_psi (z - g_psi) - F_theta` stays at ~2e-7 -- float32 noise --\nso the represented plant is provably untouched by control training. The kappa\nbound on `f_psi` holds before and after training, because the parameterization\ncannot express a violating operator.\n\n**Training the splitting does real work.** Trained `psi` regulates far better\nthan untrained `psi`, which is the same operator class with no control design in\nit. That gap is attributable entirely to the gauge choice.\n\n**It does not beat LQR from every initial condition.** With `z0_spread = 0.35`\nthe initial conditions are drawn tightly around `z*`, so `psi` is trained for\n*regulation*, not global capture -- and the numbers show exactly that shape:\nnear-target initial conditions land around 0.05, the far ones (|q0| = 0.9) do\nmarkedly worse, and the static latent LQR is better on average over this mixed\nset. Read the table as a statement about the training distribution, not as a\nverdict on the method.\n\nTo probe that directly, set `z0_spread=None` (uniform over the latent box) and\nretrain, or lengthen `T_LIST`."
   }
  ],
  "metadata": {
@@ -658,4 +595,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

experiments/duffing/sym_jg_full.yaml

@@ -0,0 +1,67 @@
+# Symmetric-J_g equilibrium map vs the incumbent, at paper settings.
+#
+# The arm under test pairs the unchanged SVD-clamped A(z) with g = grad_z Phi, whose
+# Jacobian is a Hessian and therefore symmetric by construction.  That makes the field a
+# generalized gradient system: V = 0.5||z||^2 - Phi decreases along every trajectory at
+# frozen u, the equilibria are the critical points of V, and their stability type is fixed
+# by V alone for ANY admissible A.  `l-ft-k-svd-clamp` is carried along unchanged as the
+# control -- same operator, same budget, same data, only `g` differs.
+#
+#   nohup uv run ftnode-train experiments/duffing/sym_jg_full.yaml > runs/sym_jg.log 2>&1 &
+#
+# NOTE ON COMPARISON: compare across seeds, never within one.  The equilibrium map is
+# constructed BEFORE the operator, and the scalar potential draws a different number of
+# parameters than the incumbent's m-output map (8769 vs 8964), so at a fixed seed the two
+# arms get structurally identical but numerically different A(z) at initialization.
+#
+# WHY THERE ARE NO `equilibrium.kwargs` HERE: they are GLOBAL across the variants in a file
+# -- ExperimentSpec.model_for replaces only `kind` -- so anything set here is also passed to
+# the `tanh_mlp` control arm, which raises `TypeError: unexpected keyword argument 'cap'`.
+# GradPotentialG's defaults are already the intended configuration (cap=True, g_cap=3.0,
+# w1_cap=10.0), so the arm needs none.  To vary them, the two arms have to live in separate
+# files until per-variant kwargs exist.
+
+name: sym_jg_full
+out: runs/sym_jg_full
+
+budget:
+  sigma_min: 0.1
+  kappa_max: 25.0
+  skew_frac: 0.6
+  m: 4
+
+model:
+  m: 4
+  q: 1
+  activation: silu
+  encoder:     {hidden: 64, depth: 2, z_scale: 2.0, tau: 8}
+  operator:    {hidden: 64, depth: 3, sigma_min: 0.1}
+  # R_g = 2.0 is the incumbent's l-infinity box amplitude, read only by tanh_mlp.
+  # grad_potential ignores it and certifies ||g||_2 <= g_bound instead, from its defaults:
+  #   g_cap = 3.0  -- two-sided window.  The learned equilibria sit at ||z*|| up to 2.41 and
+  #                   z* = g(z*,u), so below ~2.5 the pitchfork is unrepresentable; and
+  #                   r_in = g_cap, so at 4.0 the Nagumo ball R = kappa_max * r_in hits 100
+  #                   and stops beating the incumbent box's 2*sqrt(m) = 4.  3.0 gives R = 75.
+  #   w1_cap = 10  -- deliberately unbalanced.  Reachable curvature scales with ||W_1|| at
+  #                   fixed product, and a saddle needs lambda_max(J_g) > 1 (the incumbent
+  #                   reaches ~2.2).  Equal caps buy a weaker model for the same certificate.
+  equilibrium: {hidden: 64, depth: 3, R_g: 2.0}
+
+train:
+  n_epochs: 200
+  lr: 3.0e-3
+  batch: 64
+  clip: 1.0
+  lam_res: 1.0e-2
+  L: 200
+  L_eval: 600
+  h: 0.05
+
+data_train: {n_traj: 512, L: 200, tau: 8, seed: 0}
+data_val:   {n_traj: 64,  L: 600, tau: 8, seed: 1}
+
+variants:
+  - l-ft-k-svd-clamp                                    # control: g = R_g tanh(MLP)
+  - {operator: svd_clamp, equilibrium: grad_potential}  # slug: svd_clamp+grad_potential
+
+seeds: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

experiments/duffing/sym_jg_quick.yaml

@@ -0,0 +1,45 @@
+# QUICK twin of sym_jg_full.yaml -- a few minutes on CPU, for smoke-testing the arm.
+#
+#   uv run ftnode-train experiments/duffing/sym_jg_quick.yaml
+#   uv run ftnode-train show runs/sym_jg_quick
+#
+# This checks that the arm RUNS and that its structural invariants hold (J_g symmetric,
+# ||g||_2 within g_bound).  It does NOT produce a result: at 15 epochs the models have not
+# learned the observed coordinate, let alone the hidden one, and the multistable structure
+# has not formed.  Do not quote its numbers -- use sym_jg_full.yaml.
+
+name: sym_jg_quick
+out: runs/sym_jg_quick
+
+budget:
+  sigma_min: 0.1
+  kappa_max: 25.0
+  skew_frac: 0.6
+  m: 4
+
+model:
+  m: 4
+  q: 1
+  activation: silu
+  encoder:     {hidden: 64, depth: 2, z_scale: 2.0, tau: 8}
+  operator:    {hidden: 64, depth: 3, sigma_min: 0.1}
+  equilibrium: {hidden: 64, depth: 3, R_g: 2.0}
+
+train:
+  n_epochs: 15
+  lr: 3.0e-3
+  batch: 64
+  clip: 1.0
+  lam_res: 1.0e-2
+  L: 100
+  L_eval: 300
+  h: 0.05
+
+data_train: {n_traj: 128, L: 100, tau: 8, seed: 0}
+data_val:   {n_traj: 32,  L: 300, tau: 8, seed: 1}
+
+variants:
+  - l-ft-k-svd-clamp
+  - {operator: svd_clamp, equilibrium: grad_potential}
+
+seeds: [0, 1]

ftnode/control/__init__.py

@@ -1,8 +1,5 @@
 """Learned-splitting control for latent FT-NODE models.
 
-Promoted from ``examples/duffing/duffing_learned_splitting_control.ipynb`` and the
-``_proto_*.py`` prototypes beside it.
-
 The idea: identification fixes the plant as a latent field ``F_theta``, but the
 splitting ``F_theta = f(z)(z - g(z, u))`` is not unique -- for any invertible
 ``f``, setting ``g = z - f^{-1} F_theta`` reproduces ``F_theta`` exactly.  The

ftnode/control/operator.py

@@ -83,15 +83,15 @@ def kappa_bound(self) -> float:
 class ControlConfig:
     """Control-stage settings.
 
-    Defaults are the values ``duffing_learned_splitting_control.ipynb`` settled
-    on, each for a documented reason:
+    Defaults are the values the control study settled on, each for a documented
+    reason:
 
     ``u_bound = 0.5``
         The admissible input set is ``[-u_bound, u_bound]``.  Deliberately
         **wider** than the identification excitation range (``0.25``): the
-        notebook's reachability study finds a saturated LQR at the saddle
-        stabilizes only 19 of 39 initial conditions at ``|u| <= 0.25`` but all 39
-        at ``|u| <= 0.5``.  0.25 is not enough authority for the task.
+        reachability study finds a saturated LQR at the saddle stabilizes only 19
+        of 39 initial conditions at ``|u| <= 0.25`` but all 39 at ``|u| <= 0.5``.
+        0.25 is not enough authority for the task.
     ``k_trunc = 10``
         Truncated-BPTT window.  Backpropagating through the full 120-step loop
         produces gradient norms spiking to 1e17-1e19, which poison Adam's moments

ftnode/control/plant.py

@@ -1,10 +1,10 @@
 """The frozen identified plant the control stage designs against.
 
-In ``examples/duffing/_proto_ctrl.py`` the identified model is a module-level
-``idm`` that is **loaded from disk at import time**, and every control function
-closes over it.  That makes ``import`` do I/O, pins the working directory, and
-makes it impossible to hold two models at once.  :class:`FrozenLatentPlant`
-replaces that global with an explicit object.
+The earlier prototype held the identified model as a module-level ``idm`` and
+**loaded it from disk at import time**, and every control function closed over
+it.  That makes ``import`` do I/O, pins the working directory, and makes it
+impossible to hold two models at once.  :class:`FrozenLatentPlant` replaces that
+global with an explicit object.
 """
 from __future__ import annotations
 

ftnode/control/policy.py

@@ -6,9 +6,9 @@
 is a gauge choice rather than a property of the plant, redesigning ``psi``
 reshapes this cost landscape without touching the dynamics.
 
-Every function here takes the frozen plant explicitly.  In
-``examples/duffing/_proto_ctrl.py`` they close over a module-global ``idm``,
-``F_theta``, ``m``, ``Z_SCALE``, ``h_dt``, ``device`` and ``u_range``.
+Every function here takes the frozen plant explicitly.  The earlier prototype
+instead closed over module-global ``idm``, ``F_theta``, ``m``, ``Z_SCALE``,
+``h_dt``, ``device`` and ``u_range``.
 """
 from __future__ import annotations
 

ftnode/diagnostics.py

@@ -21,6 +21,7 @@
     "pca_2d",
     "linear_recovery_r2",
     "g_image",
+    "jg_stats",
     "empirical_lipschitz",
 ]
 
@@ -168,7 +169,45 @@ def g_image(dyn, Z, U):
 
     Thin wrapper, but it names the thing being measured: decoding this and
     comparing against the plant's true equilibrium branches is what shows whether
-    the learned ``g`` found the actual pitchfork, and ``|g|_inf <= R_g`` is the
-    box bound the ``tanh`` imposes by construction.
+    the learned ``g`` found the actual pitchfork.
+
+    **The bound to check against is the map's, not a universal one.**
+    :class:`~ftnode.latent.BoundedTanhG` gives ``|g|_inf <= R_g`` -- a box, imposed by its
+    ``tanh``.  :class:`~ftnode.latent.GradPotentialG` gives ``||g||_2 <= g_bound`` -- a
+    ball, imposed by spectral caps on its potential's weights.  Read the bound off the
+    module rather than assuming either.
     """
     return dyn.g(Z, U)
+
+
+def jg_stats(dyn, Z, U, chunk=4096):
+    """Per-sample ``(skew fraction, lambda_max(sym J_g))`` of the equilibrium map's Jacobian.
+
+    Two numbers that decide whether a symmetric-``J_g`` map is doing what it claims:
+
+    * ``||skew J_g||_F / ||J_g||_F`` -- zero by construction for a gradient map, and the
+      check that it really is a gradient.  For reference, an *unstructured* ``m x m``
+      Jacobian sits at ``sqrt((m-1)/2m)`` = 0.612 at ``m=4``, which is where the
+      ``tanh_mlp`` maps measure; the statistic is only meaningful against that null.
+    * ``lambda_max(sym J_g)`` -- exceeds 1 exactly at saddles of the potential, so it is
+      how you confirm multistability survived a bound on ``g``.  Never constrain it below
+      1: that would make ``V`` strictly convex and delete the pitchfork.
+
+    Not ``@torch.no_grad()``: the Jacobian is taken with ``torch.func``, which needs to
+    differentiate.  Results are detached.
+    """
+    from torch.func import jacrev, vmap
+
+    def g_single(zi, ui):
+        return dyn.g(zi.unsqueeze(0), ui.reshape(1, -1)).squeeze(0)
+
+    fracs, lams = [], []
+    U2 = U.unsqueeze(-1) if U.dim() == Z.dim() - 1 else U
+    for i in range(0, Z.shape[0], chunk):
+        J = vmap(jacrev(g_single, argnums=0))(Z[i : i + chunk], U2[i : i + chunk]).detach()
+        skew = 0.5 * (J - J.transpose(-1, -2))
+        denom = J.flatten(1).norm(dim=1).clamp_min(1e-12)
+        fracs.append((skew.flatten(1).norm(dim=1) / denom).cpu().numpy())
+        sym = 0.5 * (J + J.transpose(-1, -2))
+        lams.append(torch.linalg.eigvalsh(sym)[:, -1].cpu().numpy())
+    return np.concatenate(fracs), np.concatenate(lams)

ftnode/latent/__init__.py

@@ -58,19 +58,22 @@
     build_unbounded,
     build_youla,
 )
-from .equilibrium import G_KINDS, BoundedTanhG, resolve_g
+from .equilibrium import G_KINDS, BoundedTanhG, GradPotentialG, resolve_g
+from .feasibility import FeasibilityResult, fit_potential
 from .model import (
     LatentFTNODE,
     LatentNODE,
     LatentSysID,
     migrate_flat_state_dict,
 )
 from .nets import (
+    ACTIVATION_LIPSCHITZ,
     ACTIVATIONS,
     MLP,
     Encoder,
     LinearDecoder,
     is_lipschitz_1,
+    lipschitz_bound,
     resolve_activation,
 )
 from .operator import (
@@ -88,8 +91,10 @@
 __all__ = [
     # nets
     "ACTIVATIONS",
+    "ACTIVATION_LIPSCHITZ",
     "resolve_activation",
     "is_lipschitz_1",
+    "lipschitz_bound",
     "MLP",
     "Encoder",
     "LinearDecoder",
@@ -105,8 +110,11 @@
     "resolve_operator",
     # equilibrium axis
     "BoundedTanhG",
+    "GradPotentialG",
     "G_KINDS",
     "resolve_g",
+    "FeasibilityResult",
+    "fit_potential",
     # models
     "LatentSysID",
     "LatentFTNODE",