Update data_gen.py

data_gen.py

@@ -1,74 +1,243 @@
 """
-data_gen.py  v5
+data_gen.py  —  Training / test data for the elastic mesh.
 
-Each sample: (A, B, C)  where A,B,C ∈ ℝ^n, all values in [0.1, 0.9]
+Each sample is a triple (A, B, C) where:
+  A  ∈ ℝ^DIM   encodes constraints   ("what must be true")
+  B  ∈ ℝ^DIM   encodes objectives    ("what we want")
+  C  ∈ ℝ^DIM   is the analytic solution — the feasibility center the mesh must learn to produce
 
-For n=1 these are plain scalars.
-For n>1 each dimension is an independent weighted combination of A[i] and B[i],
-so the mesh must learn to route each channel correctly through the bulge.
+Five problem families, each with a geometrically distinct C:
 
-SEEN during training : heavy_a | avg | diff
-OOD  (test only)     : heavy_b   ← has never been seen, tests geometric generalisation
+  1. box_proj      — clamp B into axis-aligned box defined by A
+  2. halfspace     — project B onto hyperplane defined by A
+  3. sphere        — project B onto sphere surface defined by A
+  4. simplex       — project B onto probability simplex (A = uniform prior signal)
+  5. elastic_bal   — per-dimension weighted balance between A-center and B
+
+These cover:
+  - Bounded feasibility   (box)
+  - Equality constraints  (halfspace)
+  - Norm constraints      (sphere)
+  - Probability/sum=1     (simplex)
+  - Soft trade-offs       (elastic)
+
+The mesh sees ONLY (A, B) during inference; C is what it must reconstruct.
 """
 
-import numpy as np, json, pathlib, random, argparse
-from collections import Counter
+import numpy as np
+import json, pathlib, argparse
+from typing import List, Dict
 
-N               = 1      # embedding dimension (1 = pure scalar, set >1 for vector)
-SAMPLES_PER_TYPE = 2500
+DIM              = 32    # embedding dimension  (set to 768 for LLM-scale)
+SAMPLES_PER_TYPE = 1000  # × 5 types = 5 000 total
+
+
+# ── UTILITIES ─────────────────────────────────────────────────────────────────
+
+def normalize(v: np.ndarray) -> np.ndarray:
+    n = np.linalg.norm(v)
+    return v / (n + 1e-12)
+
+def pack(*arrays: np.ndarray, dim: int) -> np.ndarray:
+    """Concatenate + trim/pad to `dim`."""
+    v = np.concatenate(arrays)
+    if len(v) >= dim:
+        return v[:dim]
+    return np.pad(v, (0, dim - len(v)))
+
+
+# ── PROBLEM TYPE 1: BOX PROJECTION ────────────────────────────────────────────
+#
+# Constraint  A : encodes per-dimension box  [lo, hi]
+#               A[:D/2] = lo[:D/2],  A[D/2:] = hi[:D/2]
+# Objective   B : unconstrained target point in ℝ^D
+# Solution    C : clip(B, lo, hi)   — nearest point in box to B
+#
+# Meaning: "stay within resource/capacity bounds while aiming for B"
+
+def gen_box(n: int, dim: int, rng: np.random.Generator) -> List[Dict]:
+    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: HALFSPACE PROJECTION ──────────────────────────────────────
+#
+# Constraint  A : encodes a hyperplane  nᵀx = b
+#               A = normal vector, A[0] carries the offset b
+# Objective   B : unconstrained point in ℝ^D
+# Solution    C : projection of B onto the hyperplane
+#               C = B − (nᵀB − b) · n
+#
+# Meaning: "satisfy one hard equality constraint at minimum cost to B"
+
+def gen_halfspace(n: int, dim: int, rng: np.random.Generator) -> List[Dict]:
+    data = []
+    for _ in range(n):
+        normal = normalize(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          # offset embedded in first slot
+        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'halfspace'})
+    return data
+
+
+# ── PROBLEM TYPE 3: SPHERE SURFACE ────────────────────────────────────────────
+#
+# Constraint  A : encodes a sphere (center, radius)
+#               A = center vector, A[0] overwritten with radius r
+# Objective   B : external point
+# Solution    C : point on sphere surface nearest to B
+#               C = center + r · (B − center) / ‖B − center‖
+#
+# Meaning: "satisfy a norm/budget constraint, move toward B as far as allowed"
+
+def gen_sphere(n: int, dim: int, rng: np.random.Generator) -> List[Dict]:
+    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          # radius in first slot
+        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'sphere'})
+    return data
+
+
+# ── PROBLEM TYPE 4: SIMPLEX PROJECTION ────────────────────────────────────────
+#
+# Constraint  A : uniform-prior signal (all ones) → encodes simplex constraint Σxᵢ=1, xᵢ≥0
+# Objective   B : unconstrained "belief" vector
+# Solution    C : nearest point on probability simplex to B
+#
+# Meaning: "find a valid probability distribution closest to unconstrained belief B"
+# Useful for softmax-like problems.
+
+def _proj_simplex(v: np.ndarray) -> np.ndarray:
+    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)
+
+def gen_simplex(n: int, dim: int, rng: np.random.Generator) -> List[Dict]:
+    data = []
+    for _ in range(n):
+        A = np.ones(dim)                     # simplex constraint signal
+        B = rng.uniform(-1.0, 3.0, dim)      # unconstrained belief
+        C = _proj_simplex(B)
+        data.append({'A': A.tolist(), 'B': B.tolist(), 'C': C.tolist(), 'type': 'simplex'})
+    return data
 
-DATASETS = {
-    '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),   # maps to [0.1, 0.9]
-    'heavy_b': (lambda a, b: 0.2*a + 0.8*b,      False),  # OOD
-}
 
-def generate(n=N, n_per=SAMPLES_PER_TYPE, seed=42):
-    rng = np.random.default_rng(seed)
+# ── PROBLEM TYPE 5: ELASTIC BALANCE ───────────────────────────────────────────
+#
+# Constraint  A : encodes soft constraint center + per-dimension tightness weight w ∈ [0,1]
+#               A[:D/2] = constraint centers,  A[D/2:] = tightness weights
+# Objective   B : desired goal point
+# Solution    C : per-dimension elastic balance
+#               C[j] = w[j] · a_center[j] + (1 − w[j]) · B[j]
+#
+# Meaning: "each dimension is pulled between constraint center and objective,
+#           with w[j] controlling how hard the constraint is in that dimension"
+# This is the natural problem for the elastic mesh.
+
+def gen_elastic(n: int, dim: int, rng: np.random.Generator) -> List[Dict]:
     data = []
-    for dtype, (fn, _) in DATASETS.items():
-        for _ in range(n_per):
-            a = rng.uniform(0.1, 0.9, n).tolist()
-            b = rng.uniform(0.1, 0.9, n).tolist()
-            c = [round(float(fn(a[i], b[i])), 4) for i in range(n)]
-            data.append({
-                'A': [round(v, 4) for v in a],
-                'B': [round(v, 4) for v in b],
-                'C': c,
-                'type': dtype,
-            })
-    random.shuffle(data)
+    for _ in range(n):
+        a_center = rng.uniform(-2, 2, dim)
+        w        = rng.uniform(0.05, 0.95, dim)   # per-dim tightness
+        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
 
+
+# ── ASSEMBLY ──────────────────────────────────────────────────────────────────
+
+GENERATORS = {
+    'box_proj':  gen_box,
+    'halfspace': gen_halfspace,
+    'sphere':    gen_sphere,
+    'simplex':   gen_simplex,
+    'elastic':   gen_elastic,
+}
+
+def generate_all(n_per_type: int = SAMPLES_PER_TYPE,
+                 dim: int       = DIM,
+                 seed: int      = 42) -> List[Dict]:
+    rng  = np.random.default_rng(seed)
+    data = []
+    for fn in GENERATORS.values():
+        data.extend(fn(n_per_type, dim, rng))
+    idx = rng.permutation(len(data))
+    return [data[i] for i in idx]
+
+
+# ── MAIN ──────────────────────────────────────────────────────────────────────
+
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--n',   type=int, default=N,               help='input dimensions')
-    parser.add_argument('--spt', type=int, default=SAMPLES_PER_TYPE,help='samples per type')
-    parser.add_argument('--out', type=str, default='data')
+    parser = argparse.ArgumentParser(description='Generate elastic mesh training data')
+    parser.add_argument('--dim', type=int, default=DIM,              help='embedding dimension')
+    parser.add_argument('--n',   type=int, default=SAMPLES_PER_TYPE, help='samples per problem type')
+    parser.add_argument('--out', type=str, default='data',           help='output directory')
     args = parser.parse_args()
 
-    data = generate(args.n, args.spt)
-
-    seen = [d for d in data if DATASETS[d['type']][1]]
-    ood  = [d for d in data if not DATASETS[d['type']][1]]
+    print(f"\n{'─'*50}")
+    print(f"  Generating {5 * args.n} samples  |  dim={args.dim}")
+    print(f"{'─'*50}")
 
-    split = int(len(seen) * 0.9)
-    train = seen[:split]
-    test  = seen[split:] + ood
-    random.shuffle(test)
+    data  = generate_all(args.n, args.dim)
+    split = int(len(data) * 0.9)
+    train, test = data[:split], data[split:]
 
     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)
 
-    tr = Counter(d['type'] for d in train)
-    te = Counter(d['type'] for d in test)
+    # Per-type statistics
+    from collections import Counter
+    train_types = Counter(d['type'] for d in train)
+    test_types  = Counter(d['type'] for d in test)
+
+    print(f"\n  Train : {len(train)}")
+    print(f"  Test  : {len(test)}\n")
+    print(f"  {'Type':<14} {'Train':>8} {'Test':>7}  C-norm (mean)")
+    print(f"  {'─'*14} {'─'*8} {'─'*7}  {'─'*14}")
+    for t in GENERATORS:
+        subset = [d for d in data if d['type'] == t]
+        norms  = [np.linalg.norm(d['C']) for d in subset]
+        print(f"  {t:<14} {train_types[t]:>8} {test_types[t]:>7}  "
+              f"{np.mean(norms):.3f} ± {np.std(norms):.3f}")
+
+    # Sanity check one sample per type
+    print(f"\n  Sanity check (first sample per type):")
+    seen = set()
+    for d in data:
+        if d['type'] in seen: continue
+        seen.add(d['type'])
+        A, B, C = map(np.array, [d['A'], d['B'], d['C']])
+        err = np.linalg.norm(A - B)
+        print(f"  [{d['type']:<12}]  "
+              f"‖A‖={np.linalg.norm(A):.2f}  ‖B‖={np.linalg.norm(B):.2f}  "
+              f"‖C‖={np.linalg.norm(C):.2f}  ‖A-B‖={err:.2f}")
 
-    print(f"\n  dim={args.n}  total={len(data)}")
-    print(f"  {'Type':<12} {'Train':>7} {'Test':>7}  Split")
-    print(f"  {'─'*12} {'─'*7} {'─'*7}  {'─'*8}")
-    for t, (_, seen_flag) in DATASETS.items():
-        print(f"  {t:<12} {tr.get(t,0):>7} {te.get(t,0):>7}  {'SEEN' if seen_flag else 'OOD ✗'}")
-    print(f"\n  → {out}/train.json  {out}/test.json\n")
+    print(f"\n  Saved → {out}/train.json  {out}/test.json\n")