From aff86c1f2274d31a2f7645a31d2e2c6bba0b81dc Mon Sep 17 00:00:00 2001 From: GernotMaier Date: Sun, 30 Aug 2026 10:20:18 +0200 Subject: [PATCH 1/3] Separate energy from direction training --- README.md | 12 +++ src/eventdisplay_ml/config.py | 43 +++++++++- src/eventdisplay_ml/evaluate.py | 27 +++--- src/eventdisplay_ml/models.py | 130 ++++++++++++++++++++--------- tests/test_config.py | 42 ++++++++++ tests/test_regression_contracts.py | 111 ++++++++++++++++++++++++ 6 files changed, 312 insertions(+), 53 deletions(-) diff --git a/README.md b/README.md index 3380ac6..123d824 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,18 @@ eventdisplay-ml-train-xgb-stereo \ --feature_profile reduced ``` +To train separate direction and energy models while keeping one model artifact: + +```bash +eventdisplay-ml-train-xgb-stereo \ + --input_file_list train_files.txt \ + --model_prefix models/stereo_model_separate \ + --separate_direction_energy +``` + +The usual stereo-application command detects this artifact automatically and +combines the direction and energy predictions. + **Output:** Joblib model file containing: - XGBoost trained model object diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py index 796479c..aacf31e 100644 --- a/src/eventdisplay_ml/config.py +++ b/src/eventdisplay_ml/config.py @@ -49,6 +49,14 @@ def configure_training(analysis_type): "and loss summaries for telescope positions 0-3." ), ) + parser.add_argument( + "--separate_direction_energy", + action="store_true", + help=( + "Train separate direction (Xoff/Yoff residual) and energy " + "(E residual) models in one model artifact." + ), + ) if analysis_type == "classification": parser.add_argument("--input_signal_file_list", help="List of input signal mscw files.") parser.add_argument( @@ -209,6 +217,10 @@ def configure_training(analysis_type): if analysis_type == "stereo_analysis": _logger.info(f"Minimum images (DispNImages): {model_configs.get('min_images')}") _logger.info(f"Regression feature profile: {model_configs.get('feature_profile')}") + _logger.info( + "Separate direction and energy models: %s", + model_configs.get("separate_direction_energy"), + ) _logger.info( "Regression weighting: energy=inverse-sqrt(count), min_bin_events=%d, " "multiplicity=DispNImages**2, max_combined_weight=%.1f, " @@ -226,6 +238,35 @@ def configure_training(analysis_type): model_configs["models"] = hyper_parameters( analysis_type, model_configs.get("hyperparameter_config") ) + model_configs["targets"] = target_features(analysis_type) + if analysis_type == "stereo_analysis" and model_configs["separate_direction_energy"]: + configured_models = model_configs["models"] + if set(configured_models) == {"xgboost"}: + base_config = configured_models["xgboost"] + model_configs["models"] = { + "direction": { + **base_config, + "hyper_parameters": dict(base_config.get("hyper_parameters", {})), + "targets": ["Xoff_residual", "Yoff_residual"], + }, + "energy": { + **base_config, + "hyper_parameters": dict(base_config.get("hyper_parameters", {})), + "targets": ["E_residual"], + }, + } + elif set(configured_models) != {"direction", "energy"}: + raise ValueError( + "Separate direction/energy training requires either one 'xgboost' " + "configuration or exactly 'direction' and 'energy' configurations." + ) + else: + model_configs["models"]["direction"]["targets"] = [ + "Xoff_residual", + "Yoff_residual", + ] + model_configs["models"]["energy"]["targets"] = ["E_residual"] + model_configs["regression_mode"] = "separate_direction_energy" for model_name, model_cfg in model_configs["models"].items(): hyper_params = model_cfg.get("hyper_parameters") if hyper_params is None: @@ -234,8 +275,6 @@ def configure_training(analysis_type): hyper_params["n_jobs"] = model_configs["max_cores"] if model_configs.get("random_state") is not None: hyper_params["random_state"] = model_configs["random_state"] - model_configs["targets"] = target_features(analysis_type) - if analysis_type == "stereo_analysis": model_configs["pre_cuts"] = pre_cuts_regression( min_images=model_configs.get("min_images", 2) diff --git a/src/eventdisplay_ml/evaluate.py b/src/eventdisplay_ml/evaluate.py index 2e1ac6b..1dbde2c 100644 --- a/src/eventdisplay_ml/evaluate.py +++ b/src/eventdisplay_ml/evaluate.py @@ -157,16 +157,23 @@ def evaluate_regression_model( if shap_per_energy: shap_feature_importance_by_energy(model, x_test, df, y_test, y_data.columns) - calculate_resolution( - y_pred, - y_test, - df, - percentiles=[68, 90, 95], - log_e_min=-2, - log_e_max=2.5, - n_bins=9, - name=name, - ) + required_resolution_targets = {"Xoff_residual", "Yoff_residual", "E_residual"} + if required_resolution_targets.issubset(y_data.columns): + calculate_resolution( + y_pred, + y_test, + df, + percentiles=[68, 90, 95], + log_e_min=-2, + log_e_max=2.5, + n_bins=9, + name=name, + ) + else: + _logger.info( + "%s resolution summary is deferred until direction and energy predictions are combined.", + name, + ) return shap_importance_dict diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index b2167c9..d12e523 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -360,12 +360,33 @@ def load_regression_models(model_prefix, model_name): _logger.info(f"Loading regression model: {model_path}") model_data = utils.load_joblib(model_path) - models = { - model_name: { - "model": model_data["models"][model_name]["model"], - "features": model_data.get("features", []), + regression_mode = model_data.get("regression_mode") + if regression_mode == "separate_direction_energy": + required_models = {"direction", "energy"} + missing_models = required_models - set(model_data.get("models", {})) + if missing_models: + raise KeyError( + "Separate direction/energy regression artifact is missing model(s): " + f"{', '.join(sorted(missing_models))}." + ) + models = { + name: { + "model": model_data["models"][name]["model"], + "features": model_data["models"][name].get( + "features", model_data.get("features", []) + ), + "targets": model_data["models"][name].get("targets", []), + } + for name in ("direction", "energy") + } + else: + models = { + model_name: { + "model": model_data["models"][model_name]["model"], + "features": model_data.get("features", []), + "targets": ["Xoff_residual", "Yoff_residual", "E_residual"], + } } - } par = {} for key in ("target_mean", "target_std", "training_parameters"): if key in model_data: @@ -374,6 +395,9 @@ def load_regression_models(model_prefix, model_name): if key in {"target_mean", "target_std"}: _logger.warning("Missing '%s' in regression model file: %s", key, model_path) + if regression_mode is not None: + par["regression_mode"] = regression_mode + _logger.info("Loaded regression model.") return models, par @@ -416,24 +440,34 @@ def apply_regression_models(df, model_configs): preview_rows=model_configs.get("preview_rows", 20), ) - def predict(models, target_mean_cfg, target_std_cfg, mask=None): - model_data = next(iter(models.values())) - model_input = flatten_data.reindex(columns=model_data["features"]) - if mask is not None: - model_input = model_input.loc[mask] - + def predict_one(model_data, target_mean_cfg, target_std_cfg, mask=None): + targets = model_data.get("targets", ["Xoff_residual", "Yoff_residual", "E_residual"]) if not target_mean_cfg or not target_std_cfg: raise ValueError( "Missing target standardization parameters (target_mean/target_std). " "Regenerate the regression model or load a model file that includes them." ) - target_mean = np.array( - [target_mean_cfg[key] for key in ("Xoff_residual", "Yoff_residual", "E_residual")] - ) - target_std = np.array( - [target_std_cfg[key] for key in ("Xoff_residual", "Yoff_residual", "E_residual")] + target_mean = np.array([target_mean_cfg[key] for key in targets]) + target_std = np.array([target_std_cfg[key] for key in targets]) + model_input = flatten_data.reindex(columns=model_data["features"]) + if mask is not None: + model_input = model_input.loc[mask] + prediction = np.asarray(model_data["model"].predict(model_input)).reshape( + len(model_input), -1 ) - return model_data["model"].predict(model_input) * target_std + target_mean + if prediction.shape[1] != len(targets): + raise ValueError( + f"Regression model predicted {prediction.shape[1]} outputs for {len(targets)} targets: " + f"{targets}." + ) + return prediction * target_std + target_mean + + def predict(models, target_mean_cfg, target_std_cfg, mask=None): + if set(models) == {"direction", "energy"}: + direction = predict_one(models["direction"], target_mean_cfg, target_std_cfg, mask) + energy = predict_one(models["energy"], target_mean_cfg, target_std_cfg, mask) + return np.column_stack((direction, energy)) + return predict_one(next(iter(models.values())), target_mean_cfg, target_std_cfg, mask) model_data = next(iter(model_configs["models"].values())) primary_model_input = flatten_data.reindex(columns=model_data["features"]) @@ -466,8 +500,6 @@ def predict(models, target_mean_cfg, target_std_cfg, mask=None): high_mask, ) - flatten_data = primary_model_input - # Model predicts residuals, so add them to DispBDT baseline # Extract DispBDT predictions from the flattened data disp_xoff = flatten_data["Xoff_weighted_bdt"].values @@ -946,11 +978,10 @@ def train_regression(df, model_configs): del train_disp_nimages utils.log_memory_checkpoint("after sample-weight calculation", enabled=memory_profile) - # Standardize targets to prevent energy from dominating direction in multi-target learning - # Compute mean and std from training data only + # Compute target scalers from training data only. Separate direction and + # energy models use the corresponding subsets below. target_indices = df.columns.get_indexer(targets) y_train = df.iloc[train_idx, target_indices] - y_test = df.iloc[test_idx, target_indices] y_mean = y_train.mean() y_std = y_train.std() @@ -958,8 +989,6 @@ def train_regression(df, model_configs): for target in model_configs["targets"]: _logger.info(f" {target}: mean={y_mean[target]:.6f}, std={y_std[target]:.6f}") - y_train_scaled = (y_train - y_mean) / y_std - # Store scalers for later use during inference model_configs["target_mean"] = y_mean.to_dict() model_configs["target_std"] = y_std.to_dict() @@ -994,13 +1023,24 @@ def train_regression(df, model_configs): _logger.info(f"XGBoost eval_set test events: {len(eval_idx)}") for name, cfg in model_configs.get("models", {}).items(): + model_targets = cfg.get("targets", targets) + unknown_targets = set(model_targets) - set(targets) + if unknown_targets: + raise ValueError(f"Model '{name}' has unknown regression targets: {unknown_targets}") + model_target_indices = df.columns.get_indexer(model_targets) + model_y_train = df.iloc[train_idx, model_target_indices] + model_y_test = df.iloc[test_idx, model_target_indices] + model_y_mean = model_y_train.mean() + model_y_std = model_y_train.std() + model_y_train_scaled = (model_y_train - model_y_mean) / model_y_std + _logger.info(f"Training {name}") x_train = _feature_array(df, train_idx, x_cols) - y_train_scaled_array = y_train_scaled.to_numpy(dtype=np.float32, copy=True) + y_train_scaled_array = model_y_train_scaled.to_numpy(dtype=np.float32, copy=True) x_eval = _feature_array(df, eval_idx, x_cols) - y_eval_scaled_array = ((df.iloc[eval_idx, target_indices] - y_mean) / y_std).to_numpy( - dtype=np.float32, copy=True - ) + y_eval_scaled_array = ( + (df.iloc[eval_idx, model_target_indices] - model_y_mean) / model_y_std + ).to_numpy(dtype=np.float32, copy=True) eval_set = [(x_eval, y_eval_scaled_array)] utils.log_memory_checkpoint("after building XGBoost fit arrays", enabled=memory_profile) @@ -1049,9 +1089,9 @@ def train_regression(df, model_configs): random_state, ) y_train_diagnostic = ( - y_train + model_y_train if diagnostic_train_idx is train_idx - else df.iloc[diagnostic_train_idx, target_indices] + else df.iloc[diagnostic_train_idx, model_target_indices] ) _logger.info( "Post-training diagnostic training events: %d", @@ -1062,9 +1102,9 @@ def train_regression(df, model_configs): df, diagnostic_train_idx, x_cols, - y_mean, - y_std, - targets, + model_y_mean, + model_y_std, + model_targets, prediction_chunk_size, ) utils.log_memory_checkpoint( @@ -1078,9 +1118,9 @@ def train_regression(df, model_configs): df, test_idx, x_cols, - y_mean, - y_std, - targets, + model_y_mean, + model_y_std, + model_targets, prediction_chunk_size, ) utils.log_memory_checkpoint( @@ -1091,15 +1131,15 @@ def train_regression(df, model_configs): generalization_metrics = diagnostic_utils.compute_generalization_metrics( y_train_diagnostic, y_train_pred, - y_test, + model_y_test, y_pred, - targets, + model_targets, ) residual_normality_stats = diagnostic_utils.compute_residual_normality_stats( - y_test, + model_y_test, y_pred, - targets, + model_targets, ) shap_idx = _sample_eval_indices( @@ -1110,12 +1150,20 @@ def train_regression(df, model_configs): 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) shap_importance = evaluate_regression_model( - model, x_test_shap, y_pred, y_test, df, x_cols, y_test, name + model, + x_test_shap, + y_pred, + model_y_test, + df, + x_cols, + model_y_test, + name, ) del x_test_shap utils.log_memory_checkpoint(f"{name}: after regression evaluation", enabled=memory_profile) cfg["model"] = model cfg["features"] = x_cols # Store feature names for later use + cfg["targets"] = list(model_targets) cfg["generalization_metrics"] = generalization_metrics cfg["residual_normality_stats"] = residual_normality_stats cfg["shap_importance"] = shap_importance # Store per-target SHAP importance from evaluation diff --git a/tests/test_config.py b/tests/test_config.py index 02f324b..3456c9b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -143,6 +143,48 @@ def test_configure_training_stereo_parses_reduced_feature_profile(monkeypatch): assert result["feature_profile"] == "reduced" +def test_configure_training_stereo_separates_direction_and_energy(monkeypatch): + """The separation flag must create two model configurations with fixed targets.""" + monkeypatch.setattr( + sys, + "argv", + [ + "prog", + "--model_prefix", + "model", + "--input_file_list", + "inputs.txt", + "--separate_direction_energy", + ], + ) + monkeypatch.setattr( + config, + "hyper_parameters", + lambda *_: {"xgboost": {"hyper_parameters": {"max_depth": 5}}}, + ) + monkeypatch.setattr( + config, + "pre_cuts_regression", + lambda min_images: f"cut_{min_images}", + ) + + result = config.configure_training("stereo_analysis") + + assert result["regression_mode"] == "separate_direction_energy" + assert result["models"]["direction"]["targets"] == ["Xoff_residual", "Yoff_residual"] + assert result["models"]["energy"]["targets"] == ["E_residual"] + assert result["models"]["direction"]["hyper_parameters"] == { + "max_depth": 5, + "n_jobs": 1, + "random_state": 42, + } + assert result["models"]["energy"]["hyper_parameters"] == { + "max_depth": 5, + "n_jobs": 1, + "random_state": 42, + } + + def test_configure_training_classification_parses_tmva_style(monkeypatch, model_parameters_file): monkeypatch.setattr( sys, diff --git a/tests/test_regression_contracts.py b/tests/test_regression_contracts.py index 2e3b742..3868b62 100644 --- a/tests/test_regression_contracts.py +++ b/tests/test_regression_contracts.py @@ -36,6 +36,17 @@ def predict(self, x_values): return np.zeros((len(x_values), 3), dtype=np.float32) +class TargetCountRegressor(CapturingRegressor): + """Minimal regressor whose prediction width matches one target group.""" + + def __init__(self, n_targets): + self.n_targets = n_targets + + def predict(self, x_values): + """Return zero residuals for the configured target group.""" + return np.zeros((len(x_values), self.n_targets), dtype=np.float32) + + class OrderedPredictionRegressor: """Serializable predictor whose output exposes the received feature order.""" @@ -49,6 +60,22 @@ def predict(self, x_values): return np.column_stack((beta, alpha, beta - alpha)) +class DirectionPredictionRegressor: + """Return two residuals in a way that exposes direction feature order.""" + + def predict(self, x_values): + """Return the first two input columns as direction residuals.""" + return np.column_stack((x_values.iloc[:, 0], x_values.iloc[:, 1])) + + +class EnergyPredictionRegressor: + """Return a one-dimensional energy residual to exercise application reshaping.""" + + def predict(self, x_values): + """Return the first input column as an energy residual.""" + return x_values.iloc[:, 0].to_numpy(dtype=float) + + @pytest.fixture def regression_frame(): """Return indexed data with values that make ordering errors observable.""" @@ -115,6 +142,7 @@ def test_regression_training_contract_excludes_targets_and_uses_train_only_scale def test_regression_training_reduced_profile_selects_requested_columns(monkeypatch): + """The reduced profile must select its documented feature set.""" n_events = 240 row_number = np.arange(n_events, dtype=float) reduced_columns = [ @@ -166,6 +194,41 @@ def test_regression_training_reduced_profile_selects_requested_columns(monkeypat assert result["models"]["xgboost"]["features"] == reduced_columns +def test_regression_training_supports_separate_direction_and_energy_models( + regression_frame, monkeypatch +): + """Direction and energy target groups must fit and persist independently.""" + direction_model = TargetCountRegressor(2) + energy_model = TargetCountRegressor(1) + created_models = [direction_model, energy_model] + monkeypatch.setattr("xgboost.XGBRegressor", lambda **_kwargs: created_models.pop(0)) + monkeypatch.setattr("eventdisplay_ml.models.evaluate_regression_model", lambda *_args: {}) + + result = models.train_regression( + regression_frame, + { + "targets": ["Xoff_residual", "Yoff_residual", "E_residual"], + "regression_mode": "separate_direction_energy", + "train_test_fraction": 0.5, + "random_state": 19, + "eval_max_events": 0, + "diagnostic_max_events": 0, + "models": { + "direction": { + "targets": ["Xoff_residual", "Yoff_residual"], + "hyper_parameters": {}, + }, + "energy": {"targets": ["E_residual"], "hyper_parameters": {}}, + }, + }, + ) + + assert result["models"]["direction"]["targets"] == ["Xoff_residual", "Yoff_residual"] + assert result["models"]["energy"]["targets"] == ["E_residual"] + assert direction_model.y_train.shape[1] == 2 + assert energy_model.y_train.shape[1] == 1 + + def test_persisted_regression_model_preserves_feature_order_and_reconstructs_truth( tmp_path, monkeypatch ): @@ -224,6 +287,54 @@ def test_persisted_regression_model_preserves_feature_order_and_reconstructs_tru np.testing.assert_allclose(pred_log_energy, [1.15, 2.65]) +def test_separate_direction_energy_artifact_combines_predictions(tmp_path, monkeypatch): + """Application must combine separately trained direction and energy models.""" + model_prefix = tmp_path / "separate_stereo_model" + joblib.dump( + { + "regression_mode": "separate_direction_energy", + "models": { + "direction": { + "model": DirectionPredictionRegressor(), + "features": ["direction_beta", "direction_alpha"], + "targets": ["Xoff_residual", "Yoff_residual"], + }, + "energy": { + "model": EnergyPredictionRegressor(), + "features": ["energy_feature"], + "targets": ["E_residual"], + }, + }, + "target_mean": {"Xoff_residual": 1.0, "Yoff_residual": -2.0, "E_residual": 0.5}, + "target_std": {"Xoff_residual": 0.5, "Yoff_residual": 2.0, "E_residual": 0.1}, + }, + tmp_path / "separate_stereo_model.joblib.gz", + ) + flattened = pd.DataFrame( + { + "energy_feature": [3.0, 11.0], + "Xoff_weighted_bdt": [0.2, -0.5], + "ErecS": [10.0, 100.0], + "direction_alpha": [4.0, 7.0], + "Yoff_weighted_bdt": [-1.0, 2.0], + "direction_beta": [3.0, 11.0], + } + ) + monkeypatch.setattr(models, "flatten_feature_data", lambda *_args, **_kwargs: flattened) + monkeypatch.setattr(models.data_processing, "print_variable_statistics", lambda *_args: None) + + loaded_models, parameters = models.load_regression_models(str(model_prefix), "xgboost") + assert parameters["regression_mode"] == "separate_direction_energy" + + pred_xoff, pred_yoff, pred_log_energy = models.apply_regression_models( + pd.DataFrame({"event": [1, 2]}), {"models": loaded_models, **parameters} + ) + + 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.8, 3.6]) + + 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() From eddaa73b8ce691ad07c0cf6b475187ff1c805f7f Mon Sep 17 00:00:00 2001 From: GernotMaier Date: Sun, 30 Aug 2026 10:21:47 +0200 Subject: [PATCH 2/3] changelog --- docs/changes/85.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 docs/changes/85.feature.md diff --git a/docs/changes/85.feature.md b/docs/changes/85.feature.md new file mode 100644 index 0000000..04bed35 --- /dev/null +++ b/docs/changes/85.feature.md @@ -0,0 +1 @@ +Separate direction from energy regression (training). From ac8469023d55112f1a5e8f1d14eaf24faa3de1c9 Mon Sep 17 00:00:00 2001 From: GernotMaier Date: Sun, 30 Aug 2026 10:33:48 +0200 Subject: [PATCH 3/3] training --- src/eventdisplay_ml/evaluate.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/eventdisplay_ml/evaluate.py b/src/eventdisplay_ml/evaluate.py index 1dbde2c..48989b2 100644 --- a/src/eventdisplay_ml/evaluate.py +++ b/src/eventdisplay_ml/evaluate.py @@ -158,7 +158,11 @@ def evaluate_regression_model( shap_feature_importance_by_energy(model, x_test, df, y_test, y_data.columns) required_resolution_targets = {"Xoff_residual", "Yoff_residual", "E_residual"} - if required_resolution_targets.issubset(y_data.columns): + target_names = set(y_data.columns) + is_partial_stereo_model = bool(target_names & required_resolution_targets) and not ( + required_resolution_targets.issubset(target_names) + ) + if not is_partial_stereo_model: calculate_resolution( y_pred, y_test,