Create cell1_generator_v10.py

cell1_generator_v10.py

@@ -0,0 +1,547 @@
+"""
+Hierarchical Shape Generator - Two-Tier Gate Version
+======================================================
+Generates grids only. Patch analysis split into:
+  - Local properties: intrinsic to each patch's voxels (no cross-patch info)
+  - Structural properties: relational, require neighborhood context
+
+Colab Cell 1 of 3 - runs first, populates shared namespace.
+"""
+
+import numpy as np
+from typing import Dict, Optional
+from itertools import combinations
+
+# === Grid Constants ===========================================================
+GZ, GY, GX = 8, 16, 16
+GRID_SHAPE = (GZ, GY, GX)
+GRID_VOLUME = GZ * GY * GX
+
+PATCH_Z, PATCH_Y, PATCH_X = 2, 4, 4
+PATCH_VOL = PATCH_Z * PATCH_Y * PATCH_X
+MACRO_Z, MACRO_Y, MACRO_X = GZ // PATCH_Z, GY // PATCH_Y, GX // PATCH_X
+MACRO_N = MACRO_Z * MACRO_Y * MACRO_X
+
+# Worker budget (A100 Colab limit)
+MAX_WORKERS = 10
+
+_COORDS = np.mgrid[0:GZ, 0:GY, 0:GX].reshape(3, -1).T.astype(np.float64)
+
+# === Classes ==================================================================
+CLASS_NAMES = [
+    "point", "line", "corner", "cross", "arc", "helix", "circle",
+    "triangle", "quad", "plane", "disc",
+    "tetrahedron", "cube", "pyramid", "prism", "octahedron", "pentachoron", "wedge",
+    "sphere", "hemisphere", "torus", "bowl", "saddle", "capsule", "cylinder", "cone", "channel"
+]
+NUM_CLASSES = len(CLASS_NAMES)
+CLASS_TO_IDX = {n: i for i, n in enumerate(CLASS_NAMES)}
+
+# === Two-Tier Gate Constants ==================================================
+
+# Local gates: intrinsic to each patch, no cross-patch info needed
+#   dims:        4 classes (0D point, 1D line, 2D surface, 3D volume)
+#   curvature:   3 classes (rigid, curved, combined)
+#   boundary:    1 binary  (partial fill = surface patch)
+#   axis_active: 3 binary  (which axes have extent > 1 voxel)
+NUM_LOCAL_DIMS = 4
+NUM_LOCAL_CURVS = 3
+NUM_LOCAL_BOUNDARY = 1
+NUM_LOCAL_AXES = 3
+LOCAL_GATE_DIM = NUM_LOCAL_DIMS + NUM_LOCAL_CURVS + NUM_LOCAL_BOUNDARY + NUM_LOCAL_AXES  # 11
+
+# Structural gates: relational, require neighborhood context (post-attention)
+#   topology:     2 classes (open / closed based on neighbor count)
+#   neighbor_ct:  1 continuous (normalized 0-1, raw count / 6)
+#   surface_role: 3 classes (isolated 0-1 neighbors, boundary 2-4, interior 5-6)
+NUM_STRUCT_TOPO = 2
+NUM_STRUCT_NEIGHBOR = 1
+NUM_STRUCT_ROLE = 3
+STRUCTURAL_GATE_DIM = NUM_STRUCT_TOPO + NUM_STRUCT_NEIGHBOR + NUM_STRUCT_ROLE  # 6
+
+TOTAL_GATE_DIM = LOCAL_GATE_DIM + STRUCTURAL_GATE_DIM  # 17
+
+# Legacy compat
+GATES = ["rigid", "curved", "combined", "open", "closed"]
+NUM_GATES = len(GATES)
+
+
+# === Rasterization ============================================================
+def rasterize_line(p1, p2):
+    p1, p2 = np.array(p1, dtype=float), np.array(p2, dtype=float)
+    n = max(int(np.max(np.abs(p2 - p1))) + 1, 2)
+    t = np.linspace(0, 1, n)[:, None]
+    pts = np.round(p1 + t * (p2 - p1)).astype(int)
+    return np.clip(pts, [0, 0, 0], [GZ-1, GY-1, GX-1])
+
+
+def rasterize_edges(verts, edges):
+    pts = []
+    for i, j in edges:
+        pts.append(rasterize_line(verts[i], verts[j]))
+    return np.concatenate(pts)
+
+
+def rasterize_faces(verts, faces, density=1.0):
+    pts = []
+    for f in faces:
+        v0, v1, v2 = [np.array(verts[i], dtype=float) for i in f[:3]]
+        e1, e2 = v1 - v0, v2 - v0
+        n = max(int(max(np.linalg.norm(e1), np.linalg.norm(e2)) * density) + 1, 3)
+        for u in np.linspace(0, 1, n):
+            for v in np.linspace(0, 1 - u, max(int(n * (1 - u)), 1)):
+                p = np.round(v0 + u * e1 + v * e2).astype(int)
+                pts.append(p)
+    return np.clip(np.array(pts), [0, 0, 0], [GZ-1, GY-1, GX-1])
+
+
+def rasterize_sphere(c, r, fill=True, half=False, zmin=None, zmax=None):
+    c = np.array(c, dtype=float)
+    pts = []
+    zr = range(max(0, int(c[0] - r)), min(GZ, int(c[0] + r) + 1))
+    for z in zr:
+        for y in range(max(0, int(c[1] - r)), min(GY, int(c[1] + r) + 1)):
+            for x in range(max(0, int(c[2] - r)), min(GX, int(c[2] + r) + 1)):
+                d = np.sqrt((z - c[0])**2 + (y - c[1])**2 + (x - c[2])**2)
+                if fill and d <= r:
+                    if zmin is not None and z < zmin: continue
+                    if zmax is not None and z > zmax: continue
+                    pts.append([z, y, x])
+                elif not fill and abs(d - r) < 0.8:
+                    if half and z < c[0]: continue
+                    pts.append([z, y, x])
+    return np.array(pts) if pts else np.zeros((0, 3), dtype=int)
+
+
+# === Shape Generators =========================================================
+class HierarchicalShapeGenerator:
+    def __init__(self, seed=42):
+        self.rng = np.random.RandomState(seed)
+
+    def _random_center(self, margin=3):
+        return [self.rng.randint(max(1, margin//2), max(2, GZ - margin//2)),
+                self.rng.randint(margin, GY - margin),
+                self.rng.randint(margin, GX - margin)]
+
+    def _to_grid(self, pts):
+        if len(pts) == 0: return None, None
+        grid = np.zeros(GRID_SHAPE, dtype=np.float32)
+        pts = np.clip(np.array(pts).astype(int), [0, 0, 0], [GZ-1, GY-1, GX-1])
+        grid[pts[:, 0], pts[:, 1], pts[:, 2]] = 1.0
+        return grid, pts
+
+    def generate(self, name):
+        r = self.rng
+        c = self._random_center()
+
+        try:
+            if name == "point":
+                pts = [c]
+            elif name == "line":
+                axis = r.randint(0, 3)
+                p1, p2 = list(c), list(c)
+                L = r.randint(4, [GZ, GY, GX][axis])
+                p1[axis] = max(0, c[axis] - L//2)
+                p2[axis] = min([GZ, GY, GX][axis] - 1, c[axis] + L//2)
+                pts = rasterize_line(p1, p2)
+            elif name == "corner":
+                L = r.randint(3, 7)
+                p1, p2 = list(c), list(c)
+                p1[1] = max(0, c[1] - L)
+                p2[2] = min(GX - 1, c[2] + L)
+                pts = np.concatenate([rasterize_line(c, p1), rasterize_line(c, p2)])
+            elif name == "cross":
+                L = r.randint(2, 5)
+                pts = []
+                for d in range(3):
+                    p1, p2 = list(c), list(c)
+                    p1[d] = max(0, c[d] - L)
+                    p2[d] = min([GZ, GY, GX][d] - 1, c[d] + L)
+                    pts.append(rasterize_line(p1, p2))
+                pts = np.concatenate(pts)
+            elif name == "arc":
+                R = r.uniform(2, 5)
+                t = np.linspace(0, np.pi * r.uniform(0.4, 0.9), 30)
+                pts = np.round(np.column_stack([c[0] + np.zeros_like(t), c[1] + R*np.cos(t), c[2] + R*np.sin(t)])).astype(int)
+            elif name == "helix":
+                R, H = r.uniform(2, 4), r.uniform(3, GZ - 2)
+                t = np.linspace(0, 4*np.pi, 60)
+                pts = np.round(np.column_stack([c[0] - H/2 + t/(4*np.pi)*H, c[1] + R*np.cos(t), c[2] + R*np.sin(t)])).astype(int)
+            elif name == "circle":
+                R = r.uniform(2, 5)
+                t = np.linspace(0, 2*np.pi, 40)
+                pts = np.round(np.column_stack([np.full_like(t, c[0]), c[1] + R*np.cos(t), c[2] + R*np.sin(t)])).astype(int)
+            elif name == "triangle":
+                s = r.uniform(3, 6)
+                v = [[c[0], c[1] - s, c[2]], [c[0], c[1] + s//2, c[2] - s], [c[0], c[1] + s//2, c[2] + s]]
+                pts = rasterize_edges(v, [(0,1),(1,2),(2,0)])
+            elif name == "quad":
+                s = r.randint(2, 5)
+                v = [[c[0], c[1]-s, c[2]-s], [c[0], c[1]-s, c[2]+s], [c[0], c[1]+s, c[2]+s], [c[0], c[1]+s, c[2]-s]]
+                pts = rasterize_edges(v, [(0,1),(1,2),(2,3),(3,0)])
+            elif name == "plane":
+                s = r.randint(2, 5)
+                pts = rasterize_faces([[c[0],c[1]-s,c[2]-s],[c[0],c[1]-s,c[2]+s],[c[0],c[1]+s,c[2]+s],[c[0],c[1]+s,c[2]-s]], [(0,1,2),(0,2,3)])
+            elif name == "disc":
+                R = r.uniform(2, 5)
+                pts = rasterize_sphere(c, R, fill=True)
+                pts = pts[pts[:, 0] == c[0]] if len(pts) > 0 else pts
+            elif name == "tetrahedron":
+                s = r.uniform(3, 5)
+                v = [[c[0]+s,c[1],c[2]], [c[0]-s//2,c[1]+s,c[2]], [c[0]-s//2,c[1]-s//2,c[2]+s], [c[0]-s//2,c[1]-s//2,c[2]-s]]
+                pts = rasterize_edges(v, [(0,1),(0,2),(0,3),(1,2),(1,3),(2,3)])
+            elif name == "cube":
+                s = r.randint(2, 4)
+                v = [[c[0]+d[0]*s, c[1]+d[1]*s, c[2]+d[2]*s] for d in [(-1,-1,-1),(-1,-1,1),(-1,1,1),(-1,1,-1),(1,-1,-1),(1,-1,1),(1,1,1),(1,1,-1)]]
+                pts = rasterize_edges(v, [(0,1),(1,2),(2,3),(3,0),(4,5),(5,6),(6,7),(7,4),(0,4),(1,5),(2,6),(3,7)])
+            elif name == "pyramid":
+                s = r.randint(2, 4)
+                base = [[c[0]-s,c[1]-s,c[2]-s],[c[0]-s,c[1]-s,c[2]+s],[c[0]-s,c[1]+s,c[2]+s],[c[0]-s,c[1]+s,c[2]-s]]
+                apex = [c[0]+s, c[1], c[2]]
+                v = base + [apex]
+                pts = rasterize_edges(v, [(0,1),(1,2),(2,3),(3,0),(0,4),(1,4),(2,4),(3,4)])
+            elif name == "prism":
+                s, h = r.randint(2, 4), r.randint(2, 4)
+                bottom = [[c[0]-h,c[1]-s,c[2]], [c[0]-h,c[1]+s//2,c[2]-s], [c[0]-h,c[1]+s//2,c[2]+s]]
+                top = [[b[0]+2*h, b[1], b[2]] for b in bottom]
+                v = bottom + top
+                pts = rasterize_edges(v, [(0,1),(1,2),(2,0),(3,4),(4,5),(5,3),(0,3),(1,4),(2,5)])
+            elif name == "octahedron":
+                s = r.uniform(2, 4)
+                v = [[c[0]+s,c[1],c[2]],[c[0]-s,c[1],c[2]],[c[0],c[1]+s,c[2]],[c[0],c[1]-s,c[2]],[c[0],c[1],c[2]+s],[c[0],c[1],c[2]-s]]
+                pts = rasterize_edges(v, [(0,2),(0,3),(0,4),(0,5),(1,2),(1,3),(1,4),(1,5),(2,4),(2,5),(3,4),(3,5)])
+            elif name == "pentachoron":
+                s = r.uniform(2, 4)
+                v = [[c[0]+s,c[1],c[2]],[c[0]-s//2,c[1]+s,c[2]],[c[0]-s//2,c[1]-s//2,c[2]+s],[c[0]-s//2,c[1]-s//2,c[2]-s],[c[0],c[1],c[2]]]
+                pts = rasterize_edges(v, [(i,j) for i in range(5) for j in range(i+1,5)])
+            elif name == "wedge":
+                s = r.randint(2, 4)
+                v = [[c[0]-s,c[1]-s,c[2]-s],[c[0]-s,c[1]+s,c[2]-s],[c[0]-s,c[1],c[2]+s],[c[0]+s,c[1]-s,c[2]-s],[c[0]+s,c[1]+s,c[2]-s],[c[0]+s,c[1],c[2]+s]]
+                pts = rasterize_edges(v, [(0,1),(1,2),(2,0),(3,4),(4,5),(5,3),(0,3),(1,4),(2,5)])
+            elif name == "sphere":
+                R = r.uniform(2, min(3.5, GZ//2 - 1))
+                pts = rasterize_sphere(c, R, fill=False)
+            elif name == "hemisphere":
+                R = r.uniform(2, min(3.5, GZ//2 - 1))
+                pts = rasterize_sphere(c, R, fill=False, half=True)
+            elif name == "torus":
+                R, rr = r.uniform(3, 5), r.uniform(1, 2)
+                t = np.linspace(0, 2*np.pi, 40)
+                p = np.linspace(0, 2*np.pi, 20)
+                T, P = np.meshgrid(t, p)
+                pts = np.round(np.column_stack([c[0] + rr*np.sin(P.ravel()), c[1] + (R+rr*np.cos(P.ravel()))*np.cos(T.ravel()), c[2] + (R+rr*np.cos(P.ravel()))*np.sin(T.ravel())])).astype(int)
+            elif name == "bowl":
+                R = r.uniform(2, 4)
+                pts = rasterize_sphere(c, R, fill=False)
+                pts = pts[pts[:, 0] >= c[0]] if len(pts) > 0 else pts
+            elif name == "saddle":
+                s = r.uniform(2, 4)
+                Y, X = np.mgrid[-s:s:0.5, -s:s:0.5]
+                Z = (Y**2 - X**2) / (2*s)
+                pts = np.round(np.column_stack([c[0] + Z.ravel(), c[1] + Y.ravel(), c[2] + X.ravel()])).astype(int)
+            elif name == "capsule":
+                R, H = r.uniform(1.5, 3), r.uniform(2, 4)
+                shell = rasterize_sphere(c, R, fill=False)
+                body = []
+                for z in range(max(0, int(c[0]-H//2)), min(GZ, int(c[0]+H//2)+1)):
+                    for y in range(GY):
+                        for x in range(GX):
+                            if abs(np.sqrt((y-c[1])**2 + (x-c[2])**2) - R) < 0.8:
+                                body.append([z, y, x])
+                pts = np.concatenate([shell, np.array(body) if body else np.zeros((0,3), dtype=int)])
+            elif name == "cylinder":
+                R, H = r.uniform(2, 4), r.uniform(3, GZ - 2)
+                pts = []
+                for z in range(max(0, int(c[0]-H/2)), min(GZ, int(c[0]+H/2)+1)):
+                    for y in range(GY):
+                        for x in range(GX):
+                            d = np.sqrt((y-c[1])**2 + (x-c[2])**2)
+                            if abs(d - R) < 0.8:
+                                pts.append([z, y, x])
+                pts = np.array(pts) if pts else np.zeros((0,3), dtype=int)
+            elif name == "cone":
+                R, H = r.uniform(2, 4), r.uniform(3, GZ - 2)
+                pts = []
+                for z in range(max(0, int(c[0]-H/2)), min(GZ, int(c[0]+H/2)+1)):
+                    frac = 1 - (z - (c[0]-H/2)) / H
+                    cr = R * frac
+                    for y in range(GY):
+                        for x in range(GX):
+                            d = np.sqrt((y-c[1])**2 + (x-c[2])**2)
+                            if abs(d - cr) < 0.8 and cr > 0.3:
+                                pts.append([z, y, x])
+                pts = np.array(pts) if pts else np.zeros((0,3), dtype=int)
+            elif name == "channel":
+                R = r.uniform(2, 4)
+                L = r.randint(6, GX - 2)
+                pts = []
+                for z in range(GZ):
+                    for x in range(max(0, c[2]-L//2), min(GX, c[2]+L//2)):
+                        for y in range(GY):
+                            d = np.sqrt((z - c[0])**2 + (y - c[1])**2)
+                            if abs(d - R) < 0.8:
+                                pts.append([z, y, x])
+                pts = np.array(pts) if pts else np.zeros((0,3), dtype=int)
+            else:
+                return None
+        except Exception:
+            return None
+
+        grid, pts = self._to_grid(pts)
+        if grid is not None and pts is not None and len(pts) > 0:
+            return {"grid": grid, "class_idx": CLASS_TO_IDX[name]}
+        return None
+
+    def generate_multi(self, n_shapes: int = None) -> Optional[Dict]:
+        if n_shapes is None:
+            n_shapes = self.rng.randint(2, 5)
+        names = list(self.rng.choice(CLASS_NAMES, size=n_shapes, replace=False))
+        shapes = [s for s in [self.generate(n) for n in names] if s is not None]
+        if len(shapes) < 2:
+            return None
+        grid = np.zeros(GRID_SHAPE, dtype=np.float32)
+        membership = np.zeros((MACRO_N, NUM_CLASSES), dtype=np.float32)
+        for s in shapes:
+            pts = np.argwhere(s["grid"] > 0.5)
+            grid[pts[:, 0], pts[:, 1], pts[:, 2]] = 1.0
+            patch_idx = (pts[:, 0]//PATCH_Z) * (MACRO_Y*MACRO_X) + (pts[:, 1]//PATCH_Y) * MACRO_X + (pts[:, 2]//PATCH_X)
+            np.add.at(membership[:, s["class_idx"]], patch_idx, 1.0)
+        return {"grid": grid, "membership": (membership > 0).astype(np.float32), "n_shapes": len(shapes)}
+
+
+def _worker(args):
+    seed, min_s, max_s = args
+    gen = HierarchicalShapeGenerator(seed)
+    return gen.generate_multi(gen.rng.randint(min_s, max_s + 1))
+
+
+def generate_dataset(n_samples: int, seed: int = 42, num_workers: int = MAX_WORKERS) -> Dict:
+    from multiprocessing import Pool
+    try:
+        from tqdm import tqdm
+        use_tqdm = True
+    except ImportError:
+        use_tqdm = False
+
+    tasks = [(seed * 10000 + i, 2, 4) for i in range(n_samples * 2)]
+    grids, memberships, n_shapes = [], [], []
+
+    with Pool(num_workers) as pool:
+        pbar = tqdm(total=n_samples, desc="Generating") if use_tqdm else None
+        for r in pool.imap_unordered(_worker, tasks):
+            if r is not None and len(grids) < n_samples:
+                grids.append(r["grid"])
+                memberships.append(r["membership"])
+                n_shapes.append(r["n_shapes"])
+                if pbar: pbar.update(1)
+                if len(grids) >= n_samples:
+                    break
+        if pbar: pbar.close()
+
+    return {"grids": np.array(grids), "memberships": np.array(memberships), "n_shapes": np.array(n_shapes)}
+
+
+# === Patch Analysis: Two-Tier =================================================
+
+def analyze_local_patches(grids):
+    """
+    Local patch properties — intrinsic to each patch's voxels.
+    No cross-patch information. Computable from raw patch data.
+
+    Returns:
+        occupancy:    (N, 64) float — mean voxel density
+        dims:         (N, 64) long — 0-3 (axis extent counting)
+        curvature:    (N, 64) long — 0=rigid, 1=curved, 2=combined
+        boundary:     (N, 64) float — 1.0 if partial fill (surface patch)
+        axis_active:  (N, 64, 3) float — which axes have extent > 1
+        fill_ratio:   (N, 64) float — voxels / bounding_box_volume
+    """
+    import torch
+
+    if isinstance(grids, np.ndarray):
+        grids = torch.from_numpy(grids).float()
+
+    device, N = grids.device, grids.shape[0]
+    patches = grids.view(N, MACRO_Z, PATCH_Z, MACRO_Y, PATCH_Y, MACRO_X, PATCH_X)
+    patches = patches.permute(0, 1, 3, 5, 2, 4, 6).contiguous().view(N, MACRO_N, PATCH_Z, PATCH_Y, PATCH_X)
+
+    occupancy = patches.sum(dim=(2, 3, 4)) / PATCH_VOL
+    occ_mask = occupancy > 0.01
+    occ = patches > 0.5
+
+    z_c = torch.arange(PATCH_Z, device=device).view(1, 1, PATCH_Z, 1, 1).float()
+    y_c = torch.arange(PATCH_Y, device=device).view(1, 1, 1, PATCH_Y, 1).float()
+    x_c = torch.arange(PATCH_X, device=device).view(1, 1, 1, 1, PATCH_X).float()
+    INF = 1000.0
+
+    z_ext = torch.where(occ, z_c.expand_as(patches), torch.full_like(patches, -INF)).amax(dim=(2,3,4)) - torch.where(occ, z_c.expand_as(patches), torch.full_like(patches, INF)).amin(dim=(2,3,4))
+    y_ext = torch.where(occ, y_c.expand_as(patches), torch.full_like(patches, -INF)).amax(dim=(2,3,4)) - torch.where(occ, y_c.expand_as(patches), torch.full_like(patches, INF)).amin(dim=(2,3,4))
+    x_ext = torch.where(occ, x_c.expand_as(patches), torch.full_like(patches, -INF)).amax(dim=(2,3,4)) - torch.where(occ, x_c.expand_as(patches), torch.full_like(patches, INF)).amin(dim=(2,3,4))
+
+    ext_sorted, _ = torch.stack([z_ext, y_ext, x_ext], dim=-1).clamp(min=0).sort(dim=-1, descending=True)
+    dims = torch.zeros(N, MACRO_N, dtype=torch.long, device=device)
+    dims = torch.where(ext_sorted[..., 0] >= 1, torch.tensor(1, device=device), dims)
+    dims = torch.where(ext_sorted[..., 1] >= 1, torch.tensor(2, device=device), dims)
+    dims = torch.where(ext_sorted[..., 2] >= 1, torch.tensor(3, device=device), dims)
+    dims = torch.where(~occ_mask, torch.tensor(-1, device=device), dims)
+
+    voxels = patches.sum(dim=(2, 3, 4))
+    bb_vol = ((z_ext + 1) * (y_ext + 1) * (x_ext + 1)).clamp(min=1)
+    fill_ratio = voxels / bb_vol
+    curvature = torch.where(fill_ratio > 0.6, 0, torch.where(fill_ratio < 0.3, 1, 2)).long()
+
+    boundary = ((occupancy > 0.01) & (occupancy < 0.9)).float()
+
+    axis_active = torch.stack([
+        (z_ext.clamp(min=0) >= 1).float(),
+        (y_ext.clamp(min=0) >= 1).float(),
+        (x_ext.clamp(min=0) >= 1).float(),
+    ], dim=-1)
+
+    return {
+        "occupancy": occupancy,
+        "dims": dims,
+        "curvature": curvature,
+        "boundary": boundary,
+        "axis_active": axis_active,
+        "fill_ratio": fill_ratio,
+    }
+
+
+def analyze_structural_patches(grids, local_data):
+    """
+    Structural patch properties — relational, require neighborhood context.
+    Ground truth targets for post-attention heads.
+
+    Returns:
+        topology:       (N, 64) long — 0=open (<= 3 neighbors), 1=closed (> 3)
+        neighbor_count: (N, 64) float — normalized 0-1 (raw count / 6)
+        surface_role:   (N, 64) long — 0=isolated (0-1), 1=boundary (2-4), 2=interior (5-6)
+    """
+    import torch
+    import torch.nn.functional as F
+
+    if isinstance(grids, np.ndarray):
+        grids = torch.from_numpy(grids).float()
+
+    device, N = grids.device, grids.shape[0]
+    occ_mask = local_data["occupancy"] > 0.01
+
+    occ_3d = occ_mask.float().view(N, 1, MACRO_Z, MACRO_Y, MACRO_X)
+    kernel = torch.zeros(1, 1, 3, 3, 3, device=device)
+    kernel[0, 0, 1, 1, 0] = kernel[0, 0, 1, 1, 2] = 1
+    kernel[0, 0, 1, 0, 1] = kernel[0, 0, 1, 2, 1] = 1
+    kernel[0, 0, 0, 1, 1] = kernel[0, 0, 2, 1, 1] = 1
+    raw_count = F.conv3d(occ_3d, kernel, padding=1).view(N, MACRO_N)
+
+    topology = (raw_count > 3).long()
+    neighbor_count = raw_count / 6.0
+
+    surface_role = torch.zeros(N, MACRO_N, dtype=torch.long, device=device)
+    surface_role = torch.where(raw_count >= 2, torch.tensor(1, device=device), surface_role)
+    surface_role = torch.where(raw_count >= 5, torch.tensor(2, device=device), surface_role)
+
+    return {
+        "topology": topology,
+        "neighbor_count": neighbor_count,
+        "surface_role": surface_role,
+    }
+
+
+def analyze_patches_torch(grids):
+    """Combined analysis — returns both local and structural properties."""
+    local_data = analyze_local_patches(grids)
+    struct_data = analyze_structural_patches(grids, local_data)
+
+    import torch
+    N = local_data["occupancy"].shape[0]
+    device = local_data["occupancy"].device
+    labels = torch.zeros(N, MACRO_N, NUM_GATES, device=device)
+    labels[..., 0] = (local_data["curvature"] == 0).float()
+    labels[..., 1] = (local_data["curvature"] == 1).float()
+    labels[..., 2] = (local_data["curvature"] == 2).float()
+    labels[..., 3] = (struct_data["topology"] == 0).float()
+    labels[..., 4] = (struct_data["topology"] == 1).float()
+
+    return {
+        # Local
+        "patch_occupancy": local_data["occupancy"],
+        "patch_dims": local_data["dims"],
+        "patch_curvature": local_data["curvature"],
+        "patch_boundary": local_data["boundary"],
+        "patch_axis_active": local_data["axis_active"],
+        "patch_fill_ratio": local_data["fill_ratio"],
+        # Structural
+        "patch_topology": struct_data["topology"],
+        "patch_neighbor_count": struct_data["neighbor_count"],
+        "patch_surface_role": struct_data["surface_role"],
+        # Legacy
+        "patch_labels": labels,
+    }
+
+
+# === Dataset ==================================================================
+import torch
+from torch.utils.data import Dataset
+
+
+class ShapeDataset(Dataset):
+    def __init__(self, grids, memberships, patch_data):
+        self.grids = grids
+        self.memberships = memberships
+
+        # Local
+        self.patch_occupancy = patch_data["patch_occupancy"]
+        self.patch_dims = patch_data["patch_dims"]
+        self.patch_curvature = patch_data["patch_curvature"]
+        self.patch_boundary = patch_data["patch_boundary"]
+        self.patch_axis_active = patch_data["patch_axis_active"]
+        self.patch_fill_ratio = patch_data["patch_fill_ratio"]
+
+        # Structural
+        self.patch_topology = patch_data["patch_topology"]
+        self.patch_neighbor_count = patch_data["patch_neighbor_count"]
+        self.patch_surface_role = patch_data["patch_surface_role"]
+
+        # Legacy
+        self.patch_labels = patch_data["patch_labels"]
+
+        # Derived global targets
+        self.patch_shape_count = (memberships > 0).sum(dim=-1).long()
+        self.global_shapes = (memberships.sum(dim=1) > 0).float()
+        occ_mask = self.patch_occupancy > 0.01
+        occ_count = occ_mask.sum(dim=1, keepdim=True).clamp(min=1)
+        self.global_gates = (self.patch_labels * occ_mask.unsqueeze(-1)).sum(dim=1) / occ_count
+
+    def __len__(self):
+        return len(self.grids)
+
+    def __getitem__(self, idx):
+        return {
+            "grid": self.grids[idx],
+            "patch_shape_membership": self.memberships[idx],
+            "patch_shape_count": self.patch_shape_count[idx],
+            # Local
+            "patch_occupancy": self.patch_occupancy[idx],
+            "patch_dims": self.patch_dims[idx],
+            "patch_curvature": self.patch_curvature[idx],
+            "patch_boundary": self.patch_boundary[idx],
+            "patch_axis_active": self.patch_axis_active[idx],
+            "patch_fill_ratio": self.patch_fill_ratio[idx],
+            # Structural
+            "patch_topology": self.patch_topology[idx],
+            "patch_neighbor_count": self.patch_neighbor_count[idx],
+            "patch_surface_role": self.patch_surface_role[idx],
+            # Legacy
+            "patch_labels": self.patch_labels[idx],
+            # Global
+            "global_shapes": self.global_shapes[idx],
+            "global_gates": self.global_gates[idx],
+        }
+
+
+def collate_fn(batch):
+    return {k: torch.stack([b[k] for b in batch]) for k in batch[0].keys()}
+
+
+print(f"✓ Generator ready | Local: {LOCAL_GATE_DIM}d | Structural: {STRUCTURAL_GATE_DIM}d | Total: {TOTAL_GATE_DIM}d")