Create data_generator.py

data_generator.py

@@ -0,0 +1,856 @@
+"""
+3D Voxel Shape Classifier — Complete Geometric Primitive Vocabulary
+5×5×5 binary voxel grid → rigid cascade → curvature analysis → classify
+
+38 shape classes covering:
+  - Rigid 0D-3D: points, lines, joints, triangles, quads, polyhedra, prisms
+  - Curved 1D: arcs, helices
+  - Curved 2D: circles, ellipses, discs
+  - Curved 3D solid: sphere, hemisphere, cylinder, cone, capsule, torus
+  - Curved 3D hollow: shell, tube
+  - Curved 3D open: bowl (concave), saddle (hyperbolic)
+
+Curvature types: none, convex, concave, cylindrical, conical, toroidal, hyperbolic, helical
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+import math
+from itertools import combinations
+
+
+# === SwiGLU Activation =======================================================
+
+class SwiGLU(nn.Module):
+    """
+    SwiGLU activation: out = (x @ W1) * SiLU(x @ W2)
+
+    SiLU(x) = x * sigmoid(x), aka Swish — the "Swi" in SwiGLU.
+    Unlike plain sigmoid gating, SiLU preserves gradient magnitude
+    through the gate branch while maintaining sharp gating behavior.
+
+    Used at geometric decision points where crisp on/off transitions
+    matter more than smooth interpolation.
+    """
+
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.w1 = nn.Linear(in_dim, out_dim)
+        self.w2 = nn.Linear(in_dim, out_dim)
+
+    def forward(self, x):
+        return self.w1(x) * F.silu(self.w2(x))
+
+
+# === Shape Catalog ===========================================================
+
+SHAPE_CATALOG = {
+    # ---- Rigid 0D ----
+    "point":            {"dim": 0, "curved": False, "curvature": "none"},
+
+    # ---- Rigid 1D: lines ----
+    "line_x":           {"dim": 1, "curved": False, "curvature": "none"},
+    "line_y":           {"dim": 1, "curved": False, "curvature": "none"},
+    "line_z":           {"dim": 1, "curved": False, "curvature": "none"},
+    "line_diag":        {"dim": 1, "curved": False, "curvature": "none"},
+
+    # ---- Rigid 1D: compounds ----
+    "cross":            {"dim": 1, "curved": False, "curvature": "none"},
+    "l_shape":          {"dim": 1, "curved": False, "curvature": "none"},
+    "collinear":        {"dim": 1, "curved": False, "curvature": "none"},
+
+    # ---- Rigid 2D: triangles ----
+    "triangle_xy":      {"dim": 2, "curved": False, "curvature": "none"},
+    "triangle_xz":      {"dim": 2, "curved": False, "curvature": "none"},
+    "triangle_3d":      {"dim": 2, "curved": False, "curvature": "none"},
+
+    # ---- Rigid 2D: quads ----
+    "square_xy":        {"dim": 2, "curved": False, "curvature": "none"},
+    "square_xz":        {"dim": 2, "curved": False, "curvature": "none"},
+    "rectangle":        {"dim": 2, "curved": False, "curvature": "none"},
+    "coplanar":         {"dim": 2, "curved": False, "curvature": "none"},
+
+    # ---- Rigid 2D: filled ----
+    "plane":            {"dim": 2, "curved": False, "curvature": "none"},
+
+    # ---- Rigid 3D: simplices ----
+    "tetrahedron":      {"dim": 3, "curved": False, "curvature": "none"},
+    "pyramid":          {"dim": 3, "curved": False, "curvature": "none"},
+    "pentachoron":      {"dim": 3, "curved": False, "curvature": "none"},
+
+    # ---- Rigid 3D: prisms/polyhedra ----
+    "cube":             {"dim": 3, "curved": False, "curvature": "none"},
+    "cuboid":           {"dim": 3, "curved": False, "curvature": "none"},
+    "triangular_prism": {"dim": 3, "curved": False, "curvature": "none"},
+    "octahedron":       {"dim": 3, "curved": False, "curvature": "none"},
+
+    # ---- Curved 1D ----
+    "arc":              {"dim": 1, "curved": True, "curvature": "convex"},
+    "helix":            {"dim": 1, "curved": True, "curvature": "helical"},
+
+    # ---- Curved 2D: outlines ----
+    "circle":           {"dim": 2, "curved": True, "curvature": "convex"},
+    "ellipse":          {"dim": 2, "curved": True, "curvature": "convex"},
+
+    # ---- Curved 2D: filled ----
+    "disc":             {"dim": 2, "curved": True, "curvature": "convex"},
+
+    # ---- Curved 3D: solid ----
+    "sphere":           {"dim": 3, "curved": True, "curvature": "convex"},
+    "hemisphere":       {"dim": 3, "curved": True, "curvature": "convex"},
+    "cylinder":         {"dim": 3, "curved": True, "curvature": "cylindrical"},
+    "cone":             {"dim": 3, "curved": True, "curvature": "conical"},
+    "capsule":          {"dim": 3, "curved": True, "curvature": "convex"},
+    "torus":            {"dim": 3, "curved": True, "curvature": "toroidal"},
+
+    # ---- Curved 3D: hollow ----
+    "shell":            {"dim": 3, "curved": True, "curvature": "convex"},
+    "tube":             {"dim": 3, "curved": True, "curvature": "cylindrical"},
+
+    # ---- Curved 3D: open surfaces ----
+    "bowl":             {"dim": 3, "curved": True, "curvature": "concave"},
+    "saddle":           {"dim": 3, "curved": True, "curvature": "hyperbolic"},
+}
+
+NUM_CLASSES = len(SHAPE_CATALOG)
+CLASS_NAMES = list(SHAPE_CATALOG.keys())
+CLASS_TO_IDX = {name: i for i, name in enumerate(CLASS_NAMES)}
+
+CURVATURE_TYPES = ["none", "convex", "concave", "cylindrical", "conical",
+                   "toroidal", "hyperbolic", "helical"]
+CURV_TO_IDX = {c: i for i, c in enumerate(CURVATURE_TYPES)}
+NUM_CURVATURES = len(CURVATURE_TYPES)
+
+GS = 5  # grid size
+
+
+# === Cayley-Menger Utilities =================================================
+
+def cayley_menger_det(points: np.ndarray) -> float:
+    n = len(points)
+    D = np.zeros((n, n))
+    for i in range(n):
+        for j in range(n):
+            D[i, j] = np.sum((points[i] - points[j]) ** 2)
+    CM = np.zeros((n + 1, n + 1))
+    CM[0, 1:] = 1
+    CM[1:, 0] = 1
+    CM[1:, 1:] = D
+    return np.linalg.det(CM)
+
+
+def simplex_volume(points: np.ndarray) -> float:
+    k = len(points)
+    if k < 2: return 0.0
+    cm = cayley_menger_det(points)
+    sign = (-1) ** k
+    denom = (2 ** (k - 1)) * (math.factorial(k - 1) ** 2)
+    v_sq = sign * cm / denom
+    return np.sqrt(max(0, v_sq))
+
+
+def effective_volume(points: np.ndarray) -> float:
+    k = len(points)
+    if k < 2: return 0.0
+    if k == 2: return np.linalg.norm(points[0] - points[1])
+    if k >= 3:
+        max_a = 0
+        for idx in combinations(range(min(k, 8)), 3):
+            max_a = max(max_a, simplex_volume(points[list(idx)]))
+        if k < 4: return max_a
+    if k >= 4:
+        max_v = 0
+        for idx in combinations(range(min(k, 8)), 4):
+            max_v = max(max_v, simplex_volume(points[list(idx)]))
+        return max_v
+    return 0.0
+
+
+# === Shape Generator =========================================================
+
+class ShapeGenerator:
+    def __init__(self, seed=42):
+        self.rng = np.random.RandomState(seed)
+
+    def generate(self, n_samples: int) -> list:
+        samples = []
+        per_class = n_samples // NUM_CLASSES
+        for name in CLASS_NAMES:
+            count = 0
+            attempts = 0
+            while count < per_class and attempts < per_class * 5:
+                s = self._make(name)
+                attempts += 1
+                if s is not None:
+                    samples.append(s)
+                    count += 1
+        while len(samples) < n_samples:
+            name = self.rng.choice(CLASS_NAMES)
+            s = self._make(name)
+            if s is not None:
+                samples.append(s)
+        self.rng.shuffle(samples)
+        return samples[:n_samples]
+
+    def _make(self, name: str) -> Optional[dict]:
+        info = SHAPE_CATALOG[name]
+        if info["curved"]:
+            voxels = self._curved(name)
+        else:
+            voxels = self._rigid(name)
+        if voxels is None: return None
+        voxels = np.clip(voxels, 0, GS - 1).astype(int)
+        voxels = np.unique(voxels, axis=0)
+        if len(voxels) < 1: return None
+        return self._build(name, info, voxels)
+
+    # === Rigid Generators ===
+
+    def _rigid(self, name):
+        rng = self.rng
+
+        if name == "point":
+            return rng.randint(0, GS, size=(1, 3))
+
+        elif name == "line_x":
+            y, z = rng.randint(0, GS, size=2)
+            x1, x2 = sorted(rng.choice(GS, 2, replace=False))
+            return np.array([[x1, y, z], [x2, y, z]])
+
+        elif name == "line_y":
+            x, z = rng.randint(0, GS, size=2)
+            y1, y2 = sorted(rng.choice(GS, 2, replace=False))
+            return np.array([[x, y1, z], [x, y2, z]])
+
+        elif name == "line_z":
+            x, y = rng.randint(0, GS, size=2)
+            z1, z2 = sorted(rng.choice(GS, 2, replace=False))
+            return np.array([[x, y, z1], [x, y, z2]])
+
+        elif name == "line_diag":
+            p1 = rng.randint(0, 3, size=3)
+            step = rng.randint(1, 3)
+            direction = rng.choice([-1, 1], size=3)
+            if np.sum(direction != 0) < 2:
+                direction[rng.randint(3)] = rng.choice([-1, 1])
+            p2 = np.clip(p1 + step * direction, 0, GS - 1)
+            if np.array_equal(p1, p2):
+                p2 = np.clip(p1 + np.array([1, 1, 0]), 0, GS - 1)
+            return np.array([p1, p2])
+
+        elif name == "cross":
+            # Two perpendicular lines intersecting at a point
+            cx, cy, cz = rng.randint(1, GS - 1, size=3)
+            length = rng.randint(1, 3)
+            axis1, axis2 = rng.choice(3, 2, replace=False)
+            pts = [[cx, cy, cz]]  # center
+            for sign in [-1, 1]:
+                p = [cx, cy, cz]
+                p[axis1] = np.clip(p[axis1] + sign * length, 0, GS - 1)
+                pts.append(list(p))
+            for sign in [-1, 1]:
+                p = [cx, cy, cz]
+                p[axis2] = np.clip(p[axis2] + sign * length, 0, GS - 1)
+                pts.append(list(p))
+            return np.array(pts)
+
+        elif name == "l_shape":
+            # Two lines meeting at a vertex (right angle)
+            corner = rng.randint(1, GS - 1, size=3)
+            axis1, axis2 = rng.choice(3, 2, replace=False)
+            len1 = rng.randint(1, 3)
+            len2 = rng.randint(1, 3)
+            dir1 = rng.choice([-1, 1])
+            dir2 = rng.choice([-1, 1])
+            pts = [list(corner)]
+            for i in range(1, len1 + 1):
+                p = list(corner)
+                p[axis1] = np.clip(p[axis1] + dir1 * i, 0, GS - 1)
+                pts.append(p)
+            for i in range(1, len2 + 1):
+                p = list(corner)
+                p[axis2] = np.clip(p[axis2] + dir2 * i, 0, GS - 1)
+                pts.append(p)
+            return np.array(pts)
+
+        elif name == "collinear":
+            axis = rng.randint(3)
+            fixed = rng.randint(0, GS, size=2)
+            vals = sorted(rng.choice(GS, 3, replace=False))
+            pts = np.zeros((3, 3), dtype=int)
+            for i, v in enumerate(vals):
+                pts[i, axis] = v
+                pts[i, (axis + 1) % 3] = fixed[0]
+                pts[i, (axis + 2) % 3] = fixed[1]
+            return pts
+
+        elif name == "triangle_xy":
+            z = rng.randint(0, GS)
+            pts = self._rand_pts_2d(3, min_dist=1)
+            if pts is None: return None
+            return np.column_stack([pts, np.full(3, z)])
+
+        elif name == "triangle_xz":
+            y = rng.randint(0, GS)
+            pts = self._rand_pts_2d(3, min_dist=1)
+            if pts is None: return None
+            return np.column_stack([pts[:, 0], np.full(3, y), pts[:, 1]])
+
+        elif name == "triangle_3d":
+            return self._rand_pts_3d(3, min_dist=1)
+
+        elif name == "square_xy":
+            z = rng.randint(0, GS)
+            x1, y1 = rng.randint(0, 3, size=2)
+            s = rng.randint(1, 3)
+            pts = np.array([[x1, y1, z], [x1 + s, y1, z],
+                           [x1, y1 + s, z], [x1 + s, y1 + s, z]])
+            return np.clip(pts, 0, GS - 1)
+
+        elif name == "square_xz":
+            y = rng.randint(0, GS)
+            x1, z1 = rng.randint(0, 3, size=2)
+            s = rng.randint(1, 3)
+            pts = np.array([[x1, y, z1], [x1 + s, y, z1],
+                           [x1, y, z1 + s], [x1 + s, y, z1 + s]])
+            return np.clip(pts, 0, GS - 1)
+
+        elif name == "rectangle":
+            axis = rng.randint(3)
+            val = rng.randint(0, GS)
+            a1, a2 = rng.randint(0, 3), rng.randint(0, 3)
+            w, h = rng.randint(1, 4), rng.randint(1, 3)
+            if w == h: w = min(GS - 1, w + 1)
+            c = np.array([[a1, a2], [a1 + w, a2], [a1, a2 + h], [a1 + w, a2 + h]])
+            c = np.clip(c, 0, GS - 1)
+            if axis == 0: return np.column_stack([np.full(4, val), c])
+            elif axis == 1: return np.column_stack([c[:, 0], np.full(4, val), c[:, 1]])
+            else: return np.column_stack([c, np.full(4, val)])
+
+        elif name == "coplanar":
+            pts = self._rand_pts_3d(4, min_dist=1)
+            if pts is None: return None
+            pts[:, rng.randint(3)] = pts[0, rng.randint(3)]
+            return pts
+
+        elif name == "plane":
+            # Filled rectangular slab, 1 voxel thick
+            axis = rng.randint(3)
+            val = rng.randint(0, GS)
+            a_start = rng.randint(0, 2)
+            b_start = rng.randint(0, 2)
+            a_size = rng.randint(2, GS - a_start + 1)
+            b_size = rng.randint(2, GS - b_start + 1)
+            pts = []
+            for a in range(a_start, min(GS, a_start + a_size)):
+                for b in range(b_start, min(GS, b_start + b_size)):
+                    p = [0, 0, 0]
+                    p[axis] = val
+                    p[(axis + 1) % 3] = a
+                    p[(axis + 2) % 3] = b
+                    pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "tetrahedron":
+            pts = self._rand_pts_3d(4, min_dist=1)
+            if pts is None: return None
+            centered = pts - pts.mean(axis=0)
+            _, s, _ = np.linalg.svd(centered.astype(float))
+            if s[-1] < 0.5:
+                pts[rng.randint(4), rng.randint(3)] = (pts[0, 0] + 2) % GS
+            return pts
+
+        elif name == "pyramid":
+            z_base = rng.randint(0, 3)
+            x1, y1 = rng.randint(0, 3), rng.randint(0, 3)
+            s = rng.randint(1, 3)
+            base = np.array([[x1, y1, z_base], [x1 + s, y1, z_base],
+                            [x1, y1 + s, z_base], [x1 + s, y1 + s, z_base]])
+            apex = np.array([[x1 + s // 2, y1 + s // 2, z_base + rng.randint(1, 3)]])
+            return np.clip(np.vstack([base, apex]), 0, GS - 1)
+
+        elif name == "pentachoron":
+            return self._rand_pts_3d(5, min_dist=1)
+
+        elif name == "cube":
+            x1, y1, z1 = rng.randint(0, 3, size=3)
+            s = rng.randint(1, 3)
+            pts = []
+            for dx in [0, s]:
+                for dy in [0, s]:
+                    for dz in [0, s]:
+                        pts.append([x1 + dx, y1 + dy, z1 + dz])
+            return np.clip(np.array(pts), 0, GS - 1)
+
+        elif name == "cuboid":
+            x1, y1, z1 = rng.randint(0, 2, size=3)
+            sx, sy, sz = rng.randint(1, 4, size=3)
+            # Ensure not a cube: at least 2 different edge lengths
+            if sx == sy == sz:
+                sx = min(GS - 1, sx + 1)
+            pts = []
+            for dx in [0, sx]:
+                for dy in [0, sy]:
+                    for dz in [0, sz]:
+                        pts.append([x1 + dx, y1 + dy, z1 + dz])
+            return np.clip(np.array(pts), 0, GS - 1)
+
+        elif name == "triangular_prism":
+            # Triangle in one plane, extruded along the other axis
+            axis = rng.randint(3)  # extrusion axis
+            ext_start = rng.randint(0, 3)
+            ext_len = rng.randint(1, 3)
+            tri = self._rand_pts_2d(3, min_dist=1)
+            if tri is None: return None
+            pts = []
+            for e in range(ext_start, min(GS, ext_start + ext_len + 1)):
+                for t in tri:
+                    p = [0, 0, 0]
+                    p[axis] = e
+                    p[(axis + 1) % 3] = t[0]
+                    p[(axis + 2) % 3] = t[1]
+                    pts.append(p)
+            return np.clip(np.array(pts), 0, GS - 1) if len(pts) >= 6 else None
+
+        elif name == "octahedron":
+            # 6 vertices: ±1 along each axis from center
+            cx, cy, cz = rng.randint(1, GS - 1, size=3)
+            s = rng.randint(1, 3)
+            pts = [[cx, cy, cz + s], [cx, cy, cz - s],
+                   [cx + s, cy, cz], [cx - s, cy, cz],
+                   [cx, cy + s, cz], [cx, cy - s, cz]]
+            return np.clip(np.array(pts), 0, GS - 1)
+
+        return None
+
+    # === Curved Generators ===
+
+    def _curved(self, name):
+        rng = self.rng
+        cx, cy, cz = rng.uniform(1.0, 3.0, size=3)
+
+        if name == "arc":
+            r = rng.uniform(1.2, 2.2)
+            plane = rng.choice(["xy", "xz", "yz"])
+            start = rng.uniform(0, 2 * np.pi)
+            span = rng.uniform(np.pi * 0.4, np.pi * 1.2)
+            n = rng.randint(6, 12)
+            angles = np.linspace(start, start + span, n)
+            pts = []
+            for a in angles:
+                if plane == "xy":
+                    pts.append([cx + r * np.cos(a), cy + r * np.sin(a), cz])
+                elif plane == "xz":
+                    pts.append([cx + r * np.cos(a), cy, cz + r * np.sin(a)])
+                else:
+                    pts.append([cx, cy + r * np.cos(a), cz + r * np.sin(a)])
+            pts = np.unique(np.round(np.clip(pts, 0, GS - 1)).astype(int), axis=0)
+            return pts if len(pts) >= 3 else None
+
+        elif name == "helix":
+            # Spiral through 3D: parametric curve
+            r = rng.uniform(0.8, 1.8)
+            axis = rng.randint(3)
+            pitch = rng.uniform(0.3, 0.8)  # rise per radian
+            n = rng.randint(15, 30)
+            t = np.linspace(0, 2 * np.pi * rng.uniform(1.0, 2.5), n)
+            pts = []
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            start_h = rng.uniform(0, 1.0)
+            for ti in t:
+                p = [0.0, 0.0, 0.0]
+                p[axes[0]] = center[axes[0]] + r * np.cos(ti)
+                p[axes[1]] = center[axes[1]] + r * np.sin(ti)
+                p[axis] = start_h + pitch * ti
+                pts.append(p)
+            pts = np.unique(np.round(np.clip(pts, 0, GS - 1)).astype(int), axis=0)
+            return pts if len(pts) >= 5 else None
+
+        elif name == "circle":
+            r = rng.uniform(1.0, 2.0)
+            plane = rng.choice(["xy", "xz", "yz"])
+            n = rng.randint(12, 20)
+            angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
+            pts = []
+            for a in angles:
+                if plane == "xy":
+                    pts.append([cx + r * np.cos(a), cy + r * np.sin(a), cz])
+                elif plane == "xz":
+                    pts.append([cx + r * np.cos(a), cy, cz + r * np.sin(a)])
+                else:
+                    pts.append([cx, cy + r * np.cos(a), cz + r * np.sin(a)])
+            pts = np.unique(np.round(np.clip(pts, 0, GS - 1)).astype(int), axis=0)
+            return pts if len(pts) >= 5 else None
+
+        elif name == "ellipse":
+            rx, ry = rng.uniform(0.8, 2.0), rng.uniform(0.8, 2.0)
+            if abs(rx - ry) < 0.3: rx *= 1.4
+            plane = rng.choice(["xy", "xz", "yz"])
+            n = rng.randint(12, 20)
+            angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
+            pts = []
+            for a in angles:
+                if plane == "xy":
+                    pts.append([cx + rx * np.cos(a), cy + ry * np.sin(a), cz])
+                elif plane == "xz":
+                    pts.append([cx + rx * np.cos(a), cy, cz + ry * np.sin(a)])
+                else:
+                    pts.append([cx, cy + rx * np.cos(a), cz + ry * np.sin(a)])
+            pts = np.unique(np.round(np.clip(pts, 0, GS - 1)).astype(int), axis=0)
+            return pts if len(pts) >= 5 else None
+
+        elif name == "disc":
+            # Filled circle in a plane (not just outline)
+            r = rng.uniform(1.0, 2.2)
+            axis = rng.randint(3)
+            val = round(rng.uniform(0.5, 3.5))
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    p = [0, 0, 0]
+                    p[axis] = val
+                    p[axes[0]] = x
+                    p[axes[1]] = y
+                    dist = np.sqrt((x - center[axes[0]])**2 + (y - center[axes[1]])**2)
+                    if dist <= r:
+                        pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "sphere":
+            r = rng.uniform(1.0, 2.2)
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        if (x - cx)**2 + (y - cy)**2 + (z - cz)**2 <= r**2:
+                            pts.append([x, y, z])
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "hemisphere":
+            r = rng.uniform(1.0, 2.2)
+            cut_axis = rng.randint(3)
+            center = [cx, cy, cz]
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        if (x - cx)**2 + (y - cy)**2 + (z - cz)**2 <= r**2:
+                            if p[cut_axis] >= center[cut_axis]:
+                                pts.append(p)
+            return np.array(pts) if len(pts) >= 3 else None
+
+        elif name == "cylinder":
+            r = rng.uniform(0.8, 1.8)
+            axis = rng.randint(3)
+            length = rng.randint(2, 5)
+            start = rng.randint(0, GS - length + 1)
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        if p[axis] < start or p[axis] >= start + length: continue
+                        dist_sq = sum((p[a] - center[a])**2 for a in axes)
+                        if dist_sq <= r**2:
+                            pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "cone":
+            r_base = rng.uniform(1.0, 2.0)
+            axis = rng.randint(3)
+            height = rng.randint(2, 5)
+            base_pos = rng.randint(0, GS - height + 1)
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        along = p[axis] - base_pos
+                        if along < 0 or along >= height: continue
+                        t = along / (height - 1 + 1e-6)
+                        r_at = r_base * (1.0 - t)
+                        dist_sq = sum((p[a] - center[a])**2 for a in axes)
+                        if dist_sq <= r_at**2:
+                            pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "capsule":
+            # Cylinder with hemispherical caps
+            r = rng.uniform(0.8, 1.5)
+            axis = rng.randint(3)
+            body_len = rng.randint(1, 3)
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            body_start = round(center[axis] - body_len / 2)
+            body_end = body_start + body_len
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        radial_sq = sum((p[a] - center[a])**2 for a in axes)
+                        along = p[axis]
+                        # Body
+                        if body_start <= along <= body_end and radial_sq <= r**2:
+                            pts.append(p)
+                        # Bottom cap
+                        elif along < body_start:
+                            cap_center = list(center)
+                            cap_center[axis] = body_start
+                            dist_sq = sum((p[i] - cap_center[i])**2 for i in range(3))
+                            if dist_sq <= r**2:
+                                pts.append(p)
+                        # Top cap
+                        elif along > body_end:
+                            cap_center = list(center)
+                            cap_center[axis] = body_end
+                            dist_sq = sum((p[i] - cap_center[i])**2 for i in range(3))
+                            if dist_sq <= r**2:
+                                pts.append(p)
+            return np.array(pts) if len(pts) >= 5 else None
+
+        elif name == "torus":
+            R = rng.uniform(1.2, 2.0)
+            r = rng.uniform(0.5, 0.9)
+            axis = rng.randint(3)
+            center = [cx, cy, cz]
+            ring_axes = [i for i in range(3) if i != axis]
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        dist_in_plane = np.sqrt(
+                            sum((p[a] - center[a])**2 for a in ring_axes))
+                        dist_from_ring = np.sqrt(
+                            (dist_in_plane - R)**2 + (p[axis] - center[axis])**2)
+                        if dist_from_ring <= r:
+                            pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "shell":
+            # Hollow sphere: outer radius - inner radius
+            r_out = rng.uniform(1.5, 2.3)
+            r_in = r_out - rng.uniform(0.4, 0.8)
+            if r_in < 0.3: r_in = 0.3
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        d_sq = (x - cx)**2 + (y - cy)**2 + (z - cz)**2
+                        if r_in**2 <= d_sq <= r_out**2:
+                            pts.append([x, y, z])
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "tube":
+            # Hollow cylinder
+            r_out = rng.uniform(1.0, 2.0)
+            r_in = r_out - rng.uniform(0.3, 0.7)
+            if r_in < 0.2: r_in = 0.2
+            axis = rng.randint(3)
+            length = rng.randint(2, 5)
+            start = rng.randint(0, GS - length + 1)
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        if p[axis] < start or p[axis] >= start + length: continue
+                        dist_sq = sum((p[a] - center[a])**2 for a in axes)
+                        if r_in**2 <= dist_sq <= r_out**2:
+                            pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "bowl":
+            # Paraboloid: concave surface, open on top
+            r = rng.uniform(1.2, 2.2)
+            axis = rng.randint(3)
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            thickness = 0.6
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        dist_planar = np.sqrt(
+                            sum((p[a] - center[a])**2 for a in axes))
+                        if dist_planar > r: continue
+                        # Paraboloid surface: h = k * dist^2
+                        k = 1.0 / (r + 1e-6)
+                        expected_h = center[axis] + k * dist_planar**2
+                        actual_h = p[axis]
+                        if abs(actual_h - expected_h) <= thickness:
+                            pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        elif name == "saddle":
+            # Hyperbolic paraboloid: z = k*(x^2 - y^2)
+            axis = rng.randint(3)
+            center = [cx, cy, cz]
+            axes = [i for i in range(3) if i != axis]
+            k = rng.uniform(0.3, 0.8)
+            thickness = 0.7
+            pts = []
+            for x in range(GS):
+                for y in range(GS):
+                    for z in range(GS):
+                        p = [x, y, z]
+                        da = p[axes[0]] - center[axes[0]]
+                        db = p[axes[1]] - center[axes[1]]
+                        expected_h = center[axis] + k * (da**2 - db**2)
+                        if abs(p[axis] - expected_h) <= thickness:
+                            # Limit radius so it doesn't fill everything
+                            dist_sq = da**2 + db**2
+                            if dist_sq <= 4.0:
+                                pts.append(p)
+            return np.array(pts) if len(pts) >= 4 else None
+
+        return None
+
+    # === Helpers ===
+
+    def _rand_pts_2d(self, n, min_dist=0):
+        for _ in range(50):
+            pts = set()
+            while len(pts) < n:
+                pts.add((self.rng.randint(0, GS), self.rng.randint(0, GS)))
+            pts = np.array(list(pts)[:n])
+            if min_dist <= 0 or self._check_dist(pts, min_dist):
+                return pts
+        return None
+
+    def _rand_pts_3d(self, n, min_dist=0):
+        for _ in range(100):
+            pts = set()
+            while len(pts) < n:
+                pts.add(tuple(self.rng.randint(0, GS, size=3)))
+            pts = np.array(list(pts)[:n])
+            if min_dist <= 0 or self._check_dist(pts, min_dist):
+                return pts
+        return None
+
+    def _check_dist(self, pts, min_dist):
+        for i in range(len(pts)):
+            for j in range(i + 1, len(pts)):
+                if np.sum(np.abs(pts[i] - pts[j])) < min_dist:
+                    return False
+        return True
+
+    def _build(self, name, info, voxels):
+        n = len(voxels)
+        sub = voxels[:6].astype(float) if n > 6 else voxels.astype(float)
+        cm_det = cayley_menger_det(sub)
+        volume = effective_volume(sub)
+
+        dim_conf = np.zeros(4, dtype=np.float32)
+        dim_conf[0] = 1.0
+        if n >= 2: dim_conf[1] = 1.0
+        if info["dim"] >= 2: dim_conf[2] = 1.0
+        if info["dim"] >= 3: dim_conf[3] = 1.0
+
+        grid = np.zeros((GS, GS, GS), dtype=np.float32)
+        for v in voxels:
+            grid[v[0], v[1], v[2]] = 1.0
+
+        return {
+            "grid": grid, "label": CLASS_TO_IDX[name], "class_name": name,
+            "n_points": n, "n_occupied": int(grid.sum()),
+            "cm_det": float(cm_det), "volume": float(volume),
+            "peak_dim": info["dim"], "dim_confidence": dim_conf,
+            "is_curved": info["curved"], "curvature": CURV_TO_IDX[info["curvature"]],
+        }
+
+
+# === Dataset =================================================================
+
+def _generate_chunk(args):
+    """Worker function for parallel shape generation."""
+    class_assignments, seed, start_idx = args
+    gen = ShapeGenerator(seed=seed)
+    samples = []
+    for ci in class_assignments:
+        name = CLASS_NAMES[ci]
+        for attempt in range(10):
+            s = gen._make(name)
+            if s is not None:
+                samples.append(s)
+                break
+        else:
+            s = gen._make("cube")
+            if s is not None:
+                samples.append(s)
+    return samples
+
+
+def generate_parallel(n_samples, seed=42, n_workers=8):
+    """Pre-generate all samples using multiprocessing."""
+    import multiprocessing as mp
+    per_class = n_samples // NUM_CLASSES
+    class_assignments = []
+    for ci in range(NUM_CLASSES):
+        class_assignments.extend([ci] * per_class)
+    rng = np.random.RandomState(seed)
+    while len(class_assignments) < n_samples:
+        class_assignments.append(rng.randint(0, NUM_CLASSES))
+    rng.shuffle(class_assignments)
+    class_assignments = class_assignments[:n_samples]
+
+    # Split into chunks per worker
+    chunk_size = (n_samples + n_workers - 1) // n_workers
+    chunks = []
+    for i in range(n_workers):
+        start = i * chunk_size
+        end = min(start + chunk_size, n_samples)
+        if start >= n_samples:
+            break
+        chunks.append((class_assignments[start:end], seed + i * 1000000, start))
+
+    print(f"Generating {n_samples} shapes across {len(chunks)} workers...")
+    import time; t0 = time.time()
+    with mp.Pool(n_workers) as pool:
+        results = pool.map(_generate_chunk, chunks)
+    samples = []
+    for r in results:
+        samples.extend(r)
+    rng.shuffle(samples)
+    dt = time.time() - t0
+    print(f"Generated {len(samples)} samples in {dt:.1f}s ({len(samples)/dt:.0f} samples/s)")
+    return samples
+
+
+class ShapeDataset(torch.utils.data.Dataset):
+    def __init__(self, samples):
+        self.grids = torch.tensor(np.stack([s["grid"] for s in samples]), dtype=torch.float32)
+        self.labels = torch.tensor([s["label"] for s in samples], dtype=torch.long)
+        self.dim_conf = torch.tensor(np.stack([s["dim_confidence"] for s in samples]), dtype=torch.float32)
+        self.peak_dim = torch.tensor([s["peak_dim"] for s in samples], dtype=torch.long)
+        self.volume = torch.tensor([s["volume"] for s in samples], dtype=torch.float32)
+        self.cm_det = torch.tensor([s["cm_det"] for s in samples], dtype=torch.float32)
+        self.is_curved = torch.tensor([s["is_curved"] for s in samples], dtype=torch.float32)
+        self.curvature = torch.tensor([s["curvature"] for s in samples], dtype=torch.long)
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+        return (self.grids[idx], self.labels[idx], self.dim_conf[idx],
+                self.peak_dim[idx], self.volume[idx], self.cm_det[idx],
+                self.is_curved[idx], self.curvature[idx])
+
+
+
+print(f'Loaded {NUM_CLASSES} shape classes, GS={GS}')