| """ |
| 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 |
|
|
| |
| |
| 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'} |
|
|
|
|
| |
| |
| |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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, :] |
|
|
| total_nbr_fea = torch.cat([ |
| atom_fea.unsqueeze(1).expand(N, M, self.atom_fea_len), |
| atom_nbr_fea, |
| bond_fea, |
| ], dim=2) |
|
|
| 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) |
|
|
|
|
| |
| |
| |
|
|
| 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] |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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 |
|
|
|
|
| |
| |
| |
|
|
| 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 = 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, |
| '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 |
| """ |
| |
| 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 |
|
|
