""" xai_engine.py - Ensemble inference and XAI attribution for PoreGCN HF Space. Frontend contract (app.py imports): from xai_engine import ( load_ensemble, ensemble_predict, compute_attributions, classify_scenario, substructure_breakdown, ) Design: - load_ensemble() caches models per dataset at module level so repeated Gradio calls do not reload from disk. - ensemble_predict() runs all 5 models (consensus and uncertainty). - compute_attributions() uses the best model only (compute-efficient XAI). - classify_scenario() implements the Scenario A/B/C/D trust framework. - substructure_breakdown() splits attributions into metal, linker, and pore. XAI method: signed occlusion. For each atom (or pore): attribution_i = base_prediction - prediction_with_feature_i_zeroed Positive attribution: node drives prediction higher than baseline. Negative attribution: node drives prediction lower than baseline. Scenario framework (mirrors PoreGCN_unified/xai.py): Consensus = ensemble CV < CV_THRESHOLD (10%) Agreement = (>= AGREEMENT_THRESHOLD fraction of atoms in expected direction) OR (mean signed attribution in expected direction) A: consensus AND agreement (Trustworthy) B: consensus AND NOT agreement (Overconfident) C: NOT consensus AND agreement (Underconfident) D: NOT consensus AND NOT agreement (Unreliable) """ from __future__ import annotations import glob import json import logging import os from typing import Dict, List, Optional, Tuple import numpy as np import torch from build_graph import graph_to_tensors from config import ( AGREEMENT_THRESHOLD, CV_THRESHOLD, DATASETS, DEVICE, MODELS_DIR, MOF_METALS, ) from model import create_inference_model logger = logging.getLogger(__name__) # Inline Normalizer class (copied from data_module.py to avoid import chain) # The normalizer is stored in each checkpoint; we just need load_state_dict. class _Normalizer: def __init__(self): self.mean = 0.0 self.std = 1.0 def denorm(self, x: torch.Tensor) -> torch.Tensor: return x * self.std + self.mean def load_state_dict(self, d: Dict): self.mean = float(d['mean']) self.std = float(d['std']) class _MultiNormalizer: """Per-property normalizer loaded from checkpoint state_dict.""" def __init__(self): self.property_names: List[str] = [] self.normalizers: Dict[str, _Normalizer] = {} def denorm(self, normed: torch.Tensor) -> torch.Tensor: if normed.dim() == 1: normed = normed.unsqueeze(0) out = torch.zeros_like(normed) for i, name in enumerate(self.property_names): if name in self.normalizers and i < normed.shape[1]: n = self.normalizers[name] out[:, i] = normed[:, i] * n.std + n.mean return out def load_state_dict(self, d: Dict): self.property_names = d.get('property_names', []) self.normalizers = {} for k, v in d.get('normalizers', {}).items(): norm = _Normalizer() norm.load_state_dict(v) self.normalizers[k] = norm # ============================================================================= # Module-level cache: one entry per dataset # ============================================================================= _CACHE: Dict[str, Dict] = {} # Format: {dataset: {models, best_model, normalizers, best_normalizer, # property_names, property_means}} # ============================================================================= # load_ensemble # ============================================================================= def load_ensemble( dataset: str, device: str = DEVICE, ) -> Tuple[List, object, object, List[str]]: """ Load 5-model ensemble + best model for a given dataset. Checks module-level cache first; loads from disk only on the first call for each dataset. This is safe for Gradio's multi-request environment because model weights are read-only after loading. Ensemble directory structure expected: models/_ensemble/ fold_0_best.ckpt (or similar pattern) fold_1_best.ckpt ... property_means.json Each checkpoint must be a dict with either: - 'model_state' + 'model_config' + 'normalizer' (CV format) - 'model_state_dict' + 'config' + 'normalizer_state_dict' (final format) Args: dataset: One of DATASETS ('hmof_gas', 'core_mof', 'hmof_geometric') device: Target device string (always 'cpu' for HF free tier) Returns: (all_models, best_model, best_normalizer, property_names) all_models: list of 5 PoreGCN instances in eval mode best_model: the model with lowest val_loss (used for XAI) best_normalizer: _MultiNormalizer for the best model's checkpoint property_names: list of property strings in prediction order """ global _CACHE if dataset in _CACHE: entry = _CACHE[dataset] return ( entry['models'], entry['best_model'], entry['best_normalizer'], entry['property_names'], ) ensemble_dir = os.path.join(MODELS_DIR, f'{dataset}_ensemble') if not os.path.isdir(ensemble_dir): raise FileNotFoundError( f'Ensemble directory not found: {ensemble_dir}. ' f'Expected: models/{dataset}_ensemble/ with checkpoint .ckpt or .pt files.' ) # Gather checkpoint files ckpt_files = sorted( glob.glob(os.path.join(ensemble_dir, '*.ckpt')) + glob.glob(os.path.join(ensemble_dir, '*.pt')) ) if not ckpt_files: raise FileNotFoundError( f'No checkpoint files (.ckpt or .pt) found in {ensemble_dir}.' ) # Load all checkpoints and rank by val_loss loaded: List[Tuple[float, Dict, str]] = [] for path in ckpt_files: try: ckpt = torch.load(path, map_location='cpu', weights_only=False) val_loss = float(ckpt.get('val_loss', ckpt.get('val_mae', float('inf')))) loaded.append((val_loss, ckpt, path)) except Exception as exc: logger.warning('Failed to load checkpoint %s: %s', path, exc) if not loaded: raise RuntimeError(f'Could not load any checkpoints from {ensemble_dir}.') loaded.sort(key=lambda x: x[0]) top5 = loaded[:5] logger.info( 'Loaded %d checkpoints for %s, using top %d.', len(loaded), dataset, len(top5), ) # Instantiate models models: List = [] normalizers: List[_MultiNormalizer] = [] property_names_from_ckpt: Optional[List[str]] = None for val_loss, ckpt, path in top5: try: model = create_inference_model(ckpt, device=device) models.append(model) # Load normalizer normalizer = _MultiNormalizer() norm_sd = ckpt.get('normalizer', ckpt.get('normalizer_state_dict', None)) if norm_sd is not None: normalizer.load_state_dict(norm_sd) normalizers.append(normalizer) # Capture property names from first successful checkpoint if property_names_from_ckpt is None: cfg = ckpt.get('model_config', ckpt.get('config', {})) property_names_from_ckpt = cfg.get('property_names', None) if property_names_from_ckpt is None: property_names_from_ckpt = normalizer.property_names or [] except Exception as exc: logger.warning('Skipping checkpoint %s: %s', path, exc) if not models: raise RuntimeError(f'All checkpoints in {ensemble_dir} failed to load.') best_model = models[0] best_normalizer = normalizers[0] property_names = property_names_from_ckpt or [] # Load property means (for XAI direction logic) means_path = os.path.join(ensemble_dir, 'property_means.json') property_means: Dict[str, float] = {} if os.path.exists(means_path): with open(means_path) as f: property_means = json.load(f) else: logger.warning( 'property_means.json not found in %s. ' 'XAI direction checks will use 0.0 as default mean.', ensemble_dir, ) _CACHE[dataset] = { 'models': models, 'normalizers': normalizers, 'best_model': best_model, 'best_normalizer': best_normalizer, 'property_names': property_names, 'property_means': property_means, } logger.info( 'Ensemble for %s ready: %d models, %d properties.', dataset, len(models), len(property_names), ) return models, best_model, best_normalizer, property_names # ============================================================================= # ensemble_predict # ============================================================================= def ensemble_predict( graph: Dict, models: List, normalizer, device: str = DEVICE, ) -> Dict: """ Run all ensemble models on a single graph. Args: graph: Graph dict returned by cif_to_graph(). models: List of PoreGCN instances (from load_ensemble()). normalizer: _MultiNormalizer for the best model; used to denorm outputs. All models in the ensemble share the same normalizer because they were trained on the same data split. device: Target device (always 'cpu'). Returns: Dict keyed by property name: {prop: {'mean': float, 'std': float, 'cv': float}} Values are in physical (denormalized) units. """ tensors = graph_to_tensors(graph, device=device) all_preds: List[np.ndarray] = [] for model in models: model.eval() with torch.no_grad(): raw = model( tensors['atom_fea'], tensors['bond_types'], tensors['nbr_fea_idx'], tensors['crystal_atom_idx'], tensors['pore_fea'], tensors['atom_pore_edges'], tensors['crystal_pore_idx'], ) # [1, n_props] denormed = normalizer.denorm(raw.cpu()) # [1, n_props] all_preds.append(denormed[0].numpy()) all_preds_arr = np.array(all_preds) # [n_models, n_props] means = all_preds_arr.mean(axis=0) stds = all_preds_arr.std(axis=0) property_names = normalizer.property_names results: Dict[str, Dict[str, float]] = {} for i, prop in enumerate(property_names): mean_val = float(means[i]) std_val = float(stds[i]) cv = abs(std_val / mean_val) if abs(mean_val) > 1e-8 else float(std_val) results[prop] = {'mean': mean_val, 'std': std_val, 'cv': cv} return results # ============================================================================= # compute_attributions # ============================================================================= def compute_attributions( graph: Dict, best_model, normalizer, property_means: Dict[str, float], target_prop: str, device: str = DEVICE, xai_method: str = 'fast', progress_callback=None, ) -> Dict: """ Signed attribution per atom AND per pore for one target property. Two methods, selectable via `xai_method`: 'fast' (default, ~5 s) — gradient x input. One forward + one backward pass; per-atom score = sum_f (atom_fea[i,f] * d pred / d atom_fea[i,f]). Approximates signed occlusion in sign for most atoms. 'full' (~3 min for 200 atoms) — signed occlusion. For each atom_i, zero atom_fea[i] and rerun the model; for each pore_j, zero pore_fea[j] and remove its edges; signed attribution = base_pred - masked_pred. This is the exact method reported in the manuscript and used to generate xcomprehensive_xai_results/. `progress_callback(step, total, message)` is called periodically when xai_method='full' so the UI can render a progress bar; ignored when xai_method='fast'. Args: graph: Graph dict from cif_to_graph(). best_model: PoreGCN instance (lowest val_loss). normalizer: _MultiNormalizer for denormalization. property_means: Dict {prop_name: training_mean} loaded from property_means.json in the ensemble dir. target_prop: Property name string to explain. device: Always 'cpu'. Returns: Dict: per_atom np.ndarray [N_atoms] signed attributions per_pore np.ndarray [N_pores] signed attributions expected_direction '+' if prediction > mean, '-' otherwise agreement_frac float in [0, 1] (fraction of attrs with expected sign) """ tensors = graph_to_tensors(graph, device=device) property_names = normalizer.property_names if target_prop not in property_names: raise ValueError( f'Property "{target_prop}" not in model property list: {property_names}' ) target_idx = property_names.index(target_prop) best_model.eval() n_atoms = graph['n_atoms'] n_pores = graph['n_pores'] def _forward() -> float: with torch.no_grad(): raw = best_model( tensors['atom_fea'], tensors['bond_types'], tensors['nbr_fea_idx'], tensors['crystal_atom_idx'], tensors['pore_fea'], tensors['atom_pore_edges'], tensors['crystal_pore_idx'], ) return float(normalizer.denorm(raw.cpu())[0, target_idx].item()) base_pred = _forward() # ----- xai_method='full' (signed occlusion, manuscript method) ----- if xai_method == 'full': atom_signed = np.zeros(n_atoms, dtype=np.float64) # Atom occlusion: mask one atom at a time, measure prediction change orig_atom = tensors['atom_fea'].clone() for i in range(n_atoms): masked = orig_atom.clone() masked[i] = 0.0 tensors['atom_fea'] = masked masked_pred = _forward() atom_signed[i] = base_pred - masked_pred if progress_callback is not None and (i + 1) % 10 == 0: progress_callback(i + 1, n_atoms + n_pores, f'Occlusion: atom {i+1}/{n_atoms}') tensors['atom_fea'] = orig_atom # Pore occlusion: zero pore feature and remove its edges pore_signed = np.zeros(n_pores, dtype=np.float64) if n_pores > 0: orig_pore = tensors['pore_fea'].clone() orig_edges = tensors['atom_pore_edges'] for j in range(n_pores): masked_pore = orig_pore.clone() masked_pore[j] = 0.0 if orig_edges.numel() > 0: keep = orig_edges[1] != j masked_edges = orig_edges[:, keep] else: masked_edges = orig_edges tensors['pore_fea'] = masked_pore tensors['atom_pore_edges'] = masked_edges masked_pred = _forward() pore_signed[j] = base_pred - masked_pred if progress_callback is not None and (j + 1) % 10 == 0: progress_callback(n_atoms + j + 1, n_atoms + n_pores, f'Occlusion: pore {j+1}/{n_pores}') tensors['pore_fea'] = orig_pore tensors['atom_pore_edges'] = orig_edges prop_mean = property_means.get(target_prop, 0.0) expected_positive = base_pred > prop_mean expected_direction = '+' if expected_positive else '-' all_signed = np.concatenate([atom_signed, pore_signed]) if n_pores > 0 else atom_signed n_total = len(all_signed) if n_total > 0: n_matching = int(np.sum(all_signed > 0)) if expected_positive \ else int(np.sum(all_signed < 0)) agreement_frac = n_matching / n_total mean_signed = float(np.mean(all_signed)) else: agreement_frac = 0.0 mean_signed = 0.0 return { 'per_atom': atom_signed.astype(np.float32), 'per_pore': pore_signed.astype(np.float32), 'expected_direction': expected_direction, 'agreement_frac': float(agreement_frac), 'mean_signed': mean_signed, 'expected_positive': bool(expected_positive), } # ----- xai_method='fast' (gradient x input surrogate, default) ----- # One forward + one backward pass. # Per-atom score = sum_f (atom_fea_{i,f} * d pred / d atom_fea_{i,f}). # Sign convention preserved: positive = atom drives prediction up. orig_atom_fea = tensors['atom_fea'].clone() atom_fea_grad = orig_atom_fea.detach().clone().requires_grad_(True) if n_pores > 0: pore_fea_grad = tensors['pore_fea'].detach().clone().requires_grad_(True) else: pore_fea_grad = tensors['pore_fea'] # Forward with grad enabled with torch.enable_grad(): raw_pred = best_model( atom_fea_grad, tensors['bond_types'], tensors['nbr_fea_idx'], tensors['crystal_atom_idx'], pore_fea_grad, tensors['atom_pore_edges'], tensors['crystal_pore_idx'], ) # Use raw (normalised) target as the scalar to backprop through. Denormalisation # is a constant scale + shift so it does not affect the relative attribution sign/rank. target_scalar = raw_pred[0, target_idx] target_scalar.backward() atom_signed = (atom_fea_grad.detach() * atom_fea_grad.grad).sum(dim=1).cpu().numpy().astype(np.float64) if n_pores > 0 and pore_fea_grad.grad is not None: pore_signed_grad = (pore_fea_grad.detach() * pore_fea_grad.grad).sum(dim=1).cpu().numpy().astype(np.float64) else: pore_signed_grad = None # Restore tensors['atom_fea'] = orig_atom_fea # --- Pore attributions (also via gradient x input from the same backward pass) --- pore_signed = np.zeros(n_pores, dtype=np.float64) if n_pores > 0 and pore_signed_grad is not None: pore_signed = pore_signed_grad # Agreement check: fraction of attributions in expected direction prop_mean = property_means.get(target_prop, 0.0) expected_positive = base_pred > prop_mean expected_direction = '+' if expected_positive else '-' all_signed = np.concatenate([atom_signed, pore_signed]) if n_pores > 0 else atom_signed n_total = len(all_signed) if n_total > 0: if expected_positive: n_matching = int(np.sum(all_signed > 0)) else: n_matching = int(np.sum(all_signed < 0)) agreement_frac = n_matching / n_total mean_signed = float(np.mean(all_signed)) else: agreement_frac = 0.0 mean_signed = 0.0 return { 'per_atom': atom_signed.astype(np.float32), 'per_pore': pore_signed.astype(np.float32), 'expected_direction': expected_direction, 'agreement_frac': float(agreement_frac), 'mean_signed': mean_signed, 'expected_positive': bool(expected_positive), } # ============================================================================= # classify_scenario # ============================================================================= def classify_scenario( prediction_mean: float, prediction_std: float, agreement_frac: float, property_mean: float, mean_signed: float = 0.0, expected_positive: bool = None, ) -> Tuple[str, str]: """ Classify prediction trustworthiness into Scenario A/B/C/D. Mirrors the dual-criteria framework in PoreGCN/xai.py exactly: Consensus: CV < CV_THRESHOLD (10%) Agreement (EITHER satisfies): Criterion 1: agreement_frac >= AGREEMENT_THRESHOLD (70%) Criterion 2: mean signed attribution has the expected direction Earlier versions of this function used the relaxed fallback `agreement_frac >= 0.5` for criterion 2, which silently flipped Scenario A/B/C/D classifications relative to the manuscript on structures whose attributions are bipolar (e.g. signed-occlusion output where some atoms increase the prediction and others decrease it). Fixed to use the actual mean-signed value. Args: prediction_mean: Ensemble mean prediction (physical units). prediction_std: Ensemble standard deviation. agreement_frac: Fraction of attributions in the expected direction. property_mean: Training-set mean of this property. mean_signed: Mean of the signed attributions (for criterion 2). expected_positive: If None, derived from prediction_mean > property_mean. Returns: (scenario_letter, human_label) """ cv = abs(prediction_std / prediction_mean) if abs(prediction_mean) > 1e-8 else abs(prediction_std) high_consensus = cv < CV_THRESHOLD if expected_positive is None: expected_positive = prediction_mean > property_mean # Manuscript dual criteria, exactly as in PoreGCN/xai.py check_xai_agreement ratio_ok = agreement_frac >= AGREEMENT_THRESHOLD if expected_positive: mean_dir_ok = mean_signed > 0.0 else: mean_dir_ok = mean_signed < 0.0 is_agreed = ratio_ok or mean_dir_ok scenario_labels = { 'A': 'Trustworthy (high consensus, XAI agrees)', 'B': 'Overconfident (high consensus, XAI disagrees)', 'C': 'Underconfident (low consensus, XAI agrees)', 'D': 'Unreliable (low consensus, XAI disagrees)', } if high_consensus and is_agreed: scenario = 'A' elif high_consensus and not is_agreed: scenario = 'B' elif not high_consensus and is_agreed: scenario = 'C' else: scenario = 'D' return scenario, scenario_labels[scenario] # ============================================================================= # substructure_breakdown # ============================================================================= def substructure_breakdown( graph: Dict, per_atom_attrs: np.ndarray, per_pore_attrs: np.ndarray, ) -> Dict: """ Split total attribution into metal, linker, and pore fractions. Metal/linker classification uses element symbols decoded from the one-hot portion of atom_fea (first 92 dimensions = element 1..92). MOF_METALS set from config.py covers all common MOF metal centres. Args: graph: Graph dict from cif_to_graph() (must contain 'structure'). per_atom_attrs: np.ndarray [N_atoms] signed attributions. per_pore_attrs: np.ndarray [N_pores] signed attributions. Returns: Dict: metal_frac float |metal attributions| / total |attributions| linker_frac float |linker attributions| / total |attributions| pore_frac float |pore attributions| / total |attributions| """ structure = graph.get('structure') n_atoms = graph['n_atoms'] if structure is not None and len(structure) == n_atoms: elements = [str(site.specie) for site in structure] else: # Fallback: decode from one-hot features ELEMENT_NAMES = [ 'H', 'He', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Ne', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'Ar', 'K', 'Ca', 'Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'I', 'Xe', 'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy', 'Ho', 'Er', 'Tm', 'Yb', 'Lu', 'Hf', 'Ta', 'W', 'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi', 'Po', 'At', 'Rn', 'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U', ] atom_fea = graph['atom_fea'] # [N_atoms, 120] elements = [] for i in range(n_atoms): one_hot = atom_fea[i, :92] idx = int(np.argmax(one_hot)) elements.append(ELEMENT_NAMES[idx] if idx < len(ELEMENT_NAMES) else 'X') abs_atom = np.abs(per_atom_attrs) abs_pore = np.abs(per_pore_attrs) metal_sum = sum(abs_atom[i] for i, e in enumerate(elements) if e in MOF_METALS) linker_sum = sum(abs_atom[i] for i, e in enumerate(elements) if e not in MOF_METALS) pore_sum = float(np.sum(abs_pore)) total = metal_sum + linker_sum + pore_sum + 1e-12 return { 'metal_frac': float(metal_sum / total), 'linker_frac': float(linker_sum / total), 'pore_frac': float(pore_sum / total), }