Update data_gen.py

data_gen.py

@@ -1,80 +1,54 @@
 """
 data_gen.py
-
-Scalar mesh test data.  Each sample is (a, b, c) — all single floats.
-
-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)
-
-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.
-
-SEEN during training : heavy_a | avg | diff
-OOD  (test only)     : heavy_b
-
-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.
+N-Dimensional Vector Data for the Triangulated Mesh.
 """
+import numpy as np
+import json, pathlib, random, argparse
 
-import numpy as np, json, pathlib, random, argparse
-
-SAMPLES_PER_TYPE = 2500  # × 4 types = 10 000 total
+DIM = 8  # N dimensions = N nodes at A, B, and C
+SAMPLES_PER_TYPE = 1000
 
-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
-}
-
-def generate(n_per=SAMPLES_PER_TYPE, seed=42):
+def generate(n_per=SAMPLES_PER_TYPE, dim=DIM, seed=42):
     rng = np.random.default_rng(seed)
     data = []
-    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})
+    
+    # We will train the mesh to physically wire three different logic gates:
+    for _ in range(n_per):
+        a = rng.uniform(0.1, 0.9, dim)
+        b = rng.uniform(0.1, 0.9, dim)
+        
+        # 1. Direct blending (Straight down)
+        c_blend = 0.7 * a + 0.3 * b
+        data.append({'a': round_list(a), 'b': round_list(b), 'c': round_list(c_blend), 'type': 'blend'})
+        
+        # 2. Subtraction logic
+        c_diff = 0.5 + 0.4 * (a - b)
+        data.append({'a': round_list(a), 'b': round_list(b), 'c': round_list(c_diff), 'type': 'diff'})
+
+        # 3. Lateral Routing (A rolls right by 1, B rolls left by 1)
+        # This forces the mesh to use its diagonal triangulated springs!
+        c_route = 0.5 * np.roll(a, 1) + 0.5 * np.roll(b, -1)
+        data.append({'a': round_list(a), 'b': round_list(b), 'c': round_list(c_route), 'type': 'route'})
+
     random.shuffle(data)
     return data
 
+def round_list(arr):
+    return [round(float(x), 4) for x in arr]
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--n',   type=int, default=SAMPLES_PER_TYPE)
-    parser.add_argument('--out', type=str, default='data')
+    parser.add_argument('--dim', type=int, default=DIM)
+    parser.add_argument('--n', type=int, default=SAMPLES_PER_TYPE)
     args = parser.parse_args()
 
-    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)
+    data = generate(args.n, args.dim)
+    
+    split = int(len(data) * 0.9)
+    train, test = data[:split], data[split:]
+    
+    out = pathlib.Path('data')
     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)
-
-    from collections import Counter
-    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")
+    print(f"Generated {len(data)} samples (Dim={args.dim}). Saved to data/train.json and data/test.json")