#!/usr/bin/env python3 """Counterfactual edit benchmark for editable superquadric programs. This version deliberately separates two things that the first draft blurred: 1. The target object is sampled from high-resolution procedural mesh components such as frusta, tubes, torus rims, and rounded plates. 2. The fitted representation is a superquadric program. 3. The comparison is between a semantically coherent editable program and a surface-matched program whose local SQ parts do not correspond cleanly to semantic controls. That gives the page a more honest visual standard: the ground-truth object is not itself a low-resolution superquadric sketch, while the experiment still tests whether surface reconstruction error alone is insufficient for evaluating editable primitive programs. """ from __future__ import annotations import argparse import json import math import random from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Dict, Iterable, List, Sequence, Tuple import numpy as np import torch PARAMS = [ "sx", "sy", "sz", "e1", "e2", "bend", "taper_y", "taper_z", "tx", "ty", "tz", "rx", "ry", "rz", ] BOUNDS = { "sx": (0.025, 2.2), "sy": (0.02, 1.6), "sz": (0.02, 1.6), "e1": (0.18, 1.65), "e2": (0.18, 1.65), "bend": (-2.2, 2.2), "taper_y": (-0.9, 0.9), "taper_z": (-0.9, 0.9), "tx": (-2.4, 2.4), "ty": (-2.4, 2.4), "tz": (-2.4, 2.4), "rx": (-math.pi, math.pi), "ry": (-math.pi, math.pi), "rz": (-math.pi, math.pi), } COLORS = { "body": "#0ea5b7", "handle": "#ef5a3c", "base": "#f6c445", "neck": "#ef5a3c", "shade": "#6d5bd0", "hook": "#ef5a3c", "bar": "#0ea5b7", "earcup": "#0ea5b7", "cushion": "#27221d", "band": "#ef5a3c", "mount": "#f6c445", "metal": "#aeb7c1", "face": "#f8ead2", "bulb": "#ffe18a", "shadow": "#27221d", "rim": "#0b8793", "surface": "#6d5bd0", } @dataclass class Primitive: name: str group: str values: Dict[str, float] trainable: Iterable[str] = field(default_factory=list) color: str = "#0ea5b7" def tensor(self, device: torch.device, jitter: float, rng: random.Random) -> Tuple[torch.nn.Parameter, torch.Tensor]: vals = [] mask = [] trainable = set(self.trainable) for key in PARAMS: value = self.values.get(key, default_value(key)) if key in trainable and jitter > 0: lo, hi = BOUNDS[key] span = hi - lo value += rng.gauss(0, jitter * span) value = min(hi, max(lo, value)) vals.append(value) mask.append(1.0 if key in trainable else 0.0) return torch.nn.Parameter(torch.tensor(vals, dtype=torch.float32, device=device)), torch.tensor(mask, dtype=torch.float32, device=device) @dataclass class FitFamily: key: str label: str primitives: List[Primitive] edit_note: str @dataclass class ObjectSpec: key: str title: str subtitle: str edit_label: str story: str target_components: List[Dict[str, Any]] edited_components: List[Dict[str, Any]] program: FitFamily surface: FitFamily edit: Callable[[List[Dict[str, float]], str], List[Dict[str, float]]] def default_value(key: str) -> float: if key in {"sx", "sy", "sz"}: return 0.2 if key in {"e1", "e2"}: return 0.65 return 0.0 def spow(x: torch.Tensor, power: torch.Tensor | float) -> torch.Tensor: return torch.sign(x) * torch.pow(torch.clamp(torch.abs(x), min=1e-6), power) def spow_np(x: np.ndarray, power: float) -> np.ndarray: return np.sign(x) * np.power(np.clip(np.abs(x), 1e-6, None), power) def make_uv(rows: int, cols: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: eta = torch.linspace(-math.pi / 2, math.pi / 2, rows, device=device) omega = torch.linspace(-math.pi, math.pi, cols + 1, device=device)[:-1] grid_eta, grid_omega = torch.meshgrid(eta, omega, indexing="ij") return grid_eta.reshape(-1), grid_omega.reshape(-1) def rotation_matrix(rx: torch.Tensor, ry: torch.Tensor, rz: torch.Tensor) -> torch.Tensor: cx, sx = torch.cos(rx), torch.sin(rx) cy, sy = torch.cos(ry), torch.sin(ry) cz, sz = torch.cos(rz), torch.sin(rz) one = torch.ones_like(rx) zero = torch.zeros_like(rx) rxm = torch.stack([ torch.stack([one, zero, zero]), torch.stack([zero, cx, -sx]), torch.stack([zero, sx, cx]), ]) rym = torch.stack([ torch.stack([cy, zero, sy]), torch.stack([zero, one, zero]), torch.stack([-sy, zero, cy]), ]) rzm = torch.stack([ torch.stack([cz, -sz, zero]), torch.stack([sz, cz, zero]), torch.stack([zero, zero, one]), ]) return rzm @ rym @ rxm def rotation_matrix_np(rx: float, ry: float, rz: float) -> np.ndarray: cx, sx = math.cos(rx), math.sin(rx) cy, sy = math.cos(ry), math.sin(ry) cz, sz = math.cos(rz), math.sin(rz) rxm = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]], dtype=np.float64) rym = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]], dtype=np.float64) rzm = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]], dtype=np.float64) return rzm @ rym @ rxm def transform_np(points: np.ndarray, item: Dict[str, Any]) -> np.ndarray: rot = rotation_matrix_np(float(item.get("rx", 0.0)), float(item.get("ry", 0.0)), float(item.get("rz", 0.0))) trans = np.array([item.get("tx", 0.0), item.get("ty", 0.0), item.get("tz", 0.0)], dtype=np.float64) return np.einsum("ij,kj->ik", points.astype(np.float64, copy=False), rot) + trans def primitive_surface(params: torch.Tensor, eta: torch.Tensor, omega: torch.Tensor) -> torch.Tensor: sx, sy, sz, e1, e2, bend, taper_y, taper_z, tx, ty, tz, rx, ry, rz = params ce = torch.cos(eta) se = torch.sin(eta) co = torch.cos(omega) so = torch.sin(omega) x = sx * spow(ce, e1) * spow(co, e2) y = sy * spow(ce, e1) * spow(so, e2) z = sz * spow(se, e1) axis = torch.clamp(x / torch.clamp(sx, min=1e-4), -1.0, 1.0) y = y * torch.clamp(1.0 + taper_y * axis, 0.25, 1.9) z = z * torch.clamp(1.0 + taper_z * axis, 0.25, 1.9) theta = bend * x safe_bend = torch.where(torch.abs(bend) < 1e-4, torch.ones_like(bend), bend) bx = torch.sin(theta) / safe_bend - y * torch.sin(theta) by = (1.0 - torch.cos(theta)) / safe_bend + y * torch.cos(theta) bx = torch.where(torch.abs(bend) < 1e-4, x, bx) by = torch.where(torch.abs(bend) < 1e-4, y, by) pts = torch.stack([bx, by, z], dim=-1) rot = rotation_matrix(rx, ry, rz) pts = pts @ rot.T pts = pts + torch.stack([tx, ty, tz]) return pts def superquadric_np(params: Dict[str, float], rows: int, cols: int) -> np.ndarray: eta = np.linspace(-math.pi / 2, math.pi / 2, rows) omega = np.linspace(-math.pi, math.pi, cols, endpoint=False) eta_grid, omega_grid = np.meshgrid(eta, omega, indexing="ij") ce = np.cos(eta_grid).reshape(-1) se = np.sin(eta_grid).reshape(-1) co = np.cos(omega_grid).reshape(-1) so = np.sin(omega_grid).reshape(-1) sx = params.get("sx", default_value("sx")) sy = params.get("sy", default_value("sy")) sz = params.get("sz", default_value("sz")) e1 = params.get("e1", default_value("e1")) e2 = params.get("e2", default_value("e2")) x = sx * spow_np(ce, e1) * spow_np(co, e2) y = sy * spow_np(ce, e1) * spow_np(so, e2) z = sz * spow_np(se, e1) axis = np.clip(x / max(sx, 1e-4), -1.0, 1.0) y = y * np.clip(1.0 + params.get("taper_y", 0.0) * axis, 0.25, 1.9) z = z * np.clip(1.0 + params.get("taper_z", 0.0) * axis, 0.25, 1.9) bend = params.get("bend", 0.0) if abs(bend) > 1e-4: theta = bend * x bx = np.sin(theta) / bend - y * np.sin(theta) by = (1.0 - np.cos(theta)) / bend + y * np.cos(theta) x, y = bx, by pts = np.stack([x, y, z], axis=-1) return transform_np(pts, params) def sample_program(params: torch.Tensor, eta: torch.Tensor, omega: torch.Tensor) -> torch.Tensor: clouds = [primitive_surface(params[i], eta, omega) for i in range(params.shape[0])] return torch.cat(clouds, dim=0) def chamfer(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: dist = torch.cdist(a, b) return 0.5 * (torch.mean(torch.min(dist, dim=1).values) + torch.mean(torch.min(dist, dim=0).values)) def directed_rms(a: np.ndarray, b: np.ndarray) -> float: diff = a[:, None, :] - b[None, :, :] dist = np.sum(diff * diff, axis=-1) return float(np.sqrt(np.mean(np.min(dist, axis=1)))) def chamfer_np(a: np.ndarray, b: np.ndarray) -> float: return 0.5 * (directed_rms(a, b) + directed_rms(b, a)) def coverage_np(a: np.ndarray, b: np.ndarray, threshold: float = 0.065) -> float: diff = a[:, None, :] - b[None, :, :] dist = np.sum(diff * diff, axis=-1) return float(np.mean(np.min(dist, axis=1) <= threshold * threshold)) def overlap_proxy(prims: Sequence[Dict[str, float]]) -> float: if len(prims) < 2: return 0.0 total = 0.0 count = 0 for i, a in enumerate(prims): for b in prims[i + 1:]: pa = np.array([a["tx"], a["ty"], a["tz"]]) pb = np.array([b["tx"], b["ty"], b["tz"]]) distance = float(np.linalg.norm(pa - pb)) radius = 0.42 * (a["sx"] + b["sx"] + a["sy"] + b["sy"] + a["sz"] + b["sz"]) / 3.0 total += max(0.0, radius - distance) / max(radius, 1e-4) count += 1 return total / max(1, count) def component(kind: str, name: str, group: str, color_key: str, **kwargs: Any) -> Dict[str, Any]: out: Dict[str, Any] = {"type": kind, "name": name, "group": group, "color": COLORS[color_key]} out.update(kwargs) return out def bezier(points: np.ndarray, t: np.ndarray) -> np.ndarray: if len(points) == 2: return (1 - t)[:, None] * points[0] + t[:, None] * points[1] if len(points) == 3: return ( ((1 - t) ** 2)[:, None] * points[0] + (2 * (1 - t) * t)[:, None] * points[1] + (t ** 2)[:, None] * points[2] ) if len(points) == 4: return ( ((1 - t) ** 3)[:, None] * points[0] + (3 * ((1 - t) ** 2) * t)[:, None] * points[1] + (3 * (1 - t) * (t ** 2))[:, None] * points[2] + (t ** 3)[:, None] * points[3] ) raise ValueError("tube components use 2, 3, or 4 control points") def bezier_tangent(points: np.ndarray, t: np.ndarray) -> np.ndarray: if len(points) == 2: return np.repeat((points[1] - points[0])[None, :], len(t), axis=0) if len(points) == 3: return 2 * (1 - t)[:, None] * (points[1] - points[0]) + 2 * t[:, None] * (points[2] - points[1]) return ( 3 * ((1 - t) ** 2)[:, None] * (points[1] - points[0]) + 6 * ((1 - t) * t)[:, None] * (points[2] - points[1]) + 3 * (t ** 2)[:, None] * (points[3] - points[2]) ) def tube_points(item: Dict[str, Any], rows: int, cols: int) -> np.ndarray: controls = np.array(item["points"], dtype=np.float64) t = np.linspace(0, 1, rows) centers = bezier(controls, t) tangents = bezier_tangent(controls, t) tangents = tangents / np.clip(np.linalg.norm(tangents, axis=1, keepdims=True), 1e-8, None) up = np.tile(np.array([0.0, 0.0, 1.0]), (rows, 1)) close = np.abs(np.sum(tangents * up, axis=1)) > 0.92 up[close] = np.array([0.0, 1.0, 0.0]) normals = np.cross(tangents, up) normals = normals / np.clip(np.linalg.norm(normals, axis=1, keepdims=True), 1e-8, None) binormals = np.cross(normals, tangents) angles = np.linspace(0, 2 * math.pi, cols, endpoint=False) circle = np.stack([np.cos(angles), np.sin(angles)], axis=-1) radius = float(item.get("radius", 0.05)) pts = [] for i in range(rows): ring = centers[i] + radius * (circle[:, 0:1] * normals[i] + circle[:, 1:2] * binormals[i]) pts.append(ring) return np.concatenate(pts, axis=0) def cylinder_points(item: Dict[str, Any], rows: int, cols: int, caps: bool = True) -> np.ndarray: z = np.linspace(-0.5, 0.5, rows) * float(item.get("height", 1.0)) theta = np.linspace(0, 2 * math.pi, cols, endpoint=False) zz, tt = np.meshgrid(z, theta, indexing="ij") alpha = (zz / float(item.get("height", 1.0))) + 0.5 r0 = float(item.get("radiusBottom", item.get("radius", 0.3))) r1 = float(item.get("radiusTop", item.get("radius", 0.3))) radius = (1 - alpha) * r0 + alpha * r1 side = np.stack([radius * np.cos(tt), radius * np.sin(tt), zz], axis=-1).reshape(-1, 3) if not caps: return transform_np(side, item) disks = [] rings = max(3, rows // 2) radial = np.linspace(0, 1, rings) rr, aa = np.meshgrid(radial, theta, indexing="ij") for sign, radius_cap in [(-1, r0), (1, r1)]: cap = np.stack([ radius_cap * rr * np.cos(aa), radius_cap * rr * np.sin(aa), np.full_like(rr, sign * float(item.get("height", 1.0)) * 0.5), ], axis=-1).reshape(-1, 3) disks.append(cap) return transform_np(np.concatenate([side, *disks], axis=0), item) def torus_points(item: Dict[str, Any], rows: int, cols: int) -> np.ndarray: u = np.linspace(0, 2 * math.pi, rows, endpoint=False) v = np.linspace(0, 2 * math.pi, cols, endpoint=False) uu, vv = np.meshgrid(u, v, indexing="ij") major = float(item.get("majorRadius", 0.3)) minor = float(item.get("tubeRadius", 0.02)) pts = np.stack([ (major + minor * np.cos(vv)) * np.cos(uu), (major + minor * np.cos(vv)) * np.sin(uu), minor * np.sin(vv), ], axis=-1).reshape(-1, 3) return transform_np(pts, item) def component_points(item: Dict[str, Any], rows: int, cols: int) -> np.ndarray: kind = item["type"] if kind == "sq": return superquadric_np(item, rows, cols) if kind == "tube": return tube_points(item, max(rows * 2, 24), max(cols, 24)) if kind == "cylinder": return cylinder_points(item, rows, cols, caps=not item.get("openEnded", False)) if kind == "frustum": return cylinder_points(item, rows, cols, caps=not item.get("openEnded", True)) if kind == "disk": disk = dict(item) disk["height"] = item.get("height", 0.035) disk["radius"] = item.get("radius", item.get("radiusTop", 0.3)) return cylinder_points(disk, max(3, rows // 2), cols, caps=True) if kind == "torus": return torus_points(item, max(rows, 20), max(cols, 28)) raise ValueError(f"unknown component type: {kind}") def sample_components(items: Sequence[Dict[str, Any]], rows: int, cols: int) -> np.ndarray: clouds = [component_points(item, rows, cols) for item in items] return np.concatenate(clouds, axis=0).astype(np.float32) def add_noise(points: np.ndarray, sigma: float, seed: int) -> np.ndarray: if sigma <= 0: return points rng = np.random.default_rng(seed) return points + rng.normal(0, sigma, size=points.shape).astype(np.float32) def primitive(name: str, group: str, color_key: str, trainable: Sequence[str] = (), **kwargs: float) -> Primitive: values = {key: default_value(key) for key in PARAMS} values.update(kwargs) return Primitive(name=name, group=group, values=values, trainable=trainable, color=COLORS[color_key]) def rod_primitive( name: str, group: str, color_key: str, start: Sequence[float], end: Sequence[float], radius: float, trainable: Sequence[str] = (), ) -> Primitive: a = np.array(start, dtype=np.float64) b = np.array(end, dtype=np.float64) center = (a + b) * 0.5 direction = b - a length = float(np.linalg.norm(direction)) if length < 1e-6: raise ValueError(f"degenerate rod primitive: {name}") if abs(direction[1]) > 1e-6: raise ValueError("rod_primitive currently expects x-z plane segments") # Local superquadrics are elongated along x. A y-axis rotation aligns local x # with the segment direction in the x-z plane. ry = -math.atan2(float(direction[2]), float(direction[0])) return primitive( name, group, color_key, trainable, sx=max(0.04, length * 0.53), sy=radius, sz=radius, e1=0.58, e2=0.82, tx=float(center[0]), ty=float(center[1]), tz=float(center[2]), ry=ry, ) def curve_rods( prefix: str, group: str, color_key: str, controls: Sequence[Sequence[float]], segments: int, radius: float, trainable: Sequence[str] = (), labels: Sequence[str] | None = None, ) -> List[Primitive]: points = bezier(np.array(controls, dtype=np.float64), np.linspace(0, 1, segments + 1)) names = list(labels) if labels is not None else ["lower", "mid-lower", "middle", "mid-upper", "upper", "tip"] rods = [] for index in range(segments): label = names[index] if index < len(names) else f"part-{index + 1}" rods.append( rod_primitive( f"{prefix} {label}", group, color_key, points[index], points[index + 1], radius, trainable, ) ) return rods POSE_SCALE = ["tx", "ty", "tz", "rx", "ry", "rz", "sx", "sy", "sz"] POSE_SHAPE = POSE_SCALE + ["e1", "e2"] DEFORMED = POSE_SHAPE + ["bend", "taper_y", "taper_z"] CURVE_DEFORMED = DEFORMED STATIC: List[str] = [] LOCAL_PART: List[str] = [] LAMP_SHADE_BASE = (0.66, 0.32, 0.46) LAMP_SHADE_EDIT = (0.80, 0.16, 0.63) LAMP_OPENING_LOCAL = (0.19, -0.045) LAMP_BULB_LOCAL = (0.115, -0.085) LAMP_BULB_SIZE = (0.145, 0.145, 0.125) LAMP_BULB_TY = -0.015 def rotate_xz(local_x: float, local_z: float, ry: float) -> tuple[float, float]: return ( local_x * math.cos(ry) + local_z * math.sin(ry), -local_x * math.sin(ry) + local_z * math.cos(ry), ) def inverse_rotate_xz(world_x: float, world_z: float, ry: float) -> tuple[float, float]: return ( world_x * math.cos(ry) - world_z * math.sin(ry), world_x * math.sin(ry) + world_z * math.cos(ry), ) def lamp_attached_pose(local_x: float, local_z: float, edited: bool, normal_axis: bool = False) -> tuple[float, float, float]: tx, tz, ry = LAMP_SHADE_EDIT if edited else LAMP_SHADE_BASE dx, dz = rotate_xz(local_x, local_z, ry) pose_ry = ry + (math.pi / 2 if normal_axis else 0.0) return tx + dx, tz + dz, pose_ry def edit_lamp_attached_pose(p: Dict[str, float]) -> None: base_tx, base_tz, base_ry = LAMP_SHADE_BASE edit_tx, edit_tz, edit_ry = LAMP_SHADE_EDIT local_x, local_z = inverse_rotate_xz(p["tx"] - base_tx, p["tz"] - base_tz, base_ry) next_x, next_z = rotate_xz(local_x, local_z, edit_ry) p["tx"] = edit_tx + next_x p["tz"] = edit_tz + next_z p["ry"] = p.get("ry", 0.0) - base_ry + edit_ry def serialize_prims(prims: Sequence[Primitive]) -> List[Dict[str, float]]: out = [] for p in prims: values = {key: p.values.get(key, default_value(key)) for key in PARAMS} values.update({"name": p.name, "group": p.group, "color": p.color}) out.append(values) return out def tensor_to_prims(params: torch.Tensor, family: FitFamily) -> List[Dict[str, float]]: arr = params.detach().cpu().numpy() out = [] for row, primitive_def in zip(arr, family.primitives): values = {key: float(row[i]) for i, key in enumerate(PARAMS)} values.update({ "name": primitive_def.name, "group": primitive_def.group, "color": primitive_def.color, }) out.append(values) return out def prims_to_tensor(prims: Sequence[Dict[str, float]], device: torch.device) -> torch.Tensor: rows = [[prim.get(key, default_value(key)) for key in PARAMS] for prim in prims] return torch.tensor(rows, dtype=torch.float32, device=device) def surface_np(prims: Sequence[Dict[str, float]], eta: torch.Tensor, omega: torch.Tensor, device: torch.device) -> np.ndarray: with torch.no_grad(): params = prims_to_tensor(prims, device) pts = sample_program(params, eta, omega).detach().cpu().numpy() return pts def clamp_params_(params: torch.Tensor, masks: torch.Tensor) -> None: with torch.no_grad(): for idx, key in enumerate(PARAMS): lo, hi = BOUNDS[key] params[:, idx].clamp_(lo, hi) def fit_family( family: FitFamily, target_points: torch.Tensor, device: torch.device, seed: int, iters: int, rows: int, cols: int, ) -> Dict[str, Any]: rng = random.Random(seed) eta, omega = make_uv(rows, cols, device) params_list = [] mask_list = [] for primitive_def in family.primitives: param, mask = primitive_def.tensor(device, jitter=0.014, rng=rng) params_list.append(param) mask_list.append(mask) params = torch.nn.Parameter(torch.stack(params_list, dim=0)) masks = torch.stack(mask_list, dim=0) initial_params = params.detach().clone() param_scales = torch.tensor( [max(hi - lo, 1e-3) for lo, hi in (BOUNDS[key] for key in PARAMS)], dtype=torch.float32, device=device, ) optimizer = torch.optim.Adam([params], lr=0.009) history = [] for step in range(iters): optimizer.zero_grad() model = sample_program(params, eta, omega) data = chamfer(target_points, model) compact = 0.00035 * len(family.primitives) taper = torch.mean(torch.relu(torch.abs(params[:, PARAMS.index("taper_y")]) - 0.72)) taper = taper + torch.mean(torch.relu(torch.abs(params[:, PARAMS.index("taper_z")]) - 0.72)) bend = torch.mean(torch.relu(torch.abs(params[:, PARAMS.index("bend")]) - 1.95)) normalized_drift = ((params - initial_params) / param_scales) * masks design_prior = torch.mean(normalized_drift * normalized_drift) loss = data + compact + 0.0015 * taper + 0.001 * bend + 0.18 * design_prior loss.backward() with torch.no_grad(): params.grad *= masks optimizer.step() clamp_params_(params, masks) if step % 25 == 0 or step == iters - 1: history.append(float(data.detach().cpu())) prims = tensor_to_prims(params, family) return { "seed": seed, "family": family.key, "label": family.label, "primitives": prims, "history": history, "editNote": family.edit_note, } def evaluate_fit( obj: ObjectSpec, fit: Dict[str, Any], target_points_eval: np.ndarray, edited_target_points_eval: np.ndarray, eta_eval: torch.Tensor, omega_eval: torch.Tensor, device: torch.device, ) -> Dict[str, Any]: model = surface_np(fit["primitives"], eta_eval, omega_eval, device) edited_fit_prims = obj.edit(fit["primitives"], fit["family"]) edited_model = surface_np(edited_fit_prims, eta_eval, omega_eval, device) surface_cd = chamfer_np(target_points_eval, model) edit_cd = chamfer_np(edited_target_points_eval, edited_model) coverage = coverage_np(target_points_eval, model) overlap = overlap_proxy(fit["primitives"]) score = surface_cd + 0.012 * len(fit["primitives"]) + 0.02 * overlap edit_amplification = edit_cd / max(surface_cd, 1e-4) fit["metrics"] = { "surfaceChamfer": surface_cd, "editChamfer": edit_cd, "editAmplification": edit_amplification, "coverage": coverage, "overlap": overlap, "score": score, "primitiveCount": len(fit["primitives"]), } fit["editedPrimitives"] = edited_fit_prims return fit def component_bounds(component_sets: Sequence[np.ndarray], primitive_sets: Sequence[np.ndarray]) -> Dict[str, Any]: pts = np.concatenate([*component_sets, *primitive_sets], axis=0) lo = pts.min(axis=0) hi = pts.max(axis=0) center = (lo + hi) * 0.5 radius = max(float(np.max(hi - lo)) * 0.62, 0.45) return {"center": [float(v) for v in center], "radius": radius} def edit_mug(prims: List[Dict[str, float]], family: str) -> List[Dict[str, float]]: out = [dict(p) for p in prims] for p in out: if p["group"] == "mount": p["tx"] += 0.06 if family == "program" and p["group"] == "handle": if "middle" in p["name"]: p["tx"] += 0.16 p["sx"] *= 1.06 elif "mid-upper" in p["name"] or "mid-lower" in p["name"]: p["tx"] += 0.11 else: p["tx"] += 0.06 if family == "surface" and p["group"] == "surface_handle": if "middle" in p["name"]: p["tx"] += 0.22 p["sx"] *= 0.92 elif "mid-upper" in p["name"] or "upper" in p["name"]: p["tx"] += 0.04 p["tz"] += 0.12 p["ry"] += 0.42 elif "mid-lower" in p["name"] or "lower" in p["name"]: p["tx"] += 0.04 p["tz"] -= 0.12 p["ry"] -= 0.42 return out def edit_drawer(prims: List[Dict[str, float]], family: str) -> List[Dict[str, float]]: out = [dict(p) for p in prims] for p in out: if family == "program": if p["group"] == "pull_post": p["ty"] -= 0.08 p["sx"] *= 1.55 if p["group"] == "pull_bar": p["ty"] -= 0.18 p["sz"] *= 1.05 if family == "surface" and p["group"] == "surface_pull": if "middle" in p["name"]: p["ty"] -= 0.22 p["rz"] += 0.08 elif "left" in p["name"]: p["ty"] -= 0.10 p["rx"] += 0.18 elif "right" in p["name"]: p["ty"] -= 0.10 p["rx"] -= 0.18 else: p["ty"] -= 0.14 return out def edit_lamp(prims: List[Dict[str, float]], family: str) -> List[Dict[str, float]]: out = [dict(p) for p in prims] for p in out: if family == "program" and p["group"] == "neck": if "lower" in p["name"]: p["tx"] += 0.02 p["tz"] -= 0.06 p["ry"] -= 0.04 elif "middle" in p["name"]: p["tx"] += 0.08 p["tz"] -= 0.11 elif "upper" in p["name"] or "tip" in p["name"]: p["tx"] += 0.16 p["tz"] -= 0.16 p["ry"] += 0.10 if family == "surface" and p["group"] == "surface_neck": p["rx"] += 0.14 if "lower" in p["name"]: p["tx"] += 0.02 p["tz"] -= 0.08 p["ry"] -= 0.12 elif "middle" in p["name"]: p["tx"] += 0.08 p["tz"] -= 0.02 elif "upper" in p["name"] or "tip" in p["name"]: p["tx"] += 0.17 p["tz"] += 0.08 p["ry"] += 0.18 if p["group"] in {"shade", "bulb"}: edit_lamp_attached_pose(p) return out def edit_hook(prims: List[Dict[str, float]], family: str) -> List[Dict[str, float]]: out = [dict(p) for p in prims] for p in out: if p["group"] == "hook": if family == "program": p["bend"] += 0.42 p["sx"] *= 1.1 p["tx"] += 0.08 else: p["tx"] += 0.12 if "tip" in p["name"] else 0.04 p["tz"] += 0.16 if "tip" in p["name"] else -0.05 p["rz"] += 0.3 return out def edit_headphones(prims: List[Dict[str, float]], family: str) -> List[Dict[str, float]]: out = [dict(p) for p in prims] for p in out: if p["group"] in {"earcup", "cushion"}: p["tx"] *= 1.18 if p["group"] == "band": if family == "program": p["sx"] *= 1.18 p["bend"] *= 0.88 else: p["tx"] *= 1.12 p["tz"] += 0.1 if "crown" in p["name"] else -0.02 p["rz"] += 0.2 if p["tx"] > 0 else -0.2 return out def edit_faucet(prims: List[Dict[str, float]], family: str) -> List[Dict[str, float]]: out = [dict(p) for p in prims] for p in out: if p["group"] == "spout": if family == "program": p["tx"] += 0.18 p["bend"] *= 0.94 p["rz"] += 0.08 else: p["tx"] += 0.2 if "tip" in p["name"] or "nozzle" in p["name"] else 0.08 p["tz"] += -0.03 if "tip" in p["name"] or "nozzle" in p["name"] else 0.04 p["rz"] += 0.18 return out def clone_components(items: Sequence[Dict[str, Any]]) -> List[Dict[str, Any]]: return [json.loads(json.dumps(item)) for item in items] def mug_components(edited: bool = False) -> List[Dict[str, Any]]: mount_x = 0.36 if edited else 0.30 outer_x = 0.9 if edited else 0.74 items = [ component("sq", "straight-sided cup body", "body", "body", sx=0.39, sy=0.385, sz=0.535, e1=0.2, e2=0.92, tx=-0.16, tz=0.0), component("disk", "recessed dark interior", "opening", "shadow", radius=0.32, height=0.014, tx=-0.16, tz=0.545), component("torus", "thick rounded lip", "rim", "rim", majorRadius=0.39, tubeRadius=0.026, tx=-0.16, tz=0.548), component("torus", "raised foot ring", "body", "rim", majorRadius=0.255, tubeRadius=0.016, tx=-0.16, tz=-0.535), component("sq", "upper rounded handle pad", "mount", "handle", sx=0.135, sy=0.095, sz=0.105, e1=0.32, e2=0.58, tx=mount_x, ty=0.0, tz=0.31), component("sq", "lower rounded handle pad", "mount", "handle", sx=0.135, sy=0.095, sz=0.105, e1=0.32, e2=0.58, tx=mount_x, ty=0.0, tz=-0.31), ] handle_points = [ [mount_x, 0.0, 0.31], [outer_x, 0.0, 0.39], [outer_x, 0.0, -0.39], [mount_x, 0.0, -0.31], ] items.append(component("tube", "continuous handle loop", "handle", "handle", points=handle_points, radius=0.062)) return items def lamp_components(edited: bool = False) -> List[Dict[str, Any]]: neck_points = ( [[-0.35, 0.0, -0.15], [-0.33, 0.0, 0.12], [-0.04, 0.0, 0.42], [0.46, 0.0, 0.36]] if not edited else [[-0.35, 0.0, -0.15], [-0.29, 0.0, 0.04], [0.08, 0.0, 0.26], [0.62, 0.0, 0.20]] ) shade_tx, shade_tz, shade_ry = LAMP_SHADE_EDIT if edited else LAMP_SHADE_BASE opening_tx, opening_tz, opening_ry = lamp_attached_pose(*LAMP_OPENING_LOCAL, edited, normal_axis=True) bulb_tx, bulb_tz, bulb_ry = lamp_attached_pose(*LAMP_BULB_LOCAL, edited) return [ component("sq", "weighted circular base", "base", "base", sx=0.44, sy=0.44, sz=0.055, e1=0.34, e2=0.92, tz=-0.64), component("sq", "raised base collar", "base", "metal", sx=0.16, sy=0.16, sz=0.055, e1=0.34, e2=0.86, tx=-0.35, tz=-0.58), component("sq", "short upright stem", "bar", "bar", sx=0.066, sy=0.066, sz=0.23, e1=0.48, e2=0.84, tx=-0.35, tz=-0.39), component("sq", "neck root joint", "neck", "metal", sx=0.11, sy=0.09, sz=0.095, e1=0.45, e2=0.72, tx=-0.35, tz=-0.15), component("tube", "continuous gooseneck", "neck", "neck", points=neck_points, radius=0.054), component("sq", "rounded lamp shade", "shade", "shade", sx=0.34, sy=0.28, sz=0.18, e1=0.34, e2=0.62, taper_y=-0.18, taper_z=-0.08, tx=shade_tx, ty=0.0, tz=shade_tz, ry=shade_ry), component("disk", "inner shade shadow", "shade", "shadow", radius=0.115, height=0.014, tx=opening_tx, tz=opening_tz, ry=opening_ry), component("sq", "warm bulb inside shade", "bulb", "bulb", sx=LAMP_BULB_SIZE[0], sy=LAMP_BULB_SIZE[1], sz=LAMP_BULB_SIZE[2], e1=0.88, e2=0.92, tx=bulb_tx, ty=LAMP_BULB_TY, tz=bulb_tz, ry=bulb_ry), ] def hook_components(edited: bool = False) -> List[Dict[str, Any]]: tip_x = 0.62 if edited else 0.44 tip_z = 0.08 if edited else -0.16 curve = [ [-0.28, -0.1, 0.18], [0.12, -0.19, 0.08], [0.55, -0.18, -0.42], [tip_x, -0.12, tip_z], ] return [ component("sq", "rounded wall plate", "base", "base", sx=0.38, sy=0.055, sz=0.62, e1=0.22, e2=0.28, tx=-0.36, ty=0.0, tz=0.0), component("disk", "upper screw head", "base", "shadow", radius=0.06, height=0.022, tx=-0.36, ty=-0.065, tz=0.28, rx=1.57), component("disk", "lower screw head", "base", "shadow", radius=0.06, height=0.022, tx=-0.36, ty=-0.065, tz=-0.28, rx=1.57), component("sq", "rounded hook boss", "hook", "hook", sx=0.14, sy=0.12, sz=0.13, e1=0.34, e2=0.6, tx=-0.23, ty=-0.08, tz=0.14), component("tube", "single curved hook", "hook", "hook", points=curve, radius=0.062), ] def headphones_components(edited: bool = False) -> List[Dict[str, Any]]: spread = 1.18 if edited else 1.0 left_x = -0.6 * spread right_x = 0.6 * spread band = [ [left_x, -0.01, 0.02], [-0.48 * spread, -0.02, 0.76], [0.48 * spread, -0.02, 0.76], [right_x, -0.01, 0.02], ] return [ component("sq", "left outer earcup", "earcup", "earcup", sx=0.26, sy=0.12, sz=0.34, e1=0.24, e2=0.58, tx=left_x, ty=0.0, tz=-0.25, rz=0.06), component("sq", "right outer earcup", "earcup", "earcup", sx=0.26, sy=0.12, sz=0.34, e1=0.24, e2=0.58, tx=right_x, ty=0.0, tz=-0.25, rz=-0.06), component("sq", "left soft cushion", "cushion", "cushion", sx=0.19, sy=0.052, sz=0.27, e1=0.32, e2=0.72, tx=left_x, ty=-0.12, tz=-0.25, rz=0.06), component("sq", "right soft cushion", "cushion", "cushion", sx=0.19, sy=0.052, sz=0.27, e1=0.32, e2=0.72, tx=right_x, ty=-0.12, tz=-0.25, rz=-0.06), component("tube", "left metal yoke", "band", "metal", points=[[left_x, -0.06, -0.03], [left_x * 0.98, -0.04, 0.24]], radius=0.036), component("tube", "right metal yoke", "band", "metal", points=[[right_x, -0.06, -0.03], [right_x * 0.98, -0.04, 0.24]], radius=0.036), component("tube", "continuous headband", "band", "band", points=band, radius=0.052), component("tube", "inner headband cushion", "cushion", "cushion", points=[[-0.42 * spread, -0.065, 0.53], [-0.14 * spread, -0.065, 0.61], [0.14 * spread, -0.065, 0.61], [0.42 * spread, -0.065, 0.53]], radius=0.034), ] def faucet_components(edited: bool = False) -> List[Dict[str, Any]]: reach = 0.2 if edited else 0.0 spout_points = [ [-0.22, -0.02, -0.04], [-0.21, -0.02, 0.56], [0.32 + reach, -0.02, 0.7], [0.54 + reach, -0.02, 0.3], ] return [ component("sq", "sink deck plate", "base", "face", sx=0.78, sy=0.34, sz=0.055, e1=0.2, e2=0.3, tx=0.02, ty=0.0, tz=-0.48), component("cylinder", "round base collar", "base", "metal", radius=0.19, height=0.12, tx=-0.22, tz=-0.38), component("cylinder", "vertical faucet post", "base", "metal", radius=0.09, height=0.72, tx=-0.22, tz=-0.02), component("tube", "continuous arched spout", "spout", "bar", points=spout_points, radius=0.054), component("cylinder", "downward nozzle", "spout", "shadow", radius=0.06, height=0.2, tx=0.54 + reach, ty=-0.02, tz=0.19), component("sq", "left handle knob", "base", "mount", sx=0.16, sy=0.08, sz=0.055, e1=0.32, e2=0.58, tx=-0.56, ty=-0.02, tz=-0.31, rz=0.24), component("sq", "right handle knob", "base", "mount", sx=0.16, sy=0.08, sz=0.055, e1=0.32, e2=0.58, tx=0.1, ty=-0.02, tz=-0.31, rz=-0.24), ] def drawer_components(edited: bool = False) -> List[Dict[str, Any]]: pull_y = -0.28 if edited else -0.12 bar_points = [ [-0.5, pull_y, 0.07], [-0.34, pull_y, 0.24], [0.34, pull_y, 0.24], [0.5, pull_y, 0.07], ] return [ component("sq", "flat drawer front", "base", "face", sx=0.94, sy=0.055, sz=0.52, e1=0.2, e2=0.24, tx=0.0, ty=0.03, tz=0.0), component("disk", "left mounting rosette", "mount", "mount", radius=0.105, height=0.035, tx=-0.5, ty=-0.04, tz=0.07, rx=1.57), component("disk", "right mounting rosette", "mount", "mount", radius=0.105, height=0.035, tx=0.5, ty=-0.04, tz=0.07, rx=1.57), component("cylinder", "left handle post", "pull", "metal", radius=0.055, height=abs(pull_y) + 0.05, tx=-0.5, ty=(pull_y - 0.04) * 0.5, tz=0.07, rx=1.57), component("cylinder", "right handle post", "pull", "metal", radius=0.055, height=abs(pull_y) + 0.05, tx=0.5, ty=(pull_y - 0.04) * 0.5, tz=0.07, rx=1.57), component("tube", "continuous pull handle", "pull", "bar", points=bar_points, radius=0.055), ] def object_specs() -> List[ObjectSpec]: mug_target = mug_components(False) mug_edit = mug_components(True) mug_handle_controls = [[0.30, 0.0, 0.31], [0.74, 0.0, 0.39], [0.74, 0.0, -0.39], [0.30, 0.0, -0.31]] mug_shared_prims = [ primitive("cup body", "body", "body", DEFORMED, sx=0.39, sy=0.385, sz=0.535, e1=0.2, e2=0.92, tx=-0.16, tz=0.0), primitive("recessed interior", "opening", "shadow", STATIC, sx=0.32, sy=0.32, sz=0.014, e1=0.36, e2=0.9, tx=-0.16, tz=0.545), primitive("top ceramic rim", "rim", "rim", STATIC, sx=0.39, sy=0.39, sz=0.026, e1=0.34, e2=0.9, tx=-0.16, tz=0.548), primitive("foot ring", "rim", "rim", STATIC, sx=0.255, sy=0.255, sz=0.016, e1=0.32, e2=0.86, tx=-0.16, tz=-0.535), primitive("upper handle pad", "mount", "handle", STATIC, sx=0.135, sy=0.095, sz=0.105, e1=0.34, e2=0.58, tx=0.30, tz=0.31), primitive("lower handle pad", "mount", "handle", STATIC, sx=0.135, sy=0.095, sz=0.105, e1=0.34, e2=0.58, tx=0.30, tz=-0.31), ] mug_program_prims = mug_shared_prims + [ *curve_rods( "handle semantic segment", "handle", "handle", mug_handle_controls, 5, 0.058, LOCAL_PART, labels=["upper", "mid-upper", "middle", "mid-lower", "lower"], ) ] mug_surface_prims = mug_shared_prims + curve_rods( "handle local part", "surface_handle", "surface", mug_handle_controls, 5, 0.058, LOCAL_PART, labels=["upper", "mid-upper", "middle", "mid-lower", "lower"], ) mug_program = FitFamily("program", "Editable SQ program", mug_program_prims, "The cup and linked handle segments are represented as two semantic controls: container and handle.") mug_surface = FitFamily("surface", "Surface-matched SQ program", mug_surface_prims, "The original handle surface is reconstructed by local SQ pieces rather than by one editable loop.") lamp_target = lamp_components(False) lamp_edit = lamp_components(True) lamp_neck_controls = [[-0.35, 0.0, -0.15], [-0.33, 0.0, 0.12], [-0.04, 0.0, 0.42], [0.46, 0.0, 0.36]] lamp_open_tx, lamp_open_tz, lamp_open_ry = lamp_attached_pose(*LAMP_OPENING_LOCAL, False, normal_axis=True) lamp_bulb_tx, lamp_bulb_tz, lamp_bulb_ry = lamp_attached_pose(*LAMP_BULB_LOCAL, False) lamp_shared_prims = [ primitive("base disk", "base", "base", STATIC, sx=0.44, sy=0.44, sz=0.055, e1=0.34, e2=0.92, tz=-0.64), primitive("base collar", "base", "metal", STATIC, sx=0.16, sy=0.16, sz=0.055, e1=0.34, e2=0.86, tx=-0.35, tz=-0.58), primitive("upright stem", "bar", "bar", STATIC, sx=0.066, sy=0.066, sz=0.23, e1=0.48, e2=0.84, tx=-0.35, tz=-0.39), primitive("neck root joint", "mount", "metal", STATIC, sx=0.11, sy=0.09, sz=0.095, e1=0.45, e2=0.72, tx=-0.35, tz=-0.15), primitive("shade shell", "shade", "shade", DEFORMED, sx=0.34, sy=0.28, sz=0.18, e1=0.34, e2=0.62, taper_y=-0.18, taper_z=-0.08, tx=0.66, tz=0.32, ry=0.46), primitive("shade opening", "shade", "shadow", STATIC, sx=0.115, sy=0.115, sz=0.014, e1=0.36, e2=0.86, tx=lamp_open_tx, tz=lamp_open_tz, ry=lamp_open_ry), primitive("bulb", "bulb", "bulb", STATIC, sx=LAMP_BULB_SIZE[0], sy=LAMP_BULB_SIZE[1], sz=LAMP_BULB_SIZE[2], e1=0.88, e2=0.92, tx=lamp_bulb_tx, ty=LAMP_BULB_TY, tz=lamp_bulb_tz, ry=lamp_bulb_ry), ] lamp_program_prims = lamp_shared_prims[:4] + [ *curve_rods( "neck semantic segment", "neck", "neck", lamp_neck_controls, 5, 0.054, LOCAL_PART, labels=["lower", "mid-lower", "middle", "mid-upper", "upper"], ), *lamp_shared_prims[4:], ] lamp_surface_prims = lamp_shared_prims[:4] + curve_rods("neck local part", "surface_neck", "surface", lamp_neck_controls, 5, 0.054, LOCAL_PART) + lamp_shared_prims[4:] lamp_program = FitFamily("program", "Editable SQ program", lamp_program_prims, "The gooseneck is represented as a deformable control that carries the shade and visible bulb.") lamp_surface = FitFamily("surface", "Surface-matched SQ program", lamp_surface_prims, "The neck surface is reconstructed by local SQ pieces, so the edit has no single neck control.") drawer_target = drawer_components(False) drawer_edit = drawer_components(True) drawer_pull_controls = [[-0.5, -0.12, 0.07], [-0.34, -0.12, 0.24], [0.34, -0.12, 0.24], [0.5, -0.12, 0.07]] drawer_shared_prims = [ primitive("flat drawer front", "base", "face", DEFORMED, sx=0.94, sy=0.055, sz=0.52, e1=0.2, e2=0.24, tx=0.0, ty=0.03, tz=0.0), primitive("left mounting rosette", "mount", "mount", STATIC, sx=0.105, sy=0.105, sz=0.025, e1=0.32, e2=0.72, tx=-0.5, ty=-0.04, tz=0.07, rx=math.pi / 2), primitive("right mounting rosette", "mount", "mount", STATIC, sx=0.105, sy=0.105, sz=0.025, e1=0.32, e2=0.72, tx=0.5, ty=-0.04, tz=0.07, rx=math.pi / 2), ] drawer_program_prims = drawer_shared_prims + [ primitive("left pull post", "pull_post", "metal", STATIC, sx=0.085, sy=0.055, sz=0.055, e1=0.5, e2=0.78, tx=-0.5, ty=-0.08, tz=0.07, rz=math.pi / 2), primitive("right pull post", "pull_post", "metal", STATIC, sx=0.085, sy=0.055, sz=0.055, e1=0.5, e2=0.78, tx=0.5, ty=-0.08, tz=0.07, rz=math.pi / 2), *curve_rods( "pull semantic segment", "pull_bar", "bar", drawer_pull_controls, 4, 0.055, LOCAL_PART, labels=["left", "mid-left", "mid-right", "right"], ), ] drawer_surface_prims = drawer_shared_prims + curve_rods( "pull local part", "surface_pull", "surface", drawer_pull_controls, 6, 0.055, LOCAL_PART, labels=["left post", "left bend", "middle left", "middle right", "right bend", "right post"], ) drawer_program = FitFamily("program", "Editable SQ program", drawer_program_prims, "The posts and pull bar are exposed as one handle control.") drawer_surface = FitFamily("surface", "Surface-matched SQ program", drawer_surface_prims, "The original pull is reconstructed by local SQ pieces whose boundaries are not the pull control.") return [ ObjectSpec( "mug", "Handled Mug", "An open cup with a readable side handle.", "Widen the handle loop", "The semantic program has a container control and a handle control. The surface-matched program can trace the same loop with local SQ parts, but widening the loop exposes that those local pieces are not one coherent handle control.", mug_target, mug_edit, mug_program, mug_surface, edit_mug, ), ObjectSpec( "lamp", "Desk Lamp", "A base, stem, flexible neck, rounded shade, and visible bulb.", "Bend the neck downward", "The visible geometry can be matched by short local SQs along the gooseneck. The edit asks for a different structure: one flexible neck control that bends and carries the shade rather than letting local pieces drift.", lamp_target, lamp_edit, lamp_program, lamp_surface, edit_lamp, ), ObjectSpec( "drawer", "Drawer Pull", "A drawer face with two mounts and one continuous pull handle.", "Pull the handle outward", "The object is visually simple but structurally unforgiving: the handle should move as a single pull, while the drawer face stays fixed. A surface-matched program can cover the same curve with local SQs, but the edit exposes whether those pieces form one control.", drawer_target, drawer_edit, drawer_program, drawer_surface, edit_drawer, ), ] def choose_best(fits: Sequence[Dict[str, Any]]) -> Dict[str, Any]: return min(fits, key=lambda fit: fit["metrics"]["score"]) def run(args: argparse.Namespace) -> None: torch.manual_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) if args.device == "auto": device = torch.device("mps" if torch.backends.mps.is_available() else "cpu") else: device = torch.device(args.device) eta_eval, omega_eval = make_uv(args.eval_rows, args.eval_cols, device) objects = [] for obj_index, obj in enumerate(object_specs()): print(f"[{obj.key}] sampling target components") target_eval = sample_components(obj.target_components, args.eval_rows, args.eval_cols) edited_eval = sample_components(obj.edited_components, args.eval_rows, args.eval_cols) target_train = sample_components(obj.target_components, args.target_rows, args.target_cols) target_train = add_noise(target_train, args.noise, args.seed + obj_index * 31) target_points = torch.tensor(target_train, dtype=torch.float32, device=device) fits = [] for family in [obj.program, obj.surface]: family_fits = [] for seed_i in range(args.seeds): family_offset = {"program": 0, "surface": 500}[family.key] seed = args.seed + obj_index * 1000 + seed_i * 37 + family_offset print(f" fitting {family.key} seed {seed_i + 1}/{args.seeds}") fit = fit_family(family, target_points, device, seed, args.iters, args.rows, args.cols) fit = evaluate_fit(obj, fit, target_eval, edited_eval, eta_eval, omega_eval, device) family_fits.append(fit) if device.type == "mps" and hasattr(torch, "mps"): torch.mps.empty_cache() fits.append(choose_best(family_fits)) program_fit = next(f for f in fits if f["family"] == "program") surface_fit = next(f for f in fits if f["family"] == "surface") delta_surface = surface_fit["metrics"]["surfaceChamfer"] - program_fit["metrics"]["surfaceChamfer"] delta_edit = surface_fit["metrics"]["editChamfer"] - program_fit["metrics"]["editChamfer"] amp_ratio = surface_fit["metrics"]["editAmplification"] / max(program_fit["metrics"]["editAmplification"], 1e-6) primitive_clouds = [] for fit in fits: primitive_clouds.append(surface_np(fit["primitives"], eta_eval, omega_eval, device)) primitive_clouds.append(surface_np(fit["editedPrimitives"], eta_eval, omega_eval, device)) objects.append({ "key": obj.key, "title": obj.title, "subtitle": obj.subtitle, "story": obj.story, "editLabel": obj.edit_label, "targetComponents": obj.target_components, "editedTargetComponents": obj.edited_components, "target": serialize_prims(obj.program.primitives), "editedTarget": obj.edit(serialize_prims(obj.program.primitives), "program"), "displayBounds": component_bounds([target_eval, edited_eval], primitive_clouds), "fits": fits, "summary": { "surfaceGapSurfaceMinusProgram": delta_surface, "editGapSurfaceMinusProgram": delta_edit, "programEditabilityGap": delta_edit - delta_surface, "surfaceAmpRatio": amp_ratio, "doppelgangerIndex": delta_edit / max(abs(delta_surface), 0.002), }, }) summary = { "objectCount": len(objects), "seedCount": args.seeds, "iterationsPerFit": args.iters, "trainSamplesPerComponent": args.target_rows * args.target_cols, "trainSamplesPerPrimitive": args.rows * args.cols, "evalSamplesPerComponent": args.eval_rows * args.eval_cols, "evalSamplesPerPrimitive": args.eval_rows * args.eval_cols, } for family in ["program", "surface"]: family_fits = [next(f for f in obj["fits"] if f["family"] == family) for obj in objects] summary[f"mean{family.title()}SurfaceChamfer"] = float(np.mean([f["metrics"]["surfaceChamfer"] for f in family_fits])) summary[f"mean{family.title()}EditChamfer"] = float(np.mean([f["metrics"]["editChamfer"] for f in family_fits])) summary[f"meanEditAmplification{family.title()}"] = float(np.mean([f["metrics"]["editAmplification"] for f in family_fits])) summary[f"mean{family.title()}Coverage"] = float(np.mean([f["metrics"]["coverage"] for f in family_fits])) summary["meanProgramEditabilityGap"] = summary["meanSurfaceEditChamfer"] - summary["meanProgramEditChamfer"] summary["meanSurfaceFitGap"] = summary["meanSurfaceSurfaceChamfer"] - summary["meanProgramSurfaceChamfer"] summary["meanAmpRatio"] = summary["meanEditAmplificationSurface"] / max(summary["meanEditAmplificationProgram"], 1e-6) payload = { "schema": "sq-doppelganger-v2", "generatedBy": "experiments/sq_doppelganger.py", "device": str(device), "summary": summary, "objects": objects, } out = Path(args.out) out.parent.mkdir(parents=True, exist_ok=True) serialized = json.dumps(payload, indent=2) out.write_text(serialized, encoding="utf-8") js_out = out.with_suffix(".js") js_out.write_text(f"window.__DOPPEL_DATA__ = {serialized};\n", encoding="utf-8") print(f"wrote {out}") print(f"wrote {js_out}") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--iters", type=int, default=900) parser.add_argument("--seeds", type=int, default=6) parser.add_argument("--rows", type=int, default=10, help="UV rows per fitted primitive during optimization.") parser.add_argument("--cols", type=int, default=22, help="UV columns per fitted primitive during optimization.") parser.add_argument("--target-rows", type=int, default=14, help="Sampling rows per target mesh component during optimization.") parser.add_argument("--target-cols", type=int, default=26, help="Sampling columns per target mesh component during optimization.") parser.add_argument("--eval-rows", type=int, default=24) parser.add_argument("--eval-cols", type=int, default=48) parser.add_argument("--noise", type=float, default=0.004) parser.add_argument("--seed", type=int, default=7) parser.add_argument("--device", choices=["auto", "cpu", "mps"], default="auto") parser.add_argument("--out", default="data/doppelganger_results.json") return parser.parse_args() if __name__ == "__main__": run(parse_args())