Update data_gen.py

data_gen.py

@@ -1,54 +1,74 @@
 """
-data_gen.py
-N-Dimensional Vector Data for the Triangulated Mesh.
+data_gen.py  v5
+
+Each sample: (A, B, C)  where A,B,C ∈ ℝ^n, all values in [0.1, 0.9]
+
+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.
+
+SEEN during training : heavy_a | avg | diff
+OOD  (test only)     : heavy_b   ← has never been seen, tests geometric generalisation
 """
-import numpy as np
-import json, pathlib, random, argparse
 
-DIM = 8  # N dimensions = N nodes at A, B, and C
-SAMPLES_PER_TYPE = 1000
+import numpy as np, json, pathlib, random, argparse
+from collections import Counter
 
-def generate(n_per=SAMPLES_PER_TYPE, dim=DIM, seed=42):
+N               = 1      # embedding dimension (1 = pure scalar, set >1 for vector)
+SAMPLES_PER_TYPE = 2500
+
+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)
     data = []
-    
-    # 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'})
-
+    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)
     return data
 
-def round_list(arr):
-    return [round(float(x), 4) for x in arr]
-
 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('--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')
     args = parser.parse_args()
 
-    data = generate(args.n, args.dim)
-    
-    split = int(len(data) * 0.9)
-    train, test = data[:split], data[split:]
-    
-    out = pathlib.Path('data')
+    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]]
+
+    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)
-    print(f"Generated {len(data)} samples (Dim={args.dim}). Saved to data/train.json and data/test.json")
+    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)
+
+    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")