from __future__ import annotations from typing import Any, Dict, Tuple import numpy as np from sklearn.base import BaseEstimator, RegressorMixin from sklearn.linear_model import Ridge from sklearn.neural_network import MLPRegressor from sklearn.pipeline import Pipeline from sklearn.preprocessing import StandardScaler class RawModel(BaseEstimator, RegressorMixin): """Placeholder for RAW; prediction is handled outside because it uses original model field.""" def fit(self, X, y): return self def predict(self, X): raise RuntimeError("RawModel prediction is handled outside train/predict pipeline.") # ----------------------------------------------------------------------------- # Model variants # ----------------------------------------------------------------------------- # The method names below are intentionally explicit so experiment results can be # interpreted without looking back into the code. # # Benchmark: # ridge # # XGBoost variants: # xgb_d2, xgb_d3, xgb_d4 : increasing tree depth # xgb_slow : smaller learning rate, more trees # xgb_regularized : stronger regularization/subsampling # # ANN/MLP variants using sklearn MLPRegressor: # mlp_small, mlp_medium, mlp_deep # ann_tanh : tanh activation # ann_lbfgs : deterministic quasi-Newton solver; useful for small data # # Optional CNN: # cnn1d_small : requires torch installed separately; disabled by default # ----------------------------------------------------------------------------- def _as_tuple(x: Any) -> Tuple[int, ...]: if isinstance(x, tuple): return x if isinstance(x, list): return tuple(int(v) for v in x) if isinstance(x, int): return (int(x),) raise ValueError(f"Cannot convert {x!r} to hidden_layer_sizes tuple") def _xgb_params_for_method(method: str, cfg: Dict[str, Any], random_seed: int) -> Dict[str, Any]: base = { "objective": "reg:squarederror", "n_jobs": 1, "random_state": random_seed, "tree_method": "hist", "verbosity": 0, } # Allow global override while preserving safe defaults. base.update(dict(cfg.get("xgboost_base_params", {}))) variants = dict(cfg.get("xgboost_variants", {})) if method in variants: base.update(dict(variants[method])) return base # Built-in defaults if the config does not define them. builtins = { "xgb_d2": dict(n_estimators=500, max_depth=2, learning_rate=0.03, subsample=0.90, colsample_bytree=0.90, reg_lambda=5.0, reg_alpha=0.0, min_child_weight=3.0), "xgb_d3": dict(n_estimators=500, max_depth=3, learning_rate=0.03, subsample=0.90, colsample_bytree=0.90, reg_lambda=5.0, reg_alpha=0.0, min_child_weight=3.0), "xgb_d4": dict(n_estimators=500, max_depth=4, learning_rate=0.03, subsample=0.85, colsample_bytree=0.85, reg_lambda=8.0, reg_alpha=0.0, min_child_weight=4.0), "xgb_slow": dict(n_estimators=900, max_depth=3, learning_rate=0.015, subsample=0.90, colsample_bytree=0.90, reg_lambda=8.0, reg_alpha=0.0, min_child_weight=3.0), "xgb_regularized": dict(n_estimators=700, max_depth=2, learning_rate=0.02, subsample=0.75, colsample_bytree=0.75, reg_lambda=15.0, reg_alpha=0.5, min_child_weight=5.0), } if method not in builtins: raise ValueError(f"Unknown XGBoost method variant: {method}") base.update(builtins[method]) return base def _mlp_params_for_method(method: str, cfg: Dict[str, Any], random_seed: int) -> Dict[str, Any]: base = { "random_state": random_seed, "max_iter": 700, "early_stopping": True, "validation_fraction": 0.15, "n_iter_no_change": 40, "learning_rate_init": 0.001, "alpha": 0.001, "solver": "adam", "activation": "relu", } base.update(dict(cfg.get("mlp_base_params", {}))) variants = dict(cfg.get("mlp_variants", {})) if method in variants: base.update(dict(variants[method])) else: builtins = { "mlp_small": dict(hidden_layer_sizes=(32,)), "mlp_medium": dict(hidden_layer_sizes=(64, 32)), "mlp_deep": dict(hidden_layer_sizes=(128, 64, 32), alpha=0.003), "ann_tanh": dict(hidden_layer_sizes=(64, 32), activation="tanh", alpha=0.001), "ann_lbfgs": dict(hidden_layer_sizes=(32,), activation="tanh", solver="lbfgs", alpha=0.01, max_iter=700), } if method not in builtins: raise ValueError(f"Unknown MLP/ANN method variant: {method}") base.update(builtins[method]) if "hidden_layer_sizes" in base: base["hidden_layer_sizes"] = _as_tuple(base["hidden_layer_sizes"]) # sklearn MLPRegressor does not use early_stopping with lbfgs. if base.get("solver") == "lbfgs": base.pop("early_stopping", None) base.pop("validation_fraction", None) base.pop("n_iter_no_change", None) base.pop("learning_rate_init", None) return base def make_model(method: str, cfg: Dict[str, Any], random_seed: int = 42): method = method.lower() if method == "ridge": alpha = float(cfg.get("ridge_alpha", 1.0)) return Pipeline([ ("scaler", StandardScaler()), ("model", Ridge(alpha=alpha, random_state=random_seed)), ]) if method in {"mlp", "mlp_small", "mlp_medium", "mlp_deep", "ann_tanh", "ann_lbfgs"}: # Keep backward compatibility: method "mlp" uses the medium variant unless mlp_params is supplied. if method == "mlp" and "mlp_params" in cfg: params = dict(cfg.get("mlp_params", {})) if "hidden_layer_sizes" in params: params["hidden_layer_sizes"] = _as_tuple(params["hidden_layer_sizes"]) params.setdefault("random_state", random_seed) else: method2 = "mlp_medium" if method == "mlp" else method params = _mlp_params_for_method(method2, cfg, random_seed) return Pipeline([ ("scaler", StandardScaler()), ("model", MLPRegressor(**params)), ]) if method in {"xgboost", "xgb_d2", "xgb_d3", "xgb_d4", "xgb_slow", "xgb_regularized"}: try: from xgboost import XGBRegressor except Exception as exc: # pragma: no cover raise ImportError( "xgboost is not installed. Install it or remove XGBoost variants from methods in the config." ) from exc if method == "xgboost" and "xgboost_params" in cfg: params = dict(cfg.get("xgboost_params", {})) params.setdefault("random_state", random_seed) params.setdefault("n_jobs", 1) else: method2 = "xgb_d3" if method == "xgboost" else method params = _xgb_params_for_method(method2, cfg, random_seed) return XGBRegressor(**params) if method == "cnn1d_small": from .torch_models import TorchCNN1DRegressor return Pipeline([ ("scaler", StandardScaler()), ("model", TorchCNN1DRegressor(random_seed=random_seed, **dict(cfg.get("cnn1d_small_params", {})))), ]) raise ValueError(f"Unknown model method: {method}")