diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index fb7e516..d5647e1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,14 +2,14 @@ repos: # Ruff - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.15.16 + rev: v0.16.1 hooks: - id: ruff args: ["--fix"] - id: ruff-format # https://pycqa.github.io/isort/docs/configuration/black_compatibility.html#integration-with-pre-commit - repo: https://github.com/pycqa/isort - rev: 9.0.0a3 + rev: 9.0.0b1 hooks: - id: isort args: ["--profile", "black", "--filter-files"] @@ -34,7 +34,7 @@ repos: # - id: actionlint # codespell - repo: https://github.com/codespell-project/codespell - rev: v2.4.2 + rev: v2.4.3 hooks: - id: codespell args: [ diff --git a/docs/changes/80.maintenance.md b/docs/changes/80.maintenance.md new file mode 100644 index 0000000..3378da4 --- /dev/null +++ b/docs/changes/80.maintenance.md @@ -0,0 +1 @@ +Add strong testing for regression path. diff --git a/pyproject.toml b/pyproject.toml index a45b480..dff7758 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,6 +93,7 @@ indent-width = 4 exclude = [ "__init__.py", "pyproject.toml", + "tests/", ] format.indent-style = "space" @@ -132,9 +133,6 @@ lint.ignore = [ lint.pydocstyle.convention = "numpy" -[tool.ruff.lint.per-file-ignores] -"tests/**.py" = ["D103"] - [tool.codespell] ignore-words-list = "chec,arrang,livetime" diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py index 34239df..313a073 100644 --- a/src/eventdisplay_ml/data_processing.py +++ b/src/eventdisplay_ml/data_processing.py @@ -824,7 +824,9 @@ def _flatten_training_chunk( disp_erec = df_flat["ErecS"].values # Compute log energies (ErecS already filtered > 0) - mc_e0_log = np.where(mc_e0 > 0, np.log10(mc_e0), np.nan) + mc_e0_log = np.full_like(mc_e0, np.nan, dtype=np.float32) + valid_mc_energy = mc_e0 > 0 + mc_e0_log[valid_mc_energy] = np.log10(mc_e0[valid_mc_energy]) disp_erec_log = np.log10(disp_erec) # Safe since already filtered > 0 new_cols = { diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index 1125df6..e37af42 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -283,7 +283,7 @@ def _validate_energy_bin_metadata(energy_bin, model_file): Validated metadata containing ``E_min`` and ``E_max`` keys. """ if not isinstance(energy_bin, dict): - raise ValueError( + raise TypeError( "Classification model file " f"'{model_file}' has invalid 'energy_bins_log10_tev' metadata: " "expected a dict with keys 'E_min' and 'E_max'." @@ -601,7 +601,7 @@ def process_file_chunked(analysis_type, model_configs): { eff for e_bin_models in model_configs["models"].values() - for eff in (e_bin_models.get("thresholds") or {}).keys() + for eff in (e_bin_models.get("thresholds") or {}) } ) diff --git a/src/eventdisplay_ml/scripts/diagnostic_shap_summary.py b/src/eventdisplay_ml/scripts/diagnostic_shap_summary.py index af03136..fdb0615 100644 --- a/src/eventdisplay_ml/scripts/diagnostic_shap_summary.py +++ b/src/eventdisplay_ml/scripts/diagnostic_shap_summary.py @@ -156,8 +156,8 @@ def main(): output_file = utils.joblib_basename(model_path) try: process_model_file(model_path, output_dir, output_file) - except Exception as e: - _logger.exception(f"Skipping {model_path}: failed to process model ({e})") + except Exception: + _logger.exception(f"Skipping {model_path}: failed to process model") continue _logger.info(f"\nPlots saved to {output_dir}") diff --git a/src/eventdisplay_ml/scripts/optimize_classification.py b/src/eventdisplay_ml/scripts/optimize_classification.py index 54af6df..b8e2588 100644 --- a/src/eventdisplay_ml/scripts/optimize_classification.py +++ b/src/eventdisplay_ml/scripts/optimize_classification.py @@ -142,7 +142,7 @@ def _extract_tgraph2d(graph): try: x, y, z = graph.values() return np.asarray(x), np.asarray(y), np.asarray(z) - except Exception: + except (AttributeError, ValueError, TypeError): x = np.asarray(graph.member("fX")) y = np.asarray(graph.member("fY")) z = np.asarray(graph.member("fZ")) diff --git a/src/eventdisplay_ml/scripts/plot_classification_gamma_efficiency.py b/src/eventdisplay_ml/scripts/plot_classification_gamma_efficiency.py index f3eecda..f76b0dc 100644 --- a/src/eventdisplay_ml/scripts/plot_classification_gamma_efficiency.py +++ b/src/eventdisplay_ml/scripts/plot_classification_gamma_efficiency.py @@ -70,7 +70,7 @@ def get_containment_data(directory): _logger.info(f"Percentiles p70: {p70}, p95: {p95}") else: _logger.warning("Nan percentiles") - except Exception as e: + except (OSError, ValueError, KeyError) as e: _logger.error(f"Failed reading {filename}: {e}") return pd.DataFrame(results) diff --git a/src/eventdisplay_ml/scripts/plot_training_evaluation.py b/src/eventdisplay_ml/scripts/plot_training_evaluation.py index aa0754c..15b7c98 100644 --- a/src/eventdisplay_ml/scripts/plot_training_evaluation.py +++ b/src/eventdisplay_ml/scripts/plot_training_evaluation.py @@ -131,7 +131,7 @@ def main(): type=str, help=( "Directory containing multiple joblib model files. " - "All *.joblib files will be processed.", + "All *.joblib files will be processed." ), ) parser.add_argument( @@ -218,8 +218,8 @@ def main(): ) plot_training_curves(evals_result, output_file) _logger.info(f"Saved plot for {model_path.name} to {output_file}") - except Exception as e: - _logger.exception(f"Skipping {model_path}: failed to process model ({e})") + except Exception: + _logger.exception(f"Skipping {model_path}: failed to process model") continue _logger.info("Batch plotting completed.") diff --git a/src/eventdisplay_ml/utils.py b/src/eventdisplay_ml/utils.py index f85d181..1c800e4 100644 --- a/src/eventdisplay_ml/utils.py +++ b/src/eventdisplay_ml/utils.py @@ -88,7 +88,11 @@ def discover_joblib_files(model_dir): if not model_dir.exists() or not model_dir.is_dir(): raise FileNotFoundError(f"Model directory not found: {model_dir}") - discovered_files = sorted(set(model_dir.glob("*.joblib")).union(model_dir.glob("*.joblib.gz"))) + discovered_files = sorted( + path + for path in set(model_dir.glob("*.joblib")).union(model_dir.glob("*.joblib.gz")) + if path.is_file() + ) files_by_name = {} for model_path in discovered_files: key = joblib_basename(model_path) diff --git a/tests/scripts/test_stereo_entrypoints.py b/tests/scripts/test_stereo_entrypoints.py new file mode 100644 index 0000000..d059748 --- /dev/null +++ b/tests/scripts/test_stereo_entrypoints.py @@ -0,0 +1,87 @@ +"""Unit tests for production stereo-regression console-script wiring.""" + +from unittest.mock import MagicMock + +import pandas as pd +import pytest + +from eventdisplay_ml.scripts import apply_xgb_stereo, train_xgb_stereo + + +def test_train_stereo_entrypoint_runs_the_complete_regression_pipeline(monkeypatch, caplog): + """The training CLI must connect configuration, loading, training, and saving unchanged.""" + caplog.set_level("INFO") + configured = {"input_file_list": "gamma_inputs.txt", "model_prefix": "stereo_model"} + loaded_data = pd.DataFrame({"feature": [1.0]}) + trained = {"model_prefix": "stereo_model", "models": {"xgboost": {"model": object()}}} + configure = MagicMock(return_value=configured) + load_data = MagicMock(return_value=loaded_data) + train = MagicMock(return_value=trained) + save = MagicMock() + monkeypatch.setattr(train_xgb_stereo, "configure_training", configure) + monkeypatch.setattr(train_xgb_stereo, "load_training_data", load_data) + monkeypatch.setattr(train_xgb_stereo, "train_regression", train) + monkeypatch.setattr(train_xgb_stereo, "save_models", save) + + train_xgb_stereo.main() + + configure.assert_called_once_with("stereo_analysis") + load_data.assert_called_once_with(configured, "gamma_inputs.txt", "stereo_analysis") + train.assert_called_once_with(loaded_data, configured) + save.assert_called_once_with(trained) + assert "stereo_analysis model trained successfully" in caplog.text + + +def test_train_stereo_entrypoint_does_not_save_when_regression_training_fails(monkeypatch): + """A failed regression fit must propagate and never create a partial artifact.""" + configured = {"input_file_list": "gamma_inputs.txt"} + save = MagicMock() + monkeypatch.setattr(train_xgb_stereo, "configure_training", lambda *_args: configured) + monkeypatch.setattr( + train_xgb_stereo, "load_training_data", lambda *_args: pd.DataFrame({"feature": [1.0]}) + ) + monkeypatch.setattr( + train_xgb_stereo, + "train_regression", + lambda *_args: (_ for _ in ()).throw(RuntimeError("fit failed")), + ) + monkeypatch.setattr(train_xgb_stereo, "save_models", save) + + with pytest.raises(RuntimeError, match="fit failed"): + train_xgb_stereo.main() + + save.assert_not_called() + + +def test_apply_stereo_entrypoint_passes_the_loaded_configuration_to_streaming(monkeypatch): + """The apply CLI must select stereo analysis and preserve loaded model metadata.""" + configured = { + "models": {"xgboost": {"model": object()}}, + "target_mean": {"Xoff_residual": 0.0}, + "target_std": {"Xoff_residual": 1.0}, + } + configure = MagicMock(return_value=configured) + process = MagicMock() + monkeypatch.setattr(apply_xgb_stereo, "configure_apply", configure) + monkeypatch.setattr(apply_xgb_stereo, "process_file_chunked", process) + + apply_xgb_stereo.main() + + configure.assert_called_once_with("stereo_analysis") + process.assert_called_once_with("stereo_analysis", configured) + + +def test_apply_stereo_entrypoint_does_not_stream_when_configuration_fails(monkeypatch): + """Invalid model configuration must stop before input ROOT data are processed.""" + process = MagicMock() + monkeypatch.setattr( + apply_xgb_stereo, + "configure_apply", + lambda *_args: (_ for _ in ()).throw(ValueError("missing target_std")), + ) + monkeypatch.setattr(apply_xgb_stereo, "process_file_chunked", process) + + with pytest.raises(ValueError, match="missing target_std"): + apply_xgb_stereo.main() + + process.assert_not_called() diff --git a/tests/test_models_helpers.py b/tests/test_models_helpers.py index 1e188dc..7428fc2 100644 --- a/tests/test_models_helpers.py +++ b/tests/test_models_helpers.py @@ -136,7 +136,9 @@ def test_load_classification_models_rejects_missing_energy_bin_metadata(tmp_path tmp_path / "model_ebin0.joblib", ) - with pytest.raises(ValueError, match=r"model_ebin0\.joblib.*energy_bins_log10_tev"): + with pytest.raises( + (ValueError, TypeError), match=r"model_ebin0\.joblib.*energy_bins_log10_tev" + ): models.load_classification_models(str(prefix), "xgboost") @@ -367,12 +369,14 @@ def test_process_file_chunked_uses_tmva_style_features_when_flag_set(): def fake_open(path): raise RuntimeError("uproot not needed for this assertion") - with patch("eventdisplay_ml.models.uproot.open", side_effect=fake_open): - with pytest.raises(RuntimeError, match="uproot not needed"): - models.process_file_chunked( - "classification", - {"tmva_style": True, "input_file": "dummy.root"}, - ) + with ( + patch("eventdisplay_ml.models.uproot.open", side_effect=fake_open), + pytest.raises(RuntimeError, match="uproot not needed"), + ): + models.process_file_chunked( + "classification", + {"tmva_style": True, "input_file": "dummy.root"}, + ) # Verify that tmva_style features differ from regular features assert set(expected) != set(regular_features) diff --git a/tests/test_regression_contracts.py b/tests/test_regression_contracts.py new file mode 100644 index 0000000..ac60f79 --- /dev/null +++ b/tests/test_regression_contracts.py @@ -0,0 +1,194 @@ +"""Regression production contracts for training, persisted models, and output. + +These tests deliberately exercise only the stereo-regression path. They pin the +contracts that must remain stable while the classification pipeline evolves: +target exclusion and order during training, scaler provenance, serialized-model +loading, feature reordering during inference, residual reconstruction, and the +one-output-row-per-input-event guarantee. +""" + +from unittest.mock import MagicMock, patch + +import joblib +import numpy as np +import pandas as pd +import pytest +from sklearn.model_selection import train_test_split + +from eventdisplay_ml import models + + +class CapturingRegressor: + """Minimal regressor that captures the exact arrays passed to ``fit``.""" + + best_iteration = 0 + best_score = 0.0 + + def fit(self, x_values, y_values, **kwargs): + """Record training inputs without fitting a model.""" + self.x_train = np.array(x_values, copy=True) + self.y_train = np.array(y_values, copy=True) + self.fit_kwargs = kwargs + return self + + def predict(self, x_values): + """Return a zero residual for every requested event.""" + return np.zeros((len(x_values), 3), dtype=np.float32) + + +class OrderedPredictionRegressor: + """Serializable predictor whose output exposes the received feature order.""" + + def predict(self, x_values): + """Generate residuals from the first two persisted feature columns.""" + # ``feature_beta`` must be first. If inference stops reindexing to the + # persisted order, these residuals (and thus final physics quantities) + # change immediately. + beta = x_values.iloc[:, 0].to_numpy(dtype=float) + alpha = x_values.iloc[:, 1].to_numpy(dtype=float) + return np.column_stack((beta, alpha, beta - alpha)) + + +@pytest.fixture +def regression_frame(): + """Return indexed data with values that make ordering errors observable.""" + n_events = 240 + row_number = np.arange(n_events, dtype=float) + return pd.DataFrame( + { + "feature_beta": 1000.0 + row_number, + "feature_alpha": -500.0 - row_number, + "ErecS": np.full(n_events, 3.0), + "DispNImages": np.full(n_events, 2), + "Xoff_weighted_bdt": 0.1 * row_number, + "Yoff_weighted_bdt": -0.2 * row_number, + "Xoff_residual": 2.0 + 0.01 * row_number, + "Yoff_residual": -3.0 - 0.02 * row_number, + "E_residual": 0.5 + 0.001 * row_number, + }, + index=10_000 + 7 * np.arange(n_events), + ) + + +def test_regression_training_contract_excludes_targets_and_uses_train_only_scalers( + regression_frame, +): + """Pin arrays, target order, split provenance, and weights sent to XGBoost.""" + targets = ["Xoff_residual", "Yoff_residual", "E_residual"] + config = { + "targets": targets, + "train_test_fraction": 0.5, + "random_state": 19, + "eval_max_events": 0, + "models": {"xgboost": {"hyper_parameters": {}}}, + } + captured_model = CapturingRegressor() + + with ( + patch("xgboost.XGBRegressor", return_value=captured_model), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(regression_frame, config) + + feature_columns = [column for column in regression_frame if column not in targets] + train_positions, test_positions = train_test_split( + np.arange(len(regression_frame)), train_size=0.5, random_state=19 + ) + expected_train_features = regression_frame.iloc[train_positions][feature_columns].to_numpy( + dtype=np.float32 + ) + expected_train_targets = regression_frame.iloc[train_positions][targets] + expected_mean = expected_train_targets.mean() + expected_std = expected_train_targets.std() + expected_scaled_targets = ((expected_train_targets - expected_mean) / expected_std).to_numpy( + dtype=np.float32 + ) + + assert result["features"] == feature_columns + assert result["models"]["xgboost"]["features"] == feature_columns + np.testing.assert_array_equal(captured_model.x_train, expected_train_features) + np.testing.assert_allclose(captured_model.y_train, expected_scaled_targets, atol=1e-7) + assert captured_model.fit_kwargs["sample_weight"].shape == (len(train_positions),) + assert captured_model.fit_kwargs["sample_weight_eval_set"][0].shape == (len(test_positions),) + assert result["target_mean"] == pytest.approx(expected_mean.to_dict(), abs=0.0) + assert result["target_std"] == pytest.approx(expected_std.to_dict(), abs=0.0) + + +def test_persisted_regression_model_preserves_feature_order_and_reconstructs_truth( + tmp_path, monkeypatch +): + """Load a model artifact and apply it with shuffled input columns. + + This is the production boundary: model features and target scalers come from + disk, while the flattened event data can be in a different column order. + """ + model_prefix = tmp_path / "stereo_model" + feature_order = [ + "feature_beta", + "feature_alpha", + "Xoff_weighted_bdt", + "Yoff_weighted_bdt", + "ErecS", + ] + target_mean = {"Xoff_residual": 1.0, "Yoff_residual": -2.0, "E_residual": 0.25} + target_std = {"Xoff_residual": 0.5, "Yoff_residual": 2.0, "E_residual": 0.1} + joblib.dump( + { + "models": {"xgboost": {"model": OrderedPredictionRegressor()}}, + "features": feature_order, + "target_mean": target_mean, + "target_std": target_std, + }, + tmp_path / "stereo_model.joblib.gz", + ) + loaded_models, loaded_parameters = models.load_regression_models(str(model_prefix), "xgboost") + + # Deliberately not in persisted feature order, and includes a column that + # must not reach the model. + flattened = pd.DataFrame( + { + "ignored_new_column": [99.0, 98.0], + "feature_alpha": [4.0, 7.0], + "ErecS": [10.0, 100.0], + "Xoff_weighted_bdt": [0.2, -0.5], + "feature_beta": [3.0, 11.0], + "Yoff_weighted_bdt": [-1.0, 2.0], + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *_args, **_kwargs: flattened) + monkeypatch.setattr(models.data_processing, "print_variable_statistics", lambda *_args: None) + + pred_xoff, pred_yoff, pred_log_energy = models.apply_regression_models( + pd.DataFrame({"event": [1, 2]}), + {"models": loaded_models, **loaded_parameters}, + ) + + # Raw predicted residuals are [beta, alpha, beta-alpha], then each target + # is inverse-standardized and added to its DispBDT baseline. + np.testing.assert_allclose(pred_xoff, [2.7, 6.0]) + np.testing.assert_allclose(pred_yoff, [5.0, 14.0]) + np.testing.assert_allclose(pred_log_energy, [1.15, 2.65]) + + +def test_stereo_output_writer_keeps_nan_energy_rows_and_converts_only_log_energy(monkeypatch): + """Pin ROOT payload names, float32 conversion, and row-preserving NaNs.""" + tree = MagicMock() + monkeypatch.setattr( + models, + "apply_regression_models", + lambda *_args: ( + np.array([1.25, np.nan]), + np.array([-2.5, 3.5]), + np.array([2.0, np.nan]), + ), + ) + + models._apply_model("stereo_analysis", pd.DataFrame({"event": [1, 2]}), {}, tree) + + payload = tree.extend.call_args.args[0] + assert list(payload) == ["Dir_Xoff", "Dir_Yoff", "Dir_Erec"] + assert all(values.dtype == np.float32 for values in payload.values()) + assert len(payload["Dir_Xoff"]) == 2 + assert payload["Dir_Erec"][0] == pytest.approx(100.0) + assert np.isnan(payload["Dir_Xoff"][1]) + assert np.isnan(payload["Dir_Erec"][1]) diff --git a/tests/test_regression_production_guard.py b/tests/test_regression_production_guard.py new file mode 100644 index 0000000..cfa3ce8 --- /dev/null +++ b/tests/test_regression_production_guard.py @@ -0,0 +1,1536 @@ +"""Production-grade guard tests for the stereo-regression path. + +These tests exist specifically to catch regressions introduced by changes to +the classification pipeline that shares infrastructure with the regression path. +Every test targets a concrete, observable contract. If any of these fail after +a classification refactoring, the regression path has been broken. + +Coverage areas +-------------- +1. ``train_regression`` – feature/target separation, standardisation, sample + weights, train/test split, determinism, empty-input guard. +2. ``apply_regression_models`` – residual inversion, feature reordering, ErecS + handling, high-multiplicity dispatch, output shape, NaN propagation. +3. ``load_regression_models`` – artifact structure, missing-key handling, + mutual independence from classification loader. +4. ``_apply_model`` dispatch – stereo tree names, float32 dtypes, 10^x energy + conversion, row-count preservation. +5. Internal helpers – ``_feature_array``, ``_predict_unscaled_chunked``, + ``_sample_eval_indices``, ``_regression_sample_weights``, + ``_log_energy_bin_counts_from_arrays``. +6. ``process_file_chunked`` – stereo path only, chunk/index invariants. +7. Feature schema – target list, analysis-type tag, pointing-offset exclusion. +8. ``_output_tree`` – stereo branch set, no classification branches present. +""" + +import math +from unittest.mock import MagicMock, patch + +import awkward as ak +import joblib +import numpy as np +import pandas as pd +import pytest +from sklearn.model_selection import train_test_split + +from eventdisplay_ml import data_processing, models +from eventdisplay_ml.models import ( + _feature_array, + _log_energy_bin_counts_from_arrays, + _output_tree, + _predict_unscaled_chunked, + _regression_sample_weights, + _sample_eval_indices, +) + +# --------------------------------------------------------------------------- +# Minimal helpers shared across tests +# --------------------------------------------------------------------------- + +TARGETS = ["Xoff_residual", "Yoff_residual", "E_residual"] + + +def _make_regression_df(n=300, seed=0): + """Return a deterministic regression-ready DataFrame.""" + rng = np.random.default_rng(seed) + df = pd.DataFrame( + { + "Xoff_residual": rng.normal(0.1, 0.4, n), + "Yoff_residual": rng.normal(-0.2, 0.5, n), + "E_residual": rng.normal(0.05, 0.2, n), + "ErecS": np.logspace(0, 2, n), + "DispNImages": rng.choice([2, 3, 4], n), + "Xoff_weighted_bdt": rng.normal(0, 0.5, n), + "Yoff_weighted_bdt": rng.normal(0, 0.5, n), + "feature_A": rng.uniform(-1, 1, n), + "feature_B": rng.uniform(0, 10, n), + } + ) + return df + + +def _make_base_config(seed=7, n_estimators=5): + return { + "targets": TARGETS, + "train_test_fraction": 0.5, + "random_state": seed, + "eval_max_events": 0, + "models": { + "xgboost": { + "hyper_parameters": { + "n_estimators": n_estimators, + "max_depth": 2, + "random_state": seed, + } + } + }, + } + + +class _ZeroResidualModel: + """Predict zero residuals for all events – simplest possible stand-in.""" + + best_iteration = 0 + best_score = 0.0 + + def fit(self, *args, **kwargs): + return self + + def predict(self, x): + return np.zeros((len(x), 3), dtype=np.float32) + + +class _RecordingModel: + """Record every ``fit`` call for later inspection.""" + + best_iteration = 0 + best_score = 0.0 + + def __init__(self): + self.fit_calls = [] + self.predict_calls = [] + + def fit(self, x, y, **kw): + self.fit_calls.append( + {"x": np.array(x, copy=True), "y": np.array(y, copy=True), "kwargs": kw} + ) + return self + + def predict(self, x): + self.predict_calls.append(np.array(x, copy=True)) + return np.zeros((len(x), 3), dtype=np.float32) + + +class _ColumnEchoModel: + """Return the first three feature columns as residuals so order is detectable.""" + + def predict(self, x_df): + arr = x_df.to_numpy(dtype=float) if hasattr(x_df, "to_numpy") else np.asarray(x_df) + return arr[:, :3] + + +# =========================================================================== +# 1. train_regression +# =========================================================================== + + +class TestTrainRegressionFeatureTargetSeparation: + """Targets must be excluded from feature columns, never passed to XGBoost.""" + + def test_targets_absent_from_feature_list(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + assert not any(t in result["features"] for t in TARGETS), ( + "One or more targets leaked into the feature list" + ) + + def test_feature_list_matches_columns_minus_targets(self): + df = _make_regression_df() + cfg = _make_base_config() + expected = [c for c in df.columns if c not in set(TARGETS)] + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + assert result["features"] == expected + + def test_xgboost_x_train_has_no_target_column(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df, cfg) + + x_cols_seen = rec.fit_calls[0]["x"].shape[1] + expected_n_features = len(df.columns) - len(TARGETS) + assert x_cols_seen == expected_n_features + + +class TestTrainRegressionTargetStandardisation: + """Scalers must be computed from training data only and stored correctly.""" + + def test_target_mean_and_std_stored_in_result(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + assert "target_mean" in result + assert "target_std" in result + assert set(result["target_mean"]) == set(TARGETS) + assert set(result["target_std"]) == set(TARGETS) + + def test_scalers_come_from_train_split_not_full_dataset(self): + df = _make_regression_df(n=200, seed=1) + cfg = _make_base_config(seed=1) + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + train_idx, _ = train_test_split(np.arange(len(df)), train_size=0.5, random_state=1) + for t in TARGETS: + expected_mean = df.iloc[train_idx][t].mean() + expected_std = df.iloc[train_idx][t].std() + assert result["target_mean"][t] == pytest.approx(expected_mean, abs=1e-9) + assert result["target_std"][t] == pytest.approx(expected_std, abs=1e-9) + + def test_y_train_passed_to_xgboost_is_standardised(self): + df = _make_regression_df(n=200, seed=2) + cfg = _make_base_config(seed=2) + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + y_passed = rec.fit_calls[0]["y"] # shape (n_train, 3) + for col_idx, t in enumerate(TARGETS): + col = y_passed[:, col_idx] + assert abs(col.mean()) < 0.15, f"Target {t}: scaled mean too far from 0" + assert abs(col.std() - 1.0) < 0.15, f"Target {t}: scaled std too far from 1" + + def test_target_std_is_never_zero(self): + # All three targets must have non-zero std in training. + df = _make_regression_df(n=300, seed=3) + cfg = _make_base_config(seed=3) + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + for t in TARGETS: + assert result["target_std"][t] > 0.0, f"Target {t} has zero std" + + def test_target_order_in_scaler_dicts_matches_targets_list(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + assert list(result["target_mean"].keys()) == TARGETS + assert list(result["target_std"].keys()) == TARGETS + + +class TestTrainRegressionSampleWeights: + """Sample weights must cover the training set, be finite, and be positive.""" + + def test_sample_weights_shape_matches_train_size(self): + n = 400 + df = _make_regression_df(n=n, seed=4) + cfg = _make_base_config(seed=4) + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df, cfg) + + weights = rec.fit_calls[0]["kwargs"]["sample_weight"] + assert weights is None or len(weights) == n // 2 + + def test_sample_weights_are_finite_and_positive(self): + df = _make_regression_df(n=400, seed=5) + cfg = _make_base_config(seed=5) + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df, cfg) + + w = rec.fit_calls[0]["kwargs"]["sample_weight"] + if w is not None: + assert np.all(np.isfinite(w)), "Sample weights contain non-finite values" + assert np.all(w >= 0), "Sample weights contain negative values" + + def test_eval_weights_shape_matches_eval_set(self): + df = _make_regression_df(n=400, seed=6) + cfg = _make_base_config(seed=6) + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df, cfg) + + kw = rec.fit_calls[0]["kwargs"] + if kw.get("sample_weight_eval_set") is not None: + eval_w = kw["sample_weight_eval_set"] + assert isinstance(eval_w, list) + assert len(eval_w) == 1 + assert eval_w[0] is None or len(eval_w[0]) == len(kw.get("eval_set", [[]])[0][0]) + + +class TestTrainRegressionSplitDeterminism: + """The same random_state must always yield the same feature array passed to fit.""" + + def test_identical_seed_gives_identical_x_train(self): + df = _make_regression_df(n=200, seed=10) + cfg_a = _make_base_config(seed=42) + cfg_b = _make_base_config(seed=42) + rec_a, rec_b = _RecordingModel(), _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec_a), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df.copy(), cfg_a) + with ( + patch("xgboost.XGBRegressor", return_value=rec_b), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df.copy(), cfg_b) + + np.testing.assert_array_equal(rec_a.fit_calls[0]["x"], rec_b.fit_calls[0]["x"]) + + def test_different_seed_gives_different_x_train(self): + df = _make_regression_df(n=200, seed=10) + rec_a, rec_b = _RecordingModel(), _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec_a), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df.copy(), _make_base_config(seed=1)) + with ( + patch("xgboost.XGBRegressor", return_value=rec_b), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df.copy(), _make_base_config(seed=2)) + + # Different seeds should produce different splits (very likely with 200 rows) + assert not np.array_equal(rec_a.fit_calls[0]["x"], rec_b.fit_calls[0]["x"]) + + +class TestTrainRegressionEdgeCases: + def test_empty_dataframe_returns_none(self): + result = models.train_regression(pd.DataFrame(), _make_base_config()) + assert result is None + + def test_result_contains_features_key_at_top_level(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + assert "features" in result + # Must also be stored inside the per-model dict + assert "features" in result["models"]["xgboost"] + + def test_model_object_stored_in_models_dict(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + result = models.train_regression(df, cfg) + + assert result["models"]["xgboost"]["model"] is rec + + def test_x_train_dtype_is_float32(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df, cfg) + + assert rec.fit_calls[0]["x"].dtype == np.float32 + + def test_y_train_dtype_is_float32(self): + df = _make_regression_df() + cfg = _make_base_config() + rec = _RecordingModel() + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + models.train_regression(df, cfg) + + assert rec.fit_calls[0]["y"].dtype == np.float32 + + +# =========================================================================== +# 2. apply_regression_models +# =========================================================================== + + +class TestApplyRegressionResidualInversion: + """Residuals must be unscaled and added to the DispBDT baseline.""" + + def _build_apply_config(self, model, features, target_mean, target_std): + return { + "models": {"xgboost": {"model": model, "features": features}}, + "target_mean": target_mean, + "target_std": target_std, + } + + def _flat_df(self, xoff_bdt, yoff_bdt, erec_s): + return pd.DataFrame( + { + "Xoff_weighted_bdt": xoff_bdt, + "Yoff_weighted_bdt": yoff_bdt, + "ErecS": erec_s, + } + ) + + def test_zero_residual_model_returns_mean_plus_baseline(self, monkeypatch): + """When scaled residual = 0, final = mean + baseline.""" + flat = self._flat_df([10.0], [20.0], [100.0]) + target_mean = {"Xoff_residual": 3.0, "Yoff_residual": -5.0, "E_residual": 0.4} + target_std = {"Xoff_residual": 2.0, "Yoff_residual": 1.5, "E_residual": 0.1} + + class _ZeroModel: + def predict(self, x): + return np.zeros((len(x), 3)) + + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + cfg = self._build_apply_config(_ZeroModel(), flat.columns.tolist(), target_mean, target_std) + pred_xoff, pred_yoff, pred_e = models.apply_regression_models( + pd.DataFrame({"DispNImages": [2]}), cfg + ) + + # xoff: 0*2 + 3 + 10 = 13 + np.testing.assert_allclose(pred_xoff, [13.0], atol=1e-7) + # yoff: 0*1.5 + (-5) + 20 = 15 + np.testing.assert_allclose(pred_yoff, [15.0], atol=1e-7) + # log10(100)=2; E: 0*0.1 + 0.4 + 2 = 2.4 + np.testing.assert_allclose(pred_e, [2.4], atol=1e-7) + + def test_unit_scaled_residual_inverts_correctly(self, monkeypatch): + """Scaled residual=1 => physical residual = std + mean.""" + flat = self._flat_df([0.0], [0.0], [10.0]) + target_mean = {"Xoff_residual": 1.0, "Yoff_residual": 2.0, "E_residual": 0.5} + target_std = {"Xoff_residual": 2.0, "Yoff_residual": 3.0, "E_residual": 0.1} + + class _OnesModel: + def predict(self, x): + return np.ones((len(x), 3)) + + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + cfg = self._build_apply_config(_OnesModel(), flat.columns.tolist(), target_mean, target_std) + pred_xoff, pred_yoff, pred_e = models.apply_regression_models( + pd.DataFrame({"DispNImages": [2]}), cfg + ) + + # xoff: 1*2 + 1 + 0 = 3 + np.testing.assert_allclose(pred_xoff, [3.0], atol=1e-7) + # yoff: 1*3 + 2 + 0 = 5 + np.testing.assert_allclose(pred_yoff, [5.0], atol=1e-7) + # log10(10)=1; E: 1*0.1 + 0.5 + 1 = 1.6 + np.testing.assert_allclose(pred_e, [1.6], atol=1e-7) + + def test_negative_scaled_residual(self, monkeypatch): + flat = self._flat_df([5.0], [-5.0], [1000.0]) + target_mean = {"Xoff_residual": 0.0, "Yoff_residual": 0.0, "E_residual": 0.0} + target_std = {"Xoff_residual": 1.0, "Yoff_residual": 1.0, "E_residual": 1.0} + + class _MinusOneModel: + def predict(self, x): + return -1.0 * np.ones((len(x), 3)) + + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + cfg = self._build_apply_config( + _MinusOneModel(), flat.columns.tolist(), target_mean, target_std + ) + pred_xoff, pred_yoff, pred_e = models.apply_regression_models( + pd.DataFrame({"DispNImages": [2]}), cfg + ) + + np.testing.assert_allclose(pred_xoff, [4.0], atol=1e-7) + np.testing.assert_allclose(pred_yoff, [-6.0], atol=1e-7) + np.testing.assert_allclose(pred_e, [2.0], atol=1e-7) # log10(1000)-1 = 2 + + def test_multiple_events_independent(self, monkeypatch): + n = 5 + xoff_bdt = np.arange(n, dtype=float) + yoff_bdt = np.arange(n, dtype=float) * -2 + erec_s = 10.0 ** np.arange(n, dtype=float) + flat = self._flat_df(xoff_bdt, yoff_bdt, erec_s) + + class _ConstantModel: + def predict(self, x): + # Always predicts scaled residual [0.5, -0.5, 1.0] + return np.tile([0.5, -0.5, 1.0], (len(x), 1)) + + target_mean = {"Xoff_residual": 1.0, "Yoff_residual": -1.0, "E_residual": 0.2} + target_std = {"Xoff_residual": 2.0, "Yoff_residual": 4.0, "E_residual": 0.5} + + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + cfg = self._build_apply_config( + _ConstantModel(), flat.columns.tolist(), target_mean, target_std + ) + pred_xoff, pred_yoff, pred_e = models.apply_regression_models( + pd.DataFrame({"DispNImages": np.full(n, 2)}), cfg + ) + + # physical residual_xoff = 0.5*2 + 1 = 2.0; final = baseline + 2.0 + np.testing.assert_allclose(pred_xoff, xoff_bdt + 2.0, atol=1e-7) + # physical residual_yoff = -0.5*4 + (-1) = -3.0 + np.testing.assert_allclose(pred_yoff, yoff_bdt - 3.0, atol=1e-7) + # physical residual_e = 1.0*0.5 + 0.2 = 0.7 + expected_e = np.arange(n, dtype=float) + 0.7 # log10(10^i) = i, then +0.7 + np.testing.assert_allclose(pred_e, expected_e, atol=1e-7) + + +class TestApplyRegressionFeatureReordering: + """The model must receive features in the persisted order, not the input order.""" + + def test_shuffled_input_columns_produce_correct_physics(self, monkeypatch, tmp_path): + feature_order = [ + "feat_alpha", + "feat_beta", + "ErecS", + "Xoff_weighted_bdt", + "Yoff_weighted_bdt", + ] + target_mean = {"Xoff_residual": 0.0, "Yoff_residual": 0.0, "E_residual": 0.0} + target_std = {"Xoff_residual": 1.0, "Yoff_residual": 1.0, "E_residual": 1.0} + + class _FirstColModel: + """Returns first column value as all three residuals, so order is detectable.""" + + def predict(self, x): + col0 = x.iloc[:, 0].to_numpy(dtype=float) if hasattr(x, "iloc") else x[:, 0] + return np.column_stack([col0, col0, col0]) + + # Build flattened frame in non-persisted order + flat_shuffled = pd.DataFrame( + { + "Yoff_weighted_bdt": [0.0, 0.0], + "ErecS": [100.0, 100.0], + "feat_beta": [99.0, 99.0], # must NOT be first after reindex + "feat_alpha": [7.0, 7.0], # must be first after reindex → residual = 7.0 + "Xoff_weighted_bdt": [5.0, 5.0], + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat_shuffled) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + cfg = { + "models": {"xgboost": {"model": _FirstColModel(), "features": feature_order}}, + "target_mean": target_mean, + "target_std": target_std, + } + pred_xoff, pred_yoff, pred_e = models.apply_regression_models( + pd.DataFrame({"DispNImages": [2, 2]}), cfg + ) + + # After reindex to persisted order, column 0 is feat_alpha = 7.0. + # xoff = 7.0 + Xoff_weighted_bdt(5.0) = 12.0 + np.testing.assert_allclose(pred_xoff, [12.0, 12.0], atol=1e-7) + + def test_extra_input_columns_not_forwarded_to_model(self, monkeypatch): + """Columns present in input but not in persisted features must be dropped.""" + feature_order = ["Xoff_weighted_bdt", "Yoff_weighted_bdt", "ErecS"] + received_shapes = [] + + class _ShapeRecordModel: + def predict(self, x): + received_shapes.append(x.shape if hasattr(x, "shape") else (len(x),)) + return np.zeros((len(x), 3)) + + flat = pd.DataFrame( + { + "Xoff_weighted_bdt": [1.0], + "Yoff_weighted_bdt": [2.0], + "ErecS": [10.0], + "extra_col_should_be_dropped": [999.0], + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + cfg = { + "models": {"xgboost": {"model": _ShapeRecordModel(), "features": feature_order}}, + "target_mean": dict.fromkeys(TARGETS, 0.0), + "target_std": dict.fromkeys(TARGETS, 1.0), + } + models.apply_regression_models(pd.DataFrame({"DispNImages": [2]}), cfg) + + n_cols = received_shapes[0][1] + assert n_cols == len(feature_order), ( + f"Model received {n_cols} columns but should have received {len(feature_order)}" + ) + + +class TestApplyRegressionErecSHandling: + """Invalid ErecS must propagate to NaN energy output without corrupting direction.""" + + @pytest.fixture(autouse=True) + def _patch(self, monkeypatch): + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + def _run(self, monkeypatch, erec_s_values, scaled_preds=None): + n = len(erec_s_values) + flat = pd.DataFrame( + { + "Xoff_weighted_bdt": np.zeros(n), + "Yoff_weighted_bdt": np.zeros(n), + "ErecS": erec_s_values, + } + ) + if scaled_preds is None: + scaled_preds = np.zeros((n, 3)) + + class _FixedModel: + def predict(self, x): + return scaled_preds + + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + cfg = { + "models": {"xgboost": {"model": _FixedModel(), "features": flat.columns.tolist()}}, + "target_mean": dict.fromkeys(TARGETS, 0.0), + "target_std": dict.fromkeys(TARGETS, 1.0), + } + return models.apply_regression_models(pd.DataFrame({"DispNImages": np.full(n, 2)}), cfg) + + def test_negative_erecs_produce_nan_energy(self, monkeypatch): + _, _, pred_e = self._run(monkeypatch, [-5.0, 100.0]) + assert np.isnan(pred_e[0]) + assert not np.isnan(pred_e[1]) + + def test_zero_erecs_produces_nan_energy(self, monkeypatch): + _, _, pred_e = self._run(monkeypatch, [0.0]) + assert np.isnan(pred_e[0]) + + def test_nan_erecs_produces_nan_energy(self, monkeypatch): + _, _, pred_e = self._run(monkeypatch, [np.nan]) + assert np.isnan(pred_e[0]) + + def test_invalid_erecs_do_not_corrupt_direction(self, monkeypatch): + flat = pd.DataFrame( + { + "Xoff_weighted_bdt": [10.0, 20.0], + "Yoff_weighted_bdt": [30.0, 40.0], + "ErecS": [np.nan, 100.0], + } + ) + + class _ConstModel: + def predict(self, x): + return np.tile([1.0, 2.0, 3.0], (len(x), 1)) + + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + cfg = { + "models": {"xgboost": {"model": _ConstModel(), "features": flat.columns.tolist()}}, + "target_mean": dict.fromkeys(TARGETS, 0.0), + "target_std": dict.fromkeys(TARGETS, 1.0), + } + pred_xoff, pred_yoff, _ = models.apply_regression_models( + pd.DataFrame({"DispNImages": [2, 2]}), cfg + ) + + # Direction must still be valid for the invalid-energy event + assert np.isfinite(pred_xoff[0]) + assert np.isfinite(pred_yoff[0]) + + def test_all_erecs_valid_no_nan_energy(self, monkeypatch): + _, _, pred_e = self._run(monkeypatch, [1.0, 10.0, 100.0, 1000.0]) + assert all(np.isfinite(pred_e)) + + def test_output_length_equals_input_length(self, monkeypatch): + n = 97 + rng = np.random.default_rng(0) + erec_s = np.where(rng.random(n) > 0.3, rng.uniform(1, 1000, n), np.nan) + pred_xoff, pred_yoff, pred_e = self._run(monkeypatch, erec_s) + assert len(pred_xoff) == n + assert len(pred_yoff) == n + assert len(pred_e) == n + + +class TestApplyRegressionHighMultiplicity: + """High-multiplicity events must route to the dedicated model when configured.""" + + @pytest.fixture(autouse=True) + def _patch(self, monkeypatch): + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + def _make_two_model_cfg(self, flat_df, low_pred, high_pred): + class _SelectableModel: + def __init__(self, val): + self._val = val + + def predict(self, x): + return np.tile(self._val, (len(x), 1)) + + zero_mean = dict.fromkeys(TARGETS, 0.0) + unit_std = dict.fromkeys(TARGETS, 1.0) + return { + "models": { + "xgboost": {"model": _SelectableModel(low_pred), "features": flat_df.columns} + }, + "models_high_multiplicity": { + "xgboost": { + "model": _SelectableModel(high_pred), + "features": flat_df.columns, + } + }, + "target_mean": zero_mean, + "target_std": unit_std, + "target_mean_high_multiplicity": zero_mean, + "target_std_high_multiplicity": unit_std, + } + + def test_multiplicity_2_uses_primary_model(self, monkeypatch): + flat = pd.DataFrame( + { + "Xoff_weighted_bdt": [0.0], + "Yoff_weighted_bdt": [0.0], + "ErecS": [10.0], + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + cfg = self._make_two_model_cfg(flat, [5.0, 5.0, 5.0], [9.0, 9.0, 9.0]) + + pred_xoff, _, _ = models.apply_regression_models(pd.DataFrame({"DispNImages": [2]}), cfg) + np.testing.assert_allclose(pred_xoff, [5.0]) # low model + + def test_multiplicity_3_uses_high_model(self, monkeypatch): + flat = pd.DataFrame( + { + "Xoff_weighted_bdt": [0.0], + "Yoff_weighted_bdt": [0.0], + "ErecS": [10.0], + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + cfg = self._make_two_model_cfg(flat, [5.0, 5.0, 5.0], [9.0, 9.0, 9.0]) + + pred_xoff, _, _ = models.apply_regression_models(pd.DataFrame({"DispNImages": [3]}), cfg) + np.testing.assert_allclose(pred_xoff, [9.0]) # high model + + def test_multiplicity_1_returns_nan(self, monkeypatch): + flat = pd.DataFrame( + { + "Xoff_weighted_bdt": [0.0], + "Yoff_weighted_bdt": [0.0], + "ErecS": [10.0], + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + cfg = self._make_two_model_cfg(flat, [5.0, 5.0, 5.0], [9.0, 9.0, 9.0]) + + pred_xoff, _, _ = models.apply_regression_models(pd.DataFrame({"DispNImages": [1]}), cfg) + assert np.isnan(pred_xoff[0]) + + def test_mixed_multiplicity_routes_correctly(self, monkeypatch): + n = 4 + flat = pd.DataFrame( + { + "Xoff_weighted_bdt": np.zeros(n), + "Yoff_weighted_bdt": np.zeros(n), + "ErecS": np.ones(n) * 10.0, + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + cfg = self._make_two_model_cfg(flat, [2.0, 2.0, 2.0], [7.0, 7.0, 7.0]) + + # mult: [1, 2, 3, 4] + pred_xoff, _, _ = models.apply_regression_models( + pd.DataFrame({"DispNImages": [1, 2, 3, 4]}), cfg + ) + assert np.isnan(pred_xoff[0]) # mult=1: unrouted → NaN + assert pred_xoff[1] == pytest.approx(2.0) # low model + assert pred_xoff[2] == pytest.approx(7.0) # high model + assert pred_xoff[3] == pytest.approx(7.0) # high model + + +class TestApplyRegressionMissingParams: + def test_missing_target_mean_raises_value_error(self, monkeypatch): + flat = pd.DataFrame( + {"Xoff_weighted_bdt": [1.0], "Yoff_weighted_bdt": [2.0], "ErecS": [10.0]} + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + class _M: + def predict(self, x): + return np.zeros((len(x), 3)) + + cfg = { + "models": {"xgboost": {"model": _M(), "features": flat.columns.tolist()}}, + # no target_mean or target_std + } + with pytest.raises(ValueError, match="target standardization"): + models.apply_regression_models(pd.DataFrame({"DispNImages": [2]}), cfg) + + +# =========================================================================== +# 3. load_regression_models +# =========================================================================== + + +class TestLoadRegressionModels: + """Artifact loading must be independent of the classification loader.""" + + def _write_artifact(self, tmp_path, feature_list=None, include_scalers=True): + if feature_list is None: + feature_list = ["f1", "f2", "ErecS", "Xoff_weighted_bdt", "Yoff_weighted_bdt"] + payload = { + "models": {"xgboost": {"model": _ZeroResidualModel()}}, + "features": feature_list, + } + if include_scalers: + payload["target_mean"] = dict.fromkeys(TARGETS, 0.1) + payload["target_std"] = dict.fromkeys(TARGETS, 1.0) + path = tmp_path / "stereo_test.joblib.gz" + joblib.dump(payload, path) + return path + + def test_load_returns_model_and_parameters(self, tmp_path): + self._write_artifact(tmp_path) + loaded_models, par = models.load_regression_models(str(tmp_path / "stereo_test"), "xgboost") + assert "xgboost" in loaded_models + assert "model" in loaded_models["xgboost"] + assert "features" in loaded_models["xgboost"] + + def test_load_preserves_feature_order(self, tmp_path): + feat = ["z_col", "a_col", "ErecS", "Xoff_weighted_bdt", "Yoff_weighted_bdt"] + self._write_artifact(tmp_path, feature_list=feat) + loaded_models, _ = models.load_regression_models(str(tmp_path / "stereo_test"), "xgboost") + assert loaded_models["xgboost"]["features"] == feat + + def test_load_restores_target_scalers(self, tmp_path): + self._write_artifact(tmp_path) + _, par = models.load_regression_models(str(tmp_path / "stereo_test"), "xgboost") + assert "target_mean" in par + assert "target_std" in par + for t in TARGETS: + assert t in par["target_mean"] + assert t in par["target_std"] + + def test_load_missing_file_raises_error(self, tmp_path): + with pytest.raises((FileNotFoundError, Exception)): + models.load_regression_models(str(tmp_path / "nonexistent"), "xgboost") + + def test_load_regression_does_not_touch_classification_loader(self, tmp_path, monkeypatch): + """load_regression_models must never call load_classification_models.""" + self._write_artifact(tmp_path) + called = [] + monkeypatch.setattr( + models, "load_classification_models", lambda *a, **k: called.append(True) + ) + models.load_regression_models(str(tmp_path / "stereo_test"), "xgboost") + assert called == [], "load_classification_models was called during regression load" + + def test_load_models_dispatcher_routes_stereo_correctly(self, tmp_path, monkeypatch): + """The top-level load_models dispatcher must route 'stereo_analysis' to regression.""" + expected = ({"xgboost": {"model": None, "features": []}}, {}) + monkeypatch.setattr(models, "load_regression_models", lambda *a: expected) + result = models.load_models("stereo_analysis", str(tmp_path / "model"), "xgboost") + assert result == expected + + def test_load_models_dispatcher_raises_for_unknown_type(self, tmp_path): + with pytest.raises(ValueError, match="Unknown analysis_type"): + models.load_models("unknown_type", str(tmp_path / "model"), "xgboost") + + +# =========================================================================== +# 4. _apply_model dispatch +# =========================================================================== + + +class TestApplyModelStereoDispatch: + """_apply_model must write Dir_Xoff, Dir_Yoff, Dir_Erec with correct types.""" + + def _tree_and_apply(self, monkeypatch, pred_xoff, pred_yoff, pred_log_e): + tree = MagicMock() + monkeypatch.setattr( + models, + "apply_regression_models", + lambda *a: ( + np.asarray(pred_xoff, dtype=np.float64), + np.asarray(pred_yoff, dtype=np.float64), + np.asarray(pred_log_e, dtype=np.float64), + ), + ) + models._apply_model( + "stereo_analysis", + pd.DataFrame({"event": range(len(pred_xoff))}), + {}, + tree, + ) + return tree.extend.call_args.args[0] + + def test_stereo_payload_keys(self, monkeypatch): + payload = self._tree_and_apply(monkeypatch, [1.0], [2.0], [3.0]) + assert set(payload) == {"Dir_Xoff", "Dir_Yoff", "Dir_Erec"} + + def test_stereo_payload_dtype_float32(self, monkeypatch): + payload = self._tree_and_apply(monkeypatch, [1.0, 2.0], [3.0, 4.0], [5.0, 6.0]) + for key, arr in payload.items(): + assert arr.dtype == np.float32, f"Branch {key} has dtype {arr.dtype}, expected float32" + + def test_energy_converted_from_log10(self, monkeypatch): + payload = self._tree_and_apply(monkeypatch, [0.0], [0.0], [2.0]) + assert payload["Dir_Erec"][0] == pytest.approx(100.0, rel=1e-5) + + def test_energy_log10_minus1_gives_01(self, monkeypatch): + payload = self._tree_and_apply(monkeypatch, [0.0], [0.0], [-1.0]) + assert payload["Dir_Erec"][0] == pytest.approx(0.1, rel=1e-5) + + def test_nan_energy_propagates_to_output(self, monkeypatch): + payload = self._tree_and_apply(monkeypatch, [1.0, 2.0], [3.0, 4.0], [np.nan, 1.0]) + assert np.isnan(payload["Dir_Erec"][0]) + assert np.isfinite(payload["Dir_Erec"][1]) + + def test_nan_direction_propagates(self, monkeypatch): + payload = self._tree_and_apply(monkeypatch, [np.nan], [1.0], [1.0]) + assert np.isnan(payload["Dir_Xoff"][0]) + + def test_row_count_preserved(self, monkeypatch): + n = 50 + payload = self._tree_and_apply( + monkeypatch, + np.ones(n), + np.zeros(n), + np.full(n, 2.0), + ) + assert len(payload["Dir_Xoff"]) == n + assert len(payload["Dir_Yoff"]) == n + assert len(payload["Dir_Erec"]) == n + + def test_no_classification_branches_written(self, monkeypatch): + payload = self._tree_and_apply(monkeypatch, [1.0], [1.0], [1.0]) + classification_keys = {"Gamma_Prediction", "Is_Gamma_80", "Is_Gamma_70"} + assert not (set(payload) & classification_keys) + + def test_unknown_analysis_type_raises(self): + with pytest.raises(ValueError, match="Unknown analysis_type"): + models._apply_model("bad_type", pd.DataFrame({"x": [1]}), {}, MagicMock()) + + +# =========================================================================== +# 5. Internal helpers +# =========================================================================== + + +class TestFeatureArrayHelper: + def test_returns_float32_array(self): + df = pd.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]}) + arr = _feature_array(df, np.array([0, 2]), ["a", "b"]) + assert arr.dtype == np.float32 + + def test_selects_correct_rows(self): + df = pd.DataFrame({"x": [10.0, 20.0, 30.0, 40.0], "y": [1.0, 2.0, 3.0, 4.0]}) + arr = _feature_array(df, np.array([1, 3]), ["x", "y"]) + np.testing.assert_array_equal(arr, [[20.0, 2.0], [40.0, 4.0]]) + + def test_selects_correct_columns(self): + df = pd.DataFrame({"a": [1.0], "b": [2.0], "c": [3.0]}) + arr = _feature_array(df, np.array([0]), ["b", "a"]) + assert arr[0, 0] == pytest.approx(2.0) + assert arr[0, 1] == pytest.approx(1.0) + + def test_empty_row_index(self): + df = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]}) + arr = _feature_array(df, np.array([], dtype=int), ["a", "b"]) + assert arr.shape == (0, 2) + + +class TestPredictUnscaledChunked: + def _run(self, n=10, chunk_size=None): + df = pd.DataFrame( + { + "f1": np.arange(n, dtype=float), + "f2": np.arange(n, dtype=float) * 2, + "Xoff_residual": np.zeros(n), + "Yoff_residual": np.zeros(n), + "E_residual": np.zeros(n), + } + ) + targets = TARGETS + x_cols = ["f1", "f2"] + y_mean = pd.Series({"Xoff_residual": 1.0, "Yoff_residual": 2.0, "E_residual": 3.0}) + y_std = pd.Series({"Xoff_residual": 2.0, "Yoff_residual": 0.5, "E_residual": 1.0}) + + class _IdentModel: + """Returns scaled input columns as residuals.""" + + def predict(self, x): + arr = np.asarray(x, dtype=np.float32) + return np.column_stack([arr[:, 0], arr[:, 1], arr[:, 0] - arr[:, 1]]) + + row_indices = np.arange(n) + return _predict_unscaled_chunked( + _IdentModel(), df, row_indices, x_cols, y_mean, y_std, targets, chunk_size + ) + + def test_output_shape(self): + result = self._run(n=15) + assert result.shape == (15, 3) + + def test_output_columns_are_targets(self): + result = self._run(n=5) + assert list(result.columns) == TARGETS + + def test_chunked_vs_single_pass_identical(self): + full = self._run(n=20, chunk_size=None) + chunked = self._run(n=20, chunk_size=4) + pd.testing.assert_frame_equal(full, chunked) + + def test_unscaling_is_applied(self): + # For event 0: f1=0, f2=0 → scaled_pred=[0,0,0] → unscaled = 0*std+mean + result = self._run(n=1) + np.testing.assert_allclose(result.iloc[0]["Xoff_residual"], 0.0 * 2.0 + 1.0) + np.testing.assert_allclose(result.iloc[0]["Yoff_residual"], 0.0 * 0.5 + 2.0) + + def test_index_matches_df_index(self): + n = 8 + df = pd.DataFrame( + {"f1": np.arange(n, dtype=float), "f2": np.zeros(n)}, index=np.arange(100, 100 + n) + ) + y_mean = pd.Series(dict.fromkeys(TARGETS, 0.0)) + y_std = pd.Series(dict.fromkeys(TARGETS, 1.0)) + + class _Zero: + def predict(self, x): + return np.zeros((len(x), 3)) + + result = _predict_unscaled_chunked( + _Zero(), df, np.arange(n), ["f1", "f2"], y_mean, y_std, TARGETS, None + ) + assert list(result.index) == list(df.index) + + +class TestSampleEvalIndices: + def test_no_cap_returns_all_indices(self): + idx = np.arange(100) + result = _sample_eval_indices(idx, max_events=None, random_state=0) + np.testing.assert_array_equal(result, idx) + + def test_zero_cap_returns_all_indices(self): + idx = np.arange(50) + result = _sample_eval_indices(idx, max_events=0, random_state=0) + np.testing.assert_array_equal(result, idx) + + def test_cap_larger_than_pool_returns_all(self): + idx = np.arange(30) + result = _sample_eval_indices(idx, max_events=100, random_state=0) + np.testing.assert_array_equal(result, idx) + + def test_cap_limits_output_size(self): + idx = np.arange(200) + result = _sample_eval_indices(idx, max_events=50, random_state=0) + assert len(result) == 50 + + def test_sampled_indices_are_subset_of_input(self): + idx = np.arange(500) + result = _sample_eval_indices(idx, max_events=100, random_state=42) + assert set(result).issubset(set(idx)) + + def test_same_seed_gives_same_subset(self): + idx = np.arange(200) + r1 = _sample_eval_indices(idx, max_events=60, random_state=7) + r2 = _sample_eval_indices(idx, max_events=60, random_state=7) + np.testing.assert_array_equal(r1, r2) + + def test_different_seeds_give_different_subsets(self): + idx = np.arange(200) + r1 = _sample_eval_indices(idx, max_events=60, random_state=1) + r2 = _sample_eval_indices(idx, max_events=60, random_state=2) + assert not np.array_equal(r1, r2) + + +class TestRegressionSampleWeights: + def _build_inputs(self, n=300, seed=0): + rng = np.random.default_rng(seed) + erec_s = np.logspace(0, 2.5, n) + e_residual = rng.normal(0, 0.1, n) + disp_nimages = rng.choice([2, 3, 4], n) + return erec_s, e_residual, disp_nimages + + def test_weights_are_finite_and_positive(self): + erec_s, e_residual, disp_nimages = self._build_inputs() + result = _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages) + weights = result[2] + assert np.all(np.isfinite(weights)), "Non-finite weights found" + assert np.all(weights >= 0), "Negative weights found" + + def test_weights_are_capped(self): + from eventdisplay_ml.models import _MAX_REGRESSION_SAMPLE_WEIGHT + + erec_s, e_residual, disp_nimages = self._build_inputs() + weights = _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages)[2] + assert np.all(weights <= _MAX_REGRESSION_SAMPLE_WEIGHT + 1e-6), ( + "Weights exceed the hard cap" + ) + + def test_weights_length_matches_input(self): + n = 150 + erec_s, e_residual, disp_nimages = self._build_inputs(n=n) + weights = _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages)[2] + assert len(weights) == n + + def test_higher_multiplicity_gets_higher_weight(self): + """Events with more images must receive proportionally higher weights.""" + # Two events with same energy, different multiplicity + erec_s = np.array([100.0, 100.0]) + e_residual = np.array([0.0, 0.0]) + disp_nimages = np.array([2, 4]) + # _regression_sample_weights directly to probe the multiplicity component + bins = np.linspace(-2, 2.5, 10) + energy_bin_weights = np.ones(9) + mult_mean_sq = float(np.mean(np.square(disp_nimages, dtype=np.float64))) + weights, _ = _regression_sample_weights( + erec_s, + e_residual, + disp_nimages, + bins=bins, + energy_bin_weights=energy_bin_weights, + multiplicity_mean_square=mult_mean_sq, + max_weight=50.0, + ) + # w_mult = n_tel^2 / mean_sq, so mult=4 must exceed mult=2 + assert weights[1] > weights[0], ( + "Multiplicity-4 event should have higher weight than multiplicity-2" + ) + + def test_invalid_erecs_get_zero_weight(self): + erec_s = np.array([-1.0, 0.0, np.nan, 100.0]) + e_residual = np.array([0.0, 0.0, 0.0, 0.0]) + disp_nimages = np.array([2, 2, 2, 2]) + bins = np.linspace(-2, 2.5, 10) + energy_bin_weights = np.ones(9) + mult_mean_sq = 4.0 + weights, _ = _regression_sample_weights( + erec_s, + e_residual, + disp_nimages, + bins=bins, + energy_bin_weights=energy_bin_weights, + multiplicity_mean_square=mult_mean_sq, + max_weight=50.0, + ) + assert weights[0] == 0.0, "Negative ErecS should yield weight=0" + assert weights[1] == 0.0, "Zero ErecS should yield weight=0" + assert weights[2] == 0.0, "NaN ErecS should yield weight=0" + assert weights[3] > 0.0, "Valid ErecS must yield positive weight" + + def test_normalization_scale_reuse_gives_same_weights(self): + erec_s, e_residual, disp_nimages = self._build_inputs(n=100) + bins = np.linspace(-2, 2.5, 10) + energy_bin_weights = np.ones(9) + mult_mean_sq = float(np.mean(np.square(disp_nimages, dtype=np.float64))) + + w1, norm_scale = _regression_sample_weights( + erec_s, + e_residual, + disp_nimages, + bins=bins, + energy_bin_weights=energy_bin_weights, + multiplicity_mean_square=mult_mean_sq, + max_weight=50.0, + ) + w2, _ = _regression_sample_weights( + erec_s, + e_residual, + disp_nimages, + bins=bins, + energy_bin_weights=energy_bin_weights, + multiplicity_mean_square=mult_mean_sq, + max_weight=50.0, + normalization_scale=norm_scale, + ) + np.testing.assert_array_equal(w1, w2) + + def test_all_invalid_erecs_raises(self): + # When every event has an invalid ErecS, no weights can be computed and + # the function must raise ValueError rather than silently returning zeros. + n = 20 + erec_s = np.full(n, np.nan) + e_residual = np.zeros(n) + disp_nimages = np.full(n, 2) + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + with pytest.raises(ValueError): + _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages) + + +class TestLogEnergyBinCountsFromArrays: + def test_returns_three_tuple(self): + erec_s = np.logspace(0, 2, 100) + e_residual = np.zeros(100) + disp_nimages = np.full(100, 2) + result = _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages) + assert len(result) == 3 + + def test_weight_config_returned_when_requested(self): + erec_s = np.logspace(0, 2, 200) + e_residual = np.zeros(200) + disp_nimages = np.full(200, 2) + result = _log_energy_bin_counts_from_arrays( + erec_s, e_residual, disp_nimages, return_weight_config=True + ) + assert len(result) == 4 + weight_config = result[3] + assert "bins" in weight_config + assert "energy_bin_weights" in weight_config + assert "multiplicity_mean_square" in weight_config + assert "max_weight" in weight_config + assert "normalization_scale" in weight_config + + def test_counts_dict_has_pd_interval_keys(self): + erec_s = np.logspace(0, 2, 100) + e_residual = np.zeros(100) + disp_nimages = np.full(100, 2) + _, counts, _ = _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages) + # All keys should be pd.Interval + assert all(isinstance(k, pd.Interval) for k in counts.keys()) + + def test_bins_span_expected_energy_range(self): + erec_s = np.logspace(0, 2, 100) + e_residual = np.zeros(100) + disp_nimages = np.full(100, 2) + bins, _, _ = _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages) + from eventdisplay_ml.models import _EVAL_LOG_E_MAX, _EVAL_LOG_E_MIN + + assert bins[0] == _EVAL_LOG_E_MIN + assert bins[-1] == _EVAL_LOG_E_MAX + + +# =========================================================================== +# 6. _output_tree stereo structure +# =========================================================================== + + +class TestOutputTreeStereo: + """Stereo output tree must have exactly three float32 branches.""" + + def test_stereo_creates_correct_tree(self): + root_file = MagicMock() + mock_tree = MagicMock() + root_file.mktree.return_value = mock_tree + + result = _output_tree("stereo_analysis", root_file) + + root_file.mktree.assert_called_once() + call_args = root_file.mktree.call_args + tree_name = call_args.args[0] + branches = call_args.args[1] + + assert tree_name == "StereoAnalysis" + assert set(branches.keys()) == {"Dir_Xoff", "Dir_Yoff", "Dir_Erec"} + assert all(v is np.float32 for v in branches.values()) + assert result is mock_tree + + def test_stereo_tree_has_no_classification_branches(self): + root_file = MagicMock() + root_file.mktree.return_value = MagicMock() + _output_tree("stereo_analysis", root_file) + branches = root_file.mktree.call_args.args[1] + assert "Gamma_Prediction" not in branches + assert not any(k.startswith("Is_Gamma") for k in branches) + + def test_unknown_analysis_type_raises(self): + with pytest.raises(ValueError, match="Unknown analysis_type"): + _output_tree("bad_type", MagicMock()) + + +# =========================================================================== +# 7. Feature schema – regression-specific invariants +# =========================================================================== + + +class TestRegressionFeatureSchema: + """These invariants must survive classification refactoring.""" + + def test_target_list_exact_order(self): + from eventdisplay_ml import features + + assert features.target_features("stereo_analysis") == [ + "Xoff_residual", + "Yoff_residual", + "E_residual", + ] + + def test_training_features_include_mc_truth(self): + from eventdisplay_ml import features + + train = features.features("stereo_analysis", training=True) + assert "MCxoff" in train + assert "MCyoff" in train + assert "MCe0" in train + + def test_inference_features_exclude_mc_truth(self): + from eventdisplay_ml import features + + infer = features.features("stereo_analysis", training=False) + assert "MCxoff" not in infer + assert "MCyoff" not in infer + assert "MCe0" not in infer + + def test_pointing_offsets_excluded_from_stereo(self): + from eventdisplay_ml import features + + excluded = features.excluded_features("stereo_analysis", ntel=2) + assert "fpointing_dx_0" in excluded + assert "fpointing_dy_0" in excluded + assert "fpointing_dx_1" in excluded + assert "fpointing_dy_1" in excluded + + def test_analysis_type_tag_is_stereo_not_classification(self): + """Passing 'stereo_analysis' to features() must never yield classification-only columns.""" + from eventdisplay_ml import features + + infer = features.features("stereo_analysis", training=False) + classification_only = {"ze_bin", "Gamma_Prediction", "Is_Gamma"} + assert not any(c in infer for c in classification_only) + + +# =========================================================================== +# 8. process_file_chunked – stereo path invariants +# =========================================================================== + + +class TestProcessFileChunkedStereo: + """Streaming must reset chunk indices and obey max_events for the stereo path.""" + + def _setup_mocks(self, monkeypatch, chunks, max_events=None): + input_root = MagicMock() + input_root.__enter__.return_value = {"data": MagicMock()} + input_root.__exit__.return_value = False + output_root = MagicMock() + output_root.__enter__.return_value = output_root + output_root.__exit__.return_value = False + tree = MagicMock() + applied = [] + + monkeypatch.setattr(models.uproot, "open", lambda *a: input_root) + monkeypatch.setattr(models.uproot, "recreate", lambda *a: output_root) + monkeypatch.setattr(models.uproot, "iterate", lambda *a, **k: chunks) + monkeypatch.setattr( + models.data_processing, "read_telescope_config", lambda *a: {"max_tel_id": 3} + ) + monkeypatch.setattr( + models.data_processing, "_resolve_branch_aliases", lambda *a: (["ErecS"], {}) + ) + monkeypatch.setattr(models.data_processing, "_ensure_fpointing_fields", lambda c: c) + monkeypatch.setattr(models.features, "features", lambda *a, **k: ["ErecS"]) + monkeypatch.setattr(models, "_output_tree", lambda *a: tree) + monkeypatch.setattr( + models, + "_apply_model", + lambda at, chunk, *a: applied.append((at, chunk.copy())), + ) + return applied + + def test_chunk_indices_reset_to_zero_based(self, monkeypatch): + chunks = [ak.Array([{"ErecS": 1.0}, {"ErecS": 2.0}])] + applied = self._setup_mocks(monkeypatch, chunks) + models.process_file_chunked( + "stereo_analysis", + {"input_file": "in.root", "output_file": "out.root", "chunk_size": 10}, + ) + assert applied[0][1].index.tolist() == [0, 1] + + def test_max_events_limits_total_processed(self, monkeypatch): + chunks = [ + ak.Array([{"ErecS": 1.0}, {"ErecS": 2.0}, {"ErecS": 3.0}]), + ak.Array([{"ErecS": 4.0}, {"ErecS": 5.0}]), + ] + applied = self._setup_mocks(monkeypatch, chunks, max_events=4) + models.process_file_chunked( + "stereo_analysis", + {"input_file": "in.root", "output_file": "out.root", "max_events": 4, "chunk_size": 3}, + ) + total = sum(len(chunk) for _, chunk in applied) + assert total == 4 + + def test_analysis_type_passed_as_stereo(self, monkeypatch): + chunks = [ak.Array([{"ErecS": 1.0}])] + applied = self._setup_mocks(monkeypatch, chunks) + models.process_file_chunked( + "stereo_analysis", + {"input_file": "in.root", "output_file": "out.root"}, + ) + assert all(at == "stereo_analysis" for at, _ in applied) + + def test_empty_chunk_is_skipped(self, monkeypatch): + chunks = [ak.Array([]), ak.Array([{"ErecS": 1.0}])] + applied = self._setup_mocks(monkeypatch, chunks) + models.process_file_chunked( + "stereo_analysis", + {"input_file": "in.root", "output_file": "out.root"}, + ) + assert len(applied) == 1 # empty chunk was skipped + + +# =========================================================================== +# 9. End-to-end regression: train → persist → load → apply round-trip +# =========================================================================== + + +class TestRegressionEndToEnd: + """Train a tiny stand-in regressor (via a patched ``xgboost.XGBRegressor``), save it, + load it, and verify predictions round-trip correctly through the production stack. + """ + + def test_full_round_trip_prediction_shape_and_dtype(self, tmp_path, monkeypatch): + """Train → joblib dump → load_regression_models → apply → correct output shape.""" + rng = np.random.default_rng(99) + n = 200 + df = pd.DataFrame( + { + "Xoff_residual": rng.normal(0, 0.3, n), + "Yoff_residual": rng.normal(0, 0.4, n), + "E_residual": rng.normal(0, 0.15, n), + "ErecS": np.logspace(0, 2, n), + "DispNImages": rng.choice([2, 3, 4], n), + "Xoff_weighted_bdt": rng.normal(0, 0.5, n), + "Yoff_weighted_bdt": rng.normal(0, 0.5, n), + "feat_A": rng.uniform(-1, 1, n), + } + ) + rec = _RecordingModel() + cfg = { + "targets": TARGETS, + "train_test_fraction": 0.5, + "random_state": 99, + "eval_max_events": 0, + "models": {"xgboost": {"hyper_parameters": {}}}, + } + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + trained = models.train_regression(df, cfg) + + # Manually store the recording model (train_regression puts it in cfg) + model_prefix = tmp_path / "rtt_model" + artifact_path = model_prefix.with_suffix(".joblib.gz") + joblib.dump(trained, artifact_path) + + loaded_models_dict, loaded_params = models.load_regression_models( + str(model_prefix), "xgboost" + ) + assert "target_mean" in loaded_params + assert "target_std" in loaded_params + + apply_df = df.head(10).copy() + flat = apply_df[[c for c in df.columns if c not in set(TARGETS)]] + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + apply_cfg = {"models": loaded_models_dict, **loaded_params} + pred_xoff, pred_yoff, pred_e = models.apply_regression_models(apply_df, apply_cfg) + + assert len(pred_xoff) == 10 + assert len(pred_yoff) == 10 + assert len(pred_e) == 10 + assert pred_xoff.dtype in (np.float32, np.float64) + + def test_scalers_from_trained_model_invert_apply_predictions(self, tmp_path, monkeypatch): + """Scalers stored during training must round-trip: apply recovers finite predictions.""" + rng = np.random.default_rng(77) + n = 200 + df = pd.DataFrame( + { + "Xoff_residual": rng.normal(0, 0.3, n), + "Yoff_residual": rng.normal(0, 0.4, n), + "E_residual": rng.normal(0, 0.15, n), + "ErecS": np.ones(n), # log10(1)=0 as baseline + "DispNImages": np.full(n, 2), + "Xoff_weighted_bdt": np.zeros(n), + "Yoff_weighted_bdt": np.zeros(n), + "feat": rng.uniform(0, 1, n), + } + ) + rec = _RecordingModel() + cfg = { + "targets": TARGETS, + "train_test_fraction": 0.5, + "random_state": 77, + "eval_max_events": 0, + "models": {"xgboost": {"hyper_parameters": {}}}, + } + with ( + patch("xgboost.XGBRegressor", return_value=rec), + patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}), + ): + trained = models.train_regression(df, cfg) + + # Scalers must be well-defined (finite mean, positive std) + for t in TARGETS: + assert math.isfinite(trained["target_mean"][t]) + assert trained["target_std"][t] > 0 + + # Predictions on 10-event apply frame should all be finite + apply_df = df.head(10).copy() + flat = apply_df[[c for c in df.columns if c not in set(TARGETS)]] + monkeypatch.setattr(models, "flatten_feature_data", lambda *a, **k: flat) + monkeypatch.setattr(data_processing, "print_variable_statistics", lambda *a: None) + + apply_cfg = { + "models": trained["models"], + **{k: trained[k] for k in ("target_mean", "target_std")}, + } + pred_xoff, pred_yoff, pred_e = models.apply_regression_models(apply_df, apply_cfg) + + assert all(np.isfinite(pred_xoff)) + assert all(np.isfinite(pred_yoff)) + assert all(np.isfinite(pred_e)) diff --git a/tests/test_regression_shared_contracts.py b/tests/test_regression_shared_contracts.py new file mode 100644 index 0000000..0ecc8bd --- /dev/null +++ b/tests/test_regression_shared_contracts.py @@ -0,0 +1,277 @@ +"""Stereo-regression contracts across infrastructure shared with classification. + +The classification rewrite must not change the shared configuration, feature, +flattening, streaming, or evaluation behavior used by production regression. +""" + +from unittest.mock import MagicMock + +import awkward as ak +import numpy as np +import pandas as pd + +from eventdisplay_ml import config, data_processing, evaluate, features, models + + +def test_stereo_feature_schema_keeps_target_order_and_excludes_pointing_offsets(): + """Pin the regression schema consumed by ROOT input, training, and apply.""" + assert features.target_features("stereo_analysis") == [ + "Xoff_residual", + "Yoff_residual", + "E_residual", + ] + assert features.features("stereo_analysis", training=True) == [ + "MCxoff", + "MCyoff", + "MCe0", + *features.features("stereo_analysis", training=False), + ] + excluded = features.excluded_features("stereo_analysis", ntel=3) + assert excluded == { + "fpointing_dx_0", + "fpointing_dx_1", + "fpointing_dx_2", + "fpointing_dy_0", + "fpointing_dy_1", + "fpointing_dy_2", + } + + +def test_configure_apply_stereo_keeps_primary_and_high_multiplicity_scalers_separate( + monkeypatch, +): + """Ensure shared CLI loading never loses or cross-wires regression scalers.""" + monkeypatch.setattr( + "sys.argv", + [ + "apply-xgb-stereo", + "--input_file", + "input.root", + "--model_prefix", + "two_image", + "--model_prefix_high_multiplicity", + "high_image", + "--output_file", + "prediction.root", + ], + ) + primary_scalers = { + "target_mean": {"Xoff_residual": 1.0, "Yoff_residual": 2.0, "E_residual": 3.0}, + "target_std": {"Xoff_residual": 4.0, "Yoff_residual": 5.0, "E_residual": 6.0}, + } + high_scalers = { + "target_mean": {"Xoff_residual": -1.0, "Yoff_residual": -2.0, "E_residual": -3.0}, + "target_std": {"Xoff_residual": 7.0, "Yoff_residual": 8.0, "E_residual": 9.0}, + } + loader = MagicMock( + side_effect=[ + ({"xgboost": {"model": "two"}}, primary_scalers), + ({"xgboost": {"model": "high"}}, high_scalers), + ] + ) + monkeypatch.setattr(config, "load_models", loader) + + result = config.configure_apply("stereo_analysis") + + assert result["target_mean"] == primary_scalers["target_mean"] + assert result["target_std"] == primary_scalers["target_std"] + assert result["target_mean_high_multiplicity"] == high_scalers["target_mean"] + assert result["target_std_high_multiplicity"] == high_scalers["target_std"] + assert result["models"]["xgboost"]["model"] == "two" + assert result["models_high_multiplicity"]["xgboost"]["model"] == "high" + + +def test_flatten_feature_data_for_regression_removes_targets_and_pointing(monkeypatch): + """Inference must never expose training truth or pointing corrections to a model.""" + flattened = pd.DataFrame( + { + "physics_feature": [1.0], + "Xoff_residual": [9.0], + "Yoff_residual": [8.0], + "E_residual": [7.0], + "fpointing_dx_0": [6.0], + "fpointing_dy_0": [5.0], + } + ) + flatten = MagicMock(return_value=flattened) + monkeypatch.setattr(data_processing, "flatten_telescope_data_vectorized", flatten) + + result = data_processing.flatten_feature_data( + pd.DataFrame({"raw": [1]}), + ntel=2, + analysis_type="stereo_analysis", + training=False, + tel_config={"max_tel_id": 0}, + preview_rows=0, + ) + + assert result.columns.tolist() == ["physics_feature"] + assert flatten.call_args.kwargs["analysis_type"] == "stereo_analysis" + assert flatten.call_args.kwargs["training"] is False + + +def test_stereo_training_chunk_filters_invalid_energy_and_keeps_truth_baseline_alignment( + monkeypatch, +): + """Regression residuals must be row-aligned after both validity filters.""" + raw = ak.Array( + [ + {"MCxoff": 11.0, "MCyoff": 21.0, "MCe0": 100.0}, + {"MCxoff": 12.0, "MCyoff": 22.0, "MCe0": 100.0}, + {"MCxoff": 13.0, "MCyoff": 23.0, "MCe0": 0.0}, + {"MCxoff": 14.0, "MCyoff": 24.0, "MCe0": 1000.0}, + ] + ) + monkeypatch.setattr( + data_processing, + "flatten_telescope_data_vectorized", + lambda *_args, **_kwargs: pd.DataFrame( + { + "Xoff_weighted_bdt": [1.0, 2.0, 3.0, 4.0], + "Yoff_weighted_bdt": [10.0, 20.0, 30.0, 40.0], + "ErecS": [10.0, 0.0, 100.0, 100.0], + } + ), + ) + + result = data_processing._flatten_training_chunk( + raw, + {"observatory": "veritas"}, + "stereo_analysis", + {"max_tel_id": 0}, + False, + "test chunk", + ) + + # Row 1 has invalid ErecS; row 2 has invalid MC energy. The remaining + # rows must still use their original corresponding truth and baseline. + assert result.index.tolist() == [0, 3] + np.testing.assert_allclose(result["Xoff_residual"], [10.0, 10.0]) + np.testing.assert_allclose(result["Yoff_residual"], [11.0, -16.0]) + np.testing.assert_allclose(result["E_residual"], [1.0, 1.0]) + np.testing.assert_allclose(result["ErecS"], [10.0, 100.0]) + + +def test_extra_columns_keeps_stereo_energy_linear_for_residual_targets(monkeypatch): + """The shared extra-column helper must not log-transform regression baselines.""" + monkeypatch.setattr( + data_processing, + "calculate_geomagnetic_angles", + lambda *_args, **_kwargs: np.array([0.1, 0.2], dtype=np.float32), + ) + raw = pd.DataFrame( + { + "Xoff": [1.0, 2.0], + "Yoff": [3.0, 4.0], + "Xoff_intersect": [0.5, 1.5], + "Yoff_intersect": [2.5, 3.5], + "DispNImages": [2, 3], + "img2_ang": [0.2, 0.3], + "Erec": [10.0, 100.0], + "ErecS": [3.0, 30.0], + "EmissionHeight": [8.0, 9.0], + "ArrayPointing_Azimuth": [0.0, 1.0], + "ArrayPointing_Elevation": [70.0, 71.0], + } + ) + + result = data_processing.extra_columns(raw, "stereo_analysis", True, raw.index) + + np.testing.assert_allclose(result["Erec"], [10.0, 100.0]) + np.testing.assert_allclose(result["ErecS"], [3.0, 30.0]) + np.testing.assert_allclose(result["Diff_Xoff"], [0.5, 0.5]) + np.testing.assert_allclose(result["Diff_Yoff"], [0.5, 0.5]) + + +def test_stereo_streaming_resets_chunk_indices_and_obeys_global_max_events(monkeypatch): + """Chunked ROOT application must preserve exactly the requested event rows.""" + input_root = MagicMock() + input_root.__enter__.return_value = {"data": MagicMock()} + input_root.__exit__.return_value = False + output_root = MagicMock() + output_root.__enter__.return_value = output_root + output_root.__exit__.return_value = False + tree = MagicMock() + applied_chunks = [] + chunks = [ + ak.Array([{"ErecS": 1.0}, {"ErecS": 2.0}]), + ak.Array([{"ErecS": 3.0}, {"ErecS": 4.0}]), + ] + + monkeypatch.setattr(models.uproot, "open", lambda *_args: input_root) + monkeypatch.setattr(models.uproot, "recreate", lambda *_args: output_root) + monkeypatch.setattr(models.uproot, "iterate", lambda *_args, **_kwargs: chunks) + monkeypatch.setattr( + models.data_processing, "read_telescope_config", lambda *_args: {"max_tel_id": 3} + ) + monkeypatch.setattr( + models.data_processing, "_resolve_branch_aliases", lambda *_args: (["ErecS"], {}) + ) + monkeypatch.setattr(models.data_processing, "_ensure_fpointing_fields", lambda chunk: chunk) + monkeypatch.setattr(models.features, "features", lambda *_args, **_kwargs: ["ErecS"]) + monkeypatch.setattr(models, "_output_tree", lambda *_args: tree) + monkeypatch.setattr( + models, + "_apply_model", + lambda analysis_type, chunk, *_args: applied_chunks.append((analysis_type, chunk.copy())), + ) + + models.process_file_chunked( + "stereo_analysis", + { + "input_file": "input.root", + "output_file": "output.root", + "max_events": 3, + "chunk_size": 2, + }, + ) + + assert [len(chunk) for _, chunk in applied_chunks] == [2, 1] + assert [chunk.index.tolist() for _, chunk in applied_chunks] == [[0, 1], [0]] + assert [chunk["ErecS"].tolist() for _, chunk in applied_chunks] == [[1.0, 2.0], [3.0]] + assert all(analysis_type == "stereo_analysis" for analysis_type, _ in applied_chunks) + + +def test_regression_resolution_uses_dataframe_indices_for_baseline_reconstruction(monkeypatch): + """Diagnostics must pair residuals with matching non-contiguous event rows.""" + y_test = pd.DataFrame( + { + "Xoff_residual": [1.0, -2.0], + "Yoff_residual": [3.0, 4.0], + "E_residual": [0.5, -0.5], + }, + index=[101, 909], + ) + df = pd.DataFrame( + { + "Xoff_weighted_bdt": [100.0, 10.0], + "Yoff_weighted_bdt": [200.0, 20.0], + "ErecS": [10.0, 100.0], + }, + index=[909, 101], + ) + captured = [] + real_dataframe = pd.DataFrame + + def capture_first_dataframe(*args, **kwargs): + frame = real_dataframe(*args, **kwargs) + if not captured: + captured.append(frame.copy()) + return frame + + monkeypatch.setattr(evaluate.pd, "DataFrame", capture_first_dataframe) + evaluate.calculate_resolution( + y_test.copy(), + y_test.copy(), + df, + percentiles=[68], + log_e_min=-2, + log_e_max=3, + n_bins=1, + name="xgboost", + ) + + reconstructed = captured[0] + np.testing.assert_allclose(reconstructed["MCxoff_true"], [11.0, 98.0]) + np.testing.assert_allclose(reconstructed["MCyoff_true"], [23.0, 204.0]) + np.testing.assert_allclose(reconstructed["MCe0"], [2.5, 0.5]) diff --git a/tests/test_utils.py b/tests/test_utils.py index 819ce51..dff19d7 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,9 +1,14 @@ """Unit tests for utils.py.""" import json +import logging +from types import SimpleNamespace +import joblib +import pandas as pd import pytest +from eventdisplay_ml import utils from eventdisplay_ml.utils import ( discover_joblib_files, joblib_basename, @@ -102,6 +107,14 @@ def test_discover_joblib_files_empty_dir_raises(tmp_path): discover_joblib_files(tmp_path) +def test_discover_joblib_files_ignores_directories_with_model_suffixes(tmp_path): + """A directory named like a model must never be returned as a model file.""" + (tmp_path / "not_a_model.joblib").mkdir() + (tmp_path / "model.joblib.gz").touch() + + assert discover_joblib_files(tmp_path) == [tmp_path / "model.joblib.gz"] + + # --------------------------------------------------------------------------- # read_input_file_list # --------------------------------------------------------------------------- @@ -170,6 +183,12 @@ def test_parse_image_selection_invalid_raises(): parse_image_selection("abc") +def test_parse_image_selection_malformed_comma_list_falls_through_to_clear_error(): + """Partially numeric comma lists must not be silently accepted.""" + with pytest.raises(ValueError, match="Invalid image_selection"): + parse_image_selection("1,invalid") + + # --------------------------------------------------------------------------- # load_model_parameters # --------------------------------------------------------------------------- @@ -213,6 +232,21 @@ def test_load_model_parameters_missing_file_raises(): load_model_parameters("/nonexistent/path.json") +def test_load_model_parameters_none_path_has_same_clear_error(): + """Argparse omissions should not leak a TypeError from open().""" + with pytest.raises(FileNotFoundError, match="Model parameters file not found: None"): + load_model_parameters(None) + + +def test_load_model_parameters_missing_energy_bins_raises_value_error(tmp_path): + """Selecting a bin requires the energy-bin metadata to exist.""" + params_file = tmp_path / "missing_bins.json" + params_file.write_text(json.dumps({"zenith_bins_deg": []})) + + with pytest.raises(ValueError, match="Invalid energy bin number 0"): + load_model_parameters(params_file, energy_bin_number=0) + + # --------------------------------------------------------------------------- # load_energy_range # --------------------------------------------------------------------------- @@ -265,3 +299,75 @@ def test_output_file_name_creates_parent_directory(tmp_path): def test_output_file_name_returns_string(tmp_path): result = output_file_name(tmp_path / "model") assert isinstance(result, str) + + +def test_output_file_name_combines_multiplicity_and_energy_bin(tmp_path): + """Keep the naming contract unique when both suffixes are supplied.""" + result = output_file_name(tmp_path / "model", n_tel=3, energy_bin_number=4) + assert result.endswith("model_ntel3_ebin4.joblib.gz") + + +# --------------------------------------------------------------------------- +# memory profiling and joblib loading +# --------------------------------------------------------------------------- + + +def test_max_rss_uses_platform_specific_units(monkeypatch): + """MacOS reports bytes while Linux reports KiB; normalize both to GB.""" + usage = SimpleNamespace(ru_maxrss=1024**3) + monkeypatch.setattr(utils.resource, "getrusage", lambda *_args: usage) + monkeypatch.setattr(utils.sys, "platform", "darwin") + assert utils._max_rss_gb() == pytest.approx(1.0) + + monkeypatch.setattr(utils.sys, "platform", "linux") + assert utils._max_rss_gb() == pytest.approx(1024.0) + + +def test_current_rss_reads_proc_statm_when_available(monkeypatch, tmp_path): + """Current RSS uses resident pages rather than the peak when /proc exists.""" + statm = tmp_path / "statm" + statm.write_text("100 512 0 0 0 0 0") + monkeypatch.setattr(utils, "Path", lambda _path: statm) + monkeypatch.setattr(utils.os, "sysconf", lambda _name: 4096) + + assert utils._current_rss_gb() == pytest.approx(512 * 4096 / 1024**3) + + +def test_current_rss_falls_back_to_peak_when_proc_is_unavailable(monkeypatch, tmp_path): + """macOS-like environments without /proc retain a meaningful RSS value.""" + missing_statm = tmp_path / "missing" + monkeypatch.setattr(utils, "Path", lambda _path: missing_statm) + monkeypatch.setattr(utils, "_max_rss_gb", lambda: 1.25) + + assert utils._current_rss_gb() == pytest.approx(1.25) + + +def test_log_memory_checkpoint_is_disabled_without_side_effects(monkeypatch): + """The profiling helper must be a no-op unless explicitly enabled.""" + monkeypatch.setattr(utils, "_current_rss_gb", lambda: pytest.fail("should not be called")) + utils.log_memory_checkpoint("disabled", enabled=False) + + +def test_log_memory_checkpoint_logs_timing_rss_and_dataframe_memory(monkeypatch, caplog): + """Enabled profiling reports both process and DataFrame memory details.""" + utils._profile_start_time = None + utils._profile_last_time = None + monkeypatch.setattr(utils.time, "perf_counter", lambda: 10.0) + monkeypatch.setattr(utils, "_current_rss_gb", lambda: 1.5) + monkeypatch.setattr(utils, "_max_rss_gb", lambda: 2.5) + + with caplog.at_level(logging.INFO): + utils.log_memory_checkpoint("after flattening", pd.DataFrame({"x": [1, 2]}), enabled=True) + + assert "Memory checkpoint [after flattening]" in caplog.text + assert "rss=1.50 GB" in caplog.text + assert "max_rss=2.50 GB" in caplog.text + assert "shape=(2, 1)" in caplog.text + + +def test_load_joblib_returns_payload_and_suppresses_only_shape_warning(monkeypatch): + """Model loading tolerates NumPy's known pickle warning without hiding failures.""" + payload = {"model": "trusted-test-payload"} + monkeypatch.setattr(joblib, "load", lambda _path: payload) + + assert utils.load_joblib("model.joblib.gz") is payload