Update app.py

app.py

@@ -1,54 +1,84 @@
-import time, collections, threading, random
+"""
+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
+"""
+
+import numpy as np
+import time, collections, threading, json, pathlib, random
 from fastapi import FastAPI
 from fastapi.responses import HTMLResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 
 app = FastAPI()
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
-
-# ── ELASTIC CONSTANTS ─────────────────────────────────────────────────────────
-FWD_K   = 2.2    # spring pull toward rest position
-DAMPING = 0.55   # velocity retention (lower = more visible oscillation)
-DT      = 0.12   # micro-step size
-MICRO   = 6      # micro-steps per physics_step (keeps UI responsive)
-SETTLE  = 0.004  # velocity threshold for early exit
-
-
-class SimEngine:
-    def __init__(self):
-        self.mode         = 'training'
-        self.architecture = 'additive'
-        self.dataset_type = 'housing'
-        self.n_upper      = 3      # nodes in upper bulge (A-side)
-        self.n_lower      = 3      # nodes in lower bulge (B-side)
-        self.back_alpha   = 0.45   # backward tension coupling
-        self.running      = False
+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:
+    """
+    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.
+    """
+
+    def __init__(self, dim: int = DIM, n_upper: int = 3, n_lower: int = 3):
+        self.dim     = dim
+        self.n_upper = n_upper
+        self.n_lower = n_lower
+        self.mode    = 'idle'          # 'training' | 'inference' | 'idle'
+        self.running = False
         self.batch_queue  = collections.deque()
         self.logs         = []
         self.iteration    = 0
-        self.current_error      = 0.0
-        self.current_prediction = 0.0
-        self.history            = []
+        self.step_count   = 0          # steps on current sample
+        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.current_type = 'unknown'
+        self.test_errors  = []         # list of {type, err, rel} — inference results
         self._init_mesh()
 
     # ── TOPOLOGY ──────────────────────────────────────────────────────────────
-    # Physical layout (top→bottom):
-    #   A  → U1..Un → C ← L1..Ln ← B
-    #
-    # Layers for display (top to bottom):
-    #   0: [A]
-    #   1: [U1..Un]   upper bulge
-    #   2: [C]        CENTER — the waist
-    #   3: [L1..Ln]   lower bulge
-    #   4: [B]
-    #
-    # Springs:
-    #   A  → each Ui   (K_aui)
-    #   Ui → C         (K_uic)
-    #   B  → each Li   (K_bli)
-    #   Li → C         (K_lic)
-
-    def _build_layers(self):
+
+    def _layers(self):
         return [
             ['A'],
             [f'U{i}' for i in range(1, self.n_upper + 1)],
@@ -58,376 +88,568 @@ class SimEngine:
         ]
 
     def _init_mesh(self):
-        self.iteration          = 0
-        self.current_error      = 0.0
-        self.current_prediction = 0.0
-        self.history            = []
-        self.layers             = self._build_layers()
-
-        self.nodes = {}
-        for layer in self.layers:
-            for nid in layer:
-                anchored = nid in ('A', 'B')
-                self.nodes[nid] = {'x': 0.0, 'vel': 0.0, 'anchored': anchored}
-        self.nodes['A']['x'] = 2.0
-        self.nodes['B']['x'] = 3.0
-
-        # Springs — one per directed edge
-        self.springs = {}
-        # A-side: A → Ui → C
+        self.iteration  = 0
+        self.step_count = 0
+        self.error_norm = 0.0
+        self.pred_norm  = 0.0
+        self.history    = []
+        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'),
+            }
+            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))
+        self.K = {}
         for i in range(1, self.n_upper + 1):
             uid = f'U{i}'
-            self.springs[('A',  uid)] = round(random.uniform(0.85, 1.15), 4)
-            self.springs[(uid, 'C')] = round(random.uniform(0.85, 1.15), 4)
-        # B-side: B → Li → C
+            self.K[('A', uid)] = np.random.normal(0, scale, (d, d))
+            self.K[(uid, 'C')] = np.random.normal(0, scale, (d, d))
         for i in range(1, self.n_lower + 1):
             lid = f'L{i}'
-            self.springs[('B',  lid)] = round(random.uniform(0.85, 1.15), 4)
-            self.springs[(lid, 'C')] = round(random.uniform(0.85, 1.15), 4)
+            self.K[('B', lid)] = np.random.normal(0, scale, (d, d))
+            self.K[(lid, 'C')] = np.random.normal(0, scale, (d, d))
 
-    def reset(self):
-        self.running = False
-        self.batch_queue.clear()
-        self.logs = []
-        self._init_mesh()
+        # EWC: Fisher diagonal (per element of each K matrix)
+        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()}
 
-    # ── LOGGING ───────────────────────────────────────────────────────────────
-    def add_log(self, msg):
-        self.logs.insert(0, f"[{self.iteration:04d}] {msg}")
-        if len(self.logs) > 40:
-            self.logs.pop()
+    # ── PROBLEM SETUP ─────────────────────────────────────────────────────────
 
-    # ── DATASET ───────────────────────────────────────────────────────────────
-    def ground_truth(self, a, b):
-        t = self.dataset_type
-        if   t == 'housing':        return round(a * 2.5 + b * 1.2, 4)
-        elif t == 'subtraction':    return round(a - b,              4)
-        elif t == 'multiplication': return round(a * b,              4)
-        elif t == 'quadratic':      return round(a * a + b,          4)
-        return round(a + b, 4)
+    def set_problem(self, a_vec, b_vec, c_target=None, ptype='unknown'):
+        d = self.dim
+        self.nodes['A']['x'] = np.asarray(a_vec, dtype=float)[:d]
+        self.nodes['B']['x'] = np.asarray(b_vec, dtype=float)[:d]
+        self.current_type    = ptype
+        self.step_count      = 0
 
-    # ── PROBLEM SETUP ─────────────────────────────────────────────────────────
-    def set_problem(self, a, b, c_target=None):
-        self.nodes['A']['x'] = float(a)
-        self.nodes['B']['x'] = float(b)
-        # Reset hidden nodes so elastic wave is visible from scratch
+        # Reset free nodes for fresh elastic oscillation
         for layer in self.layers[1:4]:
             for nid in layer:
                 if nid != 'C':
-                    self.nodes[nid]['x']   = 0.0
-                    self.nodes[nid]['vel'] = 0.0
+                    self.nodes[nid]['x']   = np.zeros(d)
+                    self.nodes[nid]['vel'] = np.zeros(d)
+
         c = self.nodes['C']
-        c['vel'] = 0.0
+        c['vel'] = np.zeros(d)
+
         if self.mode == 'training' and c_target is not None:
-            c['x']        = float(c_target)
+            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']        = 0.0
+            c['x']        = np.zeros(d)
+            self.c_target = (np.asarray(c_target, dtype=float)[:d]
+                             if c_target is not None else None)
 
-    # ── FEEDFORWARD REST POSITION ─────────────────────────────────────────────
-    def _rest_upper(self, uid):
-        """Rest position of an upper hidden node — driven by A."""
-        k = self.springs[('A', uid)]
-        xa = self.nodes['A']['x']
-        return k * xa  # additive or used as scale for mult
+    # ── FEEDFORWARD ───────────────────────────────────────────────────────────
 
-    def _rest_lower(self, lid):
-        """Rest position of a lower hidden node — driven by B."""
-        k = self.springs[('B', lid)]
-        xb = self.nodes['B']['x']
-        return k * xb
+    def _forward(self):
+        """
+        Exact feedforward pass (used for learning).
+        Returns (C_pred, hidden_activations).
+        """
+        xa, xb = self.nodes['A']['x'], self.nodes['B']['x']
+        hid    = {}
 
-    def _rest_c(self):
-        """Rest position of C — sum of contributions from both bulges."""
-        total = 0.0
         for i in range(1, self.n_upper + 1):
             uid = f'U{i}'
-            hu  = self.nodes[uid]['x']
-            total += self.springs[(uid, 'C')] * hu
+            hid[uid] = self.K[('A', uid)] @ xa      # ℝ^d
+
         for i in range(1, self.n_lower + 1):
             lid = f'L{i}'
-            hl  = self.nodes[lid]['x']
-            total += self.springs[(lid, 'C')] * hl
-        return total
+            hid[lid] = self.K[('B', lid)] @ xb      # ℝ^d
+
+        pred = np.zeros(self.dim)
+        for i in range(1, self.n_upper + 1):
+            pred += self.K[(f'U{i}', 'C')] @ hid[f'U{i}']
+        for i in range(1, self.n_lower + 1):
+            pred += self.K[(f'L{i}', 'C')] @ hid[f'L{i}']
+
+        return pred, hid
 
-    # ── ELASTIC RELAXATION (DISPLAY PHYSICS) ─────────────────────────────────
-    def _elastic_step(self, n_steps):
+    # ── ELASTIC DISPLAY PHYSICS ───────────────────────────────────────────────
+
+    def _elastic_step(self, n_steps: int = MICRO):
         """
         Damped-oscillator spring dynamics for visualisation.
 
-        Upper hidden nodes Ui are pulled toward K(A,Ui)·A by a forward spring.
-        Lower hidden nodes Li are pulled toward K(B,Li)·B by a forward spring.
-        C is pulled toward Σ K(Xi,C)·Xi from both sides.
-
-        Back-tension (alpha): if alpha>0, each node also feels a restoring pull
-        from downstream — i.e. C's anchored position creates tension that travels
-        back up through Ui and down through Li, making the elastic wave visible.
+        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.
         """
-        alpha = self.back_alpha
-        for _ in range(n_steps):
-            max_v = 0.0
+        xa, xb = self.nodes['A']['x'], self.nodes['B']['x']
 
-            # Upper hidden nodes
+        for _ in range(n_steps):
             for i in range(1, self.n_upper + 1):
                 uid = f'U{i}'
                 n   = self.nodes[uid]
-                f   = FWD_K * (self._rest_upper(uid) - n['x'])
-                if alpha > 0:
-                    kuc = self.springs[(uid, 'C')]
-                    f  += alpha * kuc * (self.nodes['C']['x'] - n['x'])
-                n['vel']  = n['vel'] * DAMPING + f * DT
-                n['x']   += n['vel'] * DT
-                max_v     = max(max_v, abs(n['vel']))
-
-            # Lower hidden nodes
+                rest = 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']))
+                n['vel'] = n['vel'] * DAMPING + f * DT
+                n['x']  += n['vel'] * DT
+
             for i in range(1, self.n_lower + 1):
                 lid = f'L{i}'
                 n   = self.nodes[lid]
-                f   = FWD_K * (self._rest_lower(lid) - n['x'])
-                if alpha > 0:
-                    klc = self.springs[(lid, 'C')]
-                    f  += alpha * klc * (self.nodes['C']['x'] - n['x'])
-                n['vel']  = n['vel'] * DAMPING + f * DT
-                n['x']   += n['vel'] * DT
-                max_v     = max(max_v, abs(n['vel']))
-
-            # C node (only moves in inference)
+                rest = 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']))
+                n['vel'] = n['vel'] * DAMPING + f * DT
+                n['x']  += n['vel'] * DT
+
             c = self.nodes['C']
             if not c['anchored']:
-                f         = FWD_K * (self._rest_c() - c['x'])
-                c['vel']  = c['vel'] * DAMPING + f * DT
-                c['x']   += c['vel'] * DT
-                max_v     = max(max_v, abs(c['vel']))
+                rest = np.zeros(self.dim)
+                for i in range(1, self.n_upper + 1):
+                    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'])
+                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):
+        """
+        Matrix LMS with joint optimal step.
 
-            if max_v < SETTLE:
-                break
+        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
 
-    # ── EXACT FEEDFORWARD (LEARNING) ─────────────────────────────────────────
-    def _feedforward(self):
-        """
-        Exact prediction, independent of elastic node positions.
-        Additive:       Ui = K(A,Ui)·A   Li = K(B,Li)·B
-        Multiplicative: same shape but U·L cross-product at C boundary.
-        Returns (prediction, hidden_values_dict).
+        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‖²
+
+        EWC mode: step size reduced by (1 + λ·F) per element, protecting
+        dimensions with high Fisher importance from past training.
         """
-        A, B = self.nodes['A']['x'], self.nodes['B']['x']
-        ff   = {}
+        eps  = 1e-8
+        xa   = self.nodes['A']['x']
+        xb   = self.nodes['B']['x']
+
+        # ── Output-layer joint update ──────────────────────────────────────
+        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}']))
 
         for i in range(1, self.n_upper + 1):
-            uid       = f'U{i}'
-            ff[uid]   = self.springs[('A', uid)] * A
+            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])
 
         for i in range(1, self.n_lower + 1):
-            lid       = f'L{i}'
-            ff[lid]   = self.springs[('B', lid)] * B
-
-        if self.architecture == 'multiplicative':
-            # Each Ui pairs with a Li (by index, wrap if unequal counts)
-            pred = 0.0
-            n    = max(self.n_upper, self.n_lower)
-            for i in range(n):
-                uid  = f'U{(i % self.n_upper) + 1}'
-                lid  = f'L{(i % self.n_lower) + 1}'
-                ku   = self.springs[(uid, 'C')]
-                kl   = self.springs[(lid, 'C')]
-                pred += ku * ff[uid] * kl * ff[lid]
-        else:
-            # Additive: both sides simply sum at C
-            pred = (
-                sum(self.springs[(f'U{i}', 'C')] * ff[f'U{i}'] for i in range(1, self.n_upper + 1)) +
-                sum(self.springs[(f'L{i}', 'C')] * ff[f'L{i}'] 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])
+
+        # ── Hidden-layer update (backprop) ────────────────────────────────
+        xa_denom = float(np.dot(xa, xa)) + eps
+        xb_denom = float(np.dot(xb, xb)) + eps
 
-        return pred, ff
+        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])
 
-    # ── LMS OPTIMAL-STEP UPDATE ───────────────────────────────────────────────
-    def _lms_update(self, error, ff):
+        for i in range(1, self.n_lower + 1):
+            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])
+
+    # ── FISHER ACCUMULATION (EWC) ─────────────────────────────────────────────
+
+    def _update_fisher(self, error: np.ndarray, hid: dict):
         """
-        LMS optimal step μ = error / ||∇pred||²
-        Gradients via backprop through feedforward values.
+        Accumulate Fisher diagonal via EMA of squared gradient elements.
+        High Fisher → this weight dimension was important for past problems.
         """
-        grads = {k: 0.0 for k in self.springs}
-        A, B  = self.nodes['A']['x'], self.nodes['B']['x']
+        xa = self.nodes['A']['x']
+        xb = self.nodes['B']['x']
 
-        if self.architecture == 'additive':
-            for i in range(1, self.n_upper + 1):
-                uid = f'U{i}'
-                kau = self.springs[('A', uid)]
-                kuc = self.springs[(uid, 'C')]
-                # ∂pred/∂K(A,Ui) = K(Ui,C)·A
-                grads[('A', uid)] = kuc * A
-                # ∂pred/∂K(Ui,C) = K(A,Ui)·A
-                grads[(uid, 'C')] = kau * A
-            for i in range(1, self.n_lower + 1):
-                lid = f'L{i}'
-                kbl = self.springs[('B', lid)]
-                klc = self.springs[(lid, 'C')]
-                grads[('B', lid)] = klc * B
-                grads[(lid, 'C')] = kbl * B
-        else:
-            n = max(self.n_upper, self.n_lower)
-            for i in range(n):
-                uid  = f'U{(i % self.n_upper) + 1}'
-                lid  = f'L{(i % self.n_lower) + 1}'
-                kau  = self.springs[('A', uid)]
-                kbl  = self.springs[('B', lid)]
-                ku   = self.springs[(uid, 'C')]
-                kl   = self.springs[(lid, 'C')]
-                Uv, Lv = ff[uid], ff[lid]
-                # ∂pred/∂K(A,Ui) = K(Ui,C)·A · K(Li,C)·Lv
-                grads[('A', uid)]  = ku * A * kl * Lv
-                grads[('B', lid)]  = kl * B * ku * Uv
-                grads[(uid, 'C')]  = Uv * kl * Lv
-                grads[(lid, 'C')]  = Lv * ku * Uv
-
-        norm_sq = sum(g * g for g in grads.values()) + 1e-10
-        mu = error / norm_sq
-        for key, g in grads.items():
-            self.springs[key] -= mu * g
-            self.springs[key]  = max(-30.0, min(30.0, self.springs[key]))
-
-    # ── MAIN PHYSICS STEP ─────────────────────────────────────────────────────
-    def physics_step(self):
+        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)
+
+        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)
+
+    # ── PHYSICS STEP ──────────────────────────────────────────────────────────
+
+    def physics_step(self) -> bool:
+        """One server tick: elastic display + LMS update."""
         self._elastic_step(MICRO)
 
-        pred, ff            = self._feedforward()
-        self.current_prediction = round(pred, 5)
+        pred, hid = self._forward()
+        self.pred_norm  = float(np.linalg.norm(pred))
+        self.step_count += 1
 
         c = self.nodes['C']
         if c['anchored']:
-            self.current_error = pred - c['x']
+            error           = pred - c['x']
+            self.error_norm = float(np.linalg.norm(error))
         else:
-            c['x']             = round(pred, 4)
-            self.current_error = 0.0
-
-        self.history.append(round(self.current_error, 4))
-        if len(self.history) > 200:
+            c['x']          = pred.copy()
+            error           = (pred - self.c_target
+                               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))
+        if len(self.history) > 300:
             self.history.pop(0)
 
-        if abs(self.current_error) < 0.02:
-            a, b = self.nodes['A']['x'], self.nodes['B']['x']
-            gt   = self.ground_truth(a, b)
-            self.add_log(
-                f"✓ A={a:.2f} B={b:.2f} → P={pred:.4f} GT={gt:.4f} Δ={abs(pred-gt):.4f}"
-            )
+        converged = self.error_norm < CONV_THRESH
+        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}")
+            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,
+                })
+            self._update_fisher(error, hid)
             return self._next_or_stop()
 
-        if self.mode == 'training' and c['anchored']:
-            self._lms_update(self.current_error, ff)
+        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):
+    def _next_or_stop(self) -> bool:
         if self.batch_queue:
             p = self.batch_queue.popleft()
-            self.set_problem(p['a'], p['b'], p.get('c'))
-            self.add_log(f"→ Next ({len(self.batch_queue)} queued)")
+            self.set_problem(p['A'], p['B'], p.get('C'), p.get('type', 'unknown'))
             return True
         self.running = False
-        self.add_log("◼ Batch complete.")
+        self.add_log("◼ Queue empty.")
         return False
 
-    def generate_batch(self, count=30):
-        self.batch_queue.clear()
-        for _ in range(count):
-            a = round(random.uniform(1.0, 10.0), 2)
-            b = round(random.uniform(1.0, 10.0), 2)
-            self.batch_queue.append({'a': a, 'b': b, 'c': self.ground_truth(a, b)})
-        p = self.batch_queue.popleft()
-        self.set_problem(p['a'], p['b'], p.get('c'))
+    # ── FAST 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.
+        """
+        self.running = False
+        self.mode    = 'training'
+        self.add_log(f"⚡ Offline training: {epochs} epoch(s)…")
+
+        for ep in range(1, epochs + 1):
+            random.shuffle(self.train_data)
+            total_err = 0.0
+            converged = 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]
+                self.nodes['A']['x'] = xa
+                self.nodes['B']['x'] = xb
+                self.nodes['C']['x'] = ct
+
+                for _ in range(MAX_STEPS):
+                    pred, hid = self._forward()
+                    err = pred - ct
+                    en  = float(np.linalg.norm(err))
+                    if en < CONV_THRESH:
+                        self._update_fisher(err, hid)
+                        converged += 1
+                        break
+                    self._lms_update(err, hid, ewc=False)
+
+                total_err += float(np.linalg.norm(self._forward()[0] - ct))
+
+            avg = total_err / max(len(self.train_data), 1)
+            pct = 100 * converged / max(len(self.train_data), 1)
+            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.mode = 'idle'
+
+    # ── DATA LOADING ──────────────────────────────────────────────────────────
+
+    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")
+
+    # ── QUEUE HELPERS ─────────────────────────────────────────────────────────
+
+    def start_training(self, n=None):
+        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.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"▶ {count} samples | {self.dataset_type} | U{self.n_upper}·L{self.n_lower}")
+        self.add_log(f"▶ Inference: {len(data)} samples")
+
+    def _fill_queue(self, data, anchor_c):
+        self.batch_queue.clear()
+        for d in data:
+            self.batch_queue.append(
+                {'A': d['A'], 'B': d['B'], 'C': d['C'], 'type': d.get('type','?')}
+            )
+        if self.batch_queue:
+            p = self.batch_queue.popleft()
+            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['C']['anchored'] = False
+                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':
+                            self.nodes[nid]['x']   = np.zeros(d)
+                            self.nodes[nid]['vel'] = np.zeros(d)
+
+    # ── LOGGING ───────────────────────────────────────────────────────────────
+
+    def add_log(self, msg):
+        self.logs.insert(0, f"[{self.iteration:06d}] {msg}")
+        if len(self.logs) > 60:
+            self.logs.pop()
+
+    # ── STATE SERIALISATION ───────────────────────────────────────────────────
+
+    def state_dict(self):
+        nodes_out = {}
+        for nid, n in self.nodes.items():
+            nodes_out[nid] = {
+                'norm':     round(float(np.linalg.norm(n['x'])),   4),
+                'vel_norm': round(float(np.linalg.norm(n['vel'])), 4),
+                'anchored': bool(n['anchored']),
+                'x_head':   [round(float(v), 3) for v in n['x'][:6]],
+            }
+
+        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),
+                'mean': round(float(np.mean(km)),        4),
+                'std':  round(float(np.std(km)),         4),
+                'fish': round(float(np.mean(self.fisher[(u, v)])), 5),
+            }
+
+        # Per-type inference accuracy
+        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),
+            }
+            for t, v in type_acc.items()
+        }
+
+        return {
+            'nodes':        nodes_out,
+            'springs':      springs_out,
+            'error':        round(self.error_norm, 5),
+            'pred_norm':    round(self.pred_norm,  5),
+            'iter':         self.iteration,
+            'step_count':   self.step_count,
+            'logs':         self.logs,
+            'history':      self.history[-120:],
+            'running':      self.running,
+            'mode':         self.mode,
+            'n_upper':      self.n_upper,
+            'n_lower':      self.n_lower,
+            'layers':       self.layers,
+            'queue_size':   len(self.batch_queue),
+            'train_size':   len(self.train_data),
+            'test_size':    len(self.test_data),
+            'type_acc':     acc_summary,
+            'n_test_done':  len(self.test_errors),
+            'current_type': self.current_type,
+            'dim':          self.dim,
+        }
 
 
 # ── SERVER ────────────────────────────────────────────────────────────────────
-engine = SimEngine()
+
+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})")
+
 
 def run_loop():
     while True:
         if engine.running:
             engine.physics_step()
-        time.sleep(0.028)
+        time.sleep(0.025)
 
 threading.Thread(target=run_loop, daemon=True).start()
 
+
 @app.get("/", response_class=HTMLResponse)
 async def get_ui():
     return FileResponse("index.html")
 
 @app.get("/state")
 async def get_state():
-    springs_out = {f"{u}→{v}": round(k, 5) for (u, v), k in engine.springs.items()}
-    return {
-        'nodes':        engine.nodes,
-        'springs':      springs_out,
-        'layers':       engine.layers,
-        'error':        engine.current_error,
-        'prediction':   engine.current_prediction,
-        'iter':         engine.iteration,
-        'logs':         engine.logs,
-        'history':      engine.history[-80:],
-        'running':      engine.running,
-        'mode':         engine.mode,
-        'architecture': engine.architecture,
-        'dataset_type': engine.dataset_type,
-        'n_upper':      engine.n_upper,
-        'n_lower':      engine.n_lower,
-        'back_alpha':   engine.back_alpha,
-        'queue_size':   len(engine.batch_queue),
-    }
-
-@app.post("/config")
-async def config(data: dict):
-    engine.mode         = data.get('mode',         engine.mode)
-    engine.architecture = data.get('architecture', engine.architecture)
-    engine.dataset_type = data.get('dataset',      engine.dataset_type)
-    engine.n_upper      = max(1, min(8, int(data.get('n_upper',    engine.n_upper))))
-    engine.n_lower      = max(1, min(8, int(data.get('n_lower',    engine.n_lower))))
-    engine.back_alpha   = max(0.0, min(1.0, float(data.get('back_alpha', engine.back_alpha))))
-    engine.running = False
-    engine._init_mesh()
-    engine.logs = []
-    engine.add_log(
-        f"Config: {engine.mode}|{engine.architecture}|U{engine.n_upper}·L{engine.n_lower}|α={engine.back_alpha:.2f}"
-    )
+    return engine.state_dict()
+
+# ── Training controls ─────────────────────────────────────────────────────────
+
+@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 ────────────────────────────────────────────────────────
+
 @app.post("/set_layer")
 async def set_layer(data: dict):
     layer = data.get('layer', '')
     delta = int(data.get('delta', 0))
-    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))
     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))
     engine._init_mesh()
-    engine.add_log(f"Topology → U{engine.n_upper}·L{engine.n_lower}")
+    engine.add_log(f"Topology → U{engine.n_upper} · L{engine.n_lower} | springs re-init")
     return {"ok": True, "n_upper": engine.n_upper, "n_lower": engine.n_lower}
 
-@app.post("/generate")
-async def generate(data: dict):
-    engine.generate_batch(int(data.get('count', 30)))
-    return {"ok": True}
-
-@app.post("/run_custom")
-async def run_custom(data: dict):
-    a = float(data['a'])
-    b = float(data['b'])
-    c = float(data['c']) if data.get('c') not in (None, '', 'null') else None
-    if c is None and engine.mode == 'training':
-        c = engine.ground_truth(a, b)
-    engine.batch_queue.clear()
-    engine.set_problem(a, b, c)
-    engine.running = True
-    engine.add_log(f"Custom: A={a} B={b} target={c}")
-    return {"ok": True}
-
 @app.post("/halt")
 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}
+
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)