Update data_gen.py

data_gen.py

@@ -1,211 +1,80 @@
 """
-data_gen.py  —  Training / test data for the elastic mesh.
+data_gen.py
 
-OOD TEST DESIGN
-───────────────
-  SEEN during training   : box_proj | halfspace | elastic
-  UNSEEN (OOD) at test   : sphere   | simplex
+Scalar mesh test data.  Each sample is (a, b, c) — all single floats.
 
-This lets us distinguish:
-  • Memorisation  → high acc on seen, low acc on unseen
-  • Geometry      → high acc on both  (the real claim)
+A ∈ [0.1, 0.9]  — constraint scalar (top input)
+B ∈ [0.1, 0.9]  — objective scalar (bottom input)
+C              — feasibility center (what the mesh must learn to produce)
 
-Each sample: (A, B, C)  where A=constraints, B=objectives, C=feasibility center.
-DIM = 64  (double from previous run, stress-tests before LLM scale).
-"""
-
-import numpy as np
-import json, pathlib, argparse
-from typing import List, Dict
-
-
-DIM              = 64
-SAMPLES_PER_TYPE = 1000   # × 5 types = 5 000 total
-
-
-# ── UTILITIES ─────────────────────────────────────────────────────────────────
-
-def norm(v: np.ndarray) -> np.ndarray:
-    return v / (np.linalg.norm(v) + 1e-12)
-
-def pack(*arrays, dim):
-    v = np.concatenate(arrays)
-    return v[:dim] if len(v) >= dim else np.pad(v, (0, dim - len(v)))
-
-
-# ── PROBLEM TYPE 1 (SEEN): BOX PROJECTION ────────────────────────────────────
-# C = clip(B, lo, hi)
-# A encodes the box bounds
-
-def gen_box(n, dim, rng):
-    data = []
-    for _ in range(n):
-        center = rng.uniform(-2, 2, dim)
-        half   = rng.uniform(0.3, 2.0, dim)
-        lo, hi = center - half, center + half
-        B = rng.uniform(-4, 4, dim)
-        C = np.clip(B, lo, hi)
-        A = pack(lo[:dim//2], hi[:dim//2], dim=dim)
-        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'box_proj'})
-    return data
-
-
-# ── PROBLEM TYPE 2 (SEEN): HALFSPACE PROJECTION ───────────────────────────────
-# C = B − (nᵀB − b)·n   (project B onto hyperplane nᵀx = b)
-
-def gen_halfspace(n, dim, rng):
-    data = []
-    for _ in range(n):
-        normal = norm(rng.standard_normal(dim))
-        b      = float(rng.uniform(-1, 1))
-        B      = rng.uniform(-3, 3, dim)
-        C      = B - (float(np.dot(normal, B)) - b) * normal
-        A      = normal.copy(); A[0] = b
-        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'halfspace'})
-    return data
-
-
-# ── PROBLEM TYPE 3 (SEEN): ELASTIC BALANCE ────────────────────────────────────
-# C[j] = w[j]·a_center[j] + (1−w[j])·B[j]   per-dimension soft trade-off
+All C values kept in [0.1, 0.9] so the scalar Hooke mesh can represent them
+without needing to amplify beyond the input range.
 
-def gen_elastic(n, dim, rng):
-    data = []
-    for _ in range(n):
-        a_center = rng.uniform(-2, 2, dim)
-        w        = rng.uniform(0.05, 0.95, dim)
-        B        = rng.uniform(-3, 3, dim)
-        C        = w * a_center + (1.0 - w) * B
-        A        = pack(a_center[:dim//2], w[:dim//2], dim=dim)
-        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'elastic'})
-    return data
-
-
-# ── PROBLEM TYPE 4 (OOD): SPHERE SURFACE ─────────────────────────────────────
-# C = center + r·(B−center)/‖B−center‖   (nearest point on sphere to B)
-
-def gen_sphere(n, dim, rng):
-    data = []
-    for _ in range(n):
-        center = rng.uniform(-1.5, 1.5, dim)
-        r      = float(rng.uniform(1.0, 3.0))
-        B      = rng.uniform(-4, 4, dim)
-        diff   = B - center
-        nd     = np.linalg.norm(diff)
-        if nd < 1e-10:
-            diff = np.ones(dim) / np.sqrt(dim); nd = 1.0
-        C    = center + r * diff / nd
-        A    = center.copy(); A[0] = r
-        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'sphere'})
-    return data
-
-
-# ── PROBLEM TYPE 5 (OOD): SIMPLEX PROJECTION ─────────────────────────────────
-# C = nearest point on probability simplex to B  (Σxᵢ=1, xᵢ≥0)
-
-def _proj_simplex(v):
-    n  = len(v)
-    u  = np.sort(v)[::-1]
-    cs = np.cumsum(u) - 1.0
-    rho = int(np.where(u * np.arange(1, n+1) > cs)[0][-1])
-    theta = cs[rho] / (rho + 1.0)
-    return np.maximum(v - theta, 0.0)
+SEEN during training : heavy_a | avg | diff
+OOD  (test only)     : heavy_b
 
-def gen_simplex(n, dim, rng):
-    data = []
-    for _ in range(n):
-        A = np.ones(dim)
-        B = rng.uniform(-1.0, 3.0, dim)
-        C = _proj_simplex(B)
-        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'simplex'})
-    return data
+If the mesh generalises:
+  It learned "weighted combination of A and B" as a concept
+  → it applies unseen weighting (B-heavy) without training on it.
+"""
 
+import numpy as np, json, pathlib, random, argparse
 
-# ── ASSEMBLY ──────────────────────────────────────────────────────────────────
+SAMPLES_PER_TYPE = 2500  # × 4 types = 10 000 total
 
-SEEN_TYPES = {
-    'box_proj':  gen_box,
-    'halfspace': gen_halfspace,
-    'elastic':   gen_elastic,
+DATASETS = {
+    # name : (lambda, seen?)
+    'heavy_a':  (lambda a, b: 0.8*a + 0.2*b,           True),
+    'avg':      (lambda a, b: 0.5*a + 0.5*b,           True),
+    'diff':     (lambda a, b: 0.5 + 0.4*(a - b),       True),   # signed diff, offset to [0.1,0.9]
+    'heavy_b':  (lambda a, b: 0.2*a + 0.8*b,           False),  # OOD
 }
-OOD_TYPES = {
-    'sphere':  gen_sphere,
-    'simplex': gen_simplex,
-}
-ALL_TYPES = {**SEEN_TYPES, **OOD_TYPES}
-
 
-def generate_all(n_per_type=SAMPLES_PER_TYPE, dim=DIM, seed=42):
-    rng  = np.random.default_rng(seed)
+def generate(n_per=SAMPLES_PER_TYPE, seed=42):
+    rng = np.random.default_rng(seed)
     data = []
-    for fn in ALL_TYPES.values():
-        data.extend(fn(n_per_type, dim, rng))
-    idx = rng.permutation(len(data))
-    return [data[i] for i in idx]
-
+    for dtype, (fn, _) in DATASETS.items():
+        for _ in range(n_per):
+            a = float(rng.uniform(0.1, 0.9))
+            b = float(rng.uniform(0.1, 0.9))
+            c = fn(a, b)
+            data.append({'a': round(a,4), 'b': round(b,4),
+                         'c': round(c,4), 'type': dtype})
+    random.shuffle(data)
+    return data
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--dim', type=int, default=DIM)
     parser.add_argument('--n',   type=int, default=SAMPLES_PER_TYPE)
     parser.add_argument('--out', type=str, default='data')
     args = parser.parse_args()
 
-    print(f"\n{'─'*55}")
-    print(f"  Generating {5 * args.n} samples  |  dim={args.dim}")
-    print(f"  SEEN  : box_proj | halfspace | elastic")
-    print(f"  OOD   : sphere   | simplex")
-    print(f"{'─'*55}")
-
-    rng = np.random.default_rng(42)
-
-    seen_data, ood_data = [], []
-    for t, fn in SEEN_TYPES.items():
-        seen_data.extend(fn(args.n, args.dim, rng))
-    for t, fn in OOD_TYPES.items():
-        ood_data.extend(fn(args.n, args.dim, rng))
-
-    # Shuffle within splits
-    si = rng.permutation(len(seen_data))
-    oi = rng.permutation(len(ood_data))
-    seen_data = [seen_data[i] for i in si]
-    ood_data  = [ood_data[i]  for i in oi]
-
-    # Train = 90% of SEEN only
-    # Test  = 10% of SEEN  +  ALL OOD  (so model never trained on OOD)
-    split    = int(len(seen_data) * 0.9)
-    train    = seen_data[:split]
-    test_seen = seen_data[split:]
-    test     = test_seen + ood_data
-
-    # Re-shuffle test so seen/OOD are interleaved
-    ti = rng.permutation(len(test))
-    test = [test[i] for i in ti]
+    data = generate(args.n)
+
+    # Split: train = SEEN only (90%), test = 10% SEEN + ALL OOD
+    seen = [d for d in data if DATASETS[d['type']][1]]
+    ood  = [d for d in data if not DATASETS[d['type']][1]]
+
+    split    = int(len(seen) * 0.9)
+    train    = seen[:split]
+    test     = seen[split:] + ood
+
+    random.shuffle(test)
 
     out = pathlib.Path(args.out)
     out.mkdir(exist_ok=True)
-    with open(out / 'train.json', 'w') as f: json.dump(train, f)
-    with open(out / 'test.json',  'w') as f: json.dump(test,  f)
+    with open(out/'train.json','w') as f: json.dump(train, f)
+    with open(out/'test.json', 'w') as f: json.dump(test,  f)
 
     from collections import Counter
-    tr_types = Counter(d['type'] for d in train)
-    te_types = Counter(d['type'] for d in test)
-
-    print(f"\n  {'Type':<14} {'Train':>7} {'Test':>7}  {'Split'}")
-    print(f"  {'─'*14} {'─'*7} {'─'*7}  {'─'*10}")
-    for t in ALL_TYPES:
-        label = 'OOD ✗' if t in OOD_TYPES else 'SEEN ✓'
-        print(f"  {t:<14} {tr_types.get(t,0):>7} {te_types.get(t,0):>7}  {label}")
-    print(f"\n  Total  train={len(train)}  test={len(test)}\n")
-
-    # Quick sanity: verify C is geometrically correct for first sample per type
-    print(f"  Sanity check:")
-    seen_set = set()
-    for d in train + test:
-        t = d['type']
-        if t in seen_set: continue
-        seen_set.add(t)
-        A, B, C = map(np.array, [d['A'], d['B'], d['C']])
-        print(f"  [{t:<12}]  ‖A‖={np.linalg.norm(A):.2f}  "
-              f"‖B‖={np.linalg.norm(B):.2f}  ‖C‖={np.linalg.norm(C):.2f}")
-
-    print(f"\n  Saved → {out}/train.json  {out}/test.json\n")
+    tr = Counter(d['type'] for d in train)
+    te = Counter(d['type'] for d in test)
+
+    print(f"\n{'─'*50}")
+    print(f"  {'Type':<12} {'Train':>7} {'Test':>7}  Split")
+    print(f"  {'─'*12} {'─'*7} {'─'*7}  {'─'*8}")
+    for t, (fn, seen_flag) in DATASETS.items():
+        label = 'SEEN' if seen_flag else 'OOD ✗'
+        print(f"  {t:<12} {tr.get(t,0):>7} {te.get(t,0):>7}  {label}")
+    print(f"\n  Train total: {len(train)}   Test total: {len(test)}")
+    print(f"  Saved → {out}/train.json  {out}/test.json\n")