diff --git a/README.md b/README.md index 3e8ffd3..3380ac6 100644 --- a/README.md +++ b/README.md @@ -62,7 +62,16 @@ eventdisplay-ml-train-xgb-stereo \ - XGBoost trained model object - Target standardization scalers (mean/std) - Feature list and SHAP importance rankings -- Training metadata (random state, hyperparameters) +- A frozen regression training record: the seed used for input sampling, + train/test splitting, validation/diagnostic/SHAP sampling, and XGBoost; + ordered input-file manifest; feature/target configuration; weighting and + event-limit settings; configured and effective XGBoost parameters; and the + best iteration/score + +Regression uses seed `42` when `--random_state` is omitted. The full record is +stored under `training_parameters` in the Joblib payload and is also returned +by the regression model loader, so a trained model can be audited without +reconstructing the command line. ### Applying Stereo Reconstruction Models diff --git a/docs/changes/83.feature.md b/docs/changes/83.feature.md new file mode 100644 index 0000000..d0dd478 --- /dev/null +++ b/docs/changes/83.feature.md @@ -0,0 +1,3 @@ +Persist a complete seed and effective training-parameter record in stereo +regression Joblib models, including the ordered input manifest and model +selection metadata. diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py index 3d4f325..796479c 100644 --- a/src/eventdisplay_ml/config.py +++ b/src/eventdisplay_ml/config.py @@ -83,11 +83,17 @@ def configure_training(analysis_type): "(signal/background in classification)." ), ) + random_state_default = ( + utils.DEFAULT_REGRESSION_RANDOM_STATE if analysis_type == "stereo_analysis" else None + ) parser.add_argument( "--random_state", type=int, - help="Random state for train/test split.", - default=None, + help=( + "Random state used for event sampling, train/test splitting, XGBoost, " + "and diagnostics. Regression defaults to 42." + ), + default=random_state_default, ) if analysis_type == "classification": diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py index 0c370ab..3ef6599 100644 --- a/src/eventdisplay_ml/data_processing.py +++ b/src/eventdisplay_ml/data_processing.py @@ -1016,6 +1016,13 @@ def load_training_data(model_configs, file_list, analysis_type): classification_mode = analysis_type == "classification" max_events = model_configs.get("max_events", None) random_state = model_configs.get("random_state", None) + if analysis_type == "stereo_analysis" and random_state is None: + random_state = utils.DEFAULT_REGRESSION_RANDOM_STATE + model_configs["random_state"] = random_state + _logger.info( + "No regression random_state supplied; using default seed %d", + random_state, + ) memory_profile = model_configs.get("memory_profile", False) _logger.info(f"--- Loading and Flattening Data for {analysis_type} ---") @@ -1029,6 +1036,11 @@ def load_training_data(model_configs, file_list, analysis_type): _logger.info(f"Adding zenith binning: {model_configs.get('zenith_bins_deg', [])}") input_files = utils.read_input_file_list(file_list) + if analysis_type == "stereo_analysis": + # Persist the ordered manifest used for training. With a capped input + # sample, changing file order changes the reservoir stream and hence the + # selected events even when the seed is unchanged. + model_configs["training_input_files"] = list(input_files) if classification_mode and not input_files: raise ValueError(f"Input file list is empty: {file_list}") diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index b1e3511..b2167c9 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -1,5 +1,6 @@ """Apply models for regression and classification tasks.""" +import copy import logging import re import subprocess @@ -36,6 +37,13 @@ _MAX_REGRESSION_SAMPLE_WEIGHT = 50.0 _MODEL_VALIDATION_MEMORY_BYTES = 4 * 1024**3 _MODEL_VALIDATION_TIMEOUT_SECONDS = 120 +_REGRESSION_RANDOM_SEED_NAMES = ( + "data_sampling", + "train_test_split", + "validation_sampling", + "diagnostic_sampling", + "shap_sampling", +) _logger = logging.getLogger(__name__) @@ -359,11 +367,12 @@ def load_regression_models(model_prefix, model_name): } } par = {} - for key in ("target_mean", "target_std"): + for key in ("target_mean", "target_std", "training_parameters"): if key in model_data: par[key] = model_data[key] else: - _logger.warning("Missing '%s' in regression model file: %s", key, model_path) + if key in {"target_mean", "target_std"}: + _logger.warning("Missing '%s' in regression model file: %s", key, model_path) _logger.info("Loaded regression model.") return models, par @@ -844,6 +853,17 @@ def train_regression(df, model_configs): _logger.warning("Skipping training due to empty data.") return None + # A regression run has several random consumers. Keep one explicit seed + # for all of them and freeze the resulting record in the artifact below. + random_state = model_configs.get("random_state") + if random_state is None: + random_state = utils.DEFAULT_REGRESSION_RANDOM_STATE + model_configs["random_state"] = random_state + _logger.info( + "No regression random_state supplied; using default seed %d", + random_state, + ) + memory_profile = model_configs.get("memory_profile", False) utils.log_memory_checkpoint("train_regression:start", df, enabled=memory_profile) @@ -857,13 +877,44 @@ def train_regression(df, model_configs): _logger.info(f"Features ({len(x_cols)}): {', '.join(list(x_cols))}") model_configs["features"] = list(x_cols) targets = model_configs["targets"] + random_seeds = dict.fromkeys(_REGRESSION_RANDOM_SEED_NAMES, random_state) + random_seeds["xgboost"] = {} + training_parameters = { + "analysis_type": "stereo_analysis", + "random_state": random_state, + "random_seeds": random_seeds, + "train_test_fraction": model_configs.get("train_test_fraction", 0.5), + "max_events": model_configs.get("max_events"), + "eval_max_events": model_configs.get("eval_max_events", 200000), + "diagnostic_max_events": model_configs.get("diagnostic_max_events", 100000), + "prediction_chunk_size": model_configs.get("prediction_chunk_size", 200000), + "feature_profile": profile, + "features": list(x_cols), + "targets": list(targets), + "pre_cuts": model_configs.get("pre_cuts"), + "input_file_list": model_configs.get("input_file_list"), + "input_files": list(model_configs.get("training_input_files", [])), + "weighting": { + "energy_weighting": "inverse_sqrt_bin_count", + "min_energy_bin_events": _MIN_WEIGHTED_ENERGY_BIN_EVENTS, + "multiplicity_weighting": "DispNImages**2", + "max_sample_weight": _MAX_REGRESSION_SAMPLE_WEIGHT, + }, + # Keep a frozen copy of any additional top-level training options so + # future options are not silently omitted from the audit record. + "configuration": copy.deepcopy( + {key: value for key, value in model_configs.items() if key != "models"} + ), + "models": {}, + } + model_configs["training_parameters"] = training_parameters utils.log_memory_checkpoint("after feature list creation", df, enabled=memory_profile) row_indices = np.arange(len(df)) train_idx, test_idx = train_test_split( row_indices, train_size=model_configs.get("train_test_fraction", 0.5), - random_state=model_configs.get("random_state", None), + random_state=random_state, ) utils.log_memory_checkpoint("after index train_test_split", enabled=memory_profile) @@ -923,7 +974,7 @@ def train_regression(df, model_configs): eval_idx = _sample_eval_indices( test_idx, model_configs.get("eval_max_events", 200000), - model_configs.get("random_state", None), + random_state, ) weights_eval = None if weight_config is not None: @@ -954,7 +1005,12 @@ def train_regression(df, model_configs): utils.log_memory_checkpoint("after building XGBoost fit arrays", enabled=memory_profile) utils.log_memory_checkpoint(f"{name}: before XGBRegressor init", enabled=memory_profile) - hyper_parameters = dict(cfg.get("hyper_parameters", {})) + configured_hyper_parameters = copy.deepcopy(cfg.get("hyper_parameters", {})) + hyper_parameters = dict(configured_hyper_parameters) + if random_state is not None: + # configure_training applies this override for CLI jobs; applying + # it here as well keeps the public train_regression() API aligned. + hyper_parameters["random_state"] = random_state early_stopping_rounds = hyper_parameters.pop("early_stopping_rounds", None) if early_stopping_rounds is not None: hyper_parameters["callbacks"] = [ @@ -990,7 +1046,7 @@ def train_regression(df, model_configs): diagnostic_train_idx = _sample_eval_indices( train_idx, model_configs.get("diagnostic_max_events", 100000), - model_configs.get("random_state", None), + random_state, ) y_train_diagnostic = ( y_train @@ -1049,7 +1105,7 @@ def train_regression(df, model_configs): shap_idx = _sample_eval_indices( test_idx, 1000, - model_configs.get("random_state", None), + random_state, ) x_test_shap = df.iloc[shap_idx, df.columns.get_indexer(x_cols)] utils.log_memory_checkpoint(f"{name}: before regression evaluation", enabled=memory_profile) @@ -1064,6 +1120,31 @@ def train_regression(df, model_configs): cfg["residual_normality_stats"] = residual_normality_stats cfg["shap_importance"] = shap_importance # Store per-target SHAP importance from evaluation + effective_hyper_parameters = copy.deepcopy(hyper_parameters) + if early_stopping_rounds is not None: + effective_hyper_parameters["early_stopping_rounds"] = early_stopping_rounds + xgboost_parameters = None + if hasattr(model, "get_params"): + candidate_parameters = model.get_params(deep=False) + if isinstance(candidate_parameters, dict): + xgboost_parameters = copy.deepcopy(candidate_parameters) + if xgboost_parameters is None: + xgboost_parameters = copy.deepcopy(effective_hyper_parameters) + model_random_state = { + key: xgboost_parameters.get(key) + for key in ("random_state", "seed") + if key in xgboost_parameters + } + random_seeds["xgboost"][name] = model_random_state + training_parameters["models"][name] = { + "hyper_parameters": configured_hyper_parameters, + "effective_hyper_parameters": effective_hyper_parameters, + "xgboost_parameters": xgboost_parameters, + "random_seeds": model_random_state, + "best_iteration": getattr(model, "best_iteration", None), + "best_score": getattr(model, "best_score", None), + } + return model_configs diff --git a/src/eventdisplay_ml/utils.py b/src/eventdisplay_ml/utils.py index 1c800e4..21507f8 100644 --- a/src/eventdisplay_ml/utils.py +++ b/src/eventdisplay_ml/utils.py @@ -15,6 +15,11 @@ _profile_start_time = None _profile_last_time = None +# Keep command-line regression jobs reproducible when no seed is supplied. The +# value is deliberately defined in a dependency-free module so both data +# loading and model training use the same fallback. +DEFAULT_REGRESSION_RANDOM_STATE = 42 + def _max_rss_gb(): """Return the process peak resident set size in GB.""" diff --git a/tests/test_config.py b/tests/test_config.py index e5d856c..02f324b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -90,6 +90,32 @@ def test_configure_training_stereo_enables_memory_profile(monkeypatch): assert result["memory_profile"] is True +def test_configure_training_stereo_defaults_random_state(monkeypatch): + monkeypatch.setattr( + sys, + "argv", + [ + "prog", + "--model_prefix", + "model", + "--input_file_list", + "inputs.txt", + ], + ) + monkeypatch.setattr( + config, + "hyper_parameters", + lambda *_: {"xgboost": {"hyper_parameters": {}}}, + ) + monkeypatch.setattr(config, "target_features", lambda *_: ["target_a"]) + monkeypatch.setattr(config, "pre_cuts_regression", lambda min_images: f"cut_{min_images}") + + result = config.configure_training("stereo_analysis") + + assert result["random_state"] == 42 + assert result["models"]["xgboost"]["hyper_parameters"]["random_state"] == 42 + + def test_configure_training_stereo_parses_reduced_feature_profile(monkeypatch): monkeypatch.setattr( sys, diff --git a/tests/test_data_processing_remaining.py b/tests/test_data_processing_remaining.py index ed21d97..7e920a3 100644 --- a/tests/test_data_processing_remaining.py +++ b/tests/test_data_processing_remaining.py @@ -1,6 +1,6 @@ """Tests for remaining data_processing helper branches.""" -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import awkward as ak import numpy as np @@ -286,11 +286,15 @@ def test_load_training_data_stereo_adds_residuals(monkeypatch, tel_config): ) monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *_: None) - result = data_processing.load_training_data({"max_cores": 1}, "inputs.txt", "stereo_analysis") + model_configs = {"max_cores": 1} + with patch("eventdisplay_ml.data_processing.np.random.default_rng") as rng_mock: + result = data_processing.load_training_data(model_configs, "inputs.txt", "stereo_analysis") assert result["Xoff_residual"].tolist() == pytest.approx([0.2, 0.4]) assert result["Yoff_residual"].tolist() == pytest.approx([0.1, 0.4]) assert result["E_residual"].tolist() == pytest.approx([np.log10(10.0) - np.log10(5.0), 1.0]) + assert model_configs["random_state"] == 42 + rng_mock.assert_called_once_with(42) def test_load_training_data_caps_iterated_chunks(monkeypatch, tel_config): diff --git a/tests/test_models_helpers.py b/tests/test_models_helpers.py index 6a12dab..9c185b8 100644 --- a/tests/test_models_helpers.py +++ b/tests/test_models_helpers.py @@ -58,6 +58,26 @@ def test_save_models_writes_expected_joblib(tmp_path): assert (tmp_path / "saved.joblib.gz").exists() +def test_save_models_persists_regression_training_record(tmp_path, monkeypatch): + """Verify the regression record is retained by the Joblib save boundary.""" + training_record = { + "random_state": 42, + "random_seeds": {"train_test_split": 42}, + "models": {"xgboost": {"effective_hyper_parameters": {"max_depth": 5}}}, + } + model_configs = { + "model_prefix": str(tmp_path / "saved"), + "models": {"xgboost": {"model": "trained-model"}}, + "training_parameters": training_record, + } + monkeypatch.setattr(models, "_validate_saved_model", lambda _path: None) + + models.save_models(model_configs) + + saved = joblib.load(tmp_path / "saved.joblib.gz") + assert saved["training_parameters"] == training_record + + def test_save_models_rejects_model_that_fails_validation(tmp_path, monkeypatch): model_configs = { "model_prefix": str(tmp_path / "invalid"), diff --git a/tests/test_regression_contracts.py b/tests/test_regression_contracts.py index 65194da..2e3b742 100644 --- a/tests/test_regression_contracts.py +++ b/tests/test_regression_contracts.py @@ -190,10 +190,12 @@ def test_persisted_regression_model_preserves_feature_order_and_reconstructs_tru "features": feature_order, "target_mean": target_mean, "target_std": target_std, + "training_parameters": {"random_state": 17}, }, tmp_path / "stereo_model.joblib.gz", ) loaded_models, loaded_parameters = models.load_regression_models(str(model_prefix), "xgboost") + assert loaded_parameters["training_parameters"]["random_state"] == 17 # Deliberately not in persisted feature order, and includes a column that # must not reach the model. diff --git a/tests/test_train_regression_standardization.py b/tests/test_train_regression_standardization.py index 1e55e53..7837611 100644 --- a/tests/test_train_regression_standardization.py +++ b/tests/test_train_regression_standardization.py @@ -1,5 +1,6 @@ """Tests for target standardization and energy-bin weighting in train_regression().""" +import copy from unittest.mock import MagicMock, patch import numpy as np @@ -8,7 +9,7 @@ import xgboost as xgb from sklearn.model_selection import train_test_split -from eventdisplay_ml import diagnostic_utils, models +from eventdisplay_ml import diagnostic_utils, models, utils @pytest.fixture @@ -119,6 +120,84 @@ def test_target_std_never_zero(self, regression_training_df, regression_model_co for target in cfg["targets"]: assert result["target_std"][target] > 0, f"{target} std should not be zero" + @pytest.mark.parametrize("seed_mode", ["missing", "none"]) + def test_training_record_contains_all_regression_seeds_and_parameters( + self, regression_training_df, regression_model_config, seed_mode + ): + """Persist the effective reproducibility and XGBoost training record.""" + cfg = copy.deepcopy(regression_model_config) + if seed_mode == "missing": + cfg.pop("random_state") + else: + cfg["random_state"] = None + cfg["models"]["xgboost"]["hyper_parameters"].pop("random_state") + cfg["eval_max_events"] = 10 + cfg["diagnostic_max_events"] = 10 + + class _CapturingRegressor: + best_iteration = 0 + best_score = 0.0 + + def __init__(self, **kwargs): + self.params = kwargs + + def fit(self, *_args, **_kwargs): + return self + + def predict(self, x_values): + return np.zeros((len(x_values), len(cfg["targets"]))) + + def get_params(self, deep=False): # noqa: ARG002 + return self.params + + with ( + patch("eventdisplay_ml.models.train_test_split", wraps=train_test_split) as split_mock, + patch( + "eventdisplay_ml.models._sample_eval_indices", + wraps=models._sample_eval_indices, + ) as sample_mock, + patch("xgboost.XGBRegressor", side_effect=_CapturingRegressor) as xgb_mock, + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(regression_training_df, cfg) + + record = result["training_parameters"] + assert record["random_state"] == utils.DEFAULT_REGRESSION_RANDOM_STATE + assert record["random_seeds"] == { + "data_sampling": utils.DEFAULT_REGRESSION_RANDOM_STATE, + "train_test_split": utils.DEFAULT_REGRESSION_RANDOM_STATE, + "validation_sampling": utils.DEFAULT_REGRESSION_RANDOM_STATE, + "diagnostic_sampling": utils.DEFAULT_REGRESSION_RANDOM_STATE, + "shap_sampling": utils.DEFAULT_REGRESSION_RANDOM_STATE, + "xgboost": { + "xgboost": {"random_state": utils.DEFAULT_REGRESSION_RANDOM_STATE}, + }, + } + assert cfg["random_state"] == utils.DEFAULT_REGRESSION_RANDOM_STATE + assert split_mock.call_args.kwargs["random_state"] == utils.DEFAULT_REGRESSION_RANDOM_STATE + assert sample_mock.call_count >= 3 + assert all( + call.args[2] == utils.DEFAULT_REGRESSION_RANDOM_STATE for call in sample_mock.call_args_list + ) + assert ( + xgb_mock.call_args.kwargs["random_state"] == utils.DEFAULT_REGRESSION_RANDOM_STATE + ) + model_record = record["models"]["xgboost"] + assert ( + model_record["hyper_parameters"] == cfg["models"]["xgboost"]["hyper_parameters"] + ) + assert model_record["effective_hyper_parameters"]["early_stopping_rounds"] == 2 + assert ( + model_record["effective_hyper_parameters"]["random_state"] + == utils.DEFAULT_REGRESSION_RANDOM_STATE + ) + assert ( + model_record["xgboost_parameters"]["random_state"] + == utils.DEFAULT_REGRESSION_RANDOM_STATE + ) + assert record["features"] == result["features"] + assert record["targets"] == cfg["targets"] + class TestEnergyBinWeighting: """Tests for energy-bin weighting (especially zeroing low-count bins)."""