model.py

"""
Fan Design Surrogate Model - Inference Module

Usage:
    from model import FanDesignSurrogate
    
    surrogate = FanDesignSurrogate.from_pretrained("harshaperla/fan-design-surrogate")
    
    results = surrogate.predict({
        'blade_inlet_angle_deg': 55.0,
        'blade_turning_angle_deg': 15.0,
        'chord_length_mm': 100.0,
        'blade_thickness_ratio': 0.06,
        'stagger_angle_deg': 45.0,
        'hub_tip_ratio': 0.5,
        'tip_clearance_ratio': 0.015,
        'num_blades': 12,
        'aspect_ratio': 2.5,
        'solidity': 1.0,
        'sweep_angle_deg': 0.0,
        'flow_coefficient': 0.5,
        'rotational_speed_rpm': 3000,
        'tip_radius_mm': 300.0,
    })
"""

import torch
import torch.nn as nn
import numpy as np
import json
import os
from pathlib import Path


class ResBlock(nn.Module):
    def __init__(self, dim, dropout=0.05):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(dim), nn.GELU(), nn.Linear(dim, dim), nn.Dropout(dropout),
            nn.LayerNorm(dim), nn.GELU(), nn.Linear(dim, dim), nn.Dropout(dropout),
        )
    
    def forward(self, x):
        return x + self.net(x)


class FanSurrogateNet(nn.Module):
    """Residual MLP for fan performance prediction."""
    
    def __init__(self, n_in=14, n_out=8, hidden=256, blocks=4, dropout=0.05):
        super().__init__()
        self.proj = nn.Sequential(nn.Linear(n_in, hidden), nn.GELU(), nn.Linear(hidden, hidden))
        self.blocks = nn.Sequential(*[ResBlock(hidden, dropout) for _ in range(blocks)])
        self.head = nn.Sequential(
            nn.LayerNorm(hidden), nn.GELU(),
            nn.Linear(hidden, hidden // 2), nn.GELU(),
            nn.Linear(hidden // 2, n_out),
        )
    
    def forward(self, x):
        return self.head(self.blocks(self.proj(x)))


class FanDesignSurrogate:
    """
    High-level interface for the fan design surrogate model.
    
    Handles loading, preprocessing, inference, and postprocessing.
    """
    
    def __init__(self, model_dir):
        self.model_dir = Path(model_dir)
        
        # Load config
        with open(self.model_dir / 'config.json') as f:
            self.config = json.load(f)
        
        # Load scalers
        with open(self.model_dir / 'scalers.json') as f:
            scalers = json.load(f)
        
        self.scaler_X_mean = np.array(scalers['scaler_X_mean'], dtype=np.float32)
        self.scaler_X_scale = np.array(scalers['scaler_X_scale'], dtype=np.float32)
        self.scaler_y_mean = np.array(scalers['scaler_y_mean'], dtype=np.float32)
        self.scaler_y_scale = np.array(scalers['scaler_y_scale'], dtype=np.float32)
        
        self.input_cols = self.config['input_columns']
        self.output_cols = self.config['output_columns']
        self.log_indices = self.config['log_output_indices']
        self.bounds = self.config['parameter_bounds']
        
        # Load model
        self.model = FanSurrogateNet(
            n_in=self.config['n_inputs'],
            n_out=self.config['n_outputs'],
            hidden=self.config['hidden_dim'],
            blocks=self.config['n_blocks'],
            dropout=self.config['dropout'],
        )
        state = torch.load(self.model_dir / 'model.pt', map_location='cpu', weights_only=True)
        self.model.load_state_dict(state)
        self.model.eval()
    
    @classmethod
    def from_pretrained(cls, repo_id, cache_dir=None):
        """Load model from Hugging Face Hub."""
        from huggingface_hub import snapshot_download
        local_dir = snapshot_download(repo_id, cache_dir=cache_dir)
        return cls(local_dir)
    
    def validate_inputs(self, params):
        """Check that all inputs are within training bounds."""
        warnings = []
        for col in self.input_cols:
            if col not in params:
                raise ValueError(f"Missing input parameter: {col}")
            lo, hi = self.bounds[col]
            val = params[col]
            if val < lo or val > hi:
                warnings.append(f"{col}={val} outside training range [{lo}, {hi}]")
        return warnings
    
    def predict(self, params, validate=True):
        """
        Predict fan performance from design parameters.
        
        Args:
            params: dict with all 14 input parameters
            validate: if True, warn about out-of-range inputs
            
        Returns:
            dict with 8 predicted performance metrics
        """
        if validate:
            warnings = self.validate_inputs(params)
            if warnings:
                import warnings as w
                for msg in warnings:
                    w.warn(f"Extrapolation warning: {msg}")
        
        # Build input array
        x = np.array([[params[col] for col in self.input_cols]], dtype=np.float32)
        
        # Standardize
        x_scaled = (x - self.scaler_X_mean) / self.scaler_X_scale
        
        # Predict
        with torch.no_grad():
            y_scaled = self.model(torch.from_numpy(x_scaled)).numpy()
        
        # Inverse transform
        y_proc = y_scaled * self.scaler_y_scale + self.scaler_y_mean
        
        # Undo log transform
        y_final = y_proc.copy()
        for idx in self.log_indices:
            y_final[0, idx] = np.expm1(y_proc[0, idx])
        
        return {col: float(y_final[0, i]) for i, col in enumerate(self.output_cols)}
    
    def predict_batch(self, params_list):
        """Predict for multiple designs at once (faster)."""
        n = len(params_list)
        x = np.zeros((n, len(self.input_cols)), dtype=np.float32)
        for i, params in enumerate(params_list):
            for j, col in enumerate(self.input_cols):
                x[i, j] = params[col]
        
        x_scaled = (x - self.scaler_X_mean) / self.scaler_X_scale
        
        with torch.no_grad():
            y_scaled = self.model(torch.from_numpy(x_scaled)).numpy()
        
        y_proc = y_scaled * self.scaler_y_scale + self.scaler_y_mean
        y_final = y_proc.copy()
        for idx in self.log_indices:
            y_final[:, idx] = np.expm1(y_proc[:, idx])
        
        results = []
        for i in range(n):
            results.append({col: float(y_final[i, j]) for j, col in enumerate(self.output_cols)})
        return results
    
    def sensitivity_analysis(self, baseline_params, param_name, n_points=20):
        """
        Perform one-at-a-time sensitivity analysis.
        
        Varies one parameter while keeping others at baseline values.
        Returns arrays of parameter values and corresponding predictions.
        """
        lo, hi = self.bounds[param_name]
        values = np.linspace(lo, hi, n_points)
        
        designs = []
        for val in values:
            d = baseline_params.copy()
            d[param_name] = val
            designs.append(d)
        
        predictions = self.predict_batch(designs)
        return values, predictions


# Convenience function
def predict_fan_performance(params, model_dir='.'):
    """Quick prediction without loading class."""
    surrogate = FanDesignSurrogate(model_dir)
    return surrogate.predict(params)


if __name__ == '__main__':
    # Example usage
    print("Fan Design Surrogate Model - Example")
    print("=" * 50)
    
    # Try to load from current directory
    model_dir = Path('.')
    if not (model_dir / 'model.pt').exists():
        print("Model files not found locally. Use from_pretrained():")
        print("  surrogate = FanDesignSurrogate.from_pretrained('harshaperla/fan-design-surrogate')")
        exit(0)
    
    surrogate = FanDesignSurrogate('.')
    
    # Example design
    design = {
        'blade_inlet_angle_deg': 55.0,
        'blade_turning_angle_deg': 15.0,
        'chord_length_mm': 100.0,
        'blade_thickness_ratio': 0.06,
        'stagger_angle_deg': 45.0,
        'hub_tip_ratio': 0.5,
        'tip_clearance_ratio': 0.015,
        'num_blades': 12,
        'aspect_ratio': 2.5,
        'solidity': 1.0,
        'sweep_angle_deg': 0.0,
        'flow_coefficient': 0.5,
        'rotational_speed_rpm': 3000,
        'tip_radius_mm': 300.0,
    }
    
    print("\nInput Design:")
    for k, v in design.items():
        print(f"  {k}: {v}")
    
    results = surrogate.predict(design)
    
    print("\nPredicted Performance:")
    for k, v in results.items():
        print(f"  {k}: {v:.4f}")