""" model.py - Inference-only PoreGCN architecture for HuggingFace Space. Stripped from PoreGCN_unified/model.py: - Training methods removed (no forward_with_porosity, no training utilities). - Porosity head kept because checkpoints may include those weights. - create_inference_model() is the single entry point used by xai_engine.py. Architecture summary: - BondEmbedding: learnable edge-type embeddings (5 types) - AtomConvLayer: CGCNN-style atom-atom message passing - AtomPoreConvLayer: heterogeneous atom<->pore message passing (KEY INNOVATION) - PoreGCN: stacks the above, pools atoms and pores, predicts MTL targets Frontend contract: this module is not imported directly by app.py. xai_engine.py calls create_inference_model() internally. """ from __future__ import annotations from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F # Properties that receive a Softplus non-negativity constraint. # log10_CO2_N2_selectivity is intentionally absent (log10 values can be negative). DEFAULT_NONNEG_PROPERTIES = { 'ASA', 'GSA', 'VF', 'LCD', 'PLD', 'density', 'CO2_0.01bar_298K', 'CO2_0.1bar_298K', 'CO2_0.15bar_298K', 'CO2_0.5bar_298K', 'CO2_1bar_298K', 'CH4_0.9bar_298K', 'CH4_2.5bar_298K', 'CH4_35bar_298K', 'N2_0.09bar_298K', 'N2_0.9bar_298K', 'H2_2bar_77K', 'H2_100bar_77K', 'Xe_1bar_273K', 'Xe_10bar_273K', 'thermal_stability', } POROSITY_GATED_PROPERTIES = {'ASA', 'GSA'} # ============================================================================= # BondEmbedding # ============================================================================= class BondEmbedding(nn.Module): """Learnable embedding for 5 bond/edge types including ATOM_PORE (type 4).""" def __init__(self, num_edge_types: int = 5, embedding_dim: int = 16): super().__init__() self.embedding_dim = embedding_dim self.embedding = nn.Embedding(num_edge_types, embedding_dim, padding_idx=0) def forward(self, edge_types: torch.LongTensor) -> torch.Tensor: return self.embedding(edge_types) # ============================================================================= # AtomConvLayer # ============================================================================= class AtomConvLayer(nn.Module): """CGCNN-style graph convolutional layer for atom-atom interactions.""" def __init__(self, atom_fea_len: int, bond_fea_len: int): super().__init__() self.atom_fea_len = atom_fea_len self.bond_fea_len = bond_fea_len self.fc_full = nn.Linear(2 * atom_fea_len + bond_fea_len, 2 * atom_fea_len) self.sigmoid = nn.Sigmoid() self.softplus1 = nn.Softplus() self.bn1 = nn.BatchNorm1d(2 * atom_fea_len) self.bn2 = nn.BatchNorm1d(atom_fea_len) self.softplus2 = nn.Softplus() def forward( self, atom_fea: torch.Tensor, bond_fea: torch.Tensor, nbr_fea_idx: torch.LongTensor, ) -> torch.Tensor: """ Args: atom_fea: [N_atoms, atom_fea_len] bond_fea: [N_atoms, M_neighbors, bond_fea_len] nbr_fea_idx: [N_atoms, M_neighbors] Returns: Updated atom features [N_atoms, atom_fea_len] """ N, M = nbr_fea_idx.shape atom_nbr_fea = atom_fea[nbr_fea_idx, :] # [N, M, atom_fea_len] total_nbr_fea = torch.cat([ atom_fea.unsqueeze(1).expand(N, M, self.atom_fea_len), atom_nbr_fea, bond_fea, ], dim=2) # [N, M, 2*atom_fea_len + bond_fea_len] total_gated_fea = self.fc_full(total_nbr_fea) total_gated_fea = self.bn1( total_gated_fea.view(-1, self.atom_fea_len * 2) ).view(N, M, self.atom_fea_len * 2) nbr_filter, nbr_core = total_gated_fea.chunk(2, dim=2) nbr_filter = self.sigmoid(nbr_filter) nbr_core = self.softplus1(nbr_core) nbr_sumed = torch.sum(nbr_filter * nbr_core, dim=1) nbr_sumed = self.bn2(nbr_sumed) return self.softplus2(atom_fea + nbr_sumed) # ============================================================================= # AtomPoreConvLayer # ============================================================================= class AtomPoreConvLayer(nn.Module): """ Heterogeneous convolution layer for atom-pore interactions. Key innovation of PoreGCN: atoms receive messages from adjacent pore nodes, and pore nodes receive messages from surrounding atoms. Attention aggregation. """ def __init__(self, atom_fea_len: int, pore_fea_len: int, edge_fea_len: int = 16): super().__init__() self.atom_fea_len = atom_fea_len self.pore_fea_len = pore_fea_len self.edge_fea_len = edge_fea_len self.atom_from_pore = nn.Sequential( nn.Linear(atom_fea_len + pore_fea_len + edge_fea_len, atom_fea_len * 2), nn.LayerNorm(atom_fea_len * 2), nn.ReLU(), nn.Linear(atom_fea_len * 2, atom_fea_len), ) self.pore_from_atom = nn.Sequential( nn.Linear(pore_fea_len + atom_fea_len + edge_fea_len, pore_fea_len * 2), nn.LayerNorm(pore_fea_len * 2), nn.ReLU(), nn.Linear(pore_fea_len * 2, pore_fea_len), ) self.attn_atom = nn.Linear(atom_fea_len + pore_fea_len, 1) self.attn_pore = nn.Linear(pore_fea_len + atom_fea_len, 1) def forward( self, atom_fea: torch.Tensor, pore_fea: torch.Tensor, atom_pore_edges: torch.LongTensor, edge_fea: torch.Tensor, atom_pore_batch: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Args: atom_fea: [N_atoms, atom_fea_len] pore_fea: [N_pores, pore_fea_len] atom_pore_edges: [2, N_edges] where [0]=atom_idx, [1]=pore_idx edge_fea: [N_edges, edge_fea_len] Returns: (updated_atom_fea, updated_pore_fea) """ if atom_pore_edges.numel() == 0: return atom_fea, pore_fea atom_idx = atom_pore_edges[0] pore_idx = atom_pore_edges[1] atom_on_edge = atom_fea[atom_idx] pore_on_edge = pore_fea[pore_idx] # Atoms receive messages from pores atom_pore_cat = torch.cat([atom_on_edge, pore_on_edge, edge_fea], dim=1) atom_messages = self.atom_from_pore(atom_pore_cat) attn_a = torch.softmax(self.attn_atom(torch.cat([atom_on_edge, pore_on_edge], dim=1)), dim=0) atom_messages = atom_messages * attn_a atom_update = torch.zeros_like(atom_fea) atom_update.scatter_add_(0, atom_idx.unsqueeze(1).expand(-1, self.atom_fea_len), atom_messages) atom_fea_new = F.relu(atom_fea + atom_update) # Pores receive messages from atoms pore_atom_cat = torch.cat([pore_on_edge, atom_on_edge, edge_fea], dim=1) pore_messages = self.pore_from_atom(pore_atom_cat) attn_p = torch.softmax(self.attn_pore(torch.cat([pore_on_edge, atom_on_edge], dim=1)), dim=0) pore_messages = pore_messages * attn_p pore_update = torch.zeros_like(pore_fea) pore_update.scatter_add_(0, pore_idx.unsqueeze(1).expand(-1, self.pore_fea_len), pore_messages) pore_fea_new = F.relu(pore_fea + pore_update) return atom_fea_new, pore_fea_new # ============================================================================= # PoreGCN # ============================================================================= class PoreGCN(nn.Module): """ Pore-Aware Graph Convolutional Network for MOF property prediction. Forward signature matches what xai_engine.py and build_graph.py produce (single-graph, not batched via DataLoader). """ def __init__( self, orig_atom_fea_len: int, orig_pore_fea_len: int = 8, num_bond_types: int = 5, bond_embedding_dim: int = 16, atom_fea_len: int = 64, pore_fea_len: int = 32, n_conv: int = 3, h_fea_len: int = 128, n_h: int = 1, dropout: float = 0.0, num_targets: int = 1, property_names: Optional[List[str]] = None, use_nonneg_output: bool = True, use_porosity_head: bool = False, porosity_gate_weight: float = 1.0, ): super().__init__() self.atom_fea_len = atom_fea_len self.pore_fea_len = pore_fea_len self.n_conv = n_conv self.num_targets = num_targets self.use_nonneg_output = use_nonneg_output self.use_porosity_head = use_porosity_head self.porosity_gate_weight = porosity_gate_weight self.property_names = property_names or [f'prop_{i}' for i in range(num_targets)] self.nonneg_mask = [name in DEFAULT_NONNEG_PROPERTIES for name in self.property_names] self.porosity_gate_mask = [name in POROSITY_GATED_PROPERTIES for name in self.property_names] self.atom_embedding = nn.Linear(orig_atom_fea_len, atom_fea_len) self.pore_embedding = nn.Linear(orig_pore_fea_len, pore_fea_len) self.bond_embedding = BondEmbedding(num_bond_types, bond_embedding_dim) self.atom_convs = nn.ModuleList([ AtomConvLayer(atom_fea_len, bond_embedding_dim) for _ in range(n_conv) ]) self.atom_pore_convs = nn.ModuleList([ AtomPoreConvLayer(atom_fea_len, pore_fea_len, bond_embedding_dim) for _ in range(n_conv) ]) combined_fea_len = atom_fea_len + pore_fea_len if n_h > 0: layers: list = [nn.Linear(combined_fea_len, h_fea_len), nn.ReLU()] if dropout > 0: layers.append(nn.Dropout(dropout)) for _ in range(n_h - 1): layers.extend([nn.Linear(h_fea_len, h_fea_len), nn.ReLU()]) if dropout > 0: layers.append(nn.Dropout(dropout)) layers.append(nn.Linear(h_fea_len, num_targets)) self.fc_out = nn.Sequential(*layers) else: self.fc_out = nn.Linear(combined_fea_len, num_targets) self.softplus = nn.Softplus(beta=1.0) if use_porosity_head: self.porosity_head = nn.Sequential( nn.Linear(combined_fea_len, h_fea_len), nn.ReLU(), nn.Dropout(dropout) if dropout > 0 else nn.Identity(), nn.Linear(h_fea_len, 1), nn.Sigmoid(), ) else: self.porosity_head = None def forward( self, atom_fea: torch.Tensor, bond_types: torch.LongTensor, nbr_fea_idx: torch.LongTensor, crystal_atom_idx: List[torch.LongTensor], pore_fea: torch.Tensor, atom_pore_edges: torch.LongTensor, crystal_pore_idx: List[torch.LongTensor], ) -> torch.Tensor: """ Args: atom_fea: [N_atoms, orig_atom_fea_len] bond_types: [N_atoms, M_neighbors] nbr_fea_idx: [N_atoms, M_neighbors] crystal_atom_idx: list of LongTensors, one per crystal in batch pore_fea: [N_pores, orig_pore_fea_len] (may be empty) atom_pore_edges: [2, N_ap_edges] crystal_pore_idx: list of LongTensors, one per crystal Returns: Predictions [batch_size, num_targets] """ atom_fea = self.atom_embedding(atom_fea) pore_fea = self.pore_embedding(pore_fea) bond_fea = self.bond_embedding(bond_types) n_ap_edges = atom_pore_edges.shape[1] if atom_pore_edges.numel() > 0 else 0 if n_ap_edges > 0: ap_types = torch.full((n_ap_edges,), 4, dtype=torch.long, device=atom_fea.device) ap_fea = self.bond_embedding(ap_types) else: ap_fea = torch.zeros((0, self.bond_embedding.embedding_dim), device=atom_fea.device) for i in range(self.n_conv): atom_fea = self.atom_convs[i](atom_fea, bond_fea, nbr_fea_idx) if pore_fea.shape[0] > 0: atom_fea, pore_fea = self.atom_pore_convs[i]( atom_fea, pore_fea, atom_pore_edges, ap_fea ) atom_pool = self._pool_crystals(atom_fea, crystal_atom_idx) if pore_fea.shape[0] > 0: pore_pool = self._pool_crystals(pore_fea, crystal_pore_idx) else: batch_size = len(crystal_atom_idx) pore_pool = torch.zeros((batch_size, self.pore_fea_len), device=atom_fea.device) combined = torch.cat([atom_pool, pore_pool], dim=1) out = self.fc_out(combined) if self.use_nonneg_output and any(self.nonneg_mask): cols = [] for i, is_nonneg in enumerate(self.nonneg_mask): cols.append(self.softplus(out[:, i:i+1]) if is_nonneg else out[:, i:i+1]) out = torch.cat(cols, dim=1) if self.porosity_head is not None: porosity_prob = self.porosity_head(combined) gate = 1.0 - self.porosity_gate_weight + self.porosity_gate_weight * porosity_prob.squeeze(-1) cols = [] for i, is_gated in enumerate(self.porosity_gate_mask): cols.append((out[:, i] * gate).unsqueeze(1) if is_gated else out[:, i:i+1]) out = torch.cat(cols, dim=1) return out def _pool_crystals( self, fea: torch.Tensor, crystal_idx: List[torch.LongTensor], ) -> torch.Tensor: batch_size = len(crystal_idx) fea_dim = fea.shape[1] pooled = torch.zeros((batch_size, fea_dim), device=fea.device) for i, idx in enumerate(crystal_idx): if len(idx) > 0: pooled[i] = fea[idx].mean(dim=0) return pooled # ============================================================================= # Public entry point # ============================================================================= def _infer_arch(state_dict: Dict) -> Dict: """ Infer all PoreGCN hyperparameters from a saved state_dict. This avoids needing to pass a config object at inference time, which is important because the checkpoint was saved with a training config that may not be importable from the HF Space environment. """ orig_atom_fea_len = state_dict['atom_embedding.weight'].shape[1] orig_pore_fea_len = state_dict['pore_embedding.weight'].shape[1] atom_fea_len = state_dict['atom_embedding.weight'].shape[0] pore_fea_len = state_dict['pore_embedding.weight'].shape[0] n_conv = sum( 1 for k in state_dict if k.startswith('atom_convs.') and k.endswith('.fc_full.weight') ) fc_indices = sorted( int(k.split('.')[1]) for k in state_dict if k.startswith('fc_out.') and k.endswith('.weight') ) if not fc_indices: raise ValueError('Cannot infer architecture: fc_out weights missing from state_dict') h_fea_len = state_dict['fc_out.0.weight'].shape[0] last_idx = fc_indices[-1] num_targets = state_dict[f'fc_out.{last_idx}.weight'].shape[0] n_h = len(fc_indices) - 1 gap = fc_indices[1] - fc_indices[0] if len(fc_indices) > 1 else 1 dropout = 0.1 if gap >= 3 else 0.0 has_porosity_head = any('porosity_head' in k for k in state_dict) # Bond embedding dim bond_embedding_dim = state_dict['bond_embedding.embedding.weight'].shape[1] return { 'orig_atom_fea_len': orig_atom_fea_len, 'orig_pore_fea_len': orig_pore_fea_len, 'atom_fea_len': atom_fea_len, 'pore_fea_len': pore_fea_len, 'n_conv': n_conv, 'h_fea_len': h_fea_len, 'n_h': n_h, 'dropout': dropout, # Must match training value to reproduce fc_out Sequential structure 'num_targets': num_targets, 'bond_embedding_dim': bond_embedding_dim, 'has_porosity_head': has_porosity_head, } def create_inference_model(checkpoint_dict: Dict, device: str = 'cpu') -> PoreGCN: """ Instantiate an eval-mode PoreGCN from a loaded checkpoint dict. Accepts two checkpoint formats produced by PoreGCN_unified training: - CV checkpoint format: keys 'model_state', 'model_config' - Final model format: keys 'model_state_dict', 'config' Args: checkpoint_dict: dict loaded by torch.load() device: target device string (always 'cpu' on HF free tier) Returns: PoreGCN in eval mode, weights loaded, moved to device """ # Resolve state_dict key (two checkpoint formats exist) if 'model_state' in checkpoint_dict: state_dict = checkpoint_dict['model_state'] model_cfg = checkpoint_dict.get('model_config', {}) elif 'model_state_dict' in checkpoint_dict: state_dict = checkpoint_dict['model_state_dict'] model_cfg = checkpoint_dict.get('config', {}) else: raise KeyError( 'Checkpoint must have key "model_state" or "model_state_dict". ' f'Got: {list(checkpoint_dict.keys())}' ) arch = _infer_arch(state_dict) property_names = model_cfg.get('property_names', None) use_nonneg = model_cfg.get('use_nonneg_output', True) model = PoreGCN( orig_atom_fea_len=arch['orig_atom_fea_len'], orig_pore_fea_len=arch['orig_pore_fea_len'], num_bond_types=5, bond_embedding_dim=arch['bond_embedding_dim'], atom_fea_len=arch['atom_fea_len'], pore_fea_len=arch['pore_fea_len'], n_conv=arch['n_conv'], h_fea_len=arch['h_fea_len'], n_h=arch['n_h'], dropout=arch['dropout'], num_targets=arch['num_targets'], property_names=property_names, use_nonneg_output=use_nonneg, use_porosity_head=arch['has_porosity_head'], ) model.load_state_dict(state_dict, strict=True) model.to(device) model.eval() return model