Update app.py

app.py

@@ -15,7 +15,6 @@ CONV_THRESH    = 0.02
 STIFF_EPS      = 0.0005
 STIFF_INTERVAL = 1.0
 
-
 # ── ATOMIC UNIT ───────────────────────────────────────────────────────────────
 
 class HourglassUnit:
@@ -160,7 +159,6 @@ class HourglassUnit:
             'n_lower':         self.n_lower,
         }
 
-
 # ── ENGINE ────────────────────────────────────────────────────────────────────
 
 class SimEngine:
@@ -196,8 +194,6 @@ class SimEngine:
         self.stiffness_history = []
         self._init_stack()
 
-    # ── TOPOLOGY ──────────────────────────────────────────────────────────────
-
     def _init_stack(self):
         self.iteration   = 0
         self.current_error = 0.0
@@ -208,6 +204,8 @@ class SimEngine:
         self.connections = []
 
         n    = max(1, self.n_inputs)
+        if self.stack_levels > 0: n = 2 ** self.stack_levels
+        
         nu   = self.n_upper
         nl   = self.n_lower
         arch = self.architecture
@@ -237,8 +235,6 @@ class SimEngine:
                 self.connections.append({'from_uid': ub, 'to_uid': uid, 'to_port': 'B'})
             self.topology.append(curr)
 
-    # ── DATASET ───────────────────────────────────────────────────────────────
-
     def _to_vec(self, val, n):
         if isinstance(val, (list, tuple)):
             v = [float(x) for x in val]
@@ -256,15 +252,18 @@ class SimEngine:
             else:                       result.append(round(a + b,          4))
         return result
 
-    # ── PROBLEM SETUP ─────────────────────────────────────────────────────────
-
     def set_problem(self, a, b, c_target=None):
         n = max(1, self.n_inputs)
+        if self.stack_levels > 0: n = 2 ** self.stack_levels
+        
         self._a_vec = self._to_vec(a, n)
         self._b_vec = self._to_vec(b, n)
         self._c_vec = list(self._to_vec(c_target, n)) if c_target is not None else []
-
-    # ── LOGGING ───────────────────────────────────────────────────────────────
+        
+        # PROPERLY RESET PHYSICS ONLY ONCE PER PROBLEM
+        for u in self.units.values():
+            u.reset_hidden()
+        self.current_error = 999.0  # Force out of convergence
 
     def add_log(self, msg):
         self.logs.insert(0, f"[{self.iteration:05d}] {msg}")
@@ -274,8 +273,6 @@ class SimEngine:
         self.running = False; self.batch_queue.clear()
         self.logs = []; self._init_stack()
 
-    # ── FORWARD PASS ──────────────────────────────────────────────────────────
-
     def _forward_pass(self, training):
         a_vec = self._a_vec; b_vec = self._b_vec; c_vec = self._c_vec
         alpha = self.back_alpha; io = {}
@@ -284,7 +281,7 @@ class SimEngine:
         for i, uid in enumerate(self.topology[0]):
             av = a_vec[i] if i < len(a_vec) else a_vec[-1]
             bv = b_vec[i] if i < len(b_vec) else b_vec[-1]
-            u  = self.units[uid]; u.reset_hidden()
+            u  = self.units[uid]
             anchor = training and bool(c_vec) and (only_level or self.individual_train)
             if anchor:
                 if self.individual_train and not only_level:
@@ -304,7 +301,7 @@ class SimEngine:
                 b_src = next((c['from_uid'] for c in conns if c['to_port'] == 'B'), None)
                 av    = io[a_src][2] if a_src in io else 0.0
                 bv    = io[b_src][2] if b_src in io else 0.0
-                u     = self.units[uid]; u.reset_hidden()
+                u     = self.units[uid]
                 is_root = (lv == len(self.topology) - 1)
                 anchor  = training and bool(c_vec) and is_root
                 c_tgt   = c_vec[0] if anchor else 0.0
@@ -318,75 +315,45 @@ class SimEngine:
         uid  = self.topology[0][0]; u = self.units[uid]
         c_in = self._c_vec[0] if self._c_vec else 0.0
         a_gt = self._a_vec[0] if self._a_vec else 0.0
-        u.reset_hidden()
         u.elastic_step(c_in, c_in, a_gt if training else 0.0, training, self.back_alpha)
         pred, ff = u.feedforward(c_in, c_in)
         if training:
             u.lms_update(pred - a_gt, c_in, c_in, ff)
         return pred, {uid: (c_in, c_in, pred, ff)}
 
-    # ── CREDIT ASSIGNMENT ─────────────────────────────────────────────────────
-
-    def _backprop(self, root_err, io):
-        errors = {self.topology[-1][-1]: root_err}
-        for lv in range(len(self.topology)-1, -1, -1):
-            for uid in self.topology[lv]:
-                if uid not in errors: continue
-                err = errors[uid]; av, bv, _, ff = io[uid]
-                self.units[uid].lms_update(err, av, bv, ff)
-                if lv > 0:
+    def _train(self, root_err, io):
+        m = self.credit_mode
+        if m == 'elastic_backprop':
+            errors = {self.topology[-1][-1]: root_err}
+            for lv in range(len(self.topology)-1, -1, -1):
+                for uid in self.topology[lv]:
+                    if uid not in errors: continue
+                    err = errors[uid]; av, bv, _, ff = io[uid]
+                    self.units[uid].lms_update(err, av, bv, ff)
+                    if lv > 0:
+                        u     = self.units[uid]
+                        conns = [c for c in self.connections if c['to_uid'] == uid]
+                        a_src = next((c['from_uid'] for c in conns if c['to_port'] == 'A'), None)
+                        b_src = next((c['from_uid'] for c in conns if c['to_port'] == 'B'), None)
+                        if a_src: errors[a_src] = errors.get(a_src, 0.0) + err * u.sensitivity_to_a()
+                        if b_src and b_src != a_src: errors[b_src] = errors.get(b_src, 0.0) + err * u.sensitivity_to_b()
+        elif m == 'independent':
+            errors = {self.topology[-1][-1]: root_err}
+            for lv in range(len(self.topology)-1, 0, -1):
+                for uid in self.topology[lv]:
+                    if uid not in errors: continue
+                    err = errors[uid]; av, bv, _, ff = io[uid]
+                    self.units[uid].lms_update(err, av, bv, ff)
                     u     = self.units[uid]
                     conns = [c for c in self.connections if c['to_uid'] == uid]
                     a_src = next((c['from_uid'] for c in conns if c['to_port'] == 'A'), None)
                     b_src = next((c['from_uid'] for c in conns if c['to_port'] == 'B'), None)
-                    if a_src:
-                        errors[a_src] = errors.get(a_src, 0.0) + err * u.sensitivity_to_a()
-                    if b_src and b_src != a_src:
-                        errors[b_src] = errors.get(b_src, 0.0) + err * u.sensitivity_to_b()
-
-    def _greedy(self, root_err, io):
-        n = len(self.topology)
-        for lv in range(n-1, -1, -1):
-            scale = 1.0 / max(1, n - lv)
-            for uid in self.topology[lv]:
-                av, bv, _, ff = io[uid]
-                self.units[uid].lms_update(root_err * scale, av, bv, ff)
-
-    def _param_hg(self, root_err, io):
-        for level in self.topology:
-            for uid in level:
-                av, bv, _, ff = io[uid]
-                self.units[uid].lms_update(root_err, av, bv, ff)
-
-    def _independent(self, root_err, io):
-        errors = {self.topology[-1][-1]: root_err}
-        for lv in range(len(self.topology)-1, 0, -1):
-            for uid in self.topology[lv]:
-                if uid not in errors: continue
-                err = errors[uid]; av, bv, _, ff = io[uid]
-                self.units[uid].lms_update(err, av, bv, ff)
-                u     = self.units[uid]
-                conns = [c for c in self.connections if c['to_uid'] == uid]
-                a_src = next((c['from_uid'] for c in conns if c['to_port'] == 'A'), None)
-                b_src = next((c['from_uid'] for c in conns if c['to_port'] == 'B'), None)
-                if a_src:
-                    errors[a_src] = errors.get(a_src, 0.0) + err * u.sensitivity_to_a()
-                if b_src and b_src != a_src:
-                    errors[b_src] = errors.get(b_src, 0.0) + err * u.sensitivity_to_b()
-        # Leaves: independent GT
-        for i, uid in enumerate(self.topology[0]):
-            av, bv, c_pred, ff = io[uid]
-            gt_i = self.ground_truth([av], [bv])[0]
-            self.units[uid].lms_update(c_pred - gt_i, av, bv, ff)
-
-    def _train(self, root_err, io):
-        m = self.credit_mode
-        if   m == 'elastic_backprop': self._backprop(root_err, io)
-        elif m == 'greedy':           self._greedy(root_err, io)
-        elif m == 'param_hg':         self._param_hg(root_err, io)
-        elif m == 'independent':      self._independent(root_err, io)
-
-    # ── STIFFNESS MONITOR ─────────────────────────────────────────────────────
+                    if a_src: errors[a_src] = errors.get(a_src, 0.0) + err * u.sensitivity_to_a()
+                    if b_src and b_src != a_src: errors[b_src] = errors.get(b_src, 0.0) + err * u.sensitivity_to_b()
+            for i, uid in enumerate(self.topology[0]):
+                av, bv, c_pred, ff = io[uid]
+                gt_i = self.ground_truth([av], [bv])[0]
+                self.units[uid].lms_update(c_pred - gt_i, av, bv, ff)
 
     def _update_stiffness(self):
         now = time.time()
@@ -399,12 +366,9 @@ class SimEngine:
         if len(self.stiffness_history) > 60:
             self.stiffness_history.pop(0)
 
-    # ── PHYSICS STEP ──────────────────────────────────────────────────────────
-
     def physics_step(self):
         training = (self.mode == 'training')
 
-        # Reverse mode: C→A reconstruction
         if self.reverse_mode:
             pred, io = self._forward_reverse(training)
             a_gt     = self._a_vec[0] if self._a_vec else 0.0
@@ -419,10 +383,8 @@ class SimEngine:
                 return self._next_or_stop()
             return True
 
-        # Normal forward/stacked pass
         pred, io = self._forward_pass(training)
 
-        # Flat multi-unit: average error across all leaf outputs vs their own GT
         if len(self.topology) == 1 and len(self.topology[0]) > 1:
             errors = []
             for i, uid in enumerate(self.topology[0]):
@@ -447,7 +409,6 @@ class SimEngine:
             self.iteration += 1
             return True
 
-        # Stacked: single root output vs single GT
         c_gt     = self._c_vec[0] if self._c_vec else None
         root_err = (pred - c_gt) if c_gt is not None else 0.0
         self.current_error      = round(abs(root_err), 5)
@@ -458,10 +419,7 @@ class SimEngine:
 
         if self.current_error < CONV_THRESH:
             n_units = len(self.units)
-            self.add_log(
-                f"✓ STACK L{self.stack_levels} [{n_units}u] "
-                f"P={pred:.4f} GT={c_gt:.4f} Δ={abs(root_err):.4f}"
-            )
+            self.add_log(f"✓ STACK L{self.stack_levels} [{n_units}u] P={pred:.4f} GT={c_gt:.4f} Δ={abs(root_err):.4f}")
             return self._next_or_stop()
 
         if training and c_gt is not None:
@@ -483,6 +441,8 @@ class SimEngine:
     def generate_batch(self, count=30):
         self.batch_queue.clear()
         n = max(1, self.n_inputs)
+        if self.stack_levels > 0: n = 2 ** self.stack_levels
+        
         for _ in range(count):
             a_vec = [round(random.uniform(1.0, 10.0), 2) for _ in range(n)]
             b_vec = [round(random.uniform(1.0, 10.0), 2) for _ in range(n)]
@@ -492,43 +452,31 @@ class SimEngine:
                 if self.stack_levels == 0:
                     c_vec = self.ground_truth(a_vec, b_vec)
                 else:
-                    # Stacked: single output GT is the function of the first pair
                     c_vec = [self.ground_truth([a_vec[0]], [b_vec[0]])[0]]
             self.batch_queue.append({'a': a_vec, 'b': b_vec, 'c': c_vec})
         p = self.batch_queue.popleft()
         self.set_problem(p['a'], p['b'], p.get('c'))
         self.running = True
         tag = f'L{self.stack_levels}' if self.stack_levels else 'flat'
-        self.add_log(
-            f"▶ {count} | {self.dataset_type} | "
-            f"D={n} U{self.n_upper}·L{self.n_lower} "
-            f"[{tag}|{self.credit_mode[:4]}]"
-        )
+        self.add_log(f"▶ {count} | D={n} | [{tag}|{self.credit_mode[:4]}]")
 
 
 # ── SERVER ────────────────────────────────────────────────────────────────────
 engine = SimEngine()
 
-
 def run_loop():
     while True:
-        if engine.running:
-            engine.physics_step()
+        if engine.running: engine.physics_step()
         time.sleep(0.028)
 
-
 threading.Thread(target=run_loop, daemon=True).start()
 
-
 @app.get("/", response_class=HTMLResponse)
-async def get_ui():
-    return FileResponse("index.html")
-
+async def get_ui(): return FileResponse("index.html")
 
 @app.get("/state")
 async def get_state():
     units_out = {uid: u.to_dict() for uid, u in engine.units.items()}
-    # Flatten all springs for display (prefixed with unit id)
     all_springs = {}
     for uid, u in engine.units.items():
         short = uid.replace('HG_', '')
@@ -560,124 +508,88 @@ async def get_state():
         'queue_size':        len(engine.batch_queue),
         'stiffness_active':  engine.stiffness_active,
         'stiffness_history': engine.stiffness_history,
-        'a_vec':             engine._a_vec,
-        'b_vec':             engine._b_vec,
-        'c_vec':             engine._c_vec,
     }
 
-
 @app.post("/set_mode")
 async def set_mode(data: dict):
-    engine.mode    = data.get('mode', engine.mode)
+    engine.mode = data.get('mode', engine.mode)
     engine.running = False
-    engine.add_log(f"Mode → {engine.mode}  (springs preserved)")
     return {"ok": True}
 
-
 @app.post("/config")
 async def config(data: dict):
-    new_ni = max(1, min(8,  int(data.get('n_inputs',      engine.n_inputs))))
-    new_nu = max(1, min(16, int(data.get('n_upper',       engine.n_upper))))
-    new_nl = max(1, min(16, int(data.get('n_lower',       engine.n_lower))))
-    new_sl = max(0, min(4,  int(data.get('stack_levels',  engine.stack_levels))))
+    engine.architecture    = data.get('architecture', engine.architecture)
+    engine.dataset_type    = data.get('dataset', engine.dataset_type)
+    engine.back_alpha      = max(0.0, min(1.0, float(data.get('back_alpha', engine.back_alpha))))
+    engine.credit_mode     = data.get('credit_mode', engine.credit_mode)
+    engine.reverse_mode    = bool(data.get('reverse_mode', engine.reverse_mode))
+    engine.individual_train= bool(data.get('individual_train', engine.individual_train))
+    
+    new_ni = max(1, min(16, int(data.get('n_inputs', engine.n_inputs))))
+    new_nu = max(1, min(16, int(data.get('n_upper', engine.n_upper))))
+    new_nl = max(1, min(16, int(data.get('n_lower', engine.n_lower))))
+    new_sl = max(0, min(4,  int(data.get('stack_levels', engine.stack_levels))))
 
     topo_changed = (new_ni != engine.n_inputs or new_nu != engine.n_upper or
                     new_nl != engine.n_lower  or new_sl != engine.stack_levels)
 
-    engine.architecture    = data.get('architecture',     engine.architecture)
-    engine.dataset_type    = data.get('dataset',          engine.dataset_type)
-    engine.back_alpha      = max(0.0, min(1.0, float(data.get('back_alpha', engine.back_alpha))))
-    engine.credit_mode     = data.get('credit_mode',      engine.credit_mode)
-    engine.reverse_mode    = bool(data.get('reverse_mode', engine.reverse_mode))
-    engine.individual_train= bool(data.get('individual_train', engine.individual_train))
-    engine.n_inputs        = new_ni
-    engine.n_upper         = new_nu
-    engine.n_lower         = new_nl
-    engine.stack_levels    = new_sl
-    if 'mode' in data:
-        engine.mode        = data['mode']
+    engine.n_inputs     = new_ni
+    engine.n_upper      = new_nu
+    engine.n_lower      = new_nl
+    engine.stack_levels = new_sl
 
     if topo_changed:
         engine.running = False
         engine._init_stack()
         engine.logs = []
-        n_units = len(engine.units)
-        engine.add_log(
-            f"Stack rebuilt: D={new_ni} U{new_nu}·L{new_nl} "
-            f"levels={new_sl} units={n_units} "
-            f"credit={engine.credit_mode}"
-        )
-    else:
-        engine.add_log(
-            f"Config: {engine.mode}|{engine.architecture}"
-            f"|α={engine.back_alpha:.2f}|{engine.credit_mode}"
-        )
-
-    return {"ok": True, "topo_changed": topo_changed, "n_units": len(engine.units)}
-
+    
+    return {"ok": True}
 
 @app.post("/set_layer")
 async def set_layer(data: dict):
     layer = data.get('layer', '')
     delta = int(data.get('delta', 0))
-    if   layer == 'inputs': engine.n_inputs      = max(1, min(8,  engine.n_inputs      + delta))
-    elif layer == 'upper':  engine.n_upper        = max(1, min(16, engine.n_upper        + delta))
-    elif layer == 'lower':  engine.n_lower        = max(1, min(16, engine.n_lower        + delta))
-    elif layer == 'stack':  engine.stack_levels   = max(0, min(4,  engine.stack_levels   + delta))
+    if   layer == 'inputs': engine.n_inputs     = max(1, min(16, engine.n_inputs + delta))
+    elif layer == 'upper':  engine.n_upper      = max(1, min(16, engine.n_upper + delta))
+    elif layer == 'lower':  engine.n_lower      = max(1, min(16, engine.n_lower + delta))
     engine.running = False
     engine._init_stack()
-    engine.add_log(
-        f"Topology → D={engine.n_inputs} U{engine.n_upper}·L{engine.n_lower} "
-        f"stack={engine.stack_levels} units={len(engine.units)}"
-    )
     return {
-        "ok":           True,
-        "n_inputs":     engine.n_inputs,
-        "n_upper":      engine.n_upper,
-        "n_lower":      engine.n_lower,
-        "stack_levels": engine.stack_levels,
-        "n_units":      len(engine.units),
+        "ok": True,
+        "n_inputs": engine.n_inputs, "n_upper": engine.n_upper,
+        "n_lower": engine.n_lower, "stack_levels": engine.stack_levels
     }
 
-
 @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):
     def parse(v):
         if v is None or v in ('', 'null'): return None
-        if isinstance(v, list):            return [float(x) for x in v]
-        s = str(v)
-        if ',' in s:
-            return [float(x.strip()) for x in s.split(',') if x.strip()]
-        try:    return float(s)
+        if isinstance(v, list): return [float(x) for x in v]
+        if ',' in str(v): return [float(x.strip()) for x in str(v).split(',') if x.strip()]
+        try: return float(v)
         except: return None
 
     a = parse(data.get('a')) or 5.0
     b = parse(data.get('b')) or 3.0
+    
+    # We now strictly trust the C target calculated and sent by the UI
     c = parse(data.get('c'))
-    if c is None and engine.mode == 'training':
-        c = engine.ground_truth(
-            engine._to_vec(a, engine.n_inputs),
-            engine._to_vec(b, engine.n_inputs),
-        )
+    
     engine.batch_queue.clear()
     engine.set_problem(a, b, c)
     engine.running = True
-    engine.add_log(f"Custom: A={a} B={b} C={c}")
     return {"ok": True}
 
-
 @app.post("/halt")
 async def halt():
     engine.running = False
     return {"ok": True}
 
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)