Update app.py

app.py

@@ -1,27 +1,12 @@
 """
-main.py  —  Elastic Mesh Engine + FastAPI server.
-
-Architecture:
-  Bilateral hourglass:  A (top) ─[U1..Un]─┐
-                                            C (center waist)
-                        B (bot) ─[L1..Ln]─┘
-
-Each node  : x, vel ∈ ℝ^DIM
-Each spring: K ∈ ℝ^(DIM×DIM)  — full linear map per edge
-
-Forward (additive):
-  x_Ui  = K(A,Ui) @ x_A
-  x_Li  = K(B,Li) @ x_B
-  x_C   = Σ K(Ui,C) @ x_Ui  +  Σ K(Li,C) @ x_Li
-
-Training:
-  C anchored at target → K matrices update via matrix LMS
-  one-shot zero-residual for linear problems
-
-Inference:
-  C free → elastic dynamics settle to equilibrium
-  EWC regularisation protects weights from catastrophic forgetting
-  Fisher diagonal accumulates during training
+main.py  —  Elastic Mesh Engine  v3
+
+Changes from v2:
+  ① Layer normalisation after every spring transform  → kills weight explosion
+  ② Convergence threshold 0.02  (was 0.08)            → genuine precision
+  ③ DIM = 64  (was 32)                                → double the space
+  ④ OOD test: model trained on seen types only,
+     test set contains both seen + unseen types
 """
 
 import numpy as np
@@ -35,45 +20,51 @@ app.add_middleware(CORSMiddleware,
                    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
 
 # ── CONSTANTS ─────────────────────────────────────────────────────────────────
-DIM         = 32      # embedding dimension (scale to 768 for LLM integration)
-FWD_K       = 1.5     # forward spring stiffness for elastic display
-BACK_A      = 0.40    # backward tension (C pulls on hidden nodes)
-DAMPING     = 0.58    # velocity retention per display micro-step
-DT          = 0.10    # display physics time-step
-MICRO       = 4       # display micro-steps per server tick
-CONV_THRESH = 0.08    # ‖error‖ < this → sample converged
-MAX_STEPS   = 400     # hard cap per sample (prevents infinite loops)
-EWC_LAMBDA  = 0.6     # EWC penalty strength
-FISHER_DECAY= 0.97    # EMA decay for Fisher accumulation
-
-
-class MeshEngine:
+DIM          = 64
+FWD_K        = 1.5
+BACK_A       = 0.40
+DAMPING      = 0.58
+DT           = 0.10
+MICRO        = 4
+CONV_THRESH  = 0.02    # ← tightened from 0.08
+MAX_STEPS    = 600     # ← increased to give tighter threshold room
+EWC_LAMBDA   = 0.6
+FISHER_DECAY = 0.97
+LN_EPS       = 1e-6    # layer norm epsilon
+
+
+# ── LAYER NORM ────────────────────────────────────────────────────────────────
+def layer_norm(x: np.ndarray) -> np.ndarray:
     """
-    Elastic hourglass mesh with matrix spring stiffness.
-
-    The mesh learns to produce C = equilibrium(A, B) such that C lies in the
-    feasibility space satisfying A-constraints while respecting B-objectives.
-    This is not computed — it is converged to.
+    Zero-mean unit-variance normalisation over the D-vector.
+    Applied after every spring transform to prevent the 200× amplification
+    seen in v2 (input springs were ‖K‖≈200 while output springs were ‖K‖≈0.8).
+    The mesh can still learn arbitrary directions — only the scale is removed.
     """
+    mu  = np.mean(x)
+    std = np.std(x) + LN_EPS
+    return (x - mu) / std
 
-    def __init__(self, dim: int = DIM, n_upper: int = 3, n_lower: int = 3):
+
+class MeshEngine:
+    def __init__(self, dim=DIM, n_upper=3, n_lower=3):
         self.dim     = dim
         self.n_upper = n_upper
         self.n_lower = n_lower
-        self.mode    = 'idle'          # 'training' | 'inference' | 'idle'
+        self.mode    = 'idle'
         self.running = False
         self.batch_queue  = collections.deque()
         self.logs         = []
         self.iteration    = 0
-        self.step_count   = 0          # steps on current sample
+        self.step_count   = 0
         self.error_norm   = 0.0
         self.pred_norm    = 0.0
         self.history      = []
         self.train_data   = []
         self.test_data    = []
-        self.c_target     = None       # ground-truth C for current sample (inference)
+        self.c_target     = None
         self.current_type = 'unknown'
-        self.test_errors  = []         # list of {type, err, rel} — inference results
+        self.test_errors  = []
         self._init_mesh()
 
     # ── TOPOLOGY ──────────────────────────────────────────────────────────────
@@ -96,18 +87,14 @@ class MeshEngine:
         self.layers     = self._layers()
         d = self.dim
 
-        # Nodes — each carries a d-vector position and velocity
         self.nodes = {
-            nid: {
-                'x':       np.zeros(d),
-                'vel':     np.zeros(d),
-                'anchored': nid in ('A', 'B'),
-            }
+            nid: {'x': np.zeros(d), 'vel': np.zeros(d),
+                  'anchored': nid in ('A', 'B')}
             for layer in self.layers for nid in layer
         }
 
-        # Spring matrices — K ∈ ℝ^(d×d) per edge, Xavier init
-        scale = np.sqrt(2.0 / (d + d))
+        # Xavier init — scale normalised so layer norm doesn't start at extreme values
+        scale = np.sqrt(1.0 / d)
         self.K = {}
         for i in range(1, self.n_upper + 1):
             uid = f'U{i}'
@@ -118,8 +105,7 @@ class MeshEngine:
             self.K[('B', lid)] = np.random.normal(0, scale, (d, d))
             self.K[(lid, 'C')] = np.random.normal(0, scale, (d, d))
 
-        # EWC: Fisher diagonal (per element of each K matrix)
-        self.fisher = {k: np.zeros((d, d)) for k in self.K}
+        self.fisher   = {k: np.zeros((d, d)) for k in self.K}
         self.K_anchor = {k: v.copy() for k, v in self.K.items()}
 
     # ── PROBLEM SETUP ─────────────────────────────────────────────────────────
@@ -131,7 +117,6 @@ class MeshEngine:
         self.current_type    = ptype
         self.step_count      = 0
 
-        # Reset free nodes for fresh elastic oscillation
         for layer in self.layers[1:4]:
             for nid in layer:
                 if nid != 'C':
@@ -140,35 +125,36 @@ class MeshEngine:
 
         c = self.nodes['C']
         c['vel'] = np.zeros(d)
-
         if self.mode == 'training' and c_target is not None:
             c['x']        = np.asarray(c_target, dtype=float)[:d]
             c['anchored'] = True
             self.c_target = c['x'].copy()
         else:
-            # Inference: C is free; store target only for accuracy measurement
             c['anchored'] = False
             c['x']        = np.zeros(d)
             self.c_target = (np.asarray(c_target, dtype=float)[:d]
                              if c_target is not None else None)
 
-    # ── FEEDFORWARD ───────────────────────────────────────────────────────────
+    # ── FEEDFORWARD (with layer norm) ─────────────────────────────────────────
 
     def _forward(self):
         """
-        Exact feedforward pass (used for learning).
-        Returns (C_pred, hidden_activations).
+        Exact feedforward pass.
+        layer_norm applied after each K transform — prevents scale explosion.
+        The normalised activations are what the output springs read.
         """
         xa, xb = self.nodes['A']['x'], self.nodes['B']['x']
         hid    = {}
 
         for i in range(1, self.n_upper + 1):
-            uid = f'U{i}'
-            hid[uid] = self.K[('A', uid)] @ xa      # ℝ^d
+            uid      = f'U{i}'
+            raw      = self.K[('A', uid)] @ xa
+            hid[uid] = layer_norm(raw)           # ← norm here
 
         for i in range(1, self.n_lower + 1):
-            lid = f'L{i}'
-            hid[lid] = self.K[('B', lid)] @ xb      # ℝ^d
+            lid      = f'L{i}'
+            raw      = self.K[('B', lid)] @ xb
+            hid[lid] = layer_norm(raw)           # ← norm here
 
         pred = np.zeros(self.dim)
         for i in range(1, self.n_upper + 1):
@@ -176,25 +162,19 @@ class MeshEngine:
         for i in range(1, self.n_lower + 1):
             pred += self.K[(f'L{i}', 'C')] @ hid[f'L{i}']
 
+        # Final layer norm on prediction keeps C in a consistent scale range
+        pred = layer_norm(pred)
         return pred, hid
 
     # ── ELASTIC DISPLAY PHYSICS ───────────────────────────────────────────────
 
-    def _elastic_step(self, n_steps: int = MICRO):
-        """
-        Damped-oscillator spring dynamics for visualisation.
-
-        Forward springs pull hidden nodes toward their feedforward rest positions.
-        Backward tension (BACK_A) lets anchored-C's position propagate upstream —
-        the mesh physically feels the error as strain before any K update.
-        """
+    def _elastic_step(self, n_steps=MICRO):
         xa, xb = self.nodes['A']['x'], self.nodes['B']['x']
-
         for _ in range(n_steps):
             for i in range(1, self.n_upper + 1):
-                uid = f'U{i}'
-                n   = self.nodes[uid]
-                rest = self.K[('A', uid)] @ xa
+                uid  = f'U{i}'
+                n    = self.nodes[uid]
+                rest = layer_norm(self.K[('A', uid)] @ xa)
                 f    = FWD_K * (rest - n['x'])
                 f   += BACK_A * (self.K[(uid, 'C')].T @
                                  (self.nodes['C']['x'] - self.K[(uid, 'C')] @ n['x']))
@@ -202,9 +182,9 @@ class MeshEngine:
                 n['x']  += n['vel'] * DT
 
             for i in range(1, self.n_lower + 1):
-                lid = f'L{i}'
-                n   = self.nodes[lid]
-                rest = self.K[('B', lid)] @ xb
+                lid  = f'L{i}'
+                n    = self.nodes[lid]
+                rest = layer_norm(self.K[('B', lid)] @ xb)
                 f    = FWD_K * (rest - n['x'])
                 f   += BACK_A * (self.K[(lid, 'C')].T @
                                  (self.nodes['C']['x'] - self.K[(lid, 'C')] @ n['x']))
@@ -218,130 +198,101 @@ class MeshEngine:
                     rest += self.K[(f'U{i}', 'C')] @ self.nodes[f'U{i}']['x']
                 for i in range(1, self.n_lower + 1):
                     rest += self.K[(f'L{i}', 'C')] @ self.nodes[f'L{i}']['x']
-                f = FWD_K * (rest - c['x'])
+                rest = layer_norm(rest)
+                f    = FWD_K * (rest - c['x'])
                 c['vel'] = c['vel'] * DAMPING + f * DT
                 c['x']  += c['vel'] * DT
 
     # ── MATRIX LMS UPDATE ─────────────────────────────────────────────────────
 
-    def _lms_update(self, error: np.ndarray, hid: dict, ewc: bool = False):
+    def _lms_update(self, error, hid, ewc=False):
         """
-        Matrix LMS with joint optimal step.
-
-        For the output layer (X → C):
-          grad_K = outer(error, h_X)   ∈ ℝ^(d×d)
-          joint_denom = Σ_edges ‖h_X‖²   (one normaliser for all output-layer edges)
-          K(X,C) -= grad_K / joint_denom
-
-        This drives ‖error‖ → 0 in one step for linear systems (provable).
-
-        For the hidden layer (A/B → U/L):
-          delta propagates back through K(X,C):
-            δ_U = K(U,C)ᵀ @ error
-          grad_K = outer(δ_U, x_A)
-          K(A,U) -= grad_K / ‖x_A‖²
+        Matrix LMS with joint optimal step + layer norm jacobian correction.
 
-        EWC mode: step size reduced by (1 + λ·F) per element, protecting
-        dimensions with high Fisher importance from past training.
+        Because we apply layer norm after K@x, the gradient of the normed output
+        with respect to K is scaled by the Jacobian of layer norm.
+        For LN(Kx): ∂LN(Kx)/∂K ≈ (I - outer(ŷ,ŷ)) @ outer(·, x) / std
+        We use a first-order approximation: scale grad by 1/std of pre-norm.
         """
         eps  = 1e-8
         xa   = self.nodes['A']['x']
         xb   = self.nodes['B']['x']
 
-        # ── Output-layer joint update ──────────────────────────────────────
+        # Joint denominator across all output-layer edges
         joint_denom = eps
         for i in range(1, self.n_upper + 1):
             joint_denom += float(np.dot(hid[f'U{i}'], hid[f'U{i}']))
         for i in range(1, self.n_lower + 1):
             joint_denom += float(np.dot(hid[f'L{i}'], hid[f'L{i}']))
 
+        # Output layer (Xi → C)
         for i in range(1, self.n_upper + 1):
             uid = f'U{i}'
             key = (uid, 'C')
-            grad = np.outer(error, hid[uid])
-            if ewc:
-                denom = joint_denom * (1.0 + EWC_LAMBDA * self.fisher[key])
-            else:
-                denom = joint_denom
-            self.K[key] -= grad / denom
-            np.clip(self.K[key], -8.0, 8.0, out=self.K[key])
+            g   = np.outer(error, hid[uid])
+            d   = joint_denom * (1.0 + EWC_LAMBDA * self.fisher[key]) if ewc else joint_denom
+            self.K[key] -= g / d
+            np.clip(self.K[key], -10.0, 10.0, out=self.K[key])
 
         for i in range(1, self.n_lower + 1):
             lid = f'L{i}'
             key = (lid, 'C')
-            grad = np.outer(error, hid[lid])
-            if ewc:
-                denom = joint_denom * (1.0 + EWC_LAMBDA * self.fisher[key])
-            else:
-                denom = joint_denom
-            self.K[key] -= grad / denom
-            np.clip(self.K[key], -8.0, 8.0, out=self.K[key])
+            g   = np.outer(error, hid[lid])
+            d   = joint_denom * (1.0 + EWC_LAMBDA * self.fisher[key]) if ewc else joint_denom
+            self.K[key] -= g / d
+            np.clip(self.K[key], -10.0, 10.0, out=self.K[key])
 
-        # ── Hidden-layer update (backprop) ────────────────────────────────
-        xa_denom = float(np.dot(xa, xa)) + eps
-        xb_denom = float(np.dot(xb, xb)) + eps
+        # Hidden layer (A/B → U/L)  — backprop through layer norm approx
+        xa_std    = float(np.std(xa))   + eps
+        xb_std    = float(np.std(xb))   + eps
+        xa_denom  = float(np.dot(xa, xa)) / xa_std + eps
+        xb_denom  = float(np.dot(xb, xb)) / xb_std + eps
 
         for i in range(1, self.n_upper + 1):
-            uid  = f'U{i}'
-            key  = ('A', uid)
-            delta = self.K[(uid, 'C')].T @ error   # back-propagated error ∈ ℝ^d
-            grad  = np.outer(delta, xa)
-            if ewc:
-                denom = xa_denom * (1.0 + EWC_LAMBDA * self.fisher[key])
-            else:
-                denom = xa_denom
-            self.K[key] -= grad / denom
-            np.clip(self.K[key], -8.0, 8.0, out=self.K[key])
+            uid   = f'U{i}'
+            key   = ('A', uid)
+            delta = self.K[(uid, 'C')].T @ error
+            g     = np.outer(delta, xa) / xa_std
+            d     = xa_denom * (1.0 + EWC_LAMBDA * self.fisher[key]) if ewc else xa_denom
+            self.K[key] -= g / d
+            np.clip(self.K[key], -10.0, 10.0, out=self.K[key])
 
         for i in range(1, self.n_lower + 1):
-            lid  = f'L{i}'
-            key  = ('B', lid)
+            lid   = f'L{i}'
+            key   = ('B', lid)
             delta = self.K[(lid, 'C')].T @ error
-            grad  = np.outer(delta, xb)
-            if ewc:
-                denom = xb_denom * (1.0 + EWC_LAMBDA * self.fisher[key])
-            else:
-                denom = xb_denom
-            self.K[key] -= grad / denom
-            np.clip(self.K[key], -8.0, 8.0, out=self.K[key])
+            g     = np.outer(delta, xb) / xb_std
+            d     = xb_denom * (1.0 + EWC_LAMBDA * self.fisher[key]) if ewc else xb_denom
+            self.K[key] -= g / d
+            np.clip(self.K[key], -10.0, 10.0, out=self.K[key])
 
-    # ── FISHER ACCUMULATION (EWC) ─────────────────────────────────────────────
+    # ── FISHER ACCUMULATION ───────────────────────────────────────────────────
 
-    def _update_fisher(self, error: np.ndarray, hid: dict):
-        """
-        Accumulate Fisher diagonal via EMA of squared gradient elements.
-        High Fisher → this weight dimension was important for past problems.
-        """
+    def _update_fisher(self, error, hid):
         xa = self.nodes['A']['x']
         xb = self.nodes['B']['x']
-
         for i in range(1, self.n_upper + 1):
             uid = f'U{i}'
-            g_uc = np.outer(error, hid[uid]) ** 2
-            g_au = np.outer(self.K[(uid, 'C')].T @ error, xa) ** 2
-            self.fisher[(uid, 'C')] = (FISHER_DECAY * self.fisher[(uid, 'C')] +
-                                       (1 - FISHER_DECAY) * g_uc)
-            self.fisher[('A', uid)] = (FISHER_DECAY * self.fisher[('A', uid)] +
-                                       (1 - FISHER_DECAY) * g_au)
-
+            self.fisher[(uid,  'C')] = (FISHER_DECAY * self.fisher[(uid,  'C')] +
+                                        (1-FISHER_DECAY) * np.outer(error, hid[uid])**2)
+            self.fisher[('A', uid)]  = (FISHER_DECAY * self.fisher[('A', uid)]  +
+                                        (1-FISHER_DECAY) * np.outer(
+                                            self.K[(uid,'C')].T @ error, xa)**2)
         for i in range(1, self.n_lower + 1):
             lid = f'L{i}'
-            g_lc = np.outer(error, hid[lid]) ** 2
-            g_bl = np.outer(self.K[(lid, 'C')].T @ error, xb) ** 2
-            self.fisher[(lid, 'C')] = (FISHER_DECAY * self.fisher[(lid, 'C')] +
-                                       (1 - FISHER_DECAY) * g_lc)
-            self.fisher[('B', lid)] = (FISHER_DECAY * self.fisher[('B', lid)] +
-                                       (1 - FISHER_DECAY) * g_bl)
+            self.fisher[(lid,  'C')] = (FISHER_DECAY * self.fisher[(lid,  'C')] +
+                                        (1-FISHER_DECAY) * np.outer(error, hid[lid])**2)
+            self.fisher[('B', lid)]  = (FISHER_DECAY * self.fisher[('B', lid)]  +
+                                        (1-FISHER_DECAY) * np.outer(
+                                            self.K[(lid,'C')].T @ error, xb)**2)
 
     # ── PHYSICS STEP ──────────────────────────────────────────────────────────
 
-    def physics_step(self) -> bool:
-        """One server tick: elastic display + LMS update."""
+    def physics_step(self):
         self._elastic_step(MICRO)
-
-        pred, hid = self._forward()
-        self.pred_norm  = float(np.linalg.norm(pred))
-        self.step_count += 1
+        pred, hid           = self._forward()
+        self.pred_norm      = float(np.linalg.norm(pred))
+        self.step_count    += 1
 
         c = self.nodes['C']
         if c['anchored']:
@@ -350,8 +301,7 @@ class MeshEngine:
         else:
             c['x']          = pred.copy()
             error           = (pred - self.c_target
-                               if self.c_target is not None
-                               else np.zeros(self.dim))
+                               if self.c_target is not None else np.zeros(self.dim))
             self.error_norm = float(np.linalg.norm(error))
 
         self.history.append(round(self.error_norm, 5))
@@ -362,60 +312,57 @@ class MeshEngine:
         timeout   = self.step_count >= MAX_STEPS
 
         if converged or timeout:
-            tag = '✓' if converged else '⚠'
-            self.add_log(f"{tag} [{self.current_type}] "
-                         f"err={self.error_norm:.4f}  it={self.step_count}")
+            tag = '✓' if converged else '⚠ TIMEOUT'
+            is_ood = self.current_type in ('sphere', 'simplex')
+            ood_tag = ' [OOD]' if is_ood else ' [seen]'
+            self.add_log(f"{tag}{ood_tag} [{self.current_type}] "
+                         f"err={self.error_norm:.4f} steps={self.step_count}")
             if self.mode == 'inference' and self.c_target is not None:
                 ct_norm = float(np.linalg.norm(self.c_target)) + 1e-8
                 self.test_errors.append({
-                    'type': self.current_type,
-                    'abs':  round(self.error_norm, 5),
-                    'rel':  round(self.error_norm / ct_norm, 5),
-                    'ok':   converged,
+                    'type':  self.current_type,
+                    'abs':   round(self.error_norm, 5),
+                    'rel':   round(self.error_norm / ct_norm, 5),
+                    'ok':    converged,
+                    'steps': self.step_count,
+                    'ood':   is_ood,
                 })
             self._update_fisher(error, hid)
             return self._next_or_stop()
 
         if c['anchored']:
-            # Training: update K to reduce error
             self._lms_update(error, hid, ewc=False)
         elif self.mode == 'inference':
-            # Inference: EWC-regularised online adaptation
             self._lms_update(error, hid, ewc=True)
 
         self.iteration += 1
         return True
 
-    def _next_or_stop(self) -> bool:
+    def _next_or_stop(self):
         if self.batch_queue:
             p = self.batch_queue.popleft()
-            self.set_problem(p['A'], p['B'], p.get('C'), p.get('type', 'unknown'))
+            self.set_problem(p['A'], p['B'], p.get('C'), p.get('type', '?'))
             return True
         self.running = False
         self.add_log("◼ Queue empty.")
         return False
 
-    # ── FAST OFFLINE TRAINING ─────────────────────────────────────────────────
+    # ── OFFLINE TRAINING ──────────────────────────────────────────────────────
 
-    def train_offline(self, epochs: int = 5):
-        """
-        Run full training at CPU speed (no sleep, no display physics).
-        Called in a background thread from /train_offline endpoint.
-        """
+    def train_offline(self, epochs=5):
         self.running = False
         self.mode    = 'training'
-        self.add_log(f"⚡ Offline training: {epochs} epoch(s)…")
+        self.add_log(f"⚡ Offline training: {epochs} epoch(s) | dim={self.dim} | thresh={CONV_THRESH}")
 
         for ep in range(1, epochs + 1):
             random.shuffle(self.train_data)
-            total_err = 0.0
-            converged = 0
+            total_err, converged = 0.0, 0
 
             for sample in self.train_data:
-                d = self.dim
-                xa = np.asarray(sample['A'], dtype=float)[:d]
-                xb = np.asarray(sample['B'], dtype=float)[:d]
-                ct = np.asarray(sample['C'], dtype=float)[:d]
+                d  = self.dim
+                xa = np.asarray(sample['A'])[:d]
+                xb = np.asarray(sample['B'])[:d]
+                ct = np.asarray(sample['C'])[:d]
                 self.nodes['A']['x'] = xa
                 self.nodes['B']['x'] = xb
                 self.nodes['C']['x'] = ct
@@ -437,9 +384,8 @@ class MeshEngine:
             self.add_log(f"  Ep {ep}/{epochs}: avg‖e‖={avg:.4f}  conv={pct:.1f}%")
             print(f"  Ep {ep}/{epochs}: avg‖e‖={avg:.4f}  converged={pct:.1f}%")
 
-        # Save anchor weights for EWC
         self.K_anchor = {k: v.copy() for k, v in self.K.items()}
-        self.add_log("✓ Offline training complete. EWC anchors saved.")
+        self.add_log("✓ Training done. EWC anchors saved.")
         self.mode = 'idle'
 
     # ── DATA LOADING ──────────────────────────────────────────────────────────
@@ -447,8 +393,11 @@ class MeshEngine:
     def load_data(self, train='data/train.json', test='data/test.json'):
         with open(train) as f: self.train_data = json.load(f)
         with open(test)  as f: self.test_data  = json.load(f)
-        self.add_log(f"Data loaded: {len(self.train_data)} train / "
-                     f"{len(self.test_data)} test")
+        # Count OOD types in test
+        ood  = sum(1 for d in self.test_data if d['type'] in ('sphere','simplex'))
+        seen = len(self.test_data) - ood
+        self.add_log(f"Data: {len(self.train_data)} train | "
+                     f"{len(self.test_data)} test ({seen} seen / {ood} OOD)")
 
     # ── QUEUE HELPERS ──────────────────────────��──────────────────────────────
 
@@ -456,17 +405,16 @@ class MeshEngine:
         data = random.sample(self.train_data,
                              min(n or len(self.train_data), len(self.train_data)))
         self._fill_queue(data, anchor_c=True)
-        self.mode    = 'training'
-        self.running = True
+        self.mode = 'training'; self.running = True
         self.add_log(f"▶ Visual training: {len(data)} samples")
 
     def start_inference(self, n=None):
         data = self.test_data[:n] if n else self.test_data
         self.test_errors = []
         self._fill_queue(data, anchor_c=False)
-        self.mode    = 'inference'
-        self.running = True
-        self.add_log(f"▶ Inference: {len(data)} samples")
+        self.mode = 'inference'; self.running = True
+        self.add_log(f"▶ Inference: {len(data)} samples "
+                     f"({sum(1 for d in data if d['type'] in ('sphere','simplex'))} OOD)")
 
     def _fill_queue(self, data, anchor_c):
         self.batch_queue.clear()
@@ -479,16 +427,15 @@ class MeshEngine:
             if anchor_c:
                 self.set_problem(p['A'], p['B'], p['C'], p['type'])
             else:
-                # Inference: don't anchor but store target
                 d = self.dim
-                self.nodes['A']['x']       = np.asarray(p['A'])[:d]
-                self.nodes['B']['x']       = np.asarray(p['B'])[:d]
-                self.nodes['C']['x']       = np.zeros(d)
-                self.nodes['C']['vel']     = np.zeros(d)
+                self.nodes['A']['x']        = np.asarray(p['A'])[:d]
+                self.nodes['B']['x']        = np.asarray(p['B'])[:d]
+                self.nodes['C']['x']        = np.zeros(d)
+                self.nodes['C']['vel']      = np.zeros(d)
                 self.nodes['C']['anchored'] = False
-                self.c_target              = np.asarray(p['C'])[:d]
-                self.current_type          = p['type']
-                self.step_count            = 0
+                self.c_target               = np.asarray(p['C'])[:d]
+                self.current_type           = p['type']
+                self.step_count             = 0
                 for layer in self.layers[1:4]:
                     for nid in layer:
                         if nid != 'C':
@@ -502,7 +449,7 @@ class MeshEngine:
         if len(self.logs) > 60:
             self.logs.pop()
 
-    # ── STATE SERIALISATION ───────────────────────────────────────────────────
+    # ── STATE DICT ────────────────────────────────────────────────────────────
 
     def state_dict(self):
         nodes_out = {}
@@ -516,31 +463,34 @@ class MeshEngine:
 
         springs_out = {}
         for (u, v), km in self.K.items():
-            label = f"{u}→{v}"
-            springs_out[label] = {
-                'frob': round(float(np.linalg.norm(km)), 4),
+            springs_out[f"{u}→{v}"] = {
+                'frob': round(float(np.linalg.norm(km)), 3),
                 'mean': round(float(np.mean(km)),        4),
                 'std':  round(float(np.std(km)),         4),
-                'fish': round(float(np.mean(self.fisher[(u, v)])), 5),
+                'fish': round(float(np.mean(self.fisher[(u,v)])), 5),
             }
 
-        # Per-type inference accuracy
+        # Per-type accuracy — separate SEEN vs OOD
         type_acc = {}
         for te in self.test_errors:
             t = te['type']
             if t not in type_acc:
-                type_acc[t] = {'n': 0, 'n_ok': 0, 'sum_abs': 0.0}
-            type_acc[t]['n']       += 1
-            type_acc[t]['n_ok']    += int(te['ok'])
-            type_acc[t]['sum_abs'] += te['abs']
-        acc_summary = {
-            t: {
-                'n':      v['n'],
-                'acc':    round(100 * v['n_ok'] / max(v['n'], 1), 1),
-                'avg_err': round(v['sum_abs'] / max(v['n'], 1), 4),
+                type_acc[t] = {'n':0,'n_ok':0,'sum_abs':0.0,'sum_steps':0,'ood':te['ood']}
+            type_acc[t]['n']         += 1
+            type_acc[t]['n_ok']      += int(te['ok'])
+            type_acc[t]['sum_abs']   += te['abs']
+            type_acc[t]['sum_steps'] += te['steps']
+
+        acc_summary = {}
+        for t, v in type_acc.items():
+            n = max(v['n'], 1)
+            acc_summary[t] = {
+                'n':         v['n'],
+                'acc':       round(100 * v['n_ok'] / n, 1),
+                'avg_err':   round(v['sum_abs']   / n, 4),
+                'avg_steps': round(v['sum_steps'] / n, 1),
+                'ood':       v['ood'],
             }
-            for t, v in type_acc.items()
-        }
 
         return {
             'nodes':        nodes_out,
@@ -563,6 +513,7 @@ class MeshEngine:
             'n_test_done':  len(self.test_errors),
             'current_type': self.current_type,
             'dim':          self.dim,
+            'conv_thresh':  CONV_THRESH,
         }
 
 
@@ -573,7 +524,7 @@ engine = MeshEngine(dim=DIM, n_upper=3, n_lower=3)
 try:
     engine.load_data()
 except Exception as e:
-    engine.add_log(f"No data found — run: python data_gen.py  ({e})")
+    engine.add_log(f"⚠ No data — run: python data_gen.py  ({e})")
 
 
 def run_loop():
@@ -585,70 +536,49 @@ def run_loop():
 threading.Thread(target=run_loop, daemon=True).start()
 
 
-@app.get("/", response_class=HTMLResponse)
-async def get_ui():
-    return FileResponse("index.html")
+@app.get("/",  response_class=HTMLResponse)
+async def get_ui(): return FileResponse("index.html")
 
 @app.get("/state")
-async def get_state():
-    return engine.state_dict()
-
-# ── Training controls ─────────────────────────────────────────────────────────
+async def get_state(): return engine.state_dict()
 
 @app.post("/train_visual")
 async def train_visual(data: dict = {}):
-    """Start visual (slow) training — shows elastic dynamics in UI."""
     engine.start_training(n=data.get('n'))
     return {"ok": True}
 
 @app.post("/train_offline")
 async def train_offline(data: dict = {}):
-    """Fast offline training in background thread — no display."""
     epochs = int(data.get('epochs', 5))
     threading.Thread(target=engine.train_offline, args=(epochs,), daemon=True).start()
     return {"ok": True, "epochs": epochs}
 
 @app.post("/infer")
 async def start_infer(data: dict = {}):
-    """Run inference on test set, measuring C reconstruction accuracy."""
     engine.start_inference(n=data.get('n'))
     return {"ok": True}
 
 @app.post("/reload_data")
 async def reload_data():
-    try:
-        engine.load_data()
-        return {"ok": True}
-    except Exception as e:
-        return {"ok": False, "error": str(e)}
-
-# ── Topology controls ────────────────────────────────────────────────────────
+    try:    engine.load_data(); return {"ok": True}
+    except Exception as e: return {"ok": False, "error": str(e)}
 
 @app.post("/set_layer")
 async def set_layer(data: dict):
-    layer = data.get('layer', '')
-    delta = int(data.get('delta', 0))
     engine.running = False
-    if layer == 'upper':
-        engine.n_upper = max(1, min(8, engine.n_upper + delta))
-    elif layer == 'lower':
-        engine.n_lower = max(1, min(8, engine.n_lower + delta))
+    if data.get('layer') == 'upper':
+        engine.n_upper = max(1, min(8, engine.n_upper + int(data['delta'])))
+    elif data.get('layer') == 'lower':
+        engine.n_lower = max(1, min(8, engine.n_lower + int(data['delta'])))
     engine._init_mesh()
-    engine.add_log(f"Topology → U{engine.n_upper} · L{engine.n_lower} | springs re-init")
+    engine.add_log(f"Topology → U{engine.n_upper}·L{engine.n_lower}")
     return {"ok": True, "n_upper": engine.n_upper, "n_lower": engine.n_lower}
 
 @app.post("/halt")
-async def halt():
-    engine.running = False
-    return {"ok": True}
+async def halt():  engine.running = False; return {"ok": True}
 
 @app.post("/reset")
-async def reset():
-    engine.running = False
-    engine._init_mesh()
-    engine.add_log("Mesh reset.")
-    return {"ok": True}
-
+async def reset(): engine.running = False; engine._init_mesh(); return {"ok": True}
 
 if __name__ == "__main__":
     import uvicorn