Create shape_generator.py

shape_generator.py

@@ -0,0 +1,722 @@
+"""
+3D Voxel Shape Generator — VAE-Matched Resolution (8×16×16)
+=============================================================
+Adapted from v10.2 for Flux 2 VAE latent geometry:
+  - GZ=8 (channel dimension), GY=16, GX=16 (spatial dimensions)
+  - 2048 voxels per patch (vs 15,625 at 25³)
+  - Aspect ratio 1:2:2 matches VAE latent structure
+  - All 38 shape classes preserved
+  - Selective rasterization (polyhedra get edges, point-classes stay sparse)
+  - Shapes generated in native aspect ratio space
+  
+Multi-scale usage:
+  Patches extracted at any scale (4×8×8, 8×16×16, 16×32×32, 32×64×64)
+  are resized to canonical 8×16×16 before classification.
+"""
+
+import numpy as np
+from itertools import combinations
+
+# === Grid dimensions (non-cubic, VAE-matched) ================================
+GZ = 8    # channel dimension (thin)
+GY = 16   # spatial height
+GX = 16   # spatial width
+GRID_SHAPE = (GZ, GY, GX)
+GRID_VOLUME = GZ * GY * GX  # 2048
+
+# Precompute coordinate grid for vectorized generation
+_COORDS = np.mgrid[0:GZ, 0:GY, 0:GX].reshape(3, -1).T.astype(np.float64)
+
+# === Shape classes ============================================================
+CLASS_META = {
+    # 0D
+    "point":            {"dim": 0, "curved": False, "curvature": "none"},
+    # 1D
+    "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"},
+    "cross":            {"dim": 1, "curved": False, "curvature": "none"},
+    "l_shape":          {"dim": 1, "curved": False, "curvature": "none"},
+    "collinear":        {"dim": 1, "curved": False, "curvature": "none"},
+    # 2D flat
+    "triangle_xy":      {"dim": 2, "curved": False, "curvature": "none"},
+    "triangle_xz":      {"dim": 2, "curved": False, "curvature": "none"},
+    "triangle_3d":      {"dim": 2, "curved": False, "curvature": "none"},
+    "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"},
+    "plane":            {"dim": 2, "curved": False, "curvature": "none"},
+    # 3D flat
+    "tetrahedron":      {"dim": 3, "curved": False, "curvature": "none"},
+    "pyramid":          {"dim": 3, "curved": False, "curvature": "none"},
+    "pentachoron":      {"dim": 3, "curved": False, "curvature": "none"},
+    "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"},
+    # 1D curved
+    "arc":              {"dim": 1, "curved": True,  "curvature": "convex"},
+    "helix":            {"dim": 1, "curved": True,  "curvature": "helical"},
+    # 2D curved
+    "circle":           {"dim": 2, "curved": True,  "curvature": "convex"},
+    "ellipse":          {"dim": 2, "curved": True,  "curvature": "convex"},
+    "disc":             {"dim": 2, "curved": True,  "curvature": "convex"},
+    # 3D curved
+    "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"},
+    "shell":            {"dim": 3, "curved": True,  "curvature": "convex"},
+    "tube":             {"dim": 3, "curved": True,  "curvature": "cylindrical"},
+    "bowl":             {"dim": 3, "curved": True,  "curvature": "concave"},
+    "saddle":           {"dim": 3, "curved": True,  "curvature": "hyperbolic"},
+}
+
+CLASS_NAMES = list(CLASS_META.keys())
+NUM_CLASSES = len(CLASS_NAMES)
+CLASS_TO_IDX = {n: i for i, n in enumerate(CLASS_NAMES)}
+
+CURVATURE_NAMES = ["none", "convex", "concave", "cylindrical",
+                   "conical", "toroidal", "hyperbolic", "helical"]
+CURV_TO_IDX = {n: i for i, n in enumerate(CURVATURE_NAMES)}
+
+# Edge topology for rasterized shapes
+TRIANGLE_EDGES = [(0,1), (1,2), (2,0)]
+QUAD_EDGES = [(0,1), (1,3), (3,2), (2,0)]
+TETRA_EDGES = list(combinations(range(4), 2))  # 6 edges
+CUBE_EDGES = [(0,1),(0,2),(0,4),(1,3),(1,5),(2,3),(2,6),(3,7),(4,5),(4,6),(5,7),(6,7)]
+PYRAMID_EDGES = [(0,1),(1,3),(3,2),(2,0),(0,4),(1,4),(2,4),(3,4)]
+PENTA_EDGES = list(combinations(range(5), 2))  # 10 edges
+OCTA_EDGES = [(0,1),(0,2),(0,3),(0,4),(5,1),(5,2),(5,3),(5,4),(1,2),(2,3),(3,4),(4,1)]
+
+
+def rasterize_line(p1, p2):
+    """Bresenham-style 3D line rasterization between two points."""
+    p1 = np.array(p1, dtype=float)
+    p2 = np.array(p2, dtype=float)
+    diff = p2 - p1
+    n_steps = max(int(np.max(np.abs(diff))) + 1, 2)
+    t = np.linspace(0, 1, n_steps)
+    pts = p1[None, :] + t[:, None] * diff[None, :]
+    pts = np.round(pts).astype(int)
+    # Clip to grid bounds
+    pts[:, 0] = np.clip(pts[:, 0], 0, GZ - 1)
+    pts[:, 1] = np.clip(pts[:, 1], 0, GY - 1)
+    pts[:, 2] = np.clip(pts[:, 2], 0, GX - 1)
+    return np.unique(pts, axis=0)
+
+
+def rasterize_edges(vertices, edges):
+    """Rasterize a complete wireframe from vertex list and edge topology."""
+    all_pts = [vertices]
+    for i, j in edges:
+        all_pts.append(rasterize_line(vertices[i], vertices[j]))
+    return np.unique(np.vstack(all_pts), axis=0)
+
+
+class ShapeGenerator:
+    def __init__(self, seed=42):
+        self.rng = np.random.RandomState(seed)
+
+    def _pts_to_result(self, pts):
+        """Convert point array to grid + metadata."""
+        pts = np.atleast_2d(pts).astype(int)
+        # Clip to grid
+        pts[:, 0] = np.clip(pts[:, 0], 0, GZ - 1)
+        pts[:, 1] = np.clip(pts[:, 1], 0, GY - 1)
+        pts[:, 2] = np.clip(pts[:, 2], 0, GX - 1)
+        pts = np.unique(pts, axis=0)
+        grid = np.zeros(GRID_SHAPE, dtype=np.float32)
+        grid[pts[:, 0], pts[:, 1], pts[:, 2]] = 1.0
+        return {"grid": grid, "n_occupied": int(pts.shape[0]), "points": pts}
+
+    def _rand_center(self, margin_z=2, margin_yx=3):
+        """Random center respecting aspect ratio margins."""
+        cz = self.rng.uniform(margin_z, GZ - margin_z)
+        cy = self.rng.uniform(margin_yx, GY - margin_yx)
+        cx = self.rng.uniform(margin_yx, GX - margin_yx)
+        return np.array([cz, cy, cx])
+
+    def _rand_pts_2d(self, n, min_dist=2):
+        """Random 2D points in YX plane."""
+        for _ in range(100):
+            pts = np.column_stack([
+                self.rng.randint(1, GY - 1, n),
+                self.rng.randint(1, GX - 1, n)])
+            dists = [np.linalg.norm(pts[i] - pts[j])
+                     for i in range(n) for j in range(i+1, n)]
+            if all(d >= min_dist for d in dists):
+                return pts
+        return None
+
+    def _rand_pts_3d(self, n, min_dist=2):
+        """Random 3D points respecting grid bounds."""
+        for _ in range(100):
+            pts = np.column_stack([
+                self.rng.randint(0, GZ, n),
+                self.rng.randint(1, GY - 1, n),
+                self.rng.randint(1, GX - 1, n)])
+            dists = [np.linalg.norm(pts[i] - pts[j])
+                     for i in range(n) for j in range(i+1, n)]
+            if all(d >= min_dist for d in dists):
+                return pts
+        return None
+
+    def _rigid(self, name):
+        """Generate rotation axis for rigid shapes."""
+        axes = [(1,0,0), (0,1,0), (0,0,1)]
+        return axes[self.rng.randint(len(axes))]
+
+    def _make(self, name):
+        rng = self.rng
+
+        # === 0D ===
+        if name == "point":
+            z = rng.randint(0, GZ)
+            y = rng.randint(0, GY)
+            x = rng.randint(0, GX)
+            return self._pts_to_result(np.array([[z, y, x]]))
+
+        # === 1D lines ===
+        elif name == "line_x":
+            z = rng.randint(0, GZ)
+            y = rng.randint(0, GY)
+            x1, x2 = sorted(rng.choice(GX, 2, replace=False))
+            pts = np.array([[z, y, x1], [z, y, x2]])
+            return self._pts_to_result(rasterize_line(pts[0], pts[1]))
+
+        elif name == "line_y":
+            z = rng.randint(0, GZ)
+            x = rng.randint(0, GX)
+            y1, y2 = sorted(rng.choice(GY, 2, replace=False))
+            pts = np.array([[z, y1, x], [z, y2, x]])
+            return self._pts_to_result(rasterize_line(pts[0], pts[1]))
+
+        elif name == "line_z":
+            y = rng.randint(0, GY)
+            x = rng.randint(0, GX)
+            z1, z2 = sorted(rng.choice(GZ, 2, replace=False))
+            pts = np.array([[z1, y, x], [z2, y, x]])
+            return self._pts_to_result(rasterize_line(pts[0], pts[1]))
+
+        elif name == "line_diag":
+            p1 = np.array([rng.randint(0, GZ), rng.randint(0, GY), rng.randint(0, GX)])
+            p2 = np.array([rng.randint(0, GZ), rng.randint(0, GY), rng.randint(0, GX)])
+            if np.linalg.norm(p1 - p2) < 3:
+                return None
+            return self._pts_to_result(rasterize_line(p1, p2))
+
+        elif name == "cross":
+            c = self._rand_center(margin_z=1, margin_yx=2)
+            arm_yx = rng.randint(2, min(6, GY // 2))
+            arm_z = rng.randint(1, min(3, GZ // 2))
+            pts = []
+            ci = np.round(c).astype(int)
+            # YX cross
+            for dy in range(-arm_yx, arm_yx + 1):
+                pts.append([ci[0], np.clip(ci[1] + dy, 0, GY-1), ci[2]])
+            for dx in range(-arm_yx, arm_yx + 1):
+                pts.append([ci[0], ci[1], np.clip(ci[2] + dx, 0, GX-1)])
+            # Z arm
+            for dz in range(-arm_z, arm_z + 1):
+                pts.append([np.clip(ci[0] + dz, 0, GZ-1), ci[1], ci[2]])
+            return self._pts_to_result(np.array(pts))
+
+        elif name == "l_shape":
+            c = self._rand_center(margin_z=1, margin_yx=2)
+            arm = rng.randint(2, min(5, GY // 2))
+            ci = np.round(c).astype(int)
+            pts = []
+            for dy in range(arm + 1):
+                pts.append([ci[0], np.clip(ci[1] + dy, 0, GY-1), ci[2]])
+            for dx in range(1, arm + 1):
+                pts.append([ci[0], ci[1], np.clip(ci[2] + dx, 0, GX-1)])
+            return self._pts_to_result(np.array(pts))
+
+        elif name == "collinear":
+            # Identity = 3 discrete points on a line, NOT a line segment
+            axis = rng.randint(3)
+            gs = [GZ, GY, GX]
+            vals = sorted(rng.choice(gs[axis], 3, replace=False))
+            fixed = [rng.randint(0, gs[(axis+1)%3]), rng.randint(0, gs[(axis+2)%3])]
+            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 self._pts_to_result(pts)
+
+        # === 2D flat ===
+        elif name == "triangle_xy":
+            z = rng.randint(0, GZ)
+            pts2d = self._rand_pts_2d(3, min_dist=3)
+            if pts2d is None: return None
+            return self._pts_to_result(np.column_stack([np.full(3, z), pts2d]))
+
+        elif name == "triangle_xz":
+            y = rng.randint(0, GY)
+            for _ in range(50):
+                pts = np.column_stack([
+                    rng.randint(0, GZ, 3),
+                    np.full(3, y),
+                    rng.randint(1, GX - 1, 3)])
+                dists = [np.linalg.norm(pts[i] - pts[j]) for i in range(3) for j in range(i+1, 3)]
+                if all(d >= 2 for d in dists):
+                    return self._pts_to_result(pts)
+            return None
+
+        elif name == "triangle_3d":
+            verts = self._rand_pts_3d(3, min_dist=3)
+            if verts is None: return None
+            return self._pts_to_result(verts)
+
+        elif name == "square_xy":
+            z = rng.randint(0, GZ)
+            s = rng.randint(3, min(7, GY - 2))
+            cy, cx = rng.randint(s, GY - s), rng.randint(s, GX - s)
+            verts = np.array([
+                [z, cy - s, cx - s], [z, cy - s, cx + s],
+                [z, cy + s, cx - s], [z, cy + s, cx + s]])
+            return self._pts_to_result(rasterize_edges(verts, QUAD_EDGES))
+
+        elif name == "square_xz":
+            y = rng.randint(0, GY)
+            s_z = rng.randint(1, min(3, GZ // 2))
+            s_x = rng.randint(2, min(6, GX // 2))
+            cz, cx = rng.randint(s_z, GZ - s_z), rng.randint(s_x, GX - s_x)
+            verts = np.array([
+                [cz - s_z, y, cx - s_x], [cz - s_z, y, cx + s_x],
+                [cz + s_z, y, cx - s_x], [cz + s_z, y, cx + s_x]])
+            return self._pts_to_result(rasterize_edges(verts, QUAD_EDGES))
+
+        elif name == "rectangle":
+            z = rng.randint(0, GZ)
+            sy = rng.randint(2, min(6, GY // 2))
+            sx = rng.randint(2, min(6, GX // 2))
+            while abs(sy - sx) < 2:
+                sy = rng.randint(2, min(6, GY // 2))
+                sx = rng.randint(2, min(6, GX // 2))
+            cy, cx = rng.randint(sy, GY - sy), rng.randint(sx, GX - sx)
+            verts = np.array([
+                [z, cy - sy, cx - sx], [z, cy - sy, cx + sx],
+                [z, cy + sy, cx - sx], [z, cy + sy, cx + sx]])
+            return self._pts_to_result(rasterize_edges(verts, QUAD_EDGES))
+
+        elif name == "coplanar":
+            # Identity = 4 discrete coplanar points, NOT a quadrilateral
+            pts = self._rand_pts_3d(4, min_dist=2)
+            if pts is None: return None
+            axis = rng.randint(3)
+            pts[:, axis] = pts[0, axis]
+            return self._pts_to_result(pts)
+
+        elif name == "plane":
+            axis = rng.randint(3)
+            gs = [GZ, GY, GX]
+            pos = rng.randint(0, gs[axis])
+            thick = rng.randint(1, max(2, gs[axis] // 4) + 1)
+            mask = np.zeros(GRID_SHAPE, dtype=np.float32)
+            for t in range(thick):
+                p = min(pos + t, gs[axis] - 1)
+                if axis == 0: mask[p, :, :] = 1
+                elif axis == 1: mask[:, p, :] = 1
+                else: mask[:, :, p] = 1
+            pts = np.argwhere(mask > 0)
+            return self._pts_to_result(pts)
+
+        # === 3D polyhedra (rasterized edges) ===
+        elif name == "tetrahedron":
+            verts = self._rand_pts_3d(4, min_dist=3)
+            if verts is None: return None
+            return self._pts_to_result(rasterize_edges(verts, TETRA_EDGES))
+
+        elif name == "pyramid":
+            base_y = rng.randint(2, GY - 2)
+            s = rng.randint(2, min(5, GX // 2))
+            cy, cx = rng.randint(s + 1, GY - s - 1), rng.randint(s + 1, GX - s - 1)
+            base_z = rng.randint(1, GZ - 2)
+            apex_z = rng.randint(0, GZ) if rng.random() < 0.5 else base_z + rng.randint(2, min(4, GZ - base_z))
+            apex_z = min(apex_z, GZ - 1)
+            verts = np.array([
+                [base_z, cy - s, cx - s], [base_z, cy - s, cx + s],
+                [base_z, cy + s, cx - s], [base_z, cy + s, cx + s],
+                [apex_z, cy, cx]])
+            return self._pts_to_result(rasterize_edges(verts, PYRAMID_EDGES))
+
+        elif name == "pentachoron":
+            verts = self._rand_pts_3d(5, min_dist=3)
+            if verts is None: return None
+            return self._pts_to_result(rasterize_edges(verts, PENTA_EDGES))
+
+        elif name == "cube":
+            s_z = rng.randint(1, min(3, GZ // 2))
+            s_yx = rng.randint(2, min(5, GY // 2))
+            c = self._rand_center(margin_z=s_z + 1, margin_yx=s_yx + 1)
+            ci = np.round(c).astype(int)
+            verts = np.array([
+                [ci[0]-s_z, ci[1]-s_yx, ci[2]-s_yx],
+                [ci[0]-s_z, ci[1]-s_yx, ci[2]+s_yx],
+                [ci[0]-s_z, ci[1]+s_yx, ci[2]-s_yx],
+                [ci[0]-s_z, ci[1]+s_yx, ci[2]+s_yx],
+                [ci[0]+s_z, ci[1]-s_yx, ci[2]-s_yx],
+                [ci[0]+s_z, ci[1]-s_yx, ci[2]+s_yx],
+                [ci[0]+s_z, ci[1]+s_yx, ci[2]-s_yx],
+                [ci[0]+s_z, ci[1]+s_yx, ci[2]+s_yx]])
+            return self._pts_to_result(rasterize_edges(verts, CUBE_EDGES))
+
+        elif name == "cuboid":
+            sz = rng.randint(1, min(3, GZ // 2))
+            sy = rng.randint(2, min(6, GY // 2))
+            sx = rng.randint(2, min(6, GX // 2))
+            # Ensure at least one dimension differs significantly
+            while abs(sy - sx) < 2 and abs(sz * 2 - sy) < 2:
+                sy = rng.randint(2, min(6, GY // 2))
+                sx = rng.randint(2, min(6, GX // 2))
+            c = self._rand_center(margin_z=sz + 1, margin_yx=max(sy, sx) + 1)
+            ci = np.round(c).astype(int)
+            verts = np.array([
+                [ci[0]-sz, ci[1]-sy, ci[2]-sx], [ci[0]-sz, ci[1]-sy, ci[2]+sx],
+                [ci[0]-sz, ci[1]+sy, ci[2]-sx], [ci[0]-sz, ci[1]+sy, ci[2]+sx],
+                [ci[0]+sz, ci[1]-sy, ci[2]-sx], [ci[0]+sz, ci[1]-sy, ci[2]+sx],
+                [ci[0]+sz, ci[1]+sy, ci[2]-sx], [ci[0]+sz, ci[1]+sy, ci[2]+sx]])
+            return self._pts_to_result(rasterize_edges(verts, CUBE_EDGES))
+
+        elif name == "triangular_prism":
+            z1, z2 = sorted(rng.choice(GZ, 2, replace=False))
+            pts2d = self._rand_pts_2d(3, min_dist=3)
+            if pts2d is None: return None
+            # Two triangular faces + connecting edges
+            top = np.column_stack([np.full(3, z1), pts2d])
+            bot = np.column_stack([np.full(3, z2), pts2d])
+            verts = np.vstack([top, bot])
+            edges = [(0,1),(1,2),(2,0),(3,4),(4,5),(5,3),(0,3),(1,4),(2,5)]
+            return self._pts_to_result(rasterize_edges(verts, edges))
+
+        elif name == "octahedron":
+            c = self._rand_center(margin_z=2, margin_yx=3)
+            ci = np.round(c).astype(int)
+            rz = rng.randint(1, min(3, GZ // 2))
+            ryx = rng.randint(2, min(5, GY // 2))
+            verts = np.array([
+                [ci[0], ci[1] + ryx, ci[2]], [ci[0], ci[1], ci[2] + ryx],
+                [ci[0], ci[1] - ryx, ci[2]], [ci[0], ci[1], ci[2] - ryx],
+                [ci[0] + rz, ci[1], ci[2]], [ci[0] - rz, ci[1], ci[2]]])
+            return self._pts_to_result(rasterize_edges(verts, OCTA_EDGES))
+
+        # === 1D curved ===
+        elif name == "arc":
+            plane = rng.randint(3)
+            c = self._rand_center(margin_z=1, margin_yx=2)
+            r_main = rng.uniform(2.0, min(5.0, GY / 2 - 1))
+            r_z = rng.uniform(1.0, min(3.0, GZ / 2 - 1)) if plane != 0 else r_main
+            angle_start = rng.uniform(0, np.pi)
+            angle_span = rng.uniform(np.pi / 3, np.pi)
+            n_pts = max(8, int(angle_span * r_main))
+            t = np.linspace(angle_start, angle_start + angle_span, n_pts)
+            if plane == 0:  # YX plane
+                pts = np.column_stack([
+                    np.full(n_pts, c[0]),
+                    c[1] + r_main * np.cos(t),
+                    c[2] + r_main * np.sin(t)])
+            elif plane == 1:  # ZX plane
+                pts = np.column_stack([
+                    c[0] + r_z * np.cos(t),
+                    np.full(n_pts, c[1]),
+                    c[2] + r_main * np.sin(t)])
+            else:  # ZY plane
+                pts = np.column_stack([
+                    c[0] + r_z * np.cos(t),
+                    c[1] + r_main * np.sin(t),
+                    np.full(n_pts, c[2])])
+            pts = np.round(pts).astype(int)
+            return self._pts_to_result(pts)
+
+        elif name == "helix":
+            c = self._rand_center(margin_z=0, margin_yx=3)
+            r = rng.uniform(1.5, min(4.0, GY / 2 - 2))
+            turns = rng.uniform(1.0, 2.5)
+            n_pts = int(turns * 20)
+            t = np.linspace(0, turns * 2 * np.pi, n_pts)
+            z_span = GZ - 1
+            pts = np.column_stack([
+                t / (turns * 2 * np.pi) * z_span,
+                c[1] + r * np.cos(t),
+                c[2] + r * np.sin(t)])
+            pts = np.round(pts).astype(int)
+            return self._pts_to_result(pts)
+
+        # === 2D curved ===
+        elif name == "circle":
+            plane = rng.randint(3)
+            c = self._rand_center(margin_z=1, margin_yx=3)
+            r = rng.uniform(2.0, min(5.0, GY / 2 - 1))
+            n_pts = max(12, int(2 * np.pi * r))
+            t = np.linspace(0, 2 * np.pi, n_pts, endpoint=False)
+            if plane == 0:
+                pts = np.column_stack([
+                    np.full(n_pts, c[0]),
+                    c[1] + r * np.cos(t),
+                    c[2] + r * np.sin(t)])
+            elif plane == 1:
+                r_z = min(r, GZ / 2 - 1)
+                pts = np.column_stack([
+                    c[0] + r_z * np.cos(t),
+                    np.full(n_pts, c[1]),
+                    c[2] + r * np.sin(t)])
+            else:
+                r_z = min(r, GZ / 2 - 1)
+                pts = np.column_stack([
+                    c[0] + r_z * np.cos(t),
+                    c[1] + r * np.sin(t),
+                    np.full(n_pts, c[2])])
+            pts = np.round(pts).astype(int)
+            return self._pts_to_result(pts)
+
+        elif name == "ellipse":
+            c = self._rand_center(margin_z=1, margin_yx=3)
+            ry = rng.uniform(2.0, min(5.0, GY / 2 - 1))
+            ratio = rng.uniform(1.6, 2.5)
+            if rng.random() < 0.5:
+                rx = ry / ratio
+            else:
+                rx = ry * ratio
+                rx = min(rx, GX / 2 - 1)
+                if rx / ry < 1.6: ry = rx / 1.6
+            n_pts = max(16, int(2 * np.pi * max(rx, ry)))
+            t = np.linspace(0, 2 * np.pi, n_pts, endpoint=False)
+            pts = np.column_stack([
+                np.full(n_pts, c[0]),
+                c[1] + ry * np.cos(t),
+                c[2] + rx * np.sin(t)])
+            pts = np.round(pts).astype(int)
+            return self._pts_to_result(pts)
+
+        elif name == "disc":
+            plane = rng.randint(3)
+            c = self._rand_center(margin_z=1, margin_yx=3)
+            r = rng.uniform(2.0, min(5.0, GY / 2 - 1))
+            if plane == 0:
+                mask = ((_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2 <= r**2) & \
+                       (np.abs(_COORDS[:, 0] - c[0]) < 0.6)
+            elif plane == 1:
+                r_z = min(r, GZ / 2 - 1)
+                mask = ((_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 +
+                        (_COORDS[:, 2] - c[2])**2 / r**2 <= 1) & \
+                       (np.abs(_COORDS[:, 1] - c[1]) < 0.6)
+            else:
+                r_z = min(r, GZ / 2 - 1)
+                mask = ((_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 +
+                        (_COORDS[:, 1] - c[1])**2 / r**2 <= 1) & \
+                       (np.abs(_COORDS[:, 2] - c[2]) < 0.6)
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 3: return None
+            return self._pts_to_result(pts)
+
+        # === 3D curved ===
+        elif name == "sphere":
+            c = self._rand_center(margin_z=2, margin_yx=3)
+            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
+            # Use ellipsoidal check respecting aspect ratio
+            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
+                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
+                 ((_COORDS[:, 2] - c[2]) / r)**2
+            mask = d2 <= 1.0
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        elif name == "hemisphere":
+            c = self._rand_center(margin_z=2, margin_yx=3)
+            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
+            cut_axis = rng.randint(3)
+            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
+                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
+                 ((_COORDS[:, 2] - c[2]) / r)**2
+            mask = d2 <= 1.0
+            if cut_axis == 0: mask &= _COORDS[:, 0] >= c[0]
+            elif cut_axis == 1: mask &= _COORDS[:, 1] >= c[1]
+            else: mask &= _COORDS[:, 2] >= c[2]
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 3: return None
+            return self._pts_to_result(pts)
+
+        elif name == "cylinder":
+            axis = rng.randint(3)
+            c = self._rand_center(margin_z=0, margin_yx=3)
+            r = rng.uniform(1.5, min(3.0, GY / 2 - 1))
+            if axis == 0:
+                d2 = (_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2
+                mask = d2 <= r**2
+            elif axis == 1:
+                r_z = min(r, GZ / 2 - 0.5)
+                d2 = (_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 + \
+                     (_COORDS[:, 2] - c[2])**2 / r**2
+                mask = d2 <= 1.0
+            else:
+                r_z = min(r, GZ / 2 - 0.5)
+                d2 = (_COORDS[:, 0] - c[0])**2 / max(r_z, 0.5)**2 + \
+                     (_COORDS[:, 1] - c[1])**2 / r**2
+                mask = d2 <= 1.0
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        elif name == "cone":
+            axis = rng.randint(3)
+            c = self._rand_center(margin_z=1, margin_yx=3)
+            r = rng.uniform(2.0, min(4.0, GY / 2 - 1))
+            gs = [GZ, GY, GX]
+            h = gs[axis] - 1
+            apex_frac = _COORDS[:, axis] / max(h, 1)
+            local_r = r * (1.0 - apex_frac)
+            if axis == 0:
+                d2 = (_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2
+            elif axis == 1:
+                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 2] - c[2])**2
+            else:
+                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 1] - c[1])**2
+            mask = d2 <= local_r**2
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        elif name == "capsule":
+            axis = rng.randint(3)
+            c = self._rand_center(margin_z=1, margin_yx=3)
+            r = rng.uniform(1.5, min(2.5, GZ / 2 - 0.5, GY / 2 - 1))
+            gs = [GZ, GY, GX]
+            half_h = rng.uniform(1.0, gs[axis] / 2 - r - 0.5)
+            # Cylinder body + spherical caps
+            dist_axis = np.abs(_COORDS[:, axis] - c[axis])
+            clamped = np.clip(dist_axis - half_h, 0, None)
+            perp_axes = [i for i in range(3) if i != axis]
+            d2 = clamped**2
+            for a in perp_axes:
+                d2 += (_COORDS[:, a] - c[a])**2
+            mask = d2 <= r**2
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        elif name == "torus":
+            c = self._rand_center(margin_z=2, margin_yx=4)
+            R = rng.uniform(2.5, min(4.0, GY / 2 - 2))
+            r = rng.uniform(0.8, min(1.5, GZ / 2 - 0.5, R * 0.5))
+            # Torus in YX plane
+            d_yx = np.sqrt((_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2)
+            d2 = (d_yx - R)**2 + (_COORDS[:, 0] - c[0])**2
+            mask = d2 <= r**2
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        elif name == "shell":
+            c = self._rand_center(margin_z=1, margin_yx=3)
+            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
+            thick = rng.uniform(0.4, 0.8)
+            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
+                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
+                 ((_COORDS[:, 2] - c[2]) / r)**2
+            mask = (d2 <= 1.0) & (d2 >= (1.0 - thick)**2)
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        elif name == "tube":
+            axis = rng.randint(3)
+            c = self._rand_center(margin_z=0, margin_yx=3)
+            r_out = rng.uniform(2.0, min(3.5, GY / 2 - 1))
+            r_in = r_out * rng.uniform(0.4, 0.7)
+            if axis == 0:
+                d2 = (_COORDS[:, 1] - c[1])**2 + (_COORDS[:, 2] - c[2])**2
+            elif axis == 1:
+                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 2] - c[2])**2
+            else:
+                d2 = (_COORDS[:, 0] - c[0])**2 + (_COORDS[:, 1] - c[1])**2
+            mask = (d2 <= r_out**2) & (d2 >= r_in**2)
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        elif name == "bowl":
+            c = self._rand_center(margin_z=1, margin_yx=3)
+            r = rng.uniform(2.0, min(3.5, GZ / 2 - 0.5, GY / 2 - 1))
+            thick = rng.uniform(0.3, 0.7)
+            d2 = ((_COORDS[:, 0] - c[0]) / r)**2 + \
+                 ((_COORDS[:, 1] - c[1]) / r)**2 + \
+                 ((_COORDS[:, 2] - c[2]) / r)**2
+            mask = (d2 <= 1.0) & (d2 >= (1.0 - thick)**2) & (_COORDS[:, 0] <= c[0])
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 3: return None
+            return self._pts_to_result(pts)
+
+        elif name == "saddle":
+            c = self._rand_center(margin_z=2, margin_yx=4)
+            scale = rng.uniform(1.5, 3.0)
+            dy = (_COORDS[:, 1] - c[1]) / scale
+            dx = (_COORDS[:, 2] - c[2]) / scale
+            z_saddle = c[0] + (dy**2 - dx**2)
+            mask = np.abs(_COORDS[:, 0] - z_saddle) < 0.8
+            pts = _COORDS[mask].astype(int)
+            if len(pts) < 4: return None
+            return self._pts_to_result(pts)
+
+        return None
+
+    def generate(self, name, max_retries=10):
+        """Generate one sample with retries."""
+        for _ in range(max_retries):
+            result = self._make(name)
+            if result is not None and result["n_occupied"] > 0:
+                return result
+        return None
+
+    def generate_dataset(self, n_per_class, seed=None):
+        """Generate balanced dataset."""
+        if seed is not None:
+            self.rng = np.random.RandomState(seed)
+        grids, labels, dims, curveds = [], [], [], []
+        for cls_idx, name in enumerate(CLASS_NAMES):
+            meta = CLASS_META[name]
+            count = 0
+            while count < n_per_class:
+                self.rng = np.random.RandomState(seed * 1000 + cls_idx * n_per_class + count if seed else None)
+                result = self.generate(name)
+                if result is not None:
+                    grids.append(result["grid"])
+                    labels.append(cls_idx)
+                    dims.append(meta["dim"])
+                    curveds.append(1 if meta["curved"] else 0)
+                    count += 1
+        return {
+            "grids": np.array(grids),
+            "labels": np.array(labels),
+            "dims": np.array(dims),
+            "curveds": np.array(curveds),
+        }
+
+
+# === Verification =============================================================
+if __name__ == "__main__":
+    gen = ShapeGenerator(seed=42)
+    print(f"Grid: {GZ}×{GY}×{GX} = {GRID_VOLUME} voxels")
+    print(f"Classes: {NUM_CLASSES}")
+    print(f"\n{'Shape':20s} {'OK':>4s} {'Avg vox':>8s}")
+    print("-" * 36)
+    for name in CLASS_NAMES:
+        ok = 0; voxels = []
+        for trial in range(20):
+            gen.rng = np.random.RandomState(trial * 100 + hash(name) % 10000)
+            s = gen.generate(name)
+            if s:
+                ok += 1
+                voxels.append(s["n_occupied"])
+        avg = np.mean(voxels) if voxels else 0
+        status = "✓" if ok >= 15 else "✗"
+        print(f"  {status} {name:20s} {ok:2d}/20 {avg:7.1f}")
+
+print(f'\nLoaded {NUM_CLASSES} shape classes, grid={GZ}×{GY}×{GX}')