Create cell2_model.py

cell2_model.py

@@ -0,0 +1,117 @@
+"""
+Superposition Patch Classifier - Architecture & Loss
+=====================================================
+Colab Cell 2 of 3 - depends on Cell 1 (generator.py) namespace.
+
+Defines model architecture and loss.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Cell 1 provides: all constants (GZ, GY, GX, MACRO_*, PATCH_*, NUM_CLASSES, NUM_GATES, CLASS_NAMES, etc.)
+#                   generate_dataset, analyze_patches_torch
+
+# === Model Architecture =======================================================
+NUM_DIMS, NUM_CURVS, NUM_TOPOS = 4, 3, 2
+
+
+class PatchEmbedding3D(nn.Module):
+    def __init__(self, patch_dim=64):
+        super().__init__()
+        self.proj = nn.Linear(PATCH_VOL, patch_dim)
+        pz = torch.arange(MACRO_Z).float() / MACRO_Z
+        py = torch.arange(MACRO_Y).float() / MACRO_Y
+        px = torch.arange(MACRO_X).float() / MACRO_X
+        pos = torch.stack(torch.meshgrid(pz, py, px, indexing='ij'), dim=-1).reshape(MACRO_N, 3)
+        self.register_buffer('pos_embed', pos)
+        self.pos_proj = nn.Linear(3, patch_dim)
+
+    def forward(self, x):
+        B = x.shape[0]
+        patches = x.view(B, MACRO_Z, PATCH_Z, MACRO_Y, PATCH_Y, MACRO_X, PATCH_X)
+        patches = patches.permute(0, 1, 3, 5, 2, 4, 6).contiguous().view(B, MACRO_N, PATCH_VOL)
+        return self.proj(patches) + self.pos_proj(self.pos_embed)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, n_heads, dropout=0.1):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
+        self.ff = nn.Sequential(
+            nn.Linear(dim, dim * 4), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim), nn.Dropout(dropout)
+        )
+        self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim)
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x), self.ln1(x), self.ln1(x))[0]
+        return x + self.ff(self.ln2(x))
+
+
+class SuperpositionPatchClassifier(nn.Module):
+    def __init__(self, embed_dim=128, patch_dim=64, n_layers=4, n_heads=4, dropout=0.1):
+        super().__init__()
+        self.patch_embed = PatchEmbedding3D(patch_dim)
+        self.proj = nn.Linear(patch_dim, embed_dim)
+        self.blocks = nn.ModuleList([TransformerBlock(embed_dim, n_heads, dropout) for _ in range(n_layers)])
+
+        # Per-patch heads
+        self.patch_dim_head = nn.Linear(embed_dim, NUM_DIMS)
+        self.patch_curv_head = nn.Linear(embed_dim, NUM_CURVS)
+        self.patch_topo_head = nn.Linear(embed_dim, NUM_TOPOS)
+        self.patch_gate_head = nn.Linear(embed_dim, NUM_GATES)
+        self.patch_shape_head = nn.Linear(embed_dim, NUM_CLASSES)
+        self.patch_conf_head = nn.Linear(embed_dim, 1)
+
+        # Global heads
+        self.global_pool = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.GELU(), nn.Linear(embed_dim, embed_dim))
+        self.global_gate_head = nn.Linear(embed_dim, NUM_GATES)
+        self.global_shape_head = nn.Linear(embed_dim, NUM_CLASSES)
+
+    def forward(self, x):
+        h = self.proj(self.patch_embed(x))
+        for blk in self.blocks:
+            h = blk(h)
+        g = self.global_pool(h.mean(dim=1))
+        return {
+            "patch_features": h,
+            "patch_dim_logits": self.patch_dim_head(h),
+            "patch_curv_logits": self.patch_curv_head(h),
+            "patch_topo_logits": self.patch_topo_head(h),
+            "patch_gate_logits": self.patch_gate_head(h),
+            "patch_gate_soft": torch.sigmoid(self.patch_gate_head(h)),
+            "patch_shape_logits": self.patch_shape_head(h),
+            "patch_confidence": torch.sigmoid(self.patch_conf_head(h)).squeeze(-1),
+            "global_features": g,
+            "global_gates": self.global_gate_head(g),
+            "global_shapes": self.global_shape_head(g),
+        }
+
+
+class SuperpositionLoss(nn.Module):
+    def __init__(self, patch_weight=1.0, global_weight=0.5, cross_weight=0.1, conf_weight=0.1):
+        super().__init__()
+        self.pw, self.gw, self.cw, self.confw = patch_weight, global_weight, cross_weight, conf_weight
+
+    def forward(self, outputs, targets):
+        occ_mask = targets["patch_occupancy"] > 0.01
+        n_occ = occ_mask.sum().clamp(min=1)
+
+        dims_loss = F.cross_entropy(outputs["patch_dim_logits"].view(-1, NUM_DIMS), targets["patch_dims"].clamp(0, NUM_DIMS-1).view(-1), reduction='none').view_as(occ_mask)
+        curv_loss = F.cross_entropy(outputs["patch_curv_logits"].view(-1, NUM_CURVS), targets["patch_curvature"].clamp(0, NUM_CURVS-1).view(-1), reduction='none').view_as(occ_mask)
+        topo_loss = F.cross_entropy(outputs["patch_topo_logits"].view(-1, NUM_TOPOS), targets["patch_topology"].clamp(0, NUM_TOPOS-1).view(-1), reduction='none').view_as(occ_mask)
+        gate_loss = F.binary_cross_entropy_with_logits(outputs["patch_gate_logits"], targets["patch_labels"], reduction='none').mean(dim=-1)
+        shape_loss = F.binary_cross_entropy_with_logits(outputs["patch_shape_logits"], targets["patch_shape_membership"], reduction='none').mean(dim=-1)
+
+        patch_loss = ((dims_loss + curv_loss + topo_loss + gate_loss + shape_loss) * occ_mask.float()).sum() / n_occ
+
+        global_gate_loss = F.binary_cross_entropy_with_logits(outputs["global_gates"], targets["global_gates"])
+        global_shape_loss = F.binary_cross_entropy_with_logits(outputs["global_shapes"], targets["global_shapes"])
+        global_loss = global_gate_loss + global_shape_loss
+
+        return {"total": self.pw * patch_loss + self.gw * global_loss, "patch": patch_loss, "global": global_loss}
+
+
+print("✓ Model ready")