""" Multi-Mechanism Normalizing Flow for Joint Mechanism Identification and Bayesian Parameter Inference from Multi-Scan-Rate CV Signals. Architecture: 1. MultiScanEncoder: per-CV CNN + Set Transformer (SAB + PMA) -> context vector 2. MechanismClassifier: MLP head -> p(mechanism | x) 3. Per-mechanism ConditionalFlow heads: p(theta_m | x, mechanism=m) The model performs two-level inference: Level 1: Mechanism probabilities p(m | x) from the classifier Level 2: Parameter posteriors p(theta_m | x, m) from mechanism-specific flows Training loss = classification CE + mechanism-specific NLL (weighted by true label). """ import torch import torch.nn as nn import torch.nn.functional as F import math from flow_model import ( SignalEncoder, ActNorm, ConditionalSplineCoupling, ConditionalAffineCoupling, MECHANISM_LIST, MECHANISM_PARAMS, ) class MechanismClassifier(nn.Module): """MLP classifier: context vector -> mechanism probabilities.""" def __init__(self, d_context=128, n_mechanisms=4, hidden_dim=128): super().__init__() self.net = nn.Sequential( nn.Linear(d_context, hidden_dim), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_dim, hidden_dim // 2), nn.GELU(), nn.Dropout(0.1), nn.Linear(hidden_dim // 2, n_mechanisms), ) def forward(self, context): """Returns raw logits [B, n_mechanisms].""" return self.net(context) class SAB(nn.Module): """Self-Attention Block (Set Transformer, Lee et al. 2019). Applies multi-head self-attention + feed-forward with residual connections and layer norm over a set of elements. """ def __init__(self, d, n_heads=4): super().__init__() self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True) self.norm1 = nn.LayerNorm(d) self.ffn = nn.Sequential(nn.Linear(d, d), nn.GELU(), nn.Linear(d, d)) self.norm2 = nn.LayerNorm(d) def forward(self, h, key_padding_mask=None): """ Args: h: [B, N, d] key_padding_mask: [B, N] True = ignore this position """ h2, _ = self.attn(h, h, h, key_padding_mask=key_padding_mask) h = self.norm1(h + h2) h = self.norm2(h + self.ffn(h)) return h class PMA(nn.Module): """Pooling by Multi-head Attention (Set Transformer, Lee et al. 2019). Uses learnable seed vectors as queries that attend to the input set, producing a fixed-size output regardless of set cardinality. """ def __init__(self, d, n_heads=4, n_seeds=1): super().__init__() self.seed = nn.Parameter(torch.randn(1, n_seeds, d)) self.attn = nn.MultiheadAttention(d, n_heads, batch_first=True) self.norm = nn.LayerNorm(d) def forward(self, h, key_padding_mask=None): """ Args: h: [B, N, d] key_padding_mask: [B, N] True = ignore this position Returns: [B, n_seeds, d] """ S = self.seed.expand(h.size(0), -1, -1) out, _ = self.attn(S, h, h, key_padding_mask=key_padding_mask) return self.norm(out + S) class MultiScanEncoder(nn.Module): """ Encode a set of multi-scan-rate CVs into a single context vector. Architecture (Set Transformer, default): 1. Shared per-CV CNN encoder -> per-CV embedding 2. Augment with [log10(sigma), log10(peak_flux)] 3. SAB: self-attention across scan rates (cross-CV interaction) 4. PMA: attention-based pooling to single vector 5. rho MLP: project to final context With aggregation='mean_pool', steps 3-4 are replaced by masked mean pooling (no learned cross-CV interaction). Input: x [B, N_sigma, 3, T], scan_mask [B, N_sigma, T], sigmas [B, N_sigma], flux_scales [B, N_sigma] Output: context [B, d_context] """ def __init__(self, in_channels=3, d_model=128, d_context=128, n_heads=4, aggregation='set_transformer'): super().__init__() self.aggregation = aggregation self.per_cv_encoder = SignalEncoder( in_channels=in_channels, d_model=d_model, d_context=d_context, ) self.cv_augment = nn.Sequential( nn.Linear(d_context + 2, d_context), nn.GELU(), ) if aggregation == 'set_transformer': self.sab = SAB(d_context, n_heads=n_heads) self.pma = PMA(d_context, n_heads=n_heads, n_seeds=1) elif aggregation == 'deepsets': self.phi = nn.Sequential( nn.Linear(d_context, d_context), nn.GELU(), ) elif aggregation == 'mean_pool': pass else: raise ValueError(f"Unknown aggregation: {aggregation!r}") self.rho = nn.Sequential( nn.Linear(d_context, d_context), nn.GELU(), nn.Linear(d_context, d_context), ) def forward(self, x, scan_mask=None, sigmas=None, flux_scales=None): """ Args: x: [B, N_sigma, 3, T] multi-scan CV waveforms scan_mask: [B, N_sigma, T] valid timestep mask sigmas: [B, N_sigma] log10 scan rates flux_scales: [B, N_sigma] log10(peak_flux) per CV Returns: context: [B, d_context] """ B, N, C, T = x.shape x_flat = x.reshape(B * N, C, T) mask_flat = scan_mask.reshape(B * N, T) if scan_mask is not None else None h_flat = self.per_cv_encoder(x_flat, mask=mask_flat) # [B*N, d_context] h = h_flat.reshape(B, N, -1) # [B, N, d_context] if sigmas is None: sigmas = torch.zeros(B, N, device=x.device) if flux_scales is None: flux_scales = torch.zeros(B, N, device=x.device) aug_features = torch.stack([sigmas, flux_scales], dim=-1) # [B, N, 2] h = self.cv_augment(torch.cat([h, aug_features], dim=-1)) # [B, N, d_context] # Build key_padding_mask: True where CV is padded (invalid) if scan_mask is not None: cv_invalid = ~scan_mask.any(dim=-1) # [B, N] True = padded else: cv_invalid = None if self.aggregation == 'set_transformer': h = self.sab(h, key_padding_mask=cv_invalid) h = self.pma(h, key_padding_mask=cv_invalid) # [B, 1, d_context] h = h.squeeze(1) # [B, d_context] elif self.aggregation in ('deepsets', 'mean_pool'): if self.aggregation == 'deepsets': h = self.phi(h) if cv_invalid is not None: cv_valid = (~cv_invalid).unsqueeze(-1).float() # [B, N, 1] h = (h * cv_valid).sum(dim=1) / cv_valid.sum(dim=1).clamp(min=1) else: h = h.mean(dim=1) # [B, d_context] context = self.rho(h) return context class MechanismFlow(nn.Module): """ Single-mechanism conditional flow: p(theta_m | context). Lightweight flow head operating on the context vector from the shared encoder. """ def __init__( self, theta_dim, d_context=128, n_coupling_layers=6, hidden_dim=96, coupling_type='spline', n_bins=8, tail_bound=5.0, ): super().__init__() self.theta_dim = theta_dim self.coupling_type = coupling_type self.tail_bound = tail_bound self.flows = nn.ModuleList() for i in range(n_coupling_layers): mask_type = 'even' if i % 2 == 0 else 'odd' self.flows.append(ActNorm(theta_dim)) if coupling_type == 'spline': self.flows.append( ConditionalSplineCoupling( dim=theta_dim, d_context=d_context, hidden_dim=hidden_dim, mask_type=mask_type, n_bins=n_bins, tail_bound=tail_bound, ) ) else: self.flows.append( ConditionalAffineCoupling( dim=theta_dim, d_context=d_context, hidden_dim=hidden_dim, mask_type=mask_type, ) ) self.register_buffer('theta_mean', torch.zeros(theta_dim)) self.register_buffer('theta_std', torch.ones(theta_dim)) def set_theta_stats(self, mean, std): self.theta_mean.copy_(torch.as_tensor(mean, dtype=torch.float32)) self.theta_std.copy_(torch.as_tensor(std, dtype=torch.float32)) def normalize_theta(self, theta): return (theta - self.theta_mean) / self.theta_std def denormalize_theta(self, theta_norm): return theta_norm * self.theta_std + self.theta_mean def forward_flow(self, z, context): total_log_det = torch.zeros(z.shape[0], device=z.device) h = z for layer in self.flows: if isinstance(layer, ActNorm): h, ld = layer(h) total_log_det += ld else: h, ld = layer(h, context) total_log_det += ld return h, total_log_det def inverse_flow(self, theta_norm, context): total_log_det = torch.zeros(theta_norm.shape[0], device=theta_norm.device) h = theta_norm for layer in reversed(self.flows): if isinstance(layer, ActNorm): h = layer.inverse(h) total_log_det -= layer.log_scale.sum() else: h, ld = layer.inverse(h, context) total_log_det += ld return h, total_log_det def log_prob(self, theta, context): """Compute log p(theta | context) for this mechanism's parameters.""" theta_norm = self.normalize_theta(theta) if self.coupling_type == 'spline': theta_norm = theta_norm.clamp(-self.tail_bound, self.tail_bound) z, log_det = self.inverse_flow(theta_norm, context) log_pz = -0.5 * (z ** 2 + math.log(2 * math.pi)).sum(dim=-1) log_det_norm = -torch.log(self.theta_std).sum() log_p = log_pz + log_det + log_det_norm return log_p.clamp(min=-50.0, max=50.0) @torch.no_grad() def sample(self, context, n_samples=100, temperature=1.0): """Sample theta from this mechanism's posterior. Args: temperature: Posterior inflation factor. Can be: - scalar: uniform scaling for all parameters - 1-D tensor of shape [theta_dim]: per-parameter scaling T > 1 broadens the posterior by scaling samples around their per-example mean in theta-space. """ B = context.shape[0] context_rep = context.unsqueeze(1).expand(-1, n_samples, -1).reshape(B * n_samples, -1) z = torch.randn(B * n_samples, self.theta_dim, device=context.device) theta_norm, _ = self.forward_flow(z, context_rep) theta = self.denormalize_theta(theta_norm) theta = theta.reshape(B, n_samples, self.theta_dim) if isinstance(temperature, torch.Tensor): T = temperature.to(theta.device).reshape(1, 1, -1) mu = theta.mean(dim=1, keepdim=True) theta = mu + T * (theta - mu) elif temperature != 1.0: mu = theta.mean(dim=1, keepdim=True) theta = mu + temperature * (theta - mu) return theta def sample_with_grad(self, context, n_samples=64): """Sample theta with gradients enabled (for calibration loss). Uses the reparameterization trick: z ~ N(0,I) is fixed noise, gradients flow through the flow's forward transform. """ B = context.shape[0] context_rep = context.unsqueeze(1).expand(-1, n_samples, -1).reshape(B * n_samples, -1) z = torch.randn(B * n_samples, self.theta_dim, device=context.device) theta_norm, _ = self.forward_flow(z, context_rep) theta = self.denormalize_theta(theta_norm) return theta.reshape(B, n_samples, self.theta_dim) SUMMARY_DIM = 21 # 3 scan rates * 6 features + 3 log10(sigma) class SummaryProjection(nn.Module): """Project hand-crafted summary statistics to context space. Replaces the Set Transformer encoder for the Summary-ECFlow ablation, keeping everything else (classifier, flow heads, training) identical. """ def __init__(self, summary_dim=SUMMARY_DIM, d_context=128): super().__init__() self.net = nn.Sequential( nn.Linear(summary_dim, d_context), nn.GELU(), nn.Linear(d_context, d_context), nn.GELU(), nn.Linear(d_context, d_context), ) def forward(self, summary): """summary: [B, summary_dim] -> [B, d_context]""" return self.net(summary) class MultiMechanismFlow(nn.Module): """ Joint mechanism identification and parameter inference model. Combines: - Multi-scan-rate signal encoder (Set Transformer over per-CV embeddings) - Mechanism classifier - Per-mechanism normalizing flow heads If use_summary_features=True, replaces the signal encoder with a simple MLP projection from hand-crafted summary statistics (21-dim) to context space, keeping all other components identical. This isolates the effect of the input representation (full signal vs. summary statistics). """ def __init__( self, d_context=128, d_model=128, n_coupling_layers=6, hidden_dim=96, coupling_type='spline', n_bins=8, tail_bound=5.0, aggregation='set_transformer', use_summary_features=False, ): super().__init__() self.n_mechanisms = len(MECHANISM_LIST) self.mechanism_list = MECHANISM_LIST self.d_context = d_context self.use_summary_features = use_summary_features if use_summary_features: self.summary_proj = SummaryProjection( summary_dim=SUMMARY_DIM, d_context=d_context, ) self.encoder = None else: self.encoder = MultiScanEncoder( in_channels=3, d_model=d_model, d_context=d_context, aggregation=aggregation, ) self.summary_proj = None self.classifier = MechanismClassifier( d_context=d_context, n_mechanisms=self.n_mechanisms, hidden_dim=hidden_dim, ) # Mechanisms whose parameters are not identifiable from the signal # (e.g. Nernst: E0_offset=0, dA=1 are constants, dB is unidentifiable). # These get a flow head for architecture uniformity but are excluded # from the NLL loss during training. self.skip_nll_mechanisms: set = set() self.flow_heads = nn.ModuleDict() for mech in MECHANISM_LIST: theta_dim = MECHANISM_PARAMS[mech]['dim'] self.flow_heads[mech] = MechanismFlow( theta_dim=theta_dim, d_context=d_context, n_coupling_layers=n_coupling_layers, hidden_dim=hidden_dim, coupling_type=coupling_type, n_bins=n_bins, tail_bound=tail_bound, ) def set_theta_stats(self, mechanism, mean, std): """Set normalization stats for a specific mechanism's flow head.""" self.flow_heads[mechanism].set_theta_stats(mean, std) def encode_signal(self, x, scan_mask=None, sigmas=None, flux_scales=None, summary=None): """ Encode input into a context vector. In full-signal mode: uses Set Transformer encoder on x. In summary mode: projects 21-dim hand-crafted features to context space. Args: x: [B, N_sigma, 3, T] (ignored in summary mode) scan_mask: [B, N_sigma, T] sigmas: [B, N_sigma] log10 scan rates flux_scales: [B, N_sigma] log10(std(flux)) per CV summary: [B, 21] hand-crafted summary statistics (summary mode only) """ if self.use_summary_features: assert summary is not None, "summary features required in summary mode" return self.summary_proj(summary) return self.encoder(x, scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales) def forward(self, x, mechanism_ids, mech_theta, mech_theta_mask=None, scan_mask=None, sigmas=None, flux_scales=None, summary=None): """ Compute classification logits and per-sample NLL for the true mechanism. Returns: dict with 'logits' [B, n_mechanisms] and 'nll' [B] """ context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales, summary=summary) logits = self.classifier(context) nll = torch.zeros(x.shape[0], device=x.device) for m_idx, mech in enumerate(MECHANISM_LIST): sel = (mechanism_ids == m_idx) if not sel.any(): continue if mech in self.skip_nll_mechanisms: continue theta_dim = MECHANISM_PARAMS[mech]['dim'] ctx_m = context[sel] theta_m = mech_theta[sel, :theta_dim] log_p = self.flow_heads[mech].log_prob(theta_m, ctx_m) bad = ~torch.isfinite(log_p) if bad.any(): log_p = torch.where(bad, torch.full_like(log_p, -10.0).detach(), log_p) nll[sel] = -log_p return {'logits': logits, 'nll': nll} def forward_with_calibration(self, x, mechanism_ids, mech_theta, mech_theta_mask=None, scan_mask=None, sigmas=None, flux_scales=None, cal_n_samples=64, cal_levels=(0.5, 0.9), cal_beta=20.0, summary=None): """ Forward pass with additional calibration loss. Extends forward() by drawing posterior samples and computing a differentiable coverage penalty that encourages the flow's credible intervals to match their nominal levels. Returns: dict with 'logits', 'nll', and 'cal_loss' (scalar) """ context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales, summary=summary) logits = self.classifier(context) nll = torch.zeros(x.shape[0], device=x.device) cal_losses = [] for m_idx, mech in enumerate(MECHANISM_LIST): sel = (mechanism_ids == m_idx) if not sel.any(): continue if mech in self.skip_nll_mechanisms: continue theta_dim = MECHANISM_PARAMS[mech]['dim'] ctx_m = context[sel] theta_m = mech_theta[sel, :theta_dim] # Standard NLL log_p = self.flow_heads[mech].log_prob(theta_m, ctx_m) bad = ~torch.isfinite(log_p) if bad.any(): log_p = torch.where(bad, torch.full_like(log_p, -10.0).detach(), log_p) nll[sel] = -log_p # Calibration: draw samples WITH gradients if ctx_m.shape[0] < 4: continue samples = self.flow_heads[mech].sample_with_grad( ctx_m, n_samples=cal_n_samples, ) # [B_m, K, D] # Inverse-spread weights: collapsed parameters (small posterior std) # get more calibration pressure than well-spread parameters. with torch.no_grad(): param_std = samples.std(dim=1).clamp(min=1e-4) # [B_m, D] inv_spread_w = 1.0 / param_std # [B_m, D] inv_spread_w = inv_spread_w / inv_spread_w.mean() for level in cal_levels: alpha = (1.0 - level) / 2.0 lower = torch.quantile(samples, alpha, dim=1) # [B_m, D] upper = torch.quantile(samples, 1 - alpha, dim=1) # [B_m, D] inside = ( torch.sigmoid(cal_beta * (theta_m - lower)) * torch.sigmoid(cal_beta * (upper - theta_m)) ) # [B_m, D] per_sample_loss = (inside - level).pow(2) # [B_m, D] cal_losses.append((per_sample_loss * inv_spread_w).mean()) if cal_losses: cal_loss = torch.stack(cal_losses).mean() else: cal_loss = torch.tensor(0.0, device=x.device) return {'logits': logits, 'nll': nll, 'cal_loss': cal_loss} @torch.no_grad() def predict(self, x, scan_mask=None, sigmas=None, flux_scales=None, n_samples=200, top_k=None, temperature=1.0, temperature_map=None, summary=None): """ Full inference: classify mechanism, then sample parameters. Args: x: [B, N_sigma, 3, T] multi-scan input signals scan_mask: [B, N_sigma, T] sigmas: [B, N_sigma] flux_scales: [B, N_sigma] n_samples: posterior samples per mechanism top_k: if set, only sample from top-k most likely mechanisms temperature: scalar fallback (>1 broadens posteriors) temperature_map: dict mapping mechanism name -> list of per-parameter temperatures. Overrides scalar temperature for mechanisms present in the map. summary: [B, 21] hand-crafted summary stats (summary mode only) Returns: dict with mechanism_probs, mechanism_pred, samples, stats """ context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales, summary=summary) logits = self.classifier(context) probs = F.softmax(logits, dim=-1) pred = probs.argmax(dim=-1) probs_clamped = probs.clamp(min=1e-7, max=1 - 1e-7) xdB = 10.0 * torch.log10(probs_clamped / (1.0 - probs_clamped)) samples_dict = {} stats_dict = {} for m_idx, mech in enumerate(MECHANISM_LIST): if top_k is not None: top_k_mechs = probs.topk(top_k, dim=-1).indices if not (top_k_mechs == m_idx).any(): samples_dict[mech] = None stats_dict[mech] = None continue T = temperature if temperature_map is not None and mech in temperature_map: T = torch.tensor(temperature_map[mech], dtype=torch.float32) s = self.flow_heads[mech].sample(context, n_samples=n_samples, temperature=T) samples_dict[mech] = s stats_dict[mech] = { 'mean': s.mean(dim=1), 'std': s.std(dim=1), 'median': s.median(dim=1).values, 'q05': s.quantile(0.05, dim=1), 'q95': s.quantile(0.95, dim=1), } return { 'mechanism_probs': probs, 'mechanism_xdB': xdB, 'mechanism_pred': pred, 'samples': samples_dict, 'stats': stats_dict, } @torch.no_grad() def predict_single_mechanism(self, x, mechanism, scan_mask=None, sigmas=None, flux_scales=None, n_samples=1000, temperature=1.0, temperature_map=None, summary=None): """Sample parameters assuming a known mechanism.""" context = self.encode_signal(x, scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales, summary=summary) T = temperature if temperature_map is not None and mechanism in temperature_map: T = torch.tensor(temperature_map[mechanism], dtype=torch.float32) samples = self.flow_heads[mechanism].sample(context, n_samples=n_samples, temperature=T) return { 'mean': samples.mean(dim=1), 'std': samples.std(dim=1), 'median': samples.median(dim=1).values, 'q05': samples.quantile(0.05, dim=1), 'q95': samples.quantile(0.95, dim=1), 'samples': samples, } def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) if __name__ == "__main__": n_mechs = len(MECHANISM_LIST) B, N_sigma, T = n_mechs, 3, 800 x = torch.randn(B, N_sigma, 3, T) scan_mask = torch.ones(B, N_sigma, T, dtype=torch.bool) sigmas = torch.randn(B, N_sigma) flux_scales = torch.randn(B, N_sigma) mechanism_ids = torch.arange(n_mechs) max_dim = max(MECHANISM_PARAMS[m]['dim'] for m in MECHANISM_LIST) mech_theta = torch.randn(B, max_dim) mech_theta_mask = torch.zeros(B, max_dim, dtype=torch.bool) for i, mid in enumerate(mechanism_ids): d = MECHANISM_PARAMS[MECHANISM_LIST[mid]]['dim'] mech_theta_mask[i, :d] = True print("=" * 60) print("Testing MultiMechanismFlow (multi-scan-rate, Set Transformer)") print("=" * 60) model = MultiMechanismFlow( d_context=128, d_model=128, n_coupling_layers=8, hidden_dim=128, coupling_type='affine', ) total_params = count_parameters(model) print(f"Total parameters: {total_params:,}") print(f" Encoder: {count_parameters(model.encoder):,}") print(f" Classifier: {count_parameters(model.classifier):,}") for mech in MECHANISM_LIST: print(f" Flow ({mech}, dim={MECHANISM_PARAMS[mech]['dim']}): " f"{count_parameters(model.flow_heads[mech]):,}") out = model(x, mechanism_ids, mech_theta, mech_theta_mask, scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales) print(f"\nForward pass:") print(f" Logits shape: {out['logits'].shape}") print(f" NLL shape: {out['nll'].shape}") print(f" NLL values: {out['nll']}") pred = model.predict(x, scan_mask=scan_mask, sigmas=sigmas, flux_scales=flux_scales, n_samples=100) print(f"\nPrediction:") print(f" Mechanism probs shape: {pred['mechanism_probs'].shape}") print(f" Mechanism xdB shape: {pred['mechanism_xdB'].shape}") print(f" Predicted mechanisms: {pred['mechanism_pred']}") print(f" xdB for predicted mechanisms:") for i in range(B): m_idx = pred['mechanism_pred'][i].item() mech = MECHANISM_LIST[m_idx] xdB_val = pred['mechanism_xdB'][i, m_idx].item() print(f" sample {i}: {mech} ({xdB_val:+.1f} dB)") for mech in MECHANISM_LIST: if pred['samples'][mech] is not None: print(f" {mech} samples shape: {pred['samples'][mech].shape}")