#!/usr/bin/env python # coding=utf-8 """ Psc_Predict MCP Server Perovskite Solar Cell Performance Prediction MCP Service Using FastMCP framework with SSE transport Designed for HuggingFace Docker deployment """ import os import re import pickle import logging from typing import Dict, List, Optional, Any import numpy as np import torch import torch.nn as nn from fastmcp import FastMCP # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Initialize FastMCP server mcp = FastMCP("Psc_Predict") # ============ CIF Parser ============ class CIFParser: """Extract crystallographic features from CIF content (93 dimensions)""" def __init__(self): self.elements = [ 'H', 'Li', 'Be', 'B', 'C', 'N', 'O', 'F', 'Na', 'Mg', 'Al', 'Si', 'P', 'S', 'Cl', 'K', 'Ca', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn', 'Ga', 'Ge', 'As', 'Se', 'Br', 'Rb', 'Sr', 'Y', 'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd', 'In', 'Sn', 'Sb', 'Te', 'I', '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' ] self.elem_to_idx = {e: i for i, e in enumerate(self.elements)} def parse(self, cif_text: str) -> np.ndarray: """Parse CIF string and return 93-dimensional feature vector""" # Handle escaped newlines if "\\n" in cif_text: cif_text = cif_text.replace("\\n", "\n") # A. Extract lattice parameters (7 dimensions) patterns = { 'a': r"_cell_length_a\s+([\d\.]+)", 'b': r"_cell_length_b\s+([\d\.]+)", 'c': r"_cell_length_c\s+([\d\.]+)", 'alpha': r"_cell_angle_alpha\s+([\d\.]+)", 'beta': r"_cell_angle_beta\s+([\d\.]+)", 'gamma': r"_cell_angle_gamma\s+([\d\.]+)", 'vol': r"_cell_volume\s+([\d\.]+)" } lattice_feats = [] for key, pat in patterns.items(): match = re.search(pat, cif_text) val = float(match.group(1)) if match else 0.0 lattice_feats.append(val) # B. Extract element composition (86 dimensions) chem_match = re.search(r"_chemical_formula_sum\s+'?([^'\n]+)'?", cif_text) elem_vec = np.zeros(len(self.elements)) if chem_match: formula = chem_match.group(1) parts = formula.replace("'", "").split() for part in parts: m = re.match(r"([A-Za-z]+)([\d\.]*)", part) if m: el = m.group(1) num = float(m.group(2)) if m.group(2) else 1.0 if el in self.elem_to_idx: elem_vec[self.elem_to_idx[el]] = num # Normalize element vector total_atoms = np.sum(elem_vec) if total_atoms > 0: elem_vec = elem_vec / total_atoms return np.concatenate([lattice_feats, elem_vec]) def get_feature_names(self) -> List[str]: return ['a', 'b', 'c', 'alpha', 'beta', 'gamma', 'vol'] + self.elements # ============ Neural Network Model ============ class MaterialNN(nn.Module): """Neural Network for material property prediction""" def __init__(self, input_dim, hidden_dims=[128, 64, 32]): super(MaterialNN, self).__init__() layers = [] in_d = input_dim for h_d in hidden_dims: layers.append(nn.Linear(in_d, h_d)) layers.append(nn.ReLU()) layers.append(nn.BatchNorm1d(h_d)) in_d = h_d layers.append(nn.Linear(in_d, 1)) self.net = nn.Sequential(*layers) def forward(self, x): return self.net(x) # ============ Model Manager ============ class ModelManager: """Manage all pretrained models (XGBoost, Random Forest, Neural Network)""" TARGETS = ['pce', 'dft_band_gap', 'energy_above_hull', 'stability_retention', 'stability_t80', 'voc', 'jsc', 'ff'] MODEL_TYPES = ['xgboost', 'random_forest', 'neural_network'] TARGET_INFO = { 'pce': {'name': 'Power Conversion Efficiency', 'unit': '%'}, 'dft_band_gap': {'name': 'DFT Band Gap', 'unit': 'eV'}, 'energy_above_hull': {'name': 'Energy Above Hull', 'unit': 'eV/atom'}, 'stability_retention': {'name': 'Stability Retention', 'unit': '%'}, 'stability_t80': {'name': 'T80 Lifetime', 'unit': 'hours'}, 'voc': {'name': 'Open Circuit Voltage', 'unit': 'V'}, 'jsc': {'name': 'Short Circuit Current Density', 'unit': 'mA/cm²'}, 'ff': {'name': 'Fill Factor', 'unit': ''} } def __init__(self, model_dir: str = "./models"): self.model_dir = model_dir self.models: Dict[str, Dict[str, Any]] = { 'xgboost': {}, 'random_forest': {}, 'neural_network': {} } self.cif_parser = CIFParser() self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self._load_all_models() def _load_all_models(self): """Load all available models""" # Load XGBoost models for target in self.TARGETS: model_path = os.path.join( self.model_dir, f"xgboost_{target}_layers-NA_seed-42_batch-32.pkl" ) if os.path.exists(model_path): try: with open(model_path, 'rb') as f: self.models['xgboost'][target] = pickle.load(f) logger.info(f"Loaded XGBoost model for {target}") except Exception as e: logger.warning(f"Failed to load XGBoost model for {target}: {e}") # Load Random Forest models for target in self.TARGETS: model_path = os.path.join( self.model_dir, f"random_forest_{target}_layers-NA_seed-42_batch-32.pkl" ) if os.path.exists(model_path): try: with open(model_path, 'rb') as f: self.models['random_forest'][target] = pickle.load(f) logger.info(f"Loaded Random Forest model for {target}") except Exception as e: logger.warning(f"Failed to load Random Forest model for {target}: {e}") # Load Neural Network models for target in self.TARGETS: model_path = os.path.join( self.model_dir, f"neural_network_{target}_layers-128-64-32_seed-42_batch-32.pth" ) if os.path.exists(model_path): try: model = MaterialNN(input_dim=93, hidden_dims=[128, 64, 32]) model.load_state_dict(torch.load(model_path, map_location=self.device)) model.to(self.device) model.eval() self.models['neural_network'][target] = model logger.info(f"Loaded Neural Network model for {target}") except Exception as e: logger.warning(f"Failed to load Neural Network model for {target}: {e}") def predict(self, cif_text: str, targets: Optional[List[str]] = None, model_type: str = 'xgboost') -> Dict[str, float]: """Predict specified targets using selected model type""" if model_type not in self.models: raise ValueError(f"Unknown model type: {model_type}. Available: {self.MODEL_TYPES}") if targets is None: targets = list(self.models[model_type].keys()) # Parse CIF features = self.cif_parser.parse(cif_text) X = features.reshape(1, -1) # Predict results = {} for target in targets: if target in self.models[model_type]: model = self.models[model_type][target] if model_type == 'neural_network': X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device) with torch.no_grad(): pred = model(X_tensor).cpu().numpy().flatten()[0] else: pred = model.predict(X)[0] results[target] = float(pred) else: results[target] = None return results def get_available_targets(self, model_type: str = 'xgboost') -> List[str]: """Return available prediction targets for a model type""" if model_type in self.models: return list(self.models[model_type].keys()) return [] def get_available_models(self) -> Dict[str, List[str]]: """Return all available models and their targets""" return { model_type: list(targets.keys()) for model_type, targets in self.models.items() if targets } # Global model manager model_manager: Optional[ModelManager] = None def get_model_manager() -> ModelManager: """Get or initialize model manager""" global model_manager if model_manager is None: model_dir = os.environ.get("MODEL_DIR", "./models") model_manager = ModelManager(model_dir) return model_manager # ============ MCP Tools ============ # Valid model types VALID_MODEL_TYPES = ['xgboost', 'random_forest', 'neural_network'] DEFAULT_MODEL_TYPE = 'xgboost' @mcp.tool() def parse_cif_features(cif: str) -> Dict[str, Any]: """ Parse a CIF file and extract features for model prediction. Extracts 93-dimensional features: - 7 lattice parameters (a, b, c, alpha, beta, gamma, volume) - 86 element composition fractions Args: cif: Crystal structure text in CIF format Returns: Dictionary containing lattice parameters and element composition """ manager = get_model_manager() features = manager.cif_parser.parse(cif) feature_names = manager.cif_parser.get_feature_names() # Separate lattice parameters and element composition lattice = dict(zip(feature_names[:7], features[:7].tolist())) # Only return non-zero elements composition = {} for i, elem in enumerate(feature_names[7:]): if features[7 + i] > 0: composition[elem] = float(features[7 + i]) return { "lattice_parameters": lattice, "composition": composition, "feature_dim": len(features), "status": "success" } @mcp.tool() def get_model_info() -> Dict[str, Any]: """ Get model information and available prediction targets. Returns information about: - Available model types (XGBoost, Random Forest, Neural Network) - All 8 prediction targets and their availability - Input feature dimensions Returns: Dictionary containing model information """ manager = get_model_manager() # Get available targets for each model type model_availability = {} for mt in VALID_MODEL_TYPES: available = manager.get_available_targets(model_type=mt) model_availability[mt] = { "available_targets": available, "count": len(available) } targets_info = [] for target in ModelManager.TARGETS: info = ModelManager.TARGET_INFO.get(target, {}) targets_info.append({ "id": target, "name": info.get('name', target), "unit": info.get('unit', ''), "xgboost": target in model_availability['xgboost']['available_targets'], "random_forest": target in model_availability['random_forest']['available_targets'], "neural_network": target in model_availability['neural_network']['available_targets'] }) return { "available_model_types": VALID_MODEL_TYPES, "default_model_type": DEFAULT_MODEL_TYPE, "recommended_model_type": "xgboost", "input_features": 93, "targets": targets_info, "model_availability": model_availability, "total_targets": len(ModelManager.TARGETS) } @mcp.tool() def list_available_models() -> Dict[str, Any]: """ List all available models and their status. Returns detailed information about which models are loaded and ready for inference. Returns: Dictionary containing model availability status for each target and model type """ manager = get_model_manager() models_status = {} for mt in VALID_MODEL_TYPES: models_status[mt] = {} for target in ModelManager.TARGETS: key = f"{mt}_{target}" is_loaded = key in manager.models models_status[mt][target] = { "loaded": is_loaded, "status": "ready" if is_loaded else "not_available" } return { "models": models_status, "model_types": VALID_MODEL_TYPES, "targets": ModelManager.TARGETS, "status": "success" } @mcp.tool() def predict_ensemble(cif: str, targets: Optional[List[str]] = None) -> Dict[str, Any]: """ Predict using all three model types and return ensemble results with comparison. Automatically calls XGBoost, Random Forest, and Neural Network models for the same input, allowing comparison of predictions across different model architectures. Also provides ensemble statistics (mean, std, min, max) for each target. Args: cif: Crystal structure text in CIF format targets: Optional list of specific targets to predict. If None, predicts all available targets. Valid targets: pce, dft_band_gap, energy_above_hull, stability_retention, stability_t80, voc, jsc, ff Returns: Dictionary containing predictions from all models and ensemble statistics """ import numpy as np manager = get_model_manager() # Determine targets to predict if targets is None: targets = ModelManager.TARGETS # Collect predictions from all models all_predictions = {} for mt in VALID_MODEL_TYPES: try: result = manager.predict(cif, list(targets), model_type=mt) all_predictions[mt] = result except Exception as e: all_predictions[mt] = {"error": str(e)} # Calculate ensemble statistics for each target ensemble_results = {} for target in targets: values = [] model_values = {} for mt in VALID_MODEL_TYPES: if mt in all_predictions and target in all_predictions[mt]: val = all_predictions[mt][target] if val is not None: values.append(val) model_values[mt] = val else: model_values[mt] = None else: model_values[mt] = None info = ModelManager.TARGET_INFO.get(target, {}) if values: ensemble_results[target] = { "name": info.get('name', target), "unit": info.get('unit', ''), "predictions": model_values, "ensemble": { "mean": float(np.mean(values)), "std": float(np.std(values)), "min": float(np.min(values)), "max": float(np.max(values)), "range": float(np.max(values) - np.min(values)), "n_models": len(values) }, "recommendation": _get_best_prediction(target, model_values) } else: ensemble_results[target] = { "name": info.get('name', target), "unit": info.get('unit', ''), "predictions": model_values, "ensemble": None, "recommendation": None } return { "targets_predicted": list(targets), "models_used": VALID_MODEL_TYPES, "results": ensemble_results, "raw_predictions": all_predictions, "status": "success" } def _get_best_prediction(target: str, model_values: Dict[str, float]) -> Dict[str, Any]: """ Provide recommendation based on model performance characteristics. XGBoost is generally recommended as it has the best overall performance. """ # XGBoost is the recommended model based on benchmark results if model_values.get('xgboost') is not None: return { "model": "xgboost", "value": model_values['xgboost'], "reason": "XGBoost recommended - best overall performance in benchmarks" } elif model_values.get('random_forest') is not None: return { "model": "random_forest", "value": model_values['random_forest'], "reason": "Random Forest - fallback when XGBoost unavailable" } elif model_values.get('neural_network') is not None: return { "model": "neural_network", "value": model_values['neural_network'], "reason": "Neural Network - fallback option" } return None # ============ MCP Resources ============ @mcp.resource("psc://info") def get_service_info() -> str: """Service information""" return """ # Psc_Predict MCP Service Perovskite Solar Cell Performance Prediction Service ## Features - Predict 8 performance metrics from CIF crystal structures - Support for single-target and multi-target prediction - Multiple model types: XGBoost (recommended), Random Forest, Neural Network ## Available Model Types 1. **XGBoost** (default, recommended) - Best overall performance 2. **Random Forest** - Good interpretability 3. **Neural Network** - 3-layer MLP (128-64-32) ## Prediction Targets 1. PCE - Power Conversion Efficiency (%) 2. DFT Band Gap - DFT calculated band gap (eV) 3. Energy Above Hull - Thermodynamic stability (eV/atom) 4. Stability Retention - Stability retention percentage (%) 5. Stability T80 - T80 lifetime (hours) 6. Voc - Open-circuit voltage (V) 7. Jsc - Short-circuit current density (mA/cm²) 8. FF - Fill factor ## Input Format CIF (Crystallographic Information File) format crystal structure text ## Usage Example Call predict_pce(cif, model_type="xgboost") to predict PCE using XGBoost model. """ @mcp.resource("psc://example-cif") def get_example_cif() -> str: """Example CIF file for testing""" return """data_CsPbI3 _symmetry_space_group_name_H-M 'P m -3 m' _cell_length_a 6.2894 _cell_length_b 6.2894 _cell_length_c 6.2894 _cell_angle_alpha 90.0 _cell_angle_beta 90.0 _cell_angle_gamma 90.0 _cell_volume 248.89 _chemical_formula_sum 'Cs1 Pb1 I3' loop_ _atom_site_label _atom_site_type_symbol _atom_site_fract_x _atom_site_fract_y _atom_site_fract_z Cs1 Cs 0.0 0.0 0.0 Pb1 Pb 0.5 0.5 0.5 I1 I 0.5 0.5 0.0 I2 I 0.5 0.0 0.5 I3 I 0.0 0.5 0.5 """ # ============ Main Entry Point ============ if __name__ == "__main__": import sys # Support command line arguments for transport selection transport = os.environ.get("MCP_TRANSPORT", "sse") host = os.environ.get("HOST", "0.0.0.0") port = int(os.environ.get("PORT", 7860)) if transport == "stdio": mcp.run() else: # SSE mode (default, for HuggingFace) mcp.run(transport="sse", host=host, port=port)