"""Small PyTorch implementations of the POC model families. These are deliberately dependency-light baselines, not claims of a production architecture or of material/defect predictive performance. """ from __future__ import annotations from typing import Optional import torch from torch import Tensor, nn import torch.nn.functional as F def _conv2(cin: int, cout: int) -> nn.Sequential: return nn.Sequential(nn.Conv2d(cin, cout, 3, padding=1), nn.GroupNorm(1, cout), nn.GELU(), nn.Conv2d(cout, cout, 3, padding=1), nn.GELU()) class CNN2DUNet(nn.Module): """2D CNN/U-Net for multi-view raster channels; output is per-pixel.""" def __init__(self, in_channels: int, out_channels: int = 1, width: int = 16, global_dim: int = 0): super().__init__() self.enc1, self.enc2 = _conv2(in_channels, width), _conv2(width, width * 2) self.pool = nn.MaxPool2d(2) self.up = nn.ConvTranspose2d(width * 2, width, 2, stride=2) self.dec, self.head = _conv2(width * 2, width), nn.Conv2d(width, out_channels, 1) self.global_dim = global_dim self.film = nn.Linear(global_dim, width * 4) if global_dim else None def forward(self, x: Tensor, process: Optional[Tensor] = None) -> Tensor: a = self.enc1(x); b = self.enc2(self.pool(a)) if self.film is not None: if process is None or process.shape != (x.shape[0], self.global_dim): raise ValueError("CNN2DUNet requires one [B,global_dim] process+mask vector per image") scale, bias = self.film(process).chunk(2, 1) b = b * (1 + scale[:, :, None, None]) + bias[:, :, None, None] u = self.up(b) if u.shape[-2:] != a.shape[-2:]: u = F.interpolate(u, size=a.shape[-2:], mode="bilinear", align_corners=False) return self.head(self.dec(torch.cat([a, u], dim=1))) def _conv3(cin: int, cout: int) -> nn.Sequential: return nn.Sequential(nn.Conv3d(cin, cout, 3, padding=1), nn.GroupNorm(1, cout), nn.GELU(), nn.Conv3d(cout, cout, 3, padding=1), nn.GELU()) class UNet3D(nn.Module): """Dense 3D U-Net; only use bounded dense occupancy/field tensors.""" def __init__(self, in_channels: int, out_channels: int = 1, width: int = 8, global_dim: int = 0): super().__init__() self.enc1, self.enc2 = _conv3(in_channels, width), _conv3(width, width * 2) self.pool = nn.MaxPool3d(2); self.up = nn.ConvTranspose3d(width * 2, width, 2, 2) self.dec, self.head = _conv3(width * 2, width), nn.Conv3d(width, out_channels, 1) self.global_dim = global_dim; self.film = nn.Linear(global_dim, width * 4) if global_dim else None def forward(self, x: Tensor, process: Optional[Tensor] = None) -> Tensor: a = self.enc1(x); b = self.enc2(self.pool(a)); u = self.up(b) if self.film is not None: if process is None or process.shape != (x.shape[0], self.global_dim): raise ValueError("UNet3D requires one [B,global_dim] process+mask vector per volume") scale, bias = self.film(process).chunk(2, 1); b = b * (1 + scale[:, :, None, None, None]) + bias[:, :, None, None, None] u = self.up(b) if u.shape[-3:] != a.shape[-3:]: u = F.interpolate(u, size=a.shape[-3:], mode="trilinear", align_corners=False) return self.head(self.dec(torch.cat([a, u], 1))) class SparseUNetOptional(nn.Module): """True sparse U-Net, available only with MinkowskiEngine (never dense fallback).""" backend = "MinkowskiEngine" def __init__(self, in_channels: int, out_channels: int = 1, width: int = 16): super().__init__() try: import MinkowskiEngine as ME # type: ignore except ImportError as exc: raise RuntimeError("Sparse U-Net unavailable: install optional MinkowskiEngine; dense tensors are not a sparse substitute.") from exc self._ME = ME self.stem = nn.Sequential(ME.MinkowskiConvolution(in_channels, width, 3, dimension=3), ME.MinkowskiBatchNorm(width), ME.MinkowskiReLU()) self.down = nn.Sequential(ME.MinkowskiConvolution(width, width * 2, 2, stride=2, dimension=3), ME.MinkowskiBatchNorm(width * 2), ME.MinkowskiReLU()) self.mid = nn.Sequential(ME.MinkowskiConvolution(width * 2, width * 2, 3, dimension=3), ME.MinkowskiBatchNorm(width * 2), ME.MinkowskiReLU()) self.up = nn.Sequential(ME.MinkowskiConvolutionTranspose(width * 2, width, 2, stride=2, dimension=3), ME.MinkowskiBatchNorm(width), ME.MinkowskiReLU()) self.head = ME.MinkowskiConvolution(width * 2, out_channels, 1, dimension=3) def forward(self, x: Tensor) -> Tensor: skip = self.stem(x); return self.head(self._ME.cat(skip, self.up(self.mid(self.down(skip))))) class PointNetSeg(nn.Module): """PointNet segmentation baseline, [B,N,C] -> [B,N,O].""" def __init__(self, in_channels: int, out_channels: int = 1, width: int = 64): super().__init__() self.local = nn.Sequential(nn.Conv1d(in_channels, width, 1), nn.GELU(), nn.Conv1d(width, width * 2, 1), nn.GELU()) self.head = nn.Sequential(nn.Conv1d(width * 4, width, 1), nn.GELU(), nn.Conv1d(width, out_channels, 1)) def forward(self, x: Tensor) -> Tensor: h = self.local(x.transpose(1, 2)); g = h.max(dim=-1, keepdim=True).values.expand_as(h) return self.head(torch.cat([h, g], 1)).transpose(1, 2) def _knn_features(x: Tensor, k: int) -> Tensor: """Dynamic feature-space kNN, bounded by caller supplied N (DGCNN smoke <=512).""" n = x.shape[1] if n == 0: raise ValueError("DGCNN needs at least one point") if n == 1: return torch.cat([x[:, :, None, :], torch.zeros_like(x[:, :, None, :])], -1) k = min(k, n - 1) dist = torch.cdist(x, x) dist.diagonal(dim1=1, dim2=2).fill_(float("inf")) idx = dist.topk(k, largest=False).indices b = torch.arange(x.shape[0], device=x.device)[:, None, None] nbr = x[b, idx] return torch.cat([x[:, :, None, :].expand_as(nbr), nbr - x[:, :, None, :]], -1) class DGCNNSeg(nn.Module): """DGCNN with recomputed feature-neighbor graphs at both EdgeConv stages.""" def __init__(self, in_channels: int, out_channels: int = 1, width: int = 32, k: int = 16): super().__init__(); self.k = k self.e1 = nn.Sequential(nn.Linear(in_channels * 2, width), nn.GELU(), nn.Linear(width, width), nn.GELU()) self.e2 = nn.Sequential(nn.Linear(width * 2, width), nn.GELU(), nn.Linear(width, width), nn.GELU()) self.head = nn.Sequential(nn.Linear(width * 3, width), nn.GELU(), nn.Linear(width, out_channels)) def _edge(self, x: Tensor, net: nn.Module) -> Tensor: return net(_knn_features(x, self.k)).max(2).values def forward(self, x: Tensor) -> Tensor: a = self._edge(x, self.e1); b = self._edge(a, self.e2) g = b.max(1, keepdim=True).values.expand_as(b) return self.head(torch.cat([a, b, g], -1)) class MeshGraphProcessor(nn.Module): """Edge-conditioned message passing with per-graph global process conditioning.""" def __init__(self, node_dim: int, edge_dim: int, global_dim: int, out_channels: int = 1, width: int = 48, steps: int = 3): super().__init__(); self.steps = steps self.node = nn.Sequential(nn.Linear(node_dim, width), nn.GELU()) self.global_enc = nn.Sequential(nn.Linear(global_dim, width), nn.GELU()) self.edge = nn.Sequential(nn.Linear(width * 2 + edge_dim + width, width), nn.GELU(), nn.Linear(width, width), nn.GELU()) self.update = nn.GRUCell(width * 2, width); self.head = nn.Linear(width, out_channels) def forward(self, node_x: Tensor, edge_index: Tensor, edge_attr: Tensor, global_x: Tensor, batch: Optional[Tensor] = None) -> Tensor: if batch is None: batch = node_x.new_zeros(node_x.shape[0], dtype=torch.long) h = self.node(node_x); g = self.global_enc(global_x)[batch] src, dst = edge_index for _ in range(self.steps): m = self.edge(torch.cat([h[src], h[dst], edge_attr, g[src]], 1)) agg = torch.zeros_like(h).index_add_(0, dst, m) h = self.update(torch.cat([h, agg], 1), h) return self.head(h) class PhysicsConstrainedRiskGNN(nn.Module): """Node-wise risk/deformation heads with a hard magnitude constraint. ``deformation_magnitude >= abs(signed_normal_deformation)`` is enforced by construction. Defect-type logits are pooled once per graph/run. """ def __init__(self, node_dim: int, edge_dim: int, global_dim: int, defect_types: int, width: int = 32, steps: int = 2): super().__init__() if defect_types < 1: raise ValueError("defect_types must be at least one") self.defect_types = defect_types self.processor = MeshGraphProcessor( node_dim, edge_dim, global_dim, out_channels=3 + defect_types, width=width, steps=steps) def forward(self, node_x: Tensor, edge_index: Tensor, edge_attr: Tensor, global_x: Tensor, batch: Tensor) -> dict[str, Tensor]: raw = self.processor(node_x, edge_index, edge_attr, global_x, batch) signed = raw[:, 1] magnitude = signed.abs() + F.softplus(raw[:, 2]) node_type = raw[:, 3:] graph_count = global_x.shape[0] type_sum = node_type.new_zeros((graph_count, self.defect_types)).index_add_(0, batch, node_type) count = torch.bincount(batch, minlength=graph_count).clamp_min(1).to(node_type.dtype)[:, None] return { "risk_logit": raw[:, 0], "signed_normal_deformation": signed, "deformation_magnitude": magnitude, "defect_type_logits": type_sum / count, } class UVNetFaceModel(nn.Module): """UV-Net-style face-grid encoder followed by face adjacency message passing.""" def __init__(self, channels: int, face_dim: int = 8, out_channels: int = 1): super().__init__(); self.grid = nn.Sequential(nn.Conv2d(channels, face_dim, 3, padding=1), nn.GELU(), nn.Conv2d(face_dim, face_dim, 3, padding=1), nn.GELU()) self.msg = nn.Sequential(nn.Linear(face_dim * 2, face_dim), nn.GELU()); self.head = nn.Linear(face_dim, out_channels) def forward(self, grids: Tensor, adjacency: Tensor) -> Tensor: # grids [F,C,H,W], adjacency [2,E], one B-Rep per call h = self.grid(grids).mean((-1, -2)); src, dst = adjacency m = self.msg(torch.cat([h[src], h[dst]], 1)); agg = torch.zeros_like(h).index_add_(0, dst, m) return self.head(h + agg) class SpectralConv2d(nn.Module): def __init__(self, cin: int, cout: int, modes: int = 8): super().__init__(); self.cin, self.cout, self.modes = cin, cout, modes self.weight_pos = nn.Parameter(torch.randn(cin, cout, modes, modes, dtype=torch.cfloat) * 0.02) self.weight_neg = nn.Parameter(torch.randn(cin, cout, modes, modes, dtype=torch.cfloat) * 0.02) def forward(self, x: Tensor) -> Tensor: xf = torch.fft.rfft2(x); out = torch.zeros(x.shape[0], self.cout, x.shape[-2], x.shape[-1] // 2 + 1, device=x.device, dtype=torch.cfloat) m2 = min(self.modes, xf.shape[-1]); pos = min(self.modes, x.shape[-2]); neg = min(self.modes, max(0, x.shape[-2] - pos)) out[:, :, :pos, :m2] = torch.einsum("bixy,ioxy->boxy", xf[:, :, :pos, :m2], self.weight_pos[:, :, :pos, :m2]) if neg: out[:, :, -neg:, :m2] = torch.einsum("bixy,ioxy->boxy", xf[:, :, -neg:, :m2], self.weight_neg[:, :, :neg, :m2]) return torch.fft.irfft2(out, s=x.shape[-2:]) class FNO2d(nn.Module): """Fourier neural operator baseline for gridded thermal/field surrogates.""" def __init__(self, in_channels: int, out_channels: int = 1, width: int = 24, modes: int = 8, global_dim: int = 0): super().__init__(); self.lift = nn.Conv2d(in_channels, width, 1) self.spec = nn.ModuleList([SpectralConv2d(width, width, modes) for _ in range(3)]) self.skip = nn.ModuleList([nn.Conv2d(width, width, 1) for _ in range(3)]); self.out = nn.Conv2d(width, out_channels, 1) self.global_dim = global_dim; self.film = nn.Linear(global_dim, width * 2) if global_dim else None def forward(self, x: Tensor, process: Optional[Tensor] = None) -> Tensor: h = self.lift(x) if self.film is not None: if process is None or process.shape != (x.shape[0], self.global_dim): raise ValueError("FNO2d requires one [B,global_dim] process+mask vector per field") s, b = self.film(process).chunk(2, 1); h = h * (1 + s[:, :, None, None]) + b[:, :, None, None] for s, p in zip(self.spec, self.skip): h = F.gelu(s(h) + p(h)) return self.out(h) class TabularSurrogate(nn.Module): """Small mask-aware tabular surrogate; feature values and masks are concatenated by the adapter.""" def __init__(self, in_channels: int, out_channels: int = 1, width: int = 64): super().__init__() self.net = nn.Sequential(nn.Linear(in_channels, width), nn.GELU(), nn.Linear(width, width), nn.GELU(), nn.Linear(width, out_channels)) def forward(self, x: Tensor) -> Tensor: return self.net(x)