train.py

"""
Fan Design Surrogate Model - Training (CPU-optimized, fast)
"""
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import json, os, time

DEVICE = torch.device('cpu')

INPUT_COLS = [
    'blade_inlet_angle_deg', 'blade_turning_angle_deg', 'chord_length_mm',
    'blade_thickness_ratio', 'stagger_angle_deg', 'hub_tip_ratio',
    'tip_clearance_ratio', 'num_blades', 'aspect_ratio', 'solidity',
    'sweep_angle_deg', 'flow_coefficient', 'rotational_speed_rpm', 'tip_radius_mm',
]
OUTPUT_COLS = [
    'total_pressure_rise_Pa', 'isentropic_efficiency', 'power_consumption_W',
    'flow_rate_m3s', 'specific_speed', 'degree_of_reaction',
    'diffusion_factor', 'noise_estimate_dBA',
]
LOG_OUTPUTS = ['total_pressure_rise_Pa', 'power_consumption_W', 'flow_rate_m3s']


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 FanSurrogate(nn.Module):
    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),
        )
        self._init()
    
    def _init(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                if m.bias is not None: nn.init.zeros_(m.bias)
    
    def forward(self, x):
        return self.head(self.blocks(self.proj(x)))


def main():
    t0 = time.time()
    print("Loading data...")
    df_train = pd.read_csv('/app/fan_design_train.csv')
    df_test = pd.read_csv('/app/fan_design_test.csv')
    
    X_all = df_train[INPUT_COLS].values.astype(np.float32)
    y_all = df_train[OUTPUT_COLS].values.astype(np.float32)
    X_test = df_test[INPUT_COLS].values.astype(np.float32)
    y_test = df_test[OUTPUT_COLS].values.astype(np.float32)
    
    X_tr, X_val, y_tr, y_val = train_test_split(X_all, y_all, test_size=0.1, random_state=42)
    
    # Normalize inputs
    sx = StandardScaler()
    X_tr_s = sx.fit_transform(X_tr).astype(np.float32)
    X_val_s = sx.transform(X_val).astype(np.float32)
    X_test_s = sx.transform(X_test).astype(np.float32)
    
    # Log + normalize outputs
    log_idx = [OUTPUT_COLS.index(c) for c in LOG_OUTPUTS]
    def apply_log(y):
        y2 = y.copy()
        for i in log_idx: y2[:, i] = np.log1p(y[:, i])
        return y2
    
    sy = StandardScaler()
    y_tr_s = sy.fit_transform(apply_log(y_tr)).astype(np.float32)
    y_val_s = sy.transform(apply_log(y_val)).astype(np.float32)
    y_test_s = sy.transform(apply_log(y_test)).astype(np.float32)
    
    # Loaders
    tr_loader = DataLoader(TensorDataset(torch.from_numpy(X_tr_s), torch.from_numpy(y_tr_s)),
                          batch_size=512, shuffle=True)
    
    # Train
    print("Training model...")
    model = FanSurrogate(14, 8, 256, 4, 0.05)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"  Model params: {n_params:,}")
    
    opt = optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-4)
    sched = optim.lr_scheduler.CosineAnnealingLR(opt, T_max=200, eta_min=1e-5)
    loss_fn = nn.MSELoss()
    
    best_val = float('inf')
    best_state = None
    patience = 0
    
    X_val_t = torch.from_numpy(X_val_s)
    y_val_t = torch.from_numpy(y_val_s)
    
    for epoch in range(200):
        model.train()
        for xb, yb in tr_loader:
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()
        sched.step()
        
        model.eval()
        with torch.no_grad():
            vl = loss_fn(model(X_val_t), y_val_t).item()
        
        if vl < best_val:
            best_val = vl
            best_state = {k: v.clone() for k, v in model.state_dict().items()}
            patience = 0
        else:
            patience += 1
        
        if (epoch+1) % 20 == 0:
            print(f"  Epoch {epoch+1}: val_loss={vl:.6f} best={best_val:.6f} "
                  f"[{time.time()-t0:.0f}s]")
        
        if patience >= 40:
            print(f"  Early stop at epoch {epoch+1}")
            break
    
    model.load_state_dict(best_state)
    
    # Evaluate
    print("\nEvaluating...")
    model.eval()
    with torch.no_grad():
        pred_s = model(torch.from_numpy(X_test_s)).numpy()
    
    pred_proc = sy.inverse_transform(pred_s)
    test_proc = sy.inverse_transform(y_test_s)
    pred_final = pred_proc.copy()
    test_final = test_proc.copy()
    for i in log_idx:
        pred_final[:, i] = np.expm1(pred_proc[:, i])
        test_final[:, i] = np.expm1(test_proc[:, i])
    
    print(f"\n{'Output':<28} {'R²':>8} {'MAPE%':>8}")
    print("=" * 48)
    results = {}
    r2s, mapes = [], []
    for i, col in enumerate(OUTPUT_COLS):
        yt, yp = test_final[:, i], pred_final[:, i]
        ss_res = np.sum((yt - yp)**2)
        ss_tot = np.sum((yt - yt.mean())**2)
        r2 = 1 - ss_res/(ss_tot + 1e-10)
        mask = np.abs(yt) > 1e-6
        mape = np.mean(np.abs((yt[mask]-yp[mask])/yt[mask])) * 100
        print(f"  {col:<28} {r2:>7.4f} {mape:>7.2f}%")
        r2s.append(r2); mapes.append(mape)
        results[col] = {'R2': float(r2), 'MAPE': float(mape)}
    print("=" * 48)
    print(f"  {'AVERAGE':<28} {np.mean(r2s):>7.4f} {np.mean(mapes):>7.2f}%")
    
    # Save
    print("\nSaving model...")
    os.makedirs('/app/model_artifacts', exist_ok=True)
    torch.save(model.state_dict(), '/app/model_artifacts/model.pt')
    np.save('/app/model_artifacts/scaler_X_mean.npy', sx.mean_)
    np.save('/app/model_artifacts/scaler_X_scale.npy', sx.scale_)
    np.save('/app/model_artifacts/scaler_y_mean.npy', sy.mean_)
    np.save('/app/model_artifacts/scaler_y_scale.npy', sy.scale_)
    
    config = {
        'architecture': 'ResidualMLP',
        'n_inputs': 14, 'n_outputs': 8,
        'hidden_dim': 256, 'n_blocks': 4, 'dropout': 0.05,
        'input_columns': INPUT_COLS,
        'output_columns': OUTPUT_COLS,
        'log_output_indices': log_idx,
        'log_output_columns': LOG_OUTPUTS,
        'test_results': results,
        'avg_r2': float(np.mean(r2s)),
        'avg_mape': float(np.mean(mapes)),
        'training_time_s': float(time.time() - t0),
    }
    with open('/app/model_artifacts/model_config.json', 'w') as f:
        json.dump(config, f, indent=2)
    
    print(f"\n✓ Complete! Avg R²={np.mean(r2s):.4f}, Avg MAPE={np.mean(mapes):.2f}%")
    print(f"  Total time: {time.time()-t0:.1f}s")


if __name__ == '__main__':
    main()