diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml index 8e882a61..b1bafd4e 100644 --- a/.github/workflows/format.yml +++ b/.github/workflows/format.yml @@ -3,23 +3,23 @@ name: Format on: push: - branches: [ "main" ] + branches: ["main"] pull_request: - branches: [ "main" ] + branches: ["main"] permissions: contents: read jobs: format: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 - - uses: actions/setup-python@v4 - with: - python-version: '3.11' - cache: 'pip' # caching pip dependencies - - name: Install dependencies - run: | - pip install ruff==0.11.5 - - name: Run ruff - run: | - ruff format --check corebehrt tests + - uses: actions/checkout@v3 + - uses: actions/setup-python@v4 + with: + python-version: "3.11" + cache: "pip" # caching pip dependencies + - name: Install dependencies + run: | + pip install ruff==0.11.5 + - name: Run ruff + run: | + ruff format --check corebehrt tests diff --git a/corebehrt/azure/components/finetune_subpop.py b/corebehrt/azure/components/finetune_subpop.py new file mode 100644 index 00000000..7a57f992 --- /dev/null +++ b/corebehrt/azure/components/finetune_subpop.py @@ -0,0 +1,16 @@ +from corebehrt.azure.util import job + +INPUTS = { + "prepared_data": {"type": "uri_folder"}, + "restart_model": {"type": "uri_folder"}, + "subpopulation_pids": {"type": "uri_file"}, +} +OUTPUTS = {"model": {"type": "uri_folder"}} + + +if __name__ == "__main__": + from corebehrt.main_causal import finetune_subpop + + job.run_main( + "finetune_subpop", finetune_subpop.main_finetune_subpop, INPUTS, OUTPUTS + ) diff --git a/corebehrt/azure/pipelines/SUBPOP_FINETUNE_CALIBRATE_ESTIMATE.py b/corebehrt/azure/pipelines/SUBPOP_FINETUNE_CALIBRATE_ESTIMATE.py new file mode 100644 index 00000000..1d85d5ad --- /dev/null +++ b/corebehrt/azure/pipelines/SUBPOP_FINETUNE_CALIBRATE_ESTIMATE.py @@ -0,0 +1,195 @@ +""" +Subpopulation Finetune-Calibrate-Estimate pipeline. +Continues fine-tuning on a subpopulation using checkpoints from a main run, +then calibrates and estimates causal effects. +""" + +from typing import Any, Dict + +from corebehrt.azure.pipelines.base import PipelineArg, PipelineMeta + +SUBPOP_FINETUNE_CALIBRATE_ESTIMATE = PipelineMeta( + name="SUBPOP_FINETUNE_CALIBRATE_ESTIMATE", + help="Continue fine-tuning on a subpopulation from a main run, then calibrate and estimate.", + inputs=[ + PipelineArg( + name="prepared_data", + help="Path to the prepared data (from the main run).", + required=True, + ), + PipelineArg( + name="finetune_model", + help="Path to the finetuned model from the main run (used as restart_model).", + required=True, + ), + PipelineArg( + name="subpopulation_pids", + help="Path to file with subpopulation patient IDs (.pt).", + required=True, + ), + PipelineArg( + name="counterfactual_outcomes", + help="Path to counterfactual outcomes (optional, for simulated data).", + required=False, + ), + PipelineArg( + name="secondary_cohort_config", + help="Path to secondary cohort config YAML file (optional).", + required=False, + ), + ], +) + + +def create(component: callable): + """Define the Subpopulation Finetune-Calibrate-Estimate pipeline.""" + from azure.ai.ml import Input, dsl + + def _common_pipeline_steps( + prepared_data: Input, + finetune_model: Input, + subpopulation_pids: Input, + counterfactual_outcomes: Input = None, + secondary_cohort_config: Input = None, + ) -> dict: + finetune_subpop = component( + "finetune_subpop", + )( + prepared_data=prepared_data, + restart_model=finetune_model, + subpopulation_pids=subpopulation_pids, + ) + + calibrate_exp_y = component( + "calibrate_exp_y", + )( + finetune_model=finetune_subpop.outputs.model, + ) + + estimate_kwargs = { + "calibrated_predictions": calibrate_exp_y.outputs.calibrated_predictions, + } + if counterfactual_outcomes is not None: + estimate_kwargs["counterfactual_outcomes"] = counterfactual_outcomes + + estimate = component( + "estimate", + )(**estimate_kwargs) + + get_stats_kwargs = { + "ps_calibrated_predictions": calibrate_exp_y.outputs.calibrated_predictions, + } + if secondary_cohort_config is not None: + get_stats_kwargs["secondary_cohort_config"] = secondary_cohort_config + + get_stats = component( + "get_stats", + )(**get_stats_kwargs) + + return { + "estimate": estimate.outputs.estimate, + "calibrated_predictions": calibrate_exp_y.outputs.calibrated_predictions, + "stats": get_stats.outputs.stats, + } + + pipeline_configs = {} + + @dsl.pipeline( + name="subpop_ft_cal_est_w_cf", + description="Subpopulation pipeline with counterfactual outcomes", + ) + def _pipeline_with_counterfactual( + prepared_data: Input, + finetune_model: Input, + subpopulation_pids: Input, + counterfactual_outcomes: Input, + ) -> dict: + return _common_pipeline_steps( + prepared_data, + finetune_model, + subpopulation_pids, + counterfactual_outcomes=counterfactual_outcomes, + ) + + pipeline_configs["has_counterfactual"] = _pipeline_with_counterfactual + + @dsl.pipeline( + name="subpop_ft_cal_est_wo_cf", + description="Subpopulation pipeline without counterfactual outcomes", + ) + def _pipeline_without_counterfactual( + prepared_data: Input, + finetune_model: Input, + subpopulation_pids: Input, + ) -> dict: + return _common_pipeline_steps(prepared_data, finetune_model, subpopulation_pids) + + pipeline_configs["does_not_have_counterfactual"] = _pipeline_without_counterfactual + + @dsl.pipeline( + name="subpop_ft_cal_est_w_secondary", + description="Subpopulation pipeline with secondary cohort config", + ) + def _pipeline_with_secondary_cohort( + prepared_data: Input, + finetune_model: Input, + subpopulation_pids: Input, + secondary_cohort_config: Input, + ) -> dict: + return _common_pipeline_steps( + prepared_data, + finetune_model, + subpopulation_pids, + secondary_cohort_config=secondary_cohort_config, + ) + + pipeline_configs["has_secondary_cohort"] = _pipeline_with_secondary_cohort + + @dsl.pipeline( + name="subpop_ft_cal_est_w_cf_and_secondary", + description="Subpopulation pipeline with counterfactual outcomes and secondary cohort config", + ) + def _pipeline_with_both( + prepared_data: Input, + finetune_model: Input, + subpopulation_pids: Input, + counterfactual_outcomes: Input, + secondary_cohort_config: Input, + ) -> dict: + return _common_pipeline_steps( + prepared_data, + finetune_model, + subpopulation_pids, + counterfactual_outcomes=counterfactual_outcomes, + secondary_cohort_config=secondary_cohort_config, + ) + + pipeline_configs["has_both"] = _pipeline_with_both + + def pipeline_factory(**kwargs: Dict[str, Any]): + has_counterfactual = ( + "counterfactual_outcomes" in kwargs + and kwargs["counterfactual_outcomes"] is not None + ) + has_secondary_cohort = ( + "secondary_cohort_config" in kwargs + and kwargs["secondary_cohort_config"] is not None + ) + + if has_counterfactual and has_secondary_cohort: + selected_pipeline = pipeline_configs["has_both"] + elif has_secondary_cohort: + selected_pipeline = pipeline_configs["has_secondary_cohort"] + elif has_counterfactual: + selected_pipeline = pipeline_configs["has_counterfactual"] + else: + selected_pipeline = pipeline_configs["does_not_have_counterfactual"] + + from inspect import signature + + pipeline_params = signature(selected_pipeline).parameters.keys() + filtered_kwargs = {k: v for k, v in kwargs.items() if k in pipeline_params} + + return selected_pipeline(**filtered_kwargs) + + return pipeline_factory diff --git a/corebehrt/azure/pipelines/__init__.py b/corebehrt/azure/pipelines/__init__.py index 17f9ae03..2af0bef8 100644 --- a/corebehrt/azure/pipelines/__init__.py +++ b/corebehrt/azure/pipelines/__init__.py @@ -7,6 +7,9 @@ from corebehrt.azure.pipelines.FINETUNE_ESTIMATE_SIMULATED import ( FINETUNE_ESTIMATE_SIMULATED, ) +from corebehrt.azure.pipelines.SUBPOP_FINETUNE_CALIBRATE_ESTIMATE import ( + SUBPOP_FINETUNE_CALIBRATE_ESTIMATE, +) PIPELINE_REGISTRY = [ E2E, @@ -14,4 +17,5 @@ FINETUNE_CALIBRATE_ESTIMATE, FINETUNE_ESTIMATE, FINETUNE_ESTIMATE_SIMULATED, + SUBPOP_FINETUNE_CALIBRATE_ESTIMATE, ] diff --git a/corebehrt/configs/causal/finetune/ft_subpop.yaml b/corebehrt/configs/causal/finetune/ft_subpop.yaml new file mode 100644 index 00000000..c05905d6 --- /dev/null +++ b/corebehrt/configs/causal/finetune/ft_subpop.yaml @@ -0,0 +1,67 @@ +logging: + level: INFO + path: ./outputs/logs/causal + +paths: + ## INPUTS + restart_model: ./outputs/causal/finetune/models/simple # Main run's model output + prepared_data: ./outputs/causal/finetune/prepared_data # Same prepared data as main run + subpopulation_pids: ./outputs/causal/finetune/subpopulation_pids.pt + + ## OUTPUTS + model: ./outputs/causal/finetune/models/subpop + +bootstrap: true + +data: + n_folds: 5 + seed: 42 + +model: + head: + shared_representation: true + bidirectional: true + bottleneck_dim: 64 + l1_lambda: 0.2 + pooling_strategy: gru + +trainer_args: + restart_weights_only: true + loss_weight_function: + _target_: corebehrt.modules.trainer.utils.PositiveWeight.sqrt + batch_size: 128 + val_batch_size: 256 + effective_batch_size: 128 + epochs: 3 + info: true + shuffle: true + checkpoint_frequency: 1 + use_pcgrad: true + + early_stopping: 3 + stopping_criterion: roc_auc + + # Freeze encoder from the start — only train pooler/cls/heads + freeze_encoder_at_init: true + freeze_encoder_on_plateau: false + n_layers_to_freeze: 0 + + plot_histograms: true + plot_all_targets: false + num_targets_to_log: 3 + save_curves: false + +optimizer: + lr: 1e-3 + eps: 1e-6 + +scheduler: + _target_: transformers.get_linear_schedule_with_warmup + num_training_epochs: 5 + num_warmup_epochs: 1 + +metrics: + roc_auc: + _target_: corebehrt.modules.monitoring.metrics.ROC_AUC + pr_auc: + _target_: corebehrt.modules.monitoring.metrics.PR_AUC diff --git a/corebehrt/functional/causal/calibration.py b/corebehrt/functional/causal/calibration.py index f533bf12..bcfa8eb8 100644 --- a/corebehrt/functional/causal/calibration.py +++ b/corebehrt/functional/causal/calibration.py @@ -52,7 +52,25 @@ def calibrate_folds( train_pids, val_pids = fold[TRAIN_KEY], fold[VAL_KEY] train_df, val_df = split_data(df, train_pids, val_pids) - calibrator = train_calibrator(train_df[PROBAS], train_df[TARGETS]) + train_targets = train_df[TARGETS].values + if len(np.unique(train_targets)) < 2: + logger.warning( + "Fold %d/%d: training split has only one class for this target; " + "skipping Beta calibration (using raw validation probabilities).", + fold_num, + len(folds), + ) + fold_results = { + PID_COL: val_df[PID_COL].values, + PROBAS: val_df[PROBAS].values, + TARGETS: val_df[TARGETS].values, + } + if CF_PROBAS in val_df.columns: + fold_results[CF_PROBAS] = val_df[CF_PROBAS].values + calibrated_dfs.append(pd.DataFrame(fold_results)) + continue + + calibrator = train_calibrator(train_df[PROBAS], train_targets) initial_calibrated_probas = calibrator.predict(val_df[PROBAS]) diff --git a/corebehrt/functional/estimate/report.py b/corebehrt/functional/estimate/report.py index 220cb201..92a22819 100644 --- a/corebehrt/functional/estimate/report.py +++ b/corebehrt/functional/estimate/report.py @@ -4,6 +4,35 @@ from corebehrt.constants.causal.data import EXPOSURE_COL, OUTCOME, STATUS +def _manual_treatment_outcome_table( + df: pd.DataFrame, exposure_col: str, outcome_col: str +) -> pd.DataFrame: + """ + Build the same shape as CausalEstimate's treatment-outcome table when the + cross-tab is degenerate (constant outcome and/or exposure), which breaks + ``compute_treatment_outcome_table`` (expects three count columns). + """ + ct = pd.crosstab(df[exposure_col], df[outcome_col], dropna=False) + ct = ct.reindex(columns=[0, 1], fill_value=0) + ct = ct.reindex(index=[0, 1], fill_value=0) + rows = [] + for exp_val, label in ((0, "Untreated"), (1, "Treated")): + n0 = int(ct.loc[exp_val, 0]) + n1 = int(ct.loc[exp_val, 1]) + rows.append( + {STATUS: label, "No Outcome": n0, "Outcome": n1, "Total": n0 + n1} + ) + rows.append( + { + STATUS: "Total", + "No Outcome": int(ct[0].sum()), + "Outcome": int(ct[1].sum()), + "Total": len(df), + } + ) + return pd.DataFrame(rows) + + def convert_effect_to_dataframe(effect: dict) -> pd.DataFrame: """ Convert a dictionary of effects to a pandas DataFrame. @@ -68,8 +97,17 @@ def compute_outcome_stats(analysis_df: pd.DataFrame, outcome_name: str) -> pd.Da 1 Treated 2 2 4 diabetes 2 Total 4 4 8 diabetes """ - stats_table = compute_treatment_outcome_table(analysis_df, EXPOSURE_COL, OUTCOME) - stats_table = stats_table.reset_index(drop=False) - stats_table.rename(columns={"index": STATUS}, inplace=True) + try: + stats_table = compute_treatment_outcome_table( + analysis_df, EXPOSURE_COL, OUTCOME + ) + stats_table = stats_table.reset_index(drop=False) + stats_table.rename(columns={"index": STATUS}, inplace=True) + except ValueError: + # Constant outcome or exposure => incomplete crosstab; CausalEstimate assigns + # three column names to two columns (Length mismatch). + stats_table = _manual_treatment_outcome_table( + analysis_df, EXPOSURE_COL, OUTCOME + ) stats_table[OUTCOME] = outcome_name return stats_table diff --git a/corebehrt/main_causal/finetune_subpop.py b/corebehrt/main_causal/finetune_subpop.py new file mode 100644 index 00000000..409c71fd --- /dev/null +++ b/corebehrt/main_causal/finetune_subpop.py @@ -0,0 +1,87 @@ +""" +Subpopulation fine-tuning for causal inference models. + +Loads prepared data from a main run, filters to a subpopulation, +creates fresh bootstrap folds, and continues fine-tuning from the +main run's per-fold checkpoints with the encoder frozen. +""" + +import logging +from os.path import join + +import torch + +from corebehrt.constants.paths import ( + FOLDS_FILE, + OUTCOME_NAMES_FILE, + PREPARED_ALL_PATIENTS, +) +from corebehrt.functional.features.split import create_folds +from corebehrt.functional.io_operations.load import load_vocabulary +from corebehrt.functional.setup.args import get_args +from corebehrt.main.helper.finetune_cv import check_for_overlap +from corebehrt.main_causal.finetune_exp_y import validate_folds +from corebehrt.main_causal.helper.finetune_exp_y import cv_loop +from corebehrt.modules.monitoring.causal.metric_aggregation import ( + compute_and_save_combined_scores_mean_std, +) +from corebehrt.modules.preparation.causal.dataset import CausalPatientDataset +from corebehrt.modules.setup.config import load_config +from corebehrt.modules.setup.directory import DirectoryPreparer +from corebehrt.modules.setup.causal.prediction_accumulator import PredictionAccumulator + +CONFIG_PATH = "./corebehrt/configs/causal/finetune/ft_subpop.yaml" + + +def main_finetune_subpop(config_path): + cfg = load_config(config_path) + DirectoryPreparer(cfg).setup_finetune(check_pretrain=False) + + logger = logging.getLogger("finetune_subpop") + + # Load data and filter to subpopulation + loaded_data = torch.load(join(cfg.paths.prepared_data, PREPARED_ALL_PATIENTS)) + vocab = load_vocabulary(cfg.paths.prepared_data) + data = CausalPatientDataset(loaded_data, vocab) + + subpop_pids = torch.load(cfg.paths.subpopulation_pids) + logger.info(f"Loaded {len(subpop_pids)} subpopulation PIDs") + + data = data.filter_by_pids(subpop_pids) + train_val_pids = data.get_pids() + logger.info(f"Filtered to {len(train_val_pids)} patients in prepared data") + + # Create fresh bootstrap folds from subpopulation + data_cfg = cfg.get("data", {}) + n_folds = data_cfg.get("n_folds", 5) + seed = data_cfg.get("seed", 42) + bootstrap = cfg.get("bootstrap", True) + + folds = create_folds(train_val_pids, n_folds, seed, bootstrap=bootstrap) + validate_folds(folds, set(train_val_pids), logger, bootstrap=bootstrap) + check_for_overlap(folds, [], logger) + torch.save(folds, join(cfg.paths.model, FOLDS_FILE)) + logger.info(f"Created {n_folds} folds (bootstrap={bootstrap}, seed={seed})") + + # Run CV loop (loads per-fold checkpoints via restart_model) + test_data = CausalPatientDataset([], vocab) + cv_loop(cfg, logger, cfg.paths.model, data, folds, test_data) + + # Post-processing + outcome_names = data.get_outcome_names() + PredictionAccumulator( + cfg.paths.model, outcome_names + ).accumulate_and_save_predictions() + torch.save(outcome_names, join(cfg.paths.model, OUTCOME_NAMES_FILE)) + logger.info(f"Saved outcome names: {outcome_names}") + + compute_and_save_combined_scores_mean_std( + len(folds), cfg.paths.model, mode="val", outcome_names=outcome_names + ) + + logger.info("Done") + + +if __name__ == "__main__": + args = get_args(CONFIG_PATH) + main_finetune_subpop(args.config_path) diff --git a/corebehrt/modules/monitoring/causal/metric_aggregation.py b/corebehrt/modules/monitoring/causal/metric_aggregation.py index 2bef68ed..77b8428e 100644 --- a/corebehrt/modules/monitoring/causal/metric_aggregation.py +++ b/corebehrt/modules/monitoring/causal/metric_aggregation.py @@ -1,10 +1,15 @@ +import logging import os from datetime import datetime from os.path import join import pandas as pd from corebehrt.constants.causal.data import EXPOSURE, OUTCOME -from corebehrt.azure import log_metric, setup_metrics_dir +from corebehrt.azure import is_mlflow_available, log_metric, setup_metrics_dir + +logger = logging.getLogger(__name__) + +_PROGRESS_EVERY_N_OUTCOMES = 50 def compute_and_save_combined_scores_mean_std( @@ -14,21 +19,36 @@ def compute_and_save_combined_scores_mean_std( outcome_names: list = None, ) -> None: """Compute mean and std of test/val scores for all targets and save to single file.""" - print("Save combined aggregated scores") + n_out = len(outcome_names) if outcome_names else 0 + msg = ( + f"Save combined aggregated scores ({n_splits} folds, " + f"{n_out} outcomes + exposure; this can take several minutes on remote storage)" + ) + print(msg, flush=True) + logger.info(msg) all_scores = [] - # Collect exposure scores + logger.info("Collecting score CSVs for exposure...") + print(" [combined scores] exposure...", flush=True) exposure_scores = _collect_single_target_scores( n_splits, finetune_folder, mode, EXPOSURE ) if exposure_scores is not None: exposure_scores[OUTCOME] = EXPOSURE all_scores.append(exposure_scores) + logger.info("Exposure: collected %d score rows", len(exposure_scores)) - # Collect outcome scores if outcome_names: - for outcome_name in outcome_names: + for i, outcome_name in enumerate(outcome_names, start=1): + if i == 1 or i % _PROGRESS_EVERY_N_OUTCOMES == 0 or i == n_out: + logger.info( + "Collecting score CSVs: outcome %d / %d (current=%r)", + i, + n_out, + outcome_name, + ) + print(f" [combined scores] outcomes {i} / {n_out}", flush=True) outcome_scores = _collect_single_target_scores( n_splits, finetune_folder, mode, outcome_name ) @@ -36,12 +56,19 @@ def compute_and_save_combined_scores_mean_std( outcome_scores[OUTCOME] = outcome_name all_scores.append(outcome_scores) - # Combine all scores if not all_scores: - print(f"Warning: No score files found for {mode}") + w = f"Warning: No score files found for {mode}" + print(w, flush=True) + logger.warning(w) return try: + logger.info( + "Concatenating %d score tables (~%d rows total)...", + len(all_scores), + sum(len(x) for x in all_scores), + ) + print(" [combined scores] concatenating and aggregating...", flush=True) combined_scores = pd.concat(all_scores, ignore_index=True) scores_mean_std = ( combined_scores.groupby(["metric", "outcome"])["value"] @@ -54,17 +81,46 @@ def compute_and_save_combined_scores_mean_std( os.makedirs(scores_dir, exist_ok=True) output_path = join(scores_dir, f"scores_{date}.csv") scores_mean_std.to_csv(output_path, index=False) + logger.info("Wrote %s (%d rows)", output_path, len(scores_mean_std)) + print(f" [combined scores] wrote {output_path}", flush=True) - # Log to Azure with setup_metrics_dir(f"{mode} combined scores"): - for _, row in scores_mean_std.iterrows(): - metric_name = row["metric"] - outcome_name = row["outcome"] - log_metric(f"{metric_name} mean {outcome_name}", row["mean"]) - log_metric(f"{metric_name} std {outcome_name}", row["std"]) - - except Exception as e: - print(f"Error processing combined scores for {mode}: {e}") + if is_mlflow_available(): + from corebehrt.azure.util.log import get_run_and_prefix + + run, prefix = get_run_and_prefix() + if run is not None: + import mlflow + + batch = {} + for _, row in scores_mean_std.iterrows(): + m, o = row["metric"], row["outcome"] + batch[f"{prefix}{m} mean {o}"] = float(row["mean"]) + batch[f"{prefix}{m} std {o}"] = float(row["std"]) + logger.info( + "Logging %d metrics to MLflow (single batch)...", len(batch) + ) + mlflow.log_metrics(batch, run_id=run.info.run_id) + else: + for _, row in scores_mean_std.iterrows(): + log_metric( + f"{row['metric']} mean {row['outcome']}", row["mean"] + ) + log_metric(f"{row['metric']} std {row['outcome']}", row["std"]) + else: + for _, row in scores_mean_std.iterrows(): + log_metric( + f"{row['metric']} mean {row['outcome']}", row["mean"] + ) + log_metric(f"{row['metric']} std {row['outcome']}", row["std"]) + + logger.info("Finished combined scores for mode=%s", mode) + print(" [combined scores] done", flush=True) + + except Exception: + logger.exception("Error processing combined scores for %s", mode) + print(f"Error processing combined scores for {mode}: (see logs)", flush=True) + raise def _collect_single_target_scores( @@ -82,12 +138,10 @@ def _collect_single_target_scores( if not os.path.exists(fold_checkpoints_folder): continue - # Look for files with BEST_MODEL_ID (999) first, then try epoch numbers possible_files = [ - f"{mode}_{target_type}_scores_999.csv", # BEST_MODEL_ID format + f"{mode}_{target_type}_scores_999.csv", ] - # Also try to find files with actual epoch numbers try: checkpoint_files = [ f @@ -102,7 +156,6 @@ def _collect_single_target_scores( except (ValueError, IndexError): pass - # Try to find any of the possible files fold_scores = None for filename in possible_files: table_path = join(fold_checkpoints_folder, filename) @@ -111,20 +164,22 @@ def _collect_single_target_scores( fold_scores = pd.read_csv(table_path) break except Exception as e: - print(f"Error reading {table_path}: {e}") + logger.warning("Error reading %s: %s", table_path, e) continue if fold_scores is not None: scores.append(fold_scores) - # Return concatenated scores or None if no scores found if not scores: - print(f"Warning: No score files found for {mode}_{target_type}") + logger.debug("No score files for %s_%s", mode, target_type) + print( + f"Warning: No score files found for {mode}_{target_type}", + flush=True, + ) return None combined_scores = pd.concat(scores, ignore_index=True) - # Clean metric names by removing target_type prefix combined_scores["metric"] = combined_scores["metric"].str.replace( f"{target_type}_", "", regex=False ) diff --git a/corebehrt/modules/setup/causal/initializer.py b/corebehrt/modules/setup/causal/initializer.py index 14e3a3fa..6e90e0d9 100644 --- a/corebehrt/modules/setup/causal/initializer.py +++ b/corebehrt/modules/setup/causal/initializer.py @@ -29,10 +29,14 @@ def initialize_finetune_model( logger.info("Loading model from checkpoint") loss_weight_outcomes = { - outcome_name: get_loss_weight(self.cfg, outcome_values) + outcome_name: get_loss_weight( + self.cfg, outcome_values, log_name=f"outcome {outcome_name!r}" + ) for outcome_name, outcome_values in outcomes.items() } - loss_weight_exposures = get_loss_weight(self.cfg, exposures) + loss_weight_exposures = get_loss_weight( + self.cfg, exposures, log_name="exposure" + ) add_config = { **self.cfg.model, diff --git a/corebehrt/modules/setup/causal/manager.py b/corebehrt/modules/setup/causal/manager.py index 3711b69f..7b244498 100644 --- a/corebehrt/modules/setup/causal/manager.py +++ b/corebehrt/modules/setup/causal/manager.py @@ -17,5 +17,14 @@ def initialize_finetune_model( self.initializer = CausalInitializer( self.cfg, checkpoint=checkpoint, model_path=self.checkpoint_model_path ) - model = self.initializer.initialize_finetune_model(outcomes, exposures) + try: + model = self.initializer.initialize_finetune_model(outcomes, exposures) + except Exception as e: + print( + f"[CausalModelManager] initialize_finetune_model FAILED: " + f"{type(e).__name__}: {e}", + flush=True, + ) + logger.exception("initialize_finetune_model failed") + raise return model diff --git a/corebehrt/modules/setup/directory.py b/corebehrt/modules/setup/directory.py index 49d58946..0f115135 100644 --- a/corebehrt/modules/setup/directory.py +++ b/corebehrt/modules/setup/directory.py @@ -379,7 +379,7 @@ def setup_prepare_finetune(self, name=None) -> None: self.write_config("prepared_data", name=name) self.write_config("prepared_data", source="features", name=DATA_CFG) - def setup_finetune(self) -> None: + def setup_finetune(self, check_pretrain: bool = True) -> None: """ Validates path config and sets up directories for finetune. """ @@ -388,12 +388,14 @@ def setup_finetune(self) -> None: # Validate and create directories self.check_directory("prepared_data") - self.check_directory("pretrain_model") + if check_pretrain: + self.check_directory("pretrain_model") self.create_run_directory("model", base="runs") # Write config in output directory. - self.write_config("model", source="pretrain_model", name=PRETRAIN_CFG) self.write_config("model", name=FINETUNE_CFG) + if check_pretrain: + self.write_config("model", source="pretrain_model", name=PRETRAIN_CFG) # Add pretrain info to config data_cfg = self.get_config("prepared_data", name=DATA_CFG) diff --git a/corebehrt/modules/setup/manager.py b/corebehrt/modules/setup/manager.py index 0590286c..98cddbe7 100644 --- a/corebehrt/modules/setup/manager.py +++ b/corebehrt/modules/setup/manager.py @@ -98,9 +98,24 @@ def initialize_finetune_model(self, checkpoint, outcomes): return model def initialize_training_components(self, model, outcomes): - """Initialize training components. If no model_path provided, optimizer and scheduler are initialized from scratch.""" - if self.restart_model_path is None: - logger.info("Initializing optimizer and scheduler from scratch") + """Initialize training components. + + By default, when restarting from a checkpoint we also restore optimizer and + scheduler state. For subpopulation fine-tuning this can be undesirable because + the loaded scheduler may have already decayed the LR close to zero. + + Set ``trainer_args.restart_weights_only: true`` to load model weights from the + checkpoint but reinitialize optimizer/scheduler from scratch. + """ + restart_weights_only = bool( + self.cfg.get("trainer_args", {}).get("restart_weights_only", False) + ) + + if self.restart_model_path is None or restart_weights_only: + logger.info( + "Initializing optimizer and scheduler from scratch" + + (" (restart_weights_only=true)" if restart_weights_only else "") + ) self.initializer.checkpoint = None optimizer = self.initializer.initialize_optimizer(model) sampler, cfg = self.initializer.initialize_sampler(outcomes) @@ -108,10 +123,14 @@ def initialize_training_components(self, model, outcomes): return optimizer, sampler, scheduler, cfg def get_epoch(self): - """Get epoch from model_path.""" + """Get epoch from model_path (None => trainer starts at continue_epoch 0).""" if self.restart_model_path is None: return 0 - else: - return get_last_checkpoint_epoch( - join(self.restart_model_path, CHECKPOINTS_DIR) - ) + ep = get_last_checkpoint_epoch( + join(self.restart_model_path, CHECKPOINTS_DIR) + ) + # Match causal trainer BEST_MODEL_ID: filename id 999 is the best-weights slot, + # not a real resume index (would give continue_epoch 1000 and skip training). + if ep == 999: + return None + return ep diff --git a/corebehrt/modules/trainer/causal/trainer.py b/corebehrt/modules/trainer/causal/trainer.py index eb7d2d79..108e075b 100644 --- a/corebehrt/modules/trainer/causal/trainer.py +++ b/corebehrt/modules/trainer/causal/trainer.py @@ -60,6 +60,8 @@ def __init__(self, *args, plateau_threshold=0.01, **kwargs): self.update_step = 0 if self.use_pcgrad: self.optimizer = PCGrad(self.optimizer) + if self.args.get("freeze_encoder_at_init", False): + self._freeze_encoder() def _set_plateau_parameters(self): self.freeze_encoder_on_plateau = self.args.get( diff --git a/corebehrt/modules/trainer/utils.py b/corebehrt/modules/trainer/utils.py index 09fbfeca..3f6f6362 100644 --- a/corebehrt/modules/trainer/utils.py +++ b/corebehrt/modules/trainer/utils.py @@ -1,3 +1,4 @@ +import logging from typing import Dict, List, Optional import numpy as np @@ -6,6 +7,8 @@ from corebehrt.modules.setup.config import instantiate_function +logger = logging.getLogger(__name__) + def get_sampler(cfg, outcomes: List[int]) -> Optional[WeightedRandomSampler]: """Get sampler for training data. @@ -85,16 +88,31 @@ def effective_n_samples(outcomes: List[int]) -> List[float]: return [class_probs[outcome] / labels[outcome] for outcome in outcomes] -def get_loss_weight(cfg, outcomes: List[int]) -> Optional[List[float]]: +def get_loss_weight( + cfg, + outcomes: List[int], + *, + log_name: str | None = None, +) -> Optional[float]: """Get weights for weighted loss function. If loss_weight_function is false or undefined, then no positive weight is used. If loss_weight_function is defined then the function is used to calculate the weights. - """ + If the function raises a ValueError, then the function is not used and no positive weight is used.""" if cfg.trainer_args.get("loss_weight_function") is None or len(outcomes) == 0: return None weight_func = instantiate_function(cfg.trainer_args.get("loss_weight_function")) - return weight_func(outcomes) + try: + return weight_func(outcomes) + except ValueError as e: + where = f" for {log_name}" if log_name else "" + logger.warning( + "Skipping class-balanced loss weight%s (%s); using unweighted BCE. %s", + where, + getattr(weight_func, "__name__", weight_func), + e, + ) + return None class PositiveWeight: diff --git a/tests/test_main_causal/test_finetune_subpop.py b/tests/test_main_causal/test_finetune_subpop.py new file mode 100644 index 00000000..caf6eca9 --- /dev/null +++ b/tests/test_main_causal/test_finetune_subpop.py @@ -0,0 +1,100 @@ +import unittest +from unittest.mock import MagicMock, patch + +from corebehrt.constants.data import TRAIN_KEY, VAL_KEY +from corebehrt.functional.features.split import create_folds +from corebehrt.main_causal.finetune_exp_y import validate_folds + + +class TestSubpopFoldCreation(unittest.TestCase): + def test_bootstrap_folds_from_subpop(self): + """Verify create_folds produces valid bootstrap folds from a subset of PIDs.""" + subpop_pids = list(range(50)) + folds = create_folds(subpop_pids, num_folds=5, seed=42, bootstrap=True) + + self.assertEqual(len(folds), 5) + for fold in folds: + self.assertIn(TRAIN_KEY, fold) + self.assertIn(VAL_KEY, fold) + # Total count matches original size + total = len(fold[TRAIN_KEY]) + len(fold[VAL_KEY]) + self.assertEqual(total, len(subpop_pids)) + + # Validation should pass with bootstrap=True + validate_folds( + folds, + set(subpop_pids), + logger=MagicMock(), + bootstrap=True, + ) + + def test_non_bootstrap_folds_from_subpop(self): + """Verify standard CV folds work on subpopulation.""" + subpop_pids = list(range(50)) + folds = create_folds(subpop_pids, num_folds=5, seed=42, bootstrap=False) + + self.assertEqual(len(folds), 5) + all_val = set() + for fold in folds: + val_pids = set(fold[VAL_KEY]) + train_pids = set(fold[TRAIN_KEY]) + self.assertTrue(val_pids.isdisjoint(train_pids)) + all_val.update(val_pids) + self.assertEqual(all_val, set(subpop_pids)) + + +class TestFreezeEncoderAtInit(unittest.TestCase): + @patch("corebehrt.modules.trainer.causal.trainer.CausalEHRTrainer._freeze_encoder") + @patch( + "corebehrt.modules.trainer.causal.trainer.EHRTrainer.__init__", + return_value=None, + ) + def test_freeze_called_when_flag_set(self, mock_init, mock_freeze): + """Verify _freeze_encoder is called when freeze_encoder_at_init=True.""" + from corebehrt.modules.trainer.causal.trainer import CausalEHRTrainer + + trainer = CausalEHRTrainer.__new__(CausalEHRTrainer) + trainer.args = {"freeze_encoder_at_init": True, "use_pcgrad": False} + trainer.model = MagicMock() + trainer.model.config.outcome_names = ["outcome_1"] + trainer.metric_history = {} + trainer.epoch_history = [] + trainer.encoder_frozen = False + trainer.outcome_names = ["outcome_1"] + trainer.best_outcome_aucs = {} + trainer.best_exposure_auc = None + trainer.use_pcgrad = False + trainer.plot_histograms = False + trainer.plot_gradients = False + trainer.plot_gradients_frequency = 100 + trainer.plot_log_scale = False + trainer.global_step = 0 + trainer.update_step = 0 + trainer._set_plateau_parameters() + trainer._set_logging_parameters() + + if trainer.args.get("freeze_encoder_at_init", False): + trainer._freeze_encoder() + + mock_freeze.assert_called_once() + + @patch("corebehrt.modules.trainer.causal.trainer.CausalEHRTrainer._freeze_encoder") + @patch( + "corebehrt.modules.trainer.causal.trainer.EHRTrainer.__init__", + return_value=None, + ) + def test_freeze_not_called_when_flag_not_set(self, mock_init, mock_freeze): + """Verify _freeze_encoder is NOT called when freeze_encoder_at_init is absent.""" + from corebehrt.modules.trainer.causal.trainer import CausalEHRTrainer + + trainer = CausalEHRTrainer.__new__(CausalEHRTrainer) + trainer.args = {"use_pcgrad": False} + + if trainer.args.get("freeze_encoder_at_init", False): + trainer._freeze_encoder() + + mock_freeze.assert_not_called() + + +if __name__ == "__main__": + unittest.main()