Spaces:
Paused
Paused
File size: 4,905 Bytes
2fd4f2f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """
Multi-Model Ensemble & Uncertainty Quantification Engine.
Executes consensus predictions across GCN, GAT, GIN, MPNN, and Baseline models
and quantifies epistemic uncertainty via standard deviation and confidence intervals.
"""
from typing import Dict, List, Any, Tuple, Optional
import os
import numpy as np
import torch
from src.data import smiles_to_graph, DATASET_CONFIGS
from src.evaluate import load_trained_model
MODEL_NAMES = ['gin', 'mpnn', 'gcn', 'gat', 'baseline']
def predict_single_model(model: Any, model_name: str, graph: Any, dataset_name: str) -> np.ndarray:
"""Run inference for a single model and return output array (shape: [num_tasks])."""
if model_name == 'baseline':
pred = model.predict([graph])[0]
if np.ndim(pred) == 0:
pred = np.array([float(pred)])
return np.asarray(pred, dtype=np.float32)
model.eval()
with torch.no_grad():
out = model(graph.x, graph.edge_index, edge_attr=graph.edge_attr)
if dataset_name == 'esol':
val = out[0].cpu().numpy()
else:
# Classification or multi-task: convert logits to probabilities
val = torch.sigmoid(out)[0].cpu().numpy()
return np.asarray(val, dtype=np.float32)
def predict_ensemble(
smiles_or_graph: Any,
dataset_name: str,
loaded_models: Optional[Dict[str, Any]] = None,
checkpoint_dir: str = 'results/checkpoints'
) -> Dict[str, Any]:
"""
Run multi-model ensemble prediction for a specific dataset endpoint.
Args:
smiles_or_graph: SMILES string or PyG Data graph.
dataset_name: One of 'esol', 'bbbp', 'tox21'.
loaded_models: Optional preloaded dict of {model_name: model}.
checkpoint_dir: Directory containing saved checkpoints.
Returns:
Dictionary containing individual predictions, ensemble mean, std, CI, and agreement.
"""
ds = dataset_name.lower()
if isinstance(smiles_or_graph, str):
graph = smiles_to_graph(smiles_or_graph)
else:
graph = smiles_or_graph
if loaded_models is None:
loaded_models = {}
for m in MODEL_NAMES:
try:
mod, _ = load_trained_model(m, ds, checkpoint_dir=checkpoint_dir)
loaded_models[m] = mod
except Exception:
pass
if not loaded_models:
raise RuntimeError(f"No trained checkpoints found for dataset '{ds}'.")
predictions = {}
vals_list = []
for m_name, model in loaded_models.items():
try:
val = predict_single_model(model, m_name, graph, ds)
predictions[m_name] = val
vals_list.append(val)
except Exception as e:
continue
if not vals_list:
raise RuntimeError(f"Failed to obtain predictions for '{ds}'.")
stacked = np.stack(vals_list, axis=0) # Shape: (num_models, num_tasks)
mean_pred = np.mean(stacked, axis=0)
std_pred = np.std(stacked, axis=0)
# 95% Confidence bounds
ci_lower = mean_pred - 1.96 * std_pred
ci_upper = mean_pred + 1.96 * std_pred
# Assess model agreement
avg_std = float(np.mean(std_pred))
if ds == 'esol':
# logS standard deviation threshold
if avg_std < 0.25:
agreement = "High Consensus (Low Uncertainty)"
elif avg_std < 0.50:
agreement = "Moderate Consensus"
else:
agreement = "Divergent Predictions (High Uncertainty)"
else:
# Probability standard deviation threshold
if avg_std < 0.10:
agreement = "High Consensus (Low Uncertainty)"
elif avg_std < 0.20:
agreement = "Moderate Consensus"
else:
agreement = "Divergent Predictions (High Uncertainty)"
return {
"dataset": ds,
"task_names": DATASET_CONFIGS[ds]['target_cols'],
"individual": predictions,
"mean": mean_pred,
"std": std_pred,
"ci_lower": ci_lower,
"ci_upper": ci_upper,
"num_models": len(vals_list),
"agreement": agreement,
"uncertainty_score": avg_std
}
def predict_all_endpoints(
smiles_or_graph: Any,
models_cache: Optional[Dict[str, Dict[str, Any]]] = None,
checkpoint_dir: str = 'results/checkpoints'
) -> Dict[str, Any]:
"""
Run full multi-endpoint prediction suite across ESOL, BBBP, and Tox21.
"""
results = {}
for ds in ['esol', 'bbbp', 'tox21']:
cache_ds = models_cache.get(ds) if models_cache else None
try:
res = predict_ensemble(
smiles_or_graph=smiles_or_graph,
dataset_name=ds,
loaded_models=cache_ds,
checkpoint_dir=checkpoint_dir
)
results[ds] = res
except Exception as e:
results[ds] = {"error": str(e)}
return results
|