Separate energy from direction training

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

docs/changes/85.feature.md

@@ -0,0 +1 @@
+Separate direction from energy regression (training).

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)

src/eventdisplay_ml/evaluate.py

@@ -157,16 +157,27 @@ 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"}
+    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,
+            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
 

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,20 +978,17 @@ 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()
 
     _logger.info("Target standardization (training set):")
     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