diff --git a/corebehrt/azure/components/train_baseline.py b/corebehrt/azure/components/train_baseline.py new file mode 100644 index 00000000..59a19cd5 --- /dev/null +++ b/corebehrt/azure/components/train_baseline.py @@ -0,0 +1,12 @@ +from corebehrt.azure.util import job + +INPUTS = { + "prepared_data": {"type": "uri_folder"}, +} +OUTPUTS = {"model": {"type": "uri_folder"}} + + +if __name__ == "__main__": + from corebehrt.main_causal import train_baseline + + job.run_main("train_baseline", train_baseline.main_baseline, INPUTS, OUTPUTS) diff --git a/corebehrt/azure/main/job.py b/corebehrt/azure/main/job.py index 0f957f1c..1f627867 100644 --- a/corebehrt/azure/main/job.py +++ b/corebehrt/azure/main/job.py @@ -57,6 +57,7 @@ def add_parser(subparsers) -> None: "get_stats", "prepare_ft_exp_y", "finetune_exp_y", + "train_baseline", "calibrate_exp_y", "xgboost_cv", "evaluate_xgboost", diff --git a/corebehrt/configs/causal/finetune/baseline.yaml b/corebehrt/configs/causal/finetune/baseline.yaml new file mode 100644 index 00000000..46de0077 --- /dev/null +++ b/corebehrt/configs/causal/finetune/baseline.yaml @@ -0,0 +1,32 @@ +logging: + level: INFO + path: ./outputs/logs/causal + +paths: + ## INPUTS + prepared_data: ./outputs/causal/finetune/prepared_data + + ## OUTPUTS + model: ./outputs/causal/finetune/models/baseline + +# Which baseline to train: logistic (L2 logistic regression) or catboost +model: logistic + +# Feature encoding of the tokenized sequences +multihot: false # false -> binary code presence; true -> code counts +include_age: true + +# Parameters set here are FIXED; the remaining tunable ones are searched by Optuna. +logistic: + max_iter: 1000 + # C: 1.0 # regularisation strength; TUNED in [1e-4, 1e2] if not set + +# catboost: +# n_estimators: 1000 +# early_stopping_rounds: 50 + +tuning: + tune_hyperparameters: true + n_trials: 10 + inner_val_size: 0.2 + reuse_hyperparameters: true diff --git a/corebehrt/configs/causal/finetune/simulated_bl.yaml b/corebehrt/configs/causal/finetune/simulated_bl.yaml index de37c3fd..65de385e 100644 --- a/corebehrt/configs/causal/finetune/simulated_bl.yaml +++ b/corebehrt/configs/causal/finetune/simulated_bl.yaml @@ -26,6 +26,7 @@ include_age: true # If true, include age-based features (mean, min, max, std, ra # ============================================================================== # These are the base parameters for the CatBoost model. # They are used directly if tuning is disabled, or as a starting point for tuning. +model: catboost catboost: n_estimators: 100 # Max number of trees. Training will likely stop early. learning_rate: 0.03 # Step size shrinkage to prevent overfitting. diff --git a/corebehrt/main_causal/helper/baseline_models.py b/corebehrt/main_causal/helper/baseline_models.py new file mode 100644 index 00000000..21a5e398 --- /dev/null +++ b/corebehrt/main_causal/helper/baseline_models.py @@ -0,0 +1,202 @@ +""" +Model factory for the tabular baselines used as a reference for the transformer. + +Two models are supported: +- `logistic`: L2-regularised logistic regression on standardised features. +- `catboost`: gradient boosting on the same features. + +Both are fitted on the one-hot/multi-hot code matrix produced by +`create_features_from_patients`, so they share the nested CV machinery in +`helper/train_baseline.py`. +""" + +import logging +from typing import Any, Dict, Optional, Tuple + +import numpy as np +import pandas as pd +import torch +from catboost import CatBoostClassifier +from sklearn.linear_model import LogisticRegression +from sklearn.pipeline import make_pipeline +from sklearn.preprocessing import StandardScaler + +LOGISTIC = "logistic" +CATBOOST = "catboost" +SUPPORTED_MODELS = (LOGISTIC, CATBOOST) + +# Parameters that are never tuned by Optuna, per model. +NON_TUNABLE_DEFAULTS = { + LOGISTIC: {"max_iter": 1000}, + CATBOOST: {"n_estimators": 1000, "early_stopping_rounds": 50}, +} + +# (type, min, max[, log_scale]) per tunable parameter. +TUNING_RANGES = { + LOGISTIC: { + "C": ("float", 1e-4, 1e2, True), + }, + CATBOOST: { + "learning_rate": ("float", 0.01, 0.3, True), + "max_depth": ("int", 4, 10), + "subsample": ("float", 0.6, 1.0, False), + "l2_leaf_reg": ("float", 1e-8, 10.0, True), + "min_data_in_leaf": ("int", 1, 100), + # colsample_bylevel is added below for CPU only (GPU does not support it). + }, +} + +# Cache for GPU detection to avoid repeated logging +_CATBOOST_DEVICE_PARAMS_CACHE = None + + +def get_model_name(cfg) -> str: + """Returns the baseline model to train, defaulting to logistic regression.""" + model_name = cfg.get("model", LOGISTIC) + if model_name not in SUPPORTED_MODELS: + raise ValueError( + f"Unknown baseline model '{model_name}'. Choose one of {SUPPORTED_MODELS}." + ) + return model_name + + +def get_base_params(cfg, model_name: str) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """ + Returns (base_params, config_params) for the selected model. + + `config_params` are the parameters explicitly set in the config; these are + held FIXED during tuning. `base_params` additionally contains the + non-tunable defaults. + """ + config_params = dict(cfg.get(model_name, {})) + base_params = {**NON_TUNABLE_DEFAULTS[model_name], **config_params} + return base_params, config_params + + +def get_tuning_ranges( + model_name: str, config_params: Dict[str, Any] +) -> Dict[str, tuple]: + """Returns the tunable parameters and their ranges for the selected model.""" + ranges = dict(TUNING_RANGES[model_name]) + if model_name == CATBOOST: + # colsample_bylevel is only supported on CPU for classification. + if _effective_device_params(config_params)["task_type"] != "GPU": + ranges["colsample_bylevel"] = ("float", 0.6, 1.0, False) + return ranges + + +def build_model( + model_name: str, + params: Dict[str, Any], + scale_pos_weight: float, + random_seed: int, +) -> Any: + """Builds an unfitted estimator with a `predict_proba` interface.""" + if model_name == LOGISTIC: + return make_pipeline( + StandardScaler(), + LogisticRegression( + class_weight="balanced", + random_state=random_seed, + **params, + ), + ) + + device_params = _effective_device_params(params) + catboost_params = _prepare_catboost_params(params, device_params) + # early_stopping_rounds is passed to fit, not to the constructor. + catboost_params.pop("early_stopping_rounds", None) + return CatBoostClassifier( + scale_pos_weight=scale_pos_weight, + random_state=random_seed, + verbose=0, + **{**device_params, **catboost_params}, + ) + + +def fit_model( + model: Any, + model_name: str, + params: Dict[str, Any], + X_train: pd.DataFrame, + y_train: np.ndarray, + X_val: Optional[pd.DataFrame] = None, + y_val: Optional[np.ndarray] = None, +) -> Any: + """Fits the estimator, using early stopping on the validation set if supported.""" + if model_name == LOGISTIC: + model.fit(X_train, y_train) + return model + + if X_val is None: + model.fit(X_train, y_train, verbose=0) + return model + + model.fit( + X_train, + y_train, + eval_set=[(X_val, y_val)], + early_stopping_rounds=params.get("early_stopping_rounds"), + verbose=0, + ) + return model + + +def _effective_device_params(params: Dict[str, Any]) -> Dict[str, Any]: + """An explicit task_type/devices in the config wins over the detected device.""" + device_params = dict(_get_catboost_device_params()) + for key in ("task_type", "devices"): + if key in params: + device_params[key] = params[key] + return device_params + + +def _get_catboost_device_params() -> Dict[str, Any]: + """ + Detect GPU availability and return appropriate CatBoost parameters. + Returns task_type and devices parameters for CatBoost. + Logs only on first call (cached). + """ + global _CATBOOST_DEVICE_PARAMS_CACHE + + if _CATBOOST_DEVICE_PARAMS_CACHE is None: + if torch.cuda.is_available(): + logging.info("GPU detected. CatBoost will use GPU for training.") + _CATBOOST_DEVICE_PARAMS_CACHE = {"task_type": "GPU", "devices": "0"} + else: + logging.info("No GPU detected. CatBoost will use CPU for training.") + _CATBOOST_DEVICE_PARAMS_CACHE = {"task_type": "CPU"} + + return _CATBOOST_DEVICE_PARAMS_CACHE + + +def _prepare_catboost_params( + params: Dict[str, Any], device_params: Dict[str, Any] +) -> Dict[str, Any]: + """ + Prepare CatBoost parameters by: + 1. Adding bootstrap_type if subsample is used (Bayesian bootstrap doesn't support subsample) + 2. Removing GPU-incompatible parameters when using GPU mode + + GPU limitations: + - colsample_bylevel (RSM) is only supported in pairwise ranking modes, not classification + """ + params_copy = params.copy() + + # Handle bootstrap type for subsample + if "subsample" in params_copy and "bootstrap_type" not in params_copy: + params_copy["bootstrap_type"] = "Bernoulli" + + # Remove GPU-incompatible parameters + if device_params.get("task_type") == "GPU": + gpu_incompatible_params = [ + "colsample_bylevel", + "colsample_bynode", + "colsample_bytree", + ] + for param in gpu_incompatible_params: + if param in params_copy: + logging.debug(f"Removing GPU-incompatible parameter: {param}") + params_copy.pop(param) + + return params_copy diff --git a/corebehrt/main_causal/helper/train_baseline.py b/corebehrt/main_causal/helper/train_baseline.py index 9cf77e10..322622b7 100644 --- a/corebehrt/main_causal/helper/train_baseline.py +++ b/corebehrt/main_causal/helper/train_baseline.py @@ -10,7 +10,6 @@ import optuna import pandas as pd import torch -from catboost import CatBoostClassifier from sklearn.metrics import roc_auc_score from sklearn.model_selection import train_test_split @@ -31,65 +30,11 @@ from corebehrt.functional.preparation.causal.one_hot import ( create_features_from_patients, ) +from corebehrt.main_causal.helper import baseline_models from corebehrt.modules.preparation.causal.dataset import CausalPatientDataset from corebehrt.modules.setup.config import Config -# Cache for GPU detection to avoid repeated logging -_CATBOOST_DEVICE_PARAMS_CACHE = None - - -def _get_catboost_device_params() -> Dict[str, Any]: - """ - Detect GPU availability and return appropriate CatBoost parameters. - Returns task_type and devices parameters for CatBoost. - Logs only on first call (cached). - """ - global _CATBOOST_DEVICE_PARAMS_CACHE - - if _CATBOOST_DEVICE_PARAMS_CACHE is None: - if torch.cuda.is_available(): - logging.info("GPU detected. CatBoost will use GPU for training.") - _CATBOOST_DEVICE_PARAMS_CACHE = {"task_type": "GPU", "devices": "0"} - else: - logging.info("No GPU detected. CatBoost will use CPU for training.") - _CATBOOST_DEVICE_PARAMS_CACHE = {"task_type": "CPU"} - - return _CATBOOST_DEVICE_PARAMS_CACHE - - -def _prepare_catboost_params( - params: Dict[str, Any], device_params: Dict[str, Any] -) -> Dict[str, Any]: - """ - Prepare CatBoost parameters by: - 1. Adding bootstrap_type if subsample is used (Bayesian bootstrap doesn't support subsample) - 2. Removing GPU-incompatible parameters when using GPU mode - - GPU limitations: - - colsample_bylevel (RSM) is only supported in pairwise ranking modes, not classification - """ - params_copy = params.copy() - - # Handle bootstrap type for subsample - if "subsample" in params_copy and "bootstrap_type" not in params_copy: - params_copy["bootstrap_type"] = "Bernoulli" - - # Remove GPU-incompatible parameters - if device_params.get("task_type") == "GPU": - gpu_incompatible_params = [ - "colsample_bylevel", - "colsample_bynode", - "colsample_bytree", - ] - for param in gpu_incompatible_params: - if param in params_copy: - logging.debug(f"Removing GPU-incompatible parameter: {param}") - params_copy.pop(param) - - return params_copy - - @dataclass class FoldPredictionData: """Container for storing predictions from each fold.""" @@ -246,6 +191,7 @@ def run_hyperparameter_tuning( config_params: Dict[str, Any], n_trials: int, scale_pos_weight: float, + model_name: str, ) -> Dict[str, Any]: """ INNER LOOP: Performs hyperparameter tuning using Optuna on a given train/val split. @@ -260,26 +206,7 @@ def run_hyperparameter_tuning( logging.info(f" Val class distribution: {np.bincount(y_val)}") logging.info(f" Scale pos weight: {scale_pos_weight:.4f}") - # Detect device type for GPU compatibility - device_params = _get_catboost_device_params() - is_gpu = device_params.get("task_type") == "GPU" - - # Define default tuning ranges for each parameter - TUNING_RANGES = { - "learning_rate": ("float", 0.01, 0.3, True), # (type, min, max, log_scale) - "max_depth": ("int", 4, 10), # (type, min, max) - "subsample": ("float", 0.6, 1.0, False), - "l2_leaf_reg": ("float", 1e-8, 10.0, True), - "min_data_in_leaf": ("int", 1, 100), - } - - # Add colsample_bylevel only if NOT using GPU (GPU doesn't support it for classification) - if not is_gpu: - TUNING_RANGES["colsample_bylevel"] = ("float", 0.6, 1.0, False) - else: - logging.info( - " GPU mode detected: skipping colsample_bylevel from tuning (GPU incompatible)" - ) + tuning_ranges = baseline_models.get_tuning_ranges(model_name, config_params) # Determine which parameters to tune vs. fix # RULE: If parameter is explicitly in CONFIG → FIXED @@ -287,7 +214,7 @@ def run_hyperparameter_tuning( params_to_tune = {} fixed_params = {} - for param_name, range_info in TUNING_RANGES.items(): + for param_name, range_info in tuning_ranges.items(): if param_name in config_params: # Parameter explicitly set in config → FIXED fixed_params[param_name] = config_params[param_name] @@ -323,27 +250,20 @@ def objective(trial: optuna.Trial): param_name, range_info[1], range_info[2] ) - # Get GPU/CPU parameters - device_params = _get_catboost_device_params() - - # Prepare trial params with proper bootstrap type and GPU compatibility - prepared_trial_params = _prepare_catboost_params(trial_params, device_params) - - model = CatBoostClassifier( - n_estimators=base_params["n_estimators"], - scale_pos_weight=scale_pos_weight, - random_state=42, - verbose=0, - **device_params, - **prepared_trial_params, + model = baseline_models.build_model( + model_name, + {**base_params, **trial_params}, + scale_pos_weight, + random_seed=42, ) - - model.fit( + baseline_models.fit_model( + model, + model_name, + base_params, X_train, y_train, - eval_set=[(X_val, y_val)], - early_stopping_rounds=base_params["early_stopping_rounds"], - verbose=0, + X_val, + y_val, ) preds = model.predict_proba(X_val)[:, 1] @@ -375,17 +295,11 @@ def objective(trial: optuna.Trial): return final_params -def _setup_model_parameters(cfg: Config) -> Tuple[Dict[str, Any], Dict[str, Any]]: - """Loads and merges CatBoost and tuning parameters from the config.""" - # Only include defaults for parameters that are NEVER tuned - NON_TUNABLE_DEFAULTS = { - "n_estimators": 1000, - "early_stopping_rounds": 50, - } - - config_params = cfg.get("catboost", {}) - # Only use config params + non-tunable defaults - base_params = {**NON_TUNABLE_DEFAULTS, **config_params} +def _setup_model_parameters( + cfg: Config, model_name: str +) -> Tuple[Dict[str, Any], Dict[str, Any]]: + """Loads and merges model and tuning parameters from the config.""" + base_params, _ = baseline_models.get_base_params(cfg, model_name) tuning_cfg = cfg.get("tuning", {}) return base_params, tuning_cfg @@ -433,6 +347,7 @@ def _get_best_params_for_fold( tuning_cfg: Dict, data: CausalPatientDataset, cfg: Config, + model_name: str, ) -> Dict[str, Any]: """Performs inner-loop splitting and hyperparameter tuning.""" logger = logging.getLogger("train_baseline") @@ -463,7 +378,7 @@ def _get_best_params_for_fold( scale_pos_weight = (y_inner_train == 0).sum() / max((y_inner_train == 1).sum(), 1) # Get config parameters for this tuning session - config_params = cfg.get("catboost", {}) + _, config_params = baseline_models.get_base_params(cfg, model_name) tuned_params = run_hyperparameter_tuning( X_inner_train, @@ -474,6 +389,7 @@ def _get_best_params_for_fold( config_params, n_trials, scale_pos_weight, + model_name, ) logger.info(" Hyperparameter tuning completed for this fold") @@ -481,7 +397,7 @@ def _get_best_params_for_fold( def _generate_counterfactual_predictions( - model: CatBoostClassifier, + model: Any, X_test: pd.DataFrame, target_name: str, logger: logging.Logger, @@ -520,6 +436,7 @@ def _train_and_evaluate_fold( target_name: str, fold_idx: int, prediction_storage: List[FoldPredictionData], + model_name: str, ) -> float: """Trains a final model and evaluates it on the holdout test set.""" logger.info(f" Training final model with outer train set size: {len(X_train)}") @@ -532,21 +449,15 @@ def _train_and_evaluate_fold( scale_pos_weight = (y_train == 0).sum() / max((y_train == 1).sum(), 1) logger.info(f" Scale pos weight for final model: {scale_pos_weight:.4f}") - # Get GPU/CPU parameters - device_params = _get_catboost_device_params() - - # Prepare best params with proper bootstrap type and GPU compatibility - prepared_best_params = _prepare_catboost_params(best_params, device_params) - - final_model = CatBoostClassifier( - scale_pos_weight=scale_pos_weight, - random_state=42, - **device_params, - **prepared_best_params, + final_model = baseline_models.build_model( + model_name, + best_params, + scale_pos_weight, + random_seed=42, ) logger.info(" Fitting final model...") - final_model.fit(X_train, y_train, verbose=0) + baseline_models.fit_model(final_model, model_name, best_params, X_train, y_train) logger.info(" Generating predictions on test set...") y_pred_proba = final_model.predict_proba(X_test)[:, 1] @@ -617,7 +528,10 @@ def nested_cv_loop( targets_to_train = cfg.get("targets", [EXPOSURE] + data.get_outcome_names()) logger.info(f"Starting Nested Cross-Validation for targets: {targets_to_train}") - base_params, tuning_cfg = _setup_model_parameters(cfg) + model_name = baseline_models.get_model_name(cfg) + logger.info(f"Baseline model: {model_name}") + + base_params, tuning_cfg = _setup_model_parameters(cfg, model_name) should_tune = tuning_cfg.get("tune_hyperparameters", True) reuse_hyperparameters = tuning_cfg.get("reuse_hyperparameters", True) @@ -656,6 +570,7 @@ def nested_cv_loop( tuning_cfg, data, cfg, + model_name, ) best_params = tuned_params if reuse_hyperparameters: @@ -688,6 +603,7 @@ def nested_cv_loop( target_name, i, prediction_storage, + model_name, ) all_unbiased_scores.append(unbiased_auc) diff --git a/experiments/causal_pipeline/base_configs/train_baseline.yaml b/experiments/causal_pipeline/base_configs/train_baseline.yaml index 1153453b..f03a190a 100644 --- a/experiments/causal_pipeline/base_configs/train_baseline.yaml +++ b/experiments/causal_pipeline/base_configs/train_baseline.yaml @@ -20,6 +20,7 @@ include_age: true # Note: Parameters specified here will be FIXED (not tuned by Optuna) # Parameters not specified will be tuned within predefined ranges # ============================================================================== +model: catboost catboost: n_estimators: 100 task_type: GPU # Enable GPU diff --git a/experiments/causal_pipeline_resample/base_configs/train_baseline.yaml b/experiments/causal_pipeline_resample/base_configs/train_baseline.yaml index 08f31041..ad579be9 100644 --- a/experiments/causal_pipeline_resample/base_configs/train_baseline.yaml +++ b/experiments/causal_pipeline_resample/base_configs/train_baseline.yaml @@ -20,6 +20,7 @@ include_age: true # Note: Parameters specified here will be FIXED (not tuned by Optuna) # Parameters not specified will be tuned within predefined ranges # ============================================================================== +model: catboost catboost: n_estimators: 100 # task_type: CPU # Default is CPU. Set to GPU if you have CUDA available diff --git a/tests/test_main_causal/test_helper/test_baseline_models.py b/tests/test_main_causal/test_helper/test_baseline_models.py new file mode 100644 index 00000000..b480b426 --- /dev/null +++ b/tests/test_main_causal/test_helper/test_baseline_models.py @@ -0,0 +1,100 @@ +import unittest + +import numpy as np +import pandas as pd + +from corebehrt.main_causal.helper import baseline_models + + +def make_separable_data(n_samples: int = 200): + rng = np.random.RandomState(0) + features = pd.DataFrame( + { + "code_a": rng.randint(0, 2, n_samples), + "code_b": rng.randint(0, 2, n_samples), + "age": rng.randint(40, 80, n_samples), + } + ) + targets = (features["code_a"] == 1).astype(int).values + return features, targets + + +class TestBaselineModels(unittest.TestCase): + def test_default_model_is_logistic(self): + self.assertEqual(baseline_models.get_model_name({}), baseline_models.LOGISTIC) + + def test_unknown_model_raises(self): + with self.assertRaises(ValueError): + baseline_models.get_model_name({"model": "randomforest"}) + + def test_base_params_merge_config_over_defaults(self): + cfg = {"logistic": {"max_iter": 50}} + base_params, config_params = baseline_models.get_base_params( + cfg, baseline_models.LOGISTIC + ) + self.assertEqual(base_params["max_iter"], 50) + self.assertEqual(config_params, {"max_iter": 50}) + + def test_tuning_ranges_are_model_specific(self): + logistic_ranges = baseline_models.get_tuning_ranges( + baseline_models.LOGISTIC, {} + ) + catboost_ranges = baseline_models.get_tuning_ranges( + baseline_models.CATBOOST, {} + ) + self.assertEqual(list(logistic_ranges), ["C"]) + self.assertIn("learning_rate", catboost_ranges) + self.assertNotIn("C", catboost_ranges) + + def test_configured_gpu_excludes_gpu_incompatible_parameters(self): + """colsample_bylevel is unsupported on GPU, also when GPU comes from the config.""" + ranges = baseline_models.get_tuning_ranges( + baseline_models.CATBOOST, {"task_type": "GPU", "devices": "0"} + ) + self.assertNotIn("colsample_bylevel", ranges) + + params = {"n_estimators": 10, "task_type": "GPU", "colsample_bylevel": 0.8} + prepared = baseline_models._prepare_catboost_params( + params, baseline_models._effective_device_params(params) + ) + self.assertNotIn("colsample_bylevel", prepared) + + def test_logistic_fits_and_predicts(self): + features, targets = make_separable_data() + params, _ = baseline_models.get_base_params({}, baseline_models.LOGISTIC) + model = baseline_models.build_model( + baseline_models.LOGISTIC, params, scale_pos_weight=1.0, random_seed=42 + ) + baseline_models.fit_model( + model, baseline_models.LOGISTIC, params, features, targets + ) + + probas = model.predict_proba(features)[:, 1] + self.assertEqual(probas.shape, targets.shape) + self.assertTrue(((probas >= 0) & (probas <= 1)).all()) + # code_a fully determines the target, so the fit should separate the classes. + self.assertGreater(probas[targets == 1].mean(), probas[targets == 0].mean()) + + def test_catboost_fits_with_early_stopping_on_validation_set(self): + features, targets = make_separable_data() + params = {"n_estimators": 10, "early_stopping_rounds": 5} + model = baseline_models.build_model( + baseline_models.CATBOOST, params, scale_pos_weight=1.0, random_seed=42 + ) + baseline_models.fit_model( + model, + baseline_models.CATBOOST, + params, + features, + targets, + features, + targets, + ) + + probas = model.predict_proba(features)[:, 1] + self.assertEqual(probas.shape, targets.shape) + self.assertGreater(probas[targets == 1].mean(), probas[targets == 0].mean()) + + +if __name__ == "__main__": + unittest.main()