""" 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()