""" Model implementations: XGBoost, Neural Network, and Weighted Ensemble. All models follow a consistent interface for training, prediction, and cross-validation to enable fair comparison. """ from __future__ import annotations import logging from copy import deepcopy from typing import Any, Dict, List, Optional, Tuple import numpy as np import torch import torch.nn as nn from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score from sklearn.model_selection import KFold, StratifiedKFold from sklearn.preprocessing import MinMaxScaler from torch.utils.data import DataLoader, TensorDataset import xgboost as xgb logger = logging.getLogger(__name__) # ============================================================================= # XGBoost Multi-Output Regressor # ============================================================================= class XGBoostMultiOutput: """ XGBoost regressor trained independently per target. This approach (separate models per target) is preferred over MultiOutputRegressor wrapper as it allows per-target hyperparameter inspection and feature importance analysis. """ def __init__(self, params: Dict[str, Any], target_names: List[str]): self.params = params self.target_names = target_names self.models: Dict[str, xgb.XGBRegressor] = {} def fit( self, X_train: np.ndarray, y_train: np.ndarray, X_val: Optional[np.ndarray] = None, y_val: Optional[np.ndarray] = None, ) -> "XGBoostMultiOutput": """Train one XGBoost model per target.""" for i, target in enumerate(self.target_names): model = xgb.XGBRegressor(**self.params) eval_set = [(X_val, y_val[:, i])] if X_val is not None else None model.fit( X_train, y_train[:, i], eval_set=eval_set, verbose=False, ) self.models[target] = model logger.debug(f"XGBoost trained for {target}") return self def predict(self, X: np.ndarray) -> np.ndarray: """Predict all targets.""" preds = np.column_stack([ self.models[t].predict(X) for t in self.target_names ]) return preds def get_feature_importance(self, feature_names: List[str]) -> Dict[str, Dict[str, float]]: """Get feature importance per target.""" importances = {} for target, model in self.models.items(): imp = model.feature_importances_ importances[target] = dict(zip(feature_names, imp.tolist())) return importances # ============================================================================= # PyTorch Neural Network # ============================================================================= class LaserEtchingNet(nn.Module): """ Multi-output regression neural network. Architecture follows Behbahani et al. (2023) with BatchNorm and Dropout for regularization, adapted to our multi-target prediction task. """ def __init__( self, n_features: int, n_outputs: int, hidden_layers: List[int] = [128, 64, 32], dropout: float = 0.2, batch_norm: bool = True, ): super().__init__() layers = [] prev_size = n_features for i, h in enumerate(hidden_layers): layers.append(nn.Linear(prev_size, h)) if batch_norm: layers.append(nn.BatchNorm1d(h)) layers.append(nn.ReLU()) # Reduce dropout in deeper layers drop_rate = dropout * (1 - i / len(hidden_layers)) if drop_rate > 0.05: layers.append(nn.Dropout(drop_rate)) prev_size = h layers.append(nn.Linear(prev_size, n_outputs)) self.network = nn.Sequential(*layers) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.network(x) class NeuralNetworkRegressor: """ Wrapper for PyTorch neural network with sklearn-like interface. Includes MinMaxScaler for input/output normalization and early stopping based on validation loss. """ def __init__( self, n_features: int, n_outputs: int, hidden_layers: List[int] = [128, 64, 32], dropout: float = 0.2, learning_rate: float = 1e-3, weight_decay: float = 1e-4, batch_size: int = 64, max_epochs: int = 500, patience: int = 30, device: Optional[str] = None, ): self.n_features = n_features self.n_outputs = n_outputs self.hidden_layers = hidden_layers self.dropout = dropout self.learning_rate = learning_rate self.weight_decay = weight_decay self.batch_size = batch_size self.max_epochs = max_epochs self.patience = patience self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.scaler_X = MinMaxScaler() self.scaler_y = MinMaxScaler() self.model: Optional[LaserEtchingNet] = None self.train_losses: List[float] = [] self.val_losses: List[float] = [] def fit( self, X_train: np.ndarray, y_train: np.ndarray, X_val: Optional[np.ndarray] = None, y_val: Optional[np.ndarray] = None, ) -> "NeuralNetworkRegressor": """Train neural network with early stopping.""" # Scale data X_train_s = self.scaler_X.fit_transform(X_train) y_train_s = self.scaler_y.fit_transform(y_train) X_val_s = self.scaler_X.transform(X_val) if X_val is not None else None y_val_s = self.scaler_y.transform(y_val) if y_val is not None else None # Create model self.model = LaserEtchingNet( n_features=self.n_features, n_outputs=self.n_outputs, hidden_layers=self.hidden_layers, dropout=self.dropout, ).to(self.device) optimizer = torch.optim.Adam( self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay, ) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode="min", factor=0.5, patience=10, min_lr=1e-6 ) criterion = nn.MSELoss() # DataLoader train_ds = TensorDataset( torch.FloatTensor(X_train_s), torch.FloatTensor(y_train_s) ) train_dl = DataLoader(train_ds, batch_size=self.batch_size, shuffle=True) if X_val_s is not None: val_X_t = torch.FloatTensor(X_val_s).to(self.device) val_y_t = torch.FloatTensor(y_val_s).to(self.device) # Training loop best_val_loss = float("inf") patience_counter = 0 best_state = None self.train_losses = [] self.val_losses = [] for epoch in range(self.max_epochs): # Train self.model.train() epoch_loss = 0 for X_batch, y_batch in train_dl: X_batch = X_batch.to(self.device) y_batch = y_batch.to(self.device) optimizer.zero_grad() pred = self.model(X_batch) loss = criterion(pred, y_batch) loss.backward() optimizer.step() epoch_loss += loss.item() avg_train_loss = epoch_loss / len(train_dl) self.train_losses.append(avg_train_loss) # Validate if X_val_s is not None: self.model.eval() with torch.no_grad(): val_pred = self.model(val_X_t) val_loss = criterion(val_pred, val_y_t).item() self.val_losses.append(val_loss) scheduler.step(val_loss) if val_loss < best_val_loss: best_val_loss = val_loss patience_counter = 0 best_state = {k: v.clone() for k, v in self.model.state_dict().items()} else: patience_counter += 1 if patience_counter >= self.patience: logger.info(f"Early stopping at epoch {epoch + 1}") break if best_state is not None: self.model.load_state_dict(best_state) logger.info(f"NN training complete. Best val loss: {best_val_loss:.6f}") return self def predict(self, X: np.ndarray) -> np.ndarray: """Predict with inverse scaling.""" self.model.eval() X_s = self.scaler_X.transform(X) with torch.no_grad(): pred_s = self.model(torch.FloatTensor(X_s).to(self.device)).cpu().numpy() return self.scaler_y.inverse_transform(pred_s) # ============================================================================= # Weighted Ensemble # ============================================================================= class WeightedEnsemble: """ Weighted ensemble combining XGBoost and Neural Network predictions. Default weights (0.6 XGB + 0.4 NN) follow findings from Petit et al. (2025) that tree-based methods slightly outperform NNs on tabular laser data, but ensemble improves robustness. """ def __init__( self, xgb_model: XGBoostMultiOutput, nn_model: NeuralNetworkRegressor, xgb_weight: float = 0.6, nn_weight: float = 0.4, ): assert abs(xgb_weight + nn_weight - 1.0) < 1e-6, "Weights must sum to 1" self.xgb_model = xgb_model self.nn_model = nn_model self.xgb_weight = xgb_weight self.nn_weight = nn_weight def predict(self, X: np.ndarray) -> np.ndarray: """Weighted average of XGBoost and NN predictions.""" xgb_pred = self.xgb_model.predict(X) nn_pred = self.nn_model.predict(X) return self.xgb_weight * xgb_pred + self.nn_weight * nn_pred # ============================================================================= # Cross-Validation Framework # ============================================================================= def cross_validate_model( model_factory, X: np.ndarray, y: np.ndarray, n_folds: int = 5, groups: Optional[np.ndarray] = None, random_state: int = 42, ) -> Dict[str, np.ndarray]: """ Perform k-fold cross-validation with a model factory. Parameters ---------- model_factory : callable Function that returns a fresh model instance X : np.ndarray Feature matrix y : np.ndarray Target matrix n_folds : int Number of CV folds groups : np.ndarray, optional Group labels for stratified splitting random_state : int Returns ------- dict CV results with keys: 'mae', 'rmse', 'r2', 'mape' — each (n_folds, n_targets) """ if groups is not None: kf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=random_state) split_iter = kf.split(X, groups) else: kf = KFold(n_splits=n_folds, shuffle=True, random_state=random_state) split_iter = kf.split(X) n_targets = y.shape[1] results = { "mae": np.zeros((n_folds, n_targets)), "rmse": np.zeros((n_folds, n_targets)), "r2": np.zeros((n_folds, n_targets)), "mape": np.zeros((n_folds, n_targets)), } for fold, (train_idx, val_idx) in enumerate(split_iter): X_train, X_val = X[train_idx], X[val_idx] y_train, y_val = y[train_idx], y[val_idx] model = model_factory() model.fit(X_train, y_train, X_val, y_val) y_pred = model.predict(X_val) for t in range(n_targets): results["mae"][fold, t] = mean_absolute_error(y_val[:, t], y_pred[:, t]) results["rmse"][fold, t] = np.sqrt(mean_squared_error(y_val[:, t], y_pred[:, t])) results["r2"][fold, t] = r2_score(y_val[:, t], y_pred[:, t]) # MAPE with epsilon to avoid division by zero mape = np.mean(np.abs((y_val[:, t] - y_pred[:, t]) / (np.abs(y_val[:, t]) + 1e-8))) * 100 results["mape"][fold, t] = mape logger.info(f"Fold {fold + 1}/{n_folds}: mean R² = {results['r2'][fold].mean():.4f}") return results