model.py

# Cell 2
# === Capacity Head ============================================================

class CapacityHead(nn.Module):
    def __init__(self, in_dim, feat_dim, init_capacity=1.0):
        super().__init__()
        self._raw_capacity = nn.Parameter(torch.tensor(math.log(math.exp(init_capacity) - 1)))
        # GELU for cascade: smooth gradients needed for overflow propagation
        self.evidence_net = nn.Sequential(
            nn.Linear(in_dim, feat_dim), nn.GELU(), nn.Linear(feat_dim, 1))
        self.feature_net = nn.Sequential(
            nn.Linear(in_dim, feat_dim), nn.GELU(), nn.Linear(feat_dim, feat_dim))
        self.retain_gate = nn.Sequential(
            nn.Linear(feat_dim + 1, feat_dim), nn.Sigmoid())
        self.overflow_gate = nn.Sequential(
            nn.Linear(feat_dim + 1, feat_dim), nn.Sigmoid())

    @property
    def capacity(self):
        return F.softplus(self._raw_capacity)

    def forward(self, x):
        cap = self.capacity
        raw_ev = F.relu(self.evidence_net(x))
        fill = torch.clamp(raw_ev / (cap + 1e-8), max=1.0)
        sat = torch.clamp((raw_ev - cap) / (cap + 1e-8), min=0.0)
        feat = self.feature_net(x)
        retained = self.retain_gate(torch.cat([feat, fill], -1)) * feat * fill
        overflow = self.overflow_gate(torch.cat([feat, sat], -1)) * feat * torch.clamp(sat, max=1.0)
        return fill, overflow, retained, cap, raw_ev


# === Differentiation Gate =====================================================

class DifferentiationGate(nn.Module):
    """
    Curvature direction analysis via occupancy field differentiation.

    Computes gradient and Laplacian of the 3D occupancy field to determine:
    - Curvature direction: convex (normals point outward) vs concave (inward)
    - Curvature alternation: where sign flips (saddle points, torus inner/outer)
    - Perturbation robustness: smoothed gradient features survive noise

    The key insight: a hemisphere and bowl occupy nearly identical voxels,
    but their occupancy gradients point in opposite directions relative
    to the center of mass. The Laplacian's sign distinguishes them.

    Outputs gate signals that modulate curvature features:
    - direction_gate: learned weighting based on gradient analysis
    - alternation_score: how much curvature sign varies spatially
    - directional_features: rich features encoding curvature orientation
    """

    def __init__(self, embed_dim=64):
        super().__init__()

        # Fixed 3D differentiation kernels — fused into single conv
        # 4 output channels: [grad_x, grad_y, grad_z, laplacian]
        diff_kernels = torch.zeros(4, 1, 3, 3, 3)
        # Sobel X
        diff_kernels[0, 0, 0, 1, 1] = -1; diff_kernels[0, 0, 2, 1, 1] = 1
        # Sobel Y
        diff_kernels[1, 0, 1, 0, 1] = -1; diff_kernels[1, 0, 1, 2, 1] = 1
        # Sobel Z
        diff_kernels[2, 0, 1, 1, 0] = -1; diff_kernels[2, 0, 1, 1, 2] = 1
        # Laplacian
        diff_kernels[3, 0, 1, 1, 1] = -6
        diff_kernels[3, 0, 0, 1, 1] = 1; diff_kernels[3, 0, 2, 1, 1] = 1
        diff_kernels[3, 0, 1, 0, 1] = 1; diff_kernels[3, 0, 1, 2, 1] = 1
        diff_kernels[3, 0, 1, 1, 0] = 1; diff_kernels[3, 0, 1, 1, 2] = 1
        self.register_buffer("diff_kernels", diff_kernels)

        # Precompute coordinate grid
        coords = torch.stack(torch.meshgrid(
            torch.arange(GS, dtype=torch.float32),
            torch.arange(GS, dtype=torch.float32),
            torch.arange(GS, dtype=torch.float32),
            indexing="ij"), dim=-1)  # (5,5,5,3)
        self.register_buffer("coords", coords)

        # Process gradient-derived features
        # Per-voxel: gradient direction, Laplacian sign, centroid-relative direction
        # Summarized as histograms and statistics

        # Gradient direction relative to centroid: 3 histogram bins per axis
        # + Laplacian sign distribution: 3 values (frac_pos, frac_neg, frac_zero)
        # + Alternation score: 1 value
        # + Per-axis gradient asymmetry: 3 values
        # + Radial gradient profile: 5 bins
        raw_feat_dim = 3 + 3 + 1 + 3 + 5  # = 15
        # Plus the 3D conv on the Laplacian field preserving spatial structure
        self.lap_conv = nn.Sequential(
            nn.Conv3d(1, 16, 3, padding=1), nn.GELU(),
            nn.Conv3d(16, 16, 3, padding=1), nn.GELU(),
            nn.AdaptiveAvgPool3d(2))  # -> (B, 16, 2, 2, 2) = 128
        lap_conv_dim = 16 * 8  # 128

        # Gradient magnitude 3D conv (encodes where boundaries are + direction)
        self.grad_conv = nn.Sequential(
            nn.Conv3d(3, 16, 3, padding=1), nn.GELU(),  # 3-channel: dx, dy, dz
            nn.Conv3d(16, 16, 3, padding=1), nn.GELU(),
            nn.AdaptiveAvgPool3d(2))  # -> (B, 16, 2, 2, 2) = 128
        grad_conv_dim = 16 * 8  # 128

        total_feat_dim = raw_feat_dim + lap_conv_dim + grad_conv_dim  # 15 + 128 + 128 = 271

        # Direction gate: SwiGLU for sharp convex/concave gating
        self.direction_net = nn.Sequential(
            SwiGLU(total_feat_dim, embed_dim),
            nn.Linear(embed_dim, embed_dim), nn.Sigmoid())

        # Directional features: SwiGLU for crisp direction encoding
        self.direction_feat_net = nn.Sequential(
            SwiGLU(total_feat_dim, embed_dim),
            nn.Linear(embed_dim, embed_dim))

    def forward(self, grid):
        """
        grid: (B, 5, 5, 5) binary occupancy

        Returns:
            direction_gate: (B, embed_dim) sigmoid gate for curvature features
            direction_feat: (B, embed_dim) additive directional features
            alternation_score: (B, 1) how much curvature alternates
        """
        B = grid.shape[0]
        device = grid.device
        vox = grid.unsqueeze(1)  # (B, 1, 5, 5, 5)

        # === Smooth occupancy before differentiation ===
        # Binary voxels produce spike gradients. Light blur creates
        # a continuous field whose derivatives are geometrically meaningful.
        vox_smooth = F.avg_pool3d(
            F.pad(vox, (1,1,1,1,1,1), mode='replicate'),
            kernel_size=3, stride=1, padding=0)  # (B, 1, 5, 5, 5)

        # === Compute gradients + Laplacian in single fused conv ===
        diff = F.conv3d(vox_smooth, self.diff_kernels, padding=1)  # (B, 4, 5, 5, 5)
        grad_field = diff[:, :3]  # (B, 3, 5, 5, 5) — gx, gy, gz
        gx, gy, gz = diff[:, 0:1], diff[:, 1:2], diff[:, 2:3]
        lap = diff[:, 3:4]  # (B, 1, 5, 5, 5)

        # === Centroid ===
        flat_grid = grid.reshape(B, -1)  # (B, 125)
        flat_coords = self.coords.reshape(-1, 3)  # (125, 3)
        total_occ = flat_grid.sum(dim=-1, keepdim=True).clamp(min=1)  # (B, 1)
        centroids = (flat_grid.unsqueeze(-1) * flat_coords.unsqueeze(0)).sum(dim=1) / total_occ  # (B, 3)

        # === Gradient direction relative to centroid ===
        grad_flat = grad_field.reshape(B, 3, -1).permute(0, 2, 1)  # (B, 125, 3)
        diff_from_center = flat_coords.unsqueeze(0) - centroids.unsqueeze(1)  # (B, 125, 3)
        diff_norm = diff_from_center / (diff_from_center.norm(dim=-1, keepdim=True) + 1e-8)
        dot_products = (grad_flat * diff_norm).sum(dim=-1)  # (B, 125)
        grad_mag = grad_flat.norm(dim=-1)  # (B, 125)
        active = (flat_grid > 0.5) & (grad_mag > 0.01)

        # Histogram of dot product signs (convex/concave/neutral fractions)
        n_active = active.float().sum(-1).clamp(min=1)
        frac_outward = ((dot_products > 0.1) & active).float().sum(-1) / n_active
        frac_inward = ((dot_products < -0.1) & active).float().sum(-1) / n_active
        frac_neutral = 1.0 - frac_outward - frac_inward
        direction_hist = torch.stack([frac_outward, frac_inward, frac_neutral], dim=-1)  # (B, 3)

        # === Laplacian sign distribution (active voxels only) ===
        lap_flat = lap.reshape(B, -1)  # (B, 125)
        lap_active = flat_grid > 0.5
        n_lap_active = lap_active.float().sum(-1).clamp(min=1)
        frac_pos_lap = ((lap_flat > 0.1) & lap_active).float().sum(-1) / n_lap_active
        frac_neg_lap = ((lap_flat < -0.1) & lap_active).float().sum(-1) / n_lap_active
        frac_zero_lap = 1.0 - frac_pos_lap - frac_neg_lap
        lap_hist = torch.stack([frac_pos_lap, frac_neg_lap, frac_zero_lap], dim=-1)  # (B, 3)

        # === Alternation score (ACTIVE VOXELS ONLY) ===
        # Only count sign flips between neighbor pairs where BOTH voxels are
        # near occupied regions. Otherwise empty space dilutes the signal.
        lap_3d = lap.squeeze(1)  # (B, 5, 5, 5)
        # Boundary mask: dilate occupancy by 1 to include immediate neighbors
        boundary_mask = F.max_pool3d(vox, kernel_size=3, stride=1, padding=1).squeeze(1)  # (B,5,5,5)

        # X-axis: both neighbors must be in boundary region
        bm_x = boundary_mask[:, 1:, :, :] * boundary_mask[:, :-1, :, :]  # (B,4,5,5)
        flip_x = (torch.sign(lap_3d[:, 1:, :, :]) * torch.sign(lap_3d[:, :-1, :, :]) < 0).float()
        active_flips_x = (flip_x * bm_x).sum(dim=(1, 2, 3))
        active_pairs_x = bm_x.sum(dim=(1, 2, 3)).clamp(min=1)

        bm_y = boundary_mask[:, :, 1:, :] * boundary_mask[:, :, :-1, :]
        flip_y = (torch.sign(lap_3d[:, :, 1:, :]) * torch.sign(lap_3d[:, :, :-1, :]) < 0).float()
        active_flips_y = (flip_y * bm_y).sum(dim=(1, 2, 3))
        active_pairs_y = bm_y.sum(dim=(1, 2, 3)).clamp(min=1)

        bm_z = boundary_mask[:, :, :, 1:] * boundary_mask[:, :, :, :-1]
        flip_z = (torch.sign(lap_3d[:, :, :, 1:]) * torch.sign(lap_3d[:, :, :, :-1]) < 0).float()
        active_flips_z = (flip_z * bm_z).sum(dim=(1, 2, 3))
        active_pairs_z = bm_z.sum(dim=(1, 2, 3)).clamp(min=1)

        alternation = ((active_flips_x / active_pairs_x +
                         active_flips_y / active_pairs_y +
                         active_flips_z / active_pairs_z) / 3.0).unsqueeze(-1)  # (B, 1)

        # === Per-axis gradient asymmetry ===
        # Asymmetry: mean gradient along each axis (nonzero = asymmetric curvature)
        gx_mean = (gx.squeeze(1) * grid).sum(dim=(1, 2, 3)) / total_occ.squeeze(-1)
        gy_mean = (gy.squeeze(1) * grid).sum(dim=(1, 2, 3)) / total_occ.squeeze(-1)
        gz_mean = (gz.squeeze(1) * grid).sum(dim=(1, 2, 3)) / total_occ.squeeze(-1)
        grad_asym = torch.stack([gx_mean, gy_mean, gz_mean], dim=-1)  # (B, 3)

        # === Radial gradient profile ===
        # How does gradient magnitude vary with distance from centroid?
        dists = diff_from_center.norm(dim=-1)  # (B, 125)
        # Arithmetic binning (Inductor-safe, no bucketize)
        # nan_to_num prevents NaN→long producing garbage indices under BF16
        bin_idx = torch.nan_to_num(dists * (5.0 / 3.5), nan=0.0).long().clamp(0, 4)
        active_mask = (flat_grid > 0.5)  # (B, 125)
        radial_grad = torch.zeros(B, 5, device=device)
        # Scatter-add: accumulate grad_mag and counts per bin
        weighted_mag = grad_mag * active_mask.float()  # zero out inactive
        one_hot = F.one_hot(bin_idx, 5).float()  # (B, 125, 5)
        active_oh = one_hot * active_mask.float().unsqueeze(-1)  # mask inactive
        counts = active_oh.sum(dim=1).clamp(min=1)  # (B, 5)
        radial_grad = (weighted_mag.unsqueeze(-1) * active_oh).sum(dim=1) / counts
        # (B, 5)

        # === Conv on Laplacian field (spatial curvature map) ===
        lap_feat = self.lap_conv(lap).reshape(B, -1)  # (B, 128)

        # === Conv on gradient field (directional boundaries) ===
        grad_feat = self.grad_conv(grad_field).reshape(B, -1)  # (B, 128)

        # === Combine all ===
        raw_feat = torch.cat([
            direction_hist,   # 3
            lap_hist,         # 3
            alternation,      # 1
            grad_asym,        # 3
            radial_grad,      # 5
        ], dim=-1)  # (B, 15)

        all_feat = torch.cat([raw_feat, lap_feat, grad_feat], dim=-1)  # (B, 271)

        direction_gate = self.direction_net(all_feat)      # (B, embed_dim) sigmoid
        direction_feat = self.direction_feat_net(all_feat)  # (B, embed_dim)

        return direction_gate, direction_feat, alternation


# === Deformation Augmentation =================================================

def deform_grid(grid, p_dropout=0.1, p_add=0.1, p_shift=0.15):
    """Fully vectorized voxel augmentation — zero CPU-GPU sync points."""
    B = grid.shape[0]
    device = grid.device
    r = torch.rand(B, 3, device=device)
    out = grid.clone()

    # --- Voxel dropout (batched, no .any() sync) ---
    drop_sel = (r[:, 0] < p_dropout).view(B, 1, 1, 1)
    keep = torch.rand_like(out) > 0.15
    out = torch.where(drop_sel, out * keep.float(), out)

    # --- Boundary addition (batched, no .any() sync) ---
    add_sel = (r[:, 1] < p_add).view(B, 1, 1, 1).float()
    dilated = F.max_pool3d(out.unsqueeze(1), kernel_size=3, stride=1, padding=1).squeeze(1)
    boundary = ((dilated > 0.5) & (out < 0.5)).float()
    add_noise = (torch.rand_like(out) < 0.3).float()
    out = (out + boundary * add_noise * add_sel).clamp(max=1.0)

    # --- Small translation (fully vectorized, no loops, no boolean indexing) ---
    shift_sel = (r[:, 2] < p_shift)  # (B,)
    axes = torch.randint(3, (B,), device=device)
    dirs = torch.randint(0, 2, (B,), device=device) * 2 - 1

    # Precompute all 6 shifted versions of full batch (cheap for 5x5x5)
    # Encode: idx = axis * 2 + (dir==1)  → [0..5], 6 = no shift
    versions = []
    for ax in range(3):
        for d in [-1, 1]:
            s = torch.roll(out, shifts=d, dims=ax + 1)  # +1 for batch dim
            # Zero wrapped edge
            if d == 1:
                if ax == 0: s[:, 0, :, :] = 0
                elif ax == 1: s[:, :, 0, :] = 0
                else: s[:, :, :, 0] = 0
            else:
                if ax == 0: s[:, -1, :, :] = 0
                elif ax == 1: s[:, :, -1, :] = 0
                else: s[:, :, :, -1] = 0
            versions.append(s)
    versions.append(out)  # index 6 = no shift (identity)
    stacked = torch.stack(versions, dim=0)  # (7, B, 5, 5, 5)

    # Per-sample assignment: which version to pick
    assign = torch.where(shift_sel, axes * 2 + (dirs == 1).long(), torch.full_like(axes, 6))
    # Gather: stacked[assign[b], b] for each b
    out = stacked[assign, torch.arange(B, device=device)]

    return out


# === Curvature Head (axis-aware) ==============================================

class CurvatureHead(nn.Module):
    """
    Axis-aware curvature detection with differentiation gating.

    1. Per-axis max projections -> 2D conv (keeps 2×2 spatial)
    2. Radial occupancy profile from centroid
    3. Axial symmetry + translation invariance scores
    4. 3D conv with spatial preservation (2×2×2)
    5. DifferentiationGate: gradient/Laplacian analysis for direction detection

    The DifferentiationGate modulates curvature features so that
    convex and concave shapes get distinct representations even when
    their occupancy patterns are nearly identical.
    """

    def __init__(self, rigid_feat_dim, fill_dim, embed_dim):
        super().__init__()

        self.plane_conv = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.GELU(),
            nn.Conv2d(16, 16, 3, padding=1), nn.GELU(),
            nn.AdaptiveAvgPool2d(2))
        plane_feat_dim = 3 * 16 * 4  # 192

        n_radial = 5
        self.radial_net = nn.Sequential(
            nn.Linear(n_radial, 32), nn.GELU(), nn.Linear(32, 16))
        radial_feat_dim = 16

        symmetry_feat_dim = 6

        self.voxel_conv = nn.Sequential(
            nn.Conv3d(1, 16, 3, padding=1), nn.GELU(),
            nn.Conv3d(16, 32, 3, padding=1), nn.GELU(),
            nn.AdaptiveAvgPool3d(2))
        voxel3d_feat_dim = 32 * 8  # 256

        # DifferentiationGate for curvature direction
        self.diff_gate = DifferentiationGate(embed_dim)

        # Pre-gate combine (without direction features)
        pre_gate_dim = (plane_feat_dim + radial_feat_dim + symmetry_feat_dim +
                        voxel3d_feat_dim + rigid_feat_dim + fill_dim)

        # Pre-gate feature projection: SwiGLU for sharp geometric feature gating
        self.pre_gate_proj = nn.Sequential(
            SwiGLU(pre_gate_dim, embed_dim * 2),
            nn.Linear(embed_dim * 2, embed_dim))

        # Post-gate: gated features + direction features + alternation + raw combine
        # = embed_dim (gated) + embed_dim (direction) + 1 (alternation) + pre_gate_dim
        post_gate_dim = embed_dim + embed_dim + 1 + pre_gate_dim

        # SwiGLU for all curvature decision heads: sharp geometric classification
        self.curved_head = nn.Sequential(
            SwiGLU(post_gate_dim, embed_dim),
            nn.Linear(embed_dim, 1), nn.Sigmoid())
        self.curv_type_head = nn.Sequential(
            SwiGLU(post_gate_dim, embed_dim),
            nn.Linear(embed_dim, NUM_CURVATURES))
        self.curv_features = nn.Sequential(
            SwiGLU(post_gate_dim, embed_dim * 2),
            nn.Linear(embed_dim * 2, embed_dim))

    def forward(self, grid, rigid_retained, fill_ratios):
        B = grid.shape[0]

        proj_x = grid.max(dim=1).values
        proj_y = grid.max(dim=2).values
        proj_z = grid.max(dim=3).values

        # Batch all 3 projections through plane_conv in single pass
        projs_batched = torch.cat([
            proj_x.unsqueeze(1), proj_y.unsqueeze(1), proj_z.unsqueeze(1)
        ], dim=0)  # (3B, 1, 5, 5)
        plane_all = self.plane_conv(projs_batched).reshape(3, B, -1)  # (3, B, 64)
        plane_feat = plane_all.permute(1, 0, 2).reshape(B, -1)  # (B, 192)

        radial = self._radial_profile(grid)
        radial_feat = self.radial_net(radial)

        sym_feat = self._symmetry_features(proj_x, proj_y, proj_z)

        vox3d_feat = self.voxel_conv(grid.unsqueeze(1)).reshape(B, -1)

        # Raw curvature features (shape-aware but direction-blind)
        raw_combined = torch.cat([
            plane_feat, radial_feat, sym_feat, vox3d_feat,
            rigid_retained, fill_ratios], dim=-1)

        # Project to gatable dimension
        pre_gate = self.pre_gate_proj(raw_combined)  # (B, embed_dim)

        # Direction analysis
        dir_gate, dir_feat, alternation = self.diff_gate(grid)

        # Apply gate: direction-modulated curvature features
        gated = pre_gate * dir_gate  # (B, embed_dim) — convex/concave differentiation

        # Full post-gate features
        combined = torch.cat([gated, dir_feat, alternation, raw_combined], dim=-1)

        is_curved = self.curved_head(combined)
        curv_logits = self.curv_type_head(combined)
        curv_feat = self.curv_features(combined)
        return is_curved, curv_logits, curv_feat, alternation

    def _radial_profile(self, grid):
        B = grid.shape[0]
        device = grid.device
        coords = torch.stack(torch.meshgrid(
            torch.arange(GS, device=device, dtype=torch.float32),
            torch.arange(GS, device=device, dtype=torch.float32),
            torch.arange(GS, device=device, dtype=torch.float32),
            indexing="ij"), dim=-1)
        flat_grid = grid.reshape(B, -1)
        flat_coords = coords.reshape(-1, 3)
        total_occ = flat_grid.sum(dim=-1, keepdim=True).clamp(min=1)
        centroids = (flat_grid.unsqueeze(-1) * flat_coords.unsqueeze(0)).sum(dim=1) / total_occ
        diffs = flat_coords.unsqueeze(0) - centroids.unsqueeze(1)
        dists = diffs.norm(dim=-1)  # (B, 125)
        max_dist = 3.5
        n_bins = 5
        # Arithmetic binning (Inductor-safe, no bucketize)
        bin_idx = torch.nan_to_num(dists * (float(n_bins) / max_dist), nan=0.0).long().clamp(0, n_bins - 1)
        one_hot = F.one_hot(bin_idx, n_bins).float()  # (B, 125, 5)
        weighted = flat_grid.unsqueeze(-1) * one_hot  # (B, 125, 5)
        profile = weighted.sum(dim=1) / total_occ  # (B, 5)
        return profile

    def _symmetry_features(self, proj_x, proj_y, proj_z):
        projs = torch.stack([proj_x, proj_y, proj_z], dim=1)  # (B, 3, H, W)
        fh = torch.flip(projs, dims=[2])
        fv = torch.flip(projs, dims=[3])
        sym = 1.0 - ((projs - fh).abs().mean(dim=(2, 3)) +
                       (projs - fv).abs().mean(dim=(2, 3))) / 2  # (B, 3)
        shift_diff = (projs[:, :, 1:, :] - projs[:, :, :-1, :]).abs().mean(dim=(2, 3))  # (B, 3)
        trans_inv = 1.0 - shift_diff
        # Interleave: [sym0, trans0, sym1, trans1, sym2, trans2]
        return torch.stack([sym[:, 0], trans_inv[:, 0],
                           sym[:, 1], trans_inv[:, 1],
                           sym[:, 2], trans_inv[:, 2]], dim=-1)  # (B, 6)


# === Confidence Computation ====================================================

def compute_confidence(logits):
    """
    Compute real calibrated confidence metrics from logits.

    Returns dict with:
        max_prob: max(softmax(logits)) — calibrated top-class probability
        margin: top1_prob - top2_prob — disambiguation strength
        entropy: -sum(p * log(p)) — total uncertainty (lower = more confident)
        confidence: margin — primary confidence signal for gating
    """
    probs = F.softmax(logits, dim=-1)
    max_prob, _ = probs.max(dim=-1)

    top2 = probs.topk(2, dim=-1).values
    margin = top2[:, 0] - top2[:, 1]

    # Entropy normalized to [0, 1] range
    log_probs = F.log_softmax(logits, dim=-1)
    entropy = -(probs * log_probs).sum(dim=-1)
    max_entropy = math.log(logits.shape[-1])
    norm_entropy = entropy / max_entropy

    return {
        "max_prob": max_prob,
        "margin": margin,
        "entropy": norm_entropy,
        "confidence": margin,  # primary signal
    }


# === Rectified Flow Arbiter ===================================================

class RectifiedFlowArbiter(nn.Module):
    """
    Rectified flow matching for ambiguous classification refinement.

    Real flow matching requires a target endpoint to define the velocity field.
    We learn class prototypes in latent space as targets: for a sample of class c,
    the target is prototype[c]. The velocity field learns to transport the
    encoded feature z0 toward the correct prototype z1 in straight lines:

        v_target = z1 - z0  (rectified: straight path from source to target)
        loss = ||v_predicted - v_target||^2  (flow matching objective)

    At inference, the arbiter integrates the learned velocity field from z0,
    landing near the correct class prototype. Classification reads off the
    nearest prototype.

    Confidence gating: velocity magnitude is scaled by (1 - margin), so
    confident first-pass predictions receive minimal correction.
    """

    def __init__(self, feat_dim, n_classes, n_steps=4, latent_dim=128, embed_dim=64):
        super().__init__()
        self.n_steps = n_steps
        self.n_classes = n_classes
        self.dt = 1.0 / n_steps
        self.latent_dim = latent_dim

        # Project features to latent space
        self.encode = nn.Sequential(
            nn.Linear(feat_dim, latent_dim * 2), nn.GELU(),
            nn.Linear(latent_dim * 2, latent_dim))

        # Learnable class prototypes — target endpoints for flow
        self.prototypes = nn.Parameter(torch.randn(n_classes, latent_dim) * 0.05)

        # Timestep embedding
        self.time_embed = nn.Sequential(
            nn.Linear(16, embed_dim), nn.GELU(),
            nn.Linear(embed_dim, embed_dim))

        # Confidence embedding
        self.conf_embed = nn.Sequential(
            nn.Linear(3, embed_dim), nn.GELU(),
            nn.Linear(embed_dim, embed_dim))

        # Velocity network: predicts flow direction in latent space
        vel_in = latent_dim + embed_dim + embed_dim
        self.velocity = nn.Sequential(
            SwiGLU(vel_in, latent_dim),
            nn.Linear(latent_dim, latent_dim),
            SwiGLU(latent_dim, latent_dim),
            nn.Linear(latent_dim, latent_dim))

        # Velocity gate: low confidence → full correction, high → minimal
        self.vel_gate = nn.Sequential(
            nn.Linear(embed_dim, latent_dim), nn.Sigmoid())

        # Classification from latent: distance to prototypes + learned head
        self.classifier_head = nn.Sequential(
            SwiGLU(latent_dim + n_classes, 96),
            nn.Linear(96, n_classes))

        # Learned confidence head for blending (differentiable, not topk)
        self.blend_head = nn.Sequential(
            nn.Linear(feat_dim, 64), nn.GELU(),
            nn.Linear(64, 1), nn.Sigmoid())

        # Post-refinement confidence
        self.refined_confidence = nn.Sequential(
            SwiGLU(latent_dim, 32),
            nn.Linear(32, 1), nn.Sigmoid())

    def _time_encoding(self, t, device):
        freqs = torch.exp(torch.linspace(0, -4, 8, device=device))
        args = t.unsqueeze(-1) * freqs.unsqueeze(0)
        return torch.cat([args.sin(), args.cos()], dim=-1)

    def _proto_logits(self, z):
        """Classify by negative distance to prototypes."""
        # (B, latent) vs (C, latent) → (B, C) distances
        dists = torch.cdist(z.unsqueeze(0), self.prototypes.unsqueeze(0)).squeeze(0)
        # Combine distance signal with learned head
        combined = torch.cat([z, -dists], dim=-1)  # (B, latent + n_classes)
        return self.classifier_head(combined)

    def forward(self, features, initial_logits, labels=None):
        """
        features: (B, feat_dim)
        initial_logits: (B, n_classes)
        labels: (B,) — only during training, for flow matching target

        Returns:
            refined_logits, refined_conf, initial_conf, trajectory_logits, flow_loss
        """
        B = features.shape[0]
        device = features.device

        # Confidence from initial logits
        initial_conf = compute_confidence(initial_logits)
        conf_input = torch.stack([
            initial_conf["max_prob"],
            initial_conf["margin"],
            initial_conf["entropy"]], dim=-1)
        conf_emb = self.conf_embed(conf_input)

        # Confidence-gated velocity magnitude
        gate = self.vel_gate(conf_emb)
        inv_conf = (1.0 - initial_conf["margin"]).unsqueeze(-1)
        adaptive_gate = gate * inv_conf

        # Encode to latent
        z0 = self.encode(features)

        # === Flow matching target ===
        flow_loss = torch.tensor(0.0, device=device)
        if labels is not None:
            # Target: class prototype for each sample
            z1 = self.prototypes[labels]  # (B, latent_dim)
            # Target velocity: straight path z0 → z1
            v_target = z1 - z0  # (B, latent_dim)

            # Sample random timestep for flow matching training
            t_rand = torch.rand(B, device=device)
            t_emb = self.time_embed(self._time_encoding(t_rand, device))

            # Interpolated position along straight path
            z_t = z0 + t_rand.unsqueeze(-1) * v_target  # (B, latent_dim)

            # Predicted velocity at this point
            vel_input = torch.cat([z_t, t_emb, conf_emb], dim=-1)
            v_pred = self.velocity(vel_input) * adaptive_gate
            v_pred = v_pred.clamp(-20, 20)

            # Flow matching loss: predicted velocity should match target
            flow_loss = F.mse_loss(v_pred, v_target.clamp(-20, 20))

        # === Inference: integrate velocity field ===
        z = z0
        trajectory_logits = []
        for step in range(self.n_steps):
            t_val = torch.full((B,), step * self.dt, device=device)
            t_emb = self.time_embed(self._time_encoding(t_val, device))

            vel_input = torch.cat([z, t_emb, conf_emb], dim=-1)
            v = self.velocity(vel_input) * adaptive_gate
            # Prevent BF16 divergence: clamp velocity magnitude
            v = v.clamp(-20, 20)

            z = z + self.dt * v
            trajectory_logits.append(self._proto_logits(z))

        refined_logits = trajectory_logits[-1]
        refined_conf = self.refined_confidence(z)

        # Learned blend weight (differentiable, from initial features)
        blend_weight = self.blend_head(features)  # (B, 1)

        return refined_logits, refined_conf, initial_conf, trajectory_logits, flow_loss, blend_weight


# === Model ====================================================================

class GeometricShapeClassifier(nn.Module):
    def __init__(self, n_classes=NUM_CLASSES, embed_dim=64, n_tracers=5):
        super().__init__()
        self.n_tracers = n_tracers
        self.embed_dim = embed_dim

        self.voxel_embed = nn.Sequential(
            nn.Linear(4, embed_dim), nn.GELU(), nn.Linear(embed_dim, embed_dim))

        coords = torch.stack(torch.meshgrid(
            torch.arange(GS, dtype=torch.float32),
            torch.arange(GS, dtype=torch.float32),
            torch.arange(GS, dtype=torch.float32),
            indexing="ij"), dim=-1) / (GS - 1)  # (5,5,5,3) normalized
        self.register_buffer("pos_grid", coords)

        self.tracer_tokens = nn.Parameter(torch.randn(n_tracers, embed_dim) * 0.02)
        self.tracer_attn = nn.MultiheadAttention(embed_dim, num_heads=4, batch_first=True)
        self.tracer_gate = nn.Sequential(nn.Linear(embed_dim * 2, embed_dim), nn.Sigmoid())
        self.tracer_interact = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim), nn.GELU(), nn.Linear(embed_dim, embed_dim))
        # SwiGLU for edge detection: sharp "edge present?" decision
        self.edge_head = nn.Sequential(
            SwiGLU(embed_dim * 2, 32), nn.Linear(32, 1))

        # Precompute all C(n_tracers, 2) pair indices for vectorized interaction
        _pi, _pj = [], []
        for i in range(n_tracers):
            for j in range(i + 1, n_tracers):
                _pi.append(i); _pj.append(j)
        self.register_buffer("_pair_i", torch.tensor(_pi, dtype=torch.long))
        self.register_buffer("_pair_j", torch.tensor(_pj, dtype=torch.long))
        self.n_pairs = len(_pi)

        pool_dim = embed_dim * n_tracers

        self.dim0 = CapacityHead(pool_dim, embed_dim, init_capacity=0.5)
        self.dim1 = CapacityHead(pool_dim + embed_dim, embed_dim, init_capacity=1.0)
        self.dim2 = CapacityHead(pool_dim + embed_dim, embed_dim, init_capacity=1.5)
        self.dim3 = CapacityHead(pool_dim + embed_dim, embed_dim, init_capacity=2.0)

        rigid_feat_dim = embed_dim * 4
        self.curvature = CurvatureHead(rigid_feat_dim, fill_dim=4, embed_dim=embed_dim)

        class_in = pool_dim + 4 + rigid_feat_dim + embed_dim + 1
        self.class_in = class_in  # Store for arbiter
        self.classifier = nn.Sequential(
            nn.Linear(class_in, 256), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(256, 128), nn.GELU(), nn.Linear(128, n_classes))

        # SwiGLU for peak dimension: sharp "which dimension?" decision
        self.peak_head = nn.Sequential(
            SwiGLU(class_in, 32), nn.Linear(32, 4))
        # Volume is continuous interpolation — keep GELU
        self.volume_head = nn.Sequential(
            nn.Linear(class_in, 64), nn.GELU(), nn.Linear(64, 1))
        # SwiGLU for CM determinant sign: sharp geometric determinant
        self.cm_head = nn.Sequential(
            SwiGLU(class_in, 64), nn.Linear(64, 1), nn.Tanh())

        # Rectified flow arbiter for ambiguous classification
        self.arbiter = RectifiedFlowArbiter(
            feat_dim=class_in, n_classes=n_classes,
            n_steps=4, latent_dim=128, embed_dim=embed_dim)

    def forward(self, grid, labels=None):
        B = grid.shape[0]
        occ = grid.reshape(B, GS**3, 1)
        pos = self.pos_grid.reshape(1, GS**3, 3).expand(B, -1, -1)
        voxel_emb = self.voxel_embed(torch.cat([occ, pos], dim=-1))

        tracers = self.tracer_tokens.unsqueeze(0).expand(B, -1, -1)
        tracers, _ = self.tracer_attn(tracers, voxel_emb, voxel_emb)

        # Vectorized pair interaction: all C(5,2)=10 pairs at once
        left = tracers[:, self._pair_i]   # (B, 10, embed_dim)
        right = tracers[:, self._pair_j]  # (B, 10, embed_dim)
        pairs = torch.cat([left, right], dim=-1)  # (B, 10, embed_dim*2)

        # Flatten to batch, run networks, reshape back
        flat_pairs = pairs.reshape(B * self.n_pairs, -1)
        gate = self.tracer_gate(flat_pairs).reshape(B, self.n_pairs, -1)
        interaction = self.tracer_interact(flat_pairs).reshape(B, self.n_pairs, -1)
        edge_lengths = self.edge_head(flat_pairs).reshape(B, self.n_pairs)

        # Scatter-add gated interactions back to both tracers in each pair
        gated = gate * interaction  # (B, 10, embed_dim)
        tracer_out = tracers.clone()
        pi_exp = self._pair_i.view(1, self.n_pairs, 1).expand(B, -1, self.embed_dim)
        pj_exp = self._pair_j.view(1, self.n_pairs, 1).expand(B, -1, self.embed_dim)
        tracer_out.scatter_add_(1, pi_exp, gated)
        tracer_out.scatter_add_(1, pj_exp, gated)
        pooled = tracer_out.reshape(B, -1)

        fill0, ovf0, ret0, cap0, _ = self.dim0(pooled)
        fill1, ovf1, ret1, cap1, _ = self.dim1(torch.cat([pooled, ovf0], -1))
        fill2, ovf2, ret2, cap2, _ = self.dim2(torch.cat([pooled, ovf1], -1))
        fill3, ovf3, ret3, cap3, _ = self.dim3(torch.cat([pooled, ovf2], -1))

        fill_ratios = torch.cat([fill0, fill1, fill2, fill3], dim=-1)
        rigid_retained = torch.cat([ret0, ret1, ret2, ret3], dim=-1)
        ovf_norms = torch.stack([
            ovf0.norm(dim=-1), ovf1.norm(dim=-1),
            ovf2.norm(dim=-1), ovf3.norm(dim=-1)], dim=-1)

        is_curved, curv_logits, curv_feat, alternation = self.curvature(grid, rigid_retained, fill_ratios)
        full = torch.cat([pooled, fill_ratios, rigid_retained, curv_feat, is_curved], dim=-1)

        # === First pass classification ===
        initial_logits = self.classifier(full)

        # === Rectified flow arbitration ===
        refined_logits, refined_conf, initial_conf, trajectory_logits, flow_loss, blend_weight = \
            self.arbiter(full, initial_logits, labels=labels)

        # === Blend: learned confidence head decides trust ===
        # blend_weight is (B, 1) sigmoid output from learned head
        final_logits = blend_weight * initial_logits + (1.0 - blend_weight) * refined_logits

        return {
            # Classification
            "class_logits": final_logits,
            "initial_logits": initial_logits,
            "refined_logits": refined_logits,
            "trajectory_logits": trajectory_logits,
            # Flow matching
            "flow_loss": flow_loss,
            # Confidence
            "confidence": initial_conf["confidence"],
            "max_prob": initial_conf["max_prob"],
            "entropy": initial_conf["entropy"],
            "refined_confidence": refined_conf,
            "blend_weight": blend_weight.squeeze(-1),
            # Auxiliary heads
            "peak_logits": self.peak_head(full),
            "volume_pred": self.volume_head(full).squeeze(-1),
            "cm_pred": self.cm_head(full).squeeze(-1),
            "edge_lengths": edge_lengths,
            "fill_ratios": fill_ratios,
            "overflows": ovf_norms,
            "capacities": torch.stack([cap0, cap1, cap2, cap3]),
            "is_curved_pred": is_curved,
            "curv_type_logits": curv_logits,
            "alternation": alternation,
            # Pre-classifier features (for cross-contrast)
            "features": full,
        }


# Quick sanity
_m = GeometricShapeClassifier()
print(f'GeometricShapeClassifier: {sum(p.numel() for p in _m.parameters()):,} params')
del _m