From 13786ec47b1f608a0c13451fd96fc061262eeb38 Mon Sep 17 00:00:00 2001
From: GernotMaier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 12:41:15 +0200
Subject: [PATCH 01/10] Classification improvement

---
 src/eventdisplay_ml/config.py           |  19 ++
 src/eventdisplay_ml/data_processing.py  | 119 ++++++--
 src/eventdisplay_ml/evaluate.py         |  47 +++-
 src/eventdisplay_ml/features.py         |  78 ++++++
 src/eventdisplay_ml/models.py           | 356 ++++++++++++++++++++----
 tests/test_classification_robustness.py |  69 +++++
 6 files changed, 620 insertions(+), 68 deletions(-)
 create mode 100644 tests/test_classification_robustness.py

diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py
index c22611b..aff214d 100644
--- a/src/eventdisplay_ml/config.py
+++ b/src/eventdisplay_ml/config.py
@@ -104,6 +104,23 @@ def configure_training(analysis_type):
             action="store_true",
             help="Remove ze_bin from gamma/hadron training features.",
         )
+        parser.add_argument(
+            "--feature_profile",
+            choices=("robust", "extended"),
+            default="robust",
+            help=(
+                "Classification feature set. 'robust' uses stable image/stereo variables "
+                "available in existing data; 'extended' retains the historical feature set."
+            ),
+        )
+        parser.add_argument(
+            "--grouped_split",
+            action=argparse.BooleanOptionalAction,
+            default=True,
+            help=(
+                "Keep events from the same source file in one split when provenance is available."
+            ),
+        )
     parser.add_argument(
         "--max_cores",
         type=int,
@@ -194,6 +211,8 @@ def configure_training(analysis_type):
             f"Balance class zenith weights: {model_configs.get('balance_class_zenith_weights')}"
         )
         _logger.info(f"Ignore ze_bin feature: {model_configs.get('ignore_ze_bin')}")
+        _logger.info(f"Classification feature profile: {model_configs.get('feature_profile')}")
+        _logger.info(f"Grouped classification split: {model_configs.get('grouped_split')}")
 
     model_configs["models"] = hyper_parameters(
         analysis_type, model_configs.get("hyperparameter_config")
diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py
index 34239df..a54f5fc 100644
--- a/src/eventdisplay_ml/data_processing.py
+++ b/src/eventdisplay_ml/data_processing.py
@@ -337,7 +337,7 @@ def _normalize_telescope_variable_to_tel_id_space(data, index_list, max_tel_id,
     row_indices, col_indices = np.where(~np.isnan(index_list))
     tel_ids = index_list[row_indices, col_indices].astype(int)
     # Filter for valid telescope IDs and valid column indices in data array
-    valid_mask = (tel_ids <= max_tel_id) & (col_indices < data.shape[1])
+    valid_mask = (tel_ids >= 0) & (tel_ids <= max_tel_id) & (col_indices < data.shape[1])
     full_matrix[row_indices[valid_mask], tel_ids[valid_mask]] = data[
         row_indices[valid_mask], col_indices[valid_mask]
     ]
@@ -483,13 +483,17 @@ def flatten_telescope_data_vectorized(
     active_mask = np.zeros((n_evt, max_tel_id + 1), dtype=bool)
     row_indices, col_indices = np.where(~np.isnan(tel_list_matrix))
     tel_ids = tel_list_matrix[row_indices, col_indices].astype(int)
-    valid_tel_mask = tel_ids <= max_tel_id
+    valid_tel_mask = (tel_ids >= 0) & (tel_ids <= max_tel_id)
     active_mask[row_indices[valid_tel_mask], tel_ids[valid_tel_mask]] = True
 
     # Pre-load and normalize size to telescope-ID space for sorting
     size_data = _normalize_telescope_variable_to_tel_id_space(
         _to_dense_array(df["size"]), index_list_for_remapping, max_tel_id, n_evt
     )
+    # A telescope absent from DispTelList_T is not a zero-sized image.  Keep it
+    # explicitly missing so sorting and the XGBoost missing-value path cannot
+    # learn a detector-slot/observing-condition proxy.
+    size_data = np.where(active_mask, size_data, np.nan)
     size_data = _clip_size_array(size_data)
 
     core_x, core_y = _get_core_arrays(df)
@@ -534,6 +538,16 @@ def flatten_telescope_data_vectorized(
                     sort_indices,
                 )
             )
+            # Geometry is useful only for an active image; inactive slots must
+            # not become a stable class/domain indicator.
+            for key in tuple(flat_features):
+                if key.startswith(f"{var}_"):
+                    sorted_tel = int(key.rsplit("_", 1)[1])
+                    flat_features[key] = np.where(
+                        active_mask[np.arange(n_evt), sort_indices[:, sorted_tel]],
+                        flat_features[key],
+                        np.nan,
+                    )
             continue
 
         data = _to_dense_array(df[var]) if _has_field(df, var) else np.full((n_evt, n_tel), np.nan)
@@ -550,6 +564,9 @@ def flatten_telescope_data_vectorized(
                 data, index_list_for_remapping, max_tel_id, n_evt
             )
 
+        if var != "tel_active":
+            data_normalized = np.where(active_mask, data_normalized, np.nan)
+
         # All variables are now in telescope-ID space; apply sorting and flatten uniformly
         data_normalized = data_normalized[np.arange(n_evt)[:, np.newaxis], sort_indices]
 
@@ -667,6 +684,7 @@ def _compute_size_area_sort_indices(size_data, active_mask, tel_config, max_tel_
         for tel_idx in range(max_tel_id + 1):
             area = mirror_lookup[tel_idx]
             size_val = sizes[evt_idx, tel_idx]
+            active = bool(active_mask[evt_idx, tel_idx])
             # Build sort key:
             #   1) valid area first (0), NaN area last (1)
             #   2) area descending via negative value
@@ -676,9 +694,14 @@ def _compute_size_area_sort_indices(size_data, active_mask, tel_config, max_tel_
             size_valid = 0 if not np.isnan(size_val) else 1
             area_key = -area if area_valid == 0 else 0.0
             size_key = -size_val if size_valid == 0 else 0.0
-            tel_entries.append((tel_idx, area_valid, area_key, size_valid, size_key))
+            # Active images always precede inactive detector slots.  This keeps
+            # slot ordering deterministic without turning missing telescopes
+            # into a large/small-image classification feature.
+            tel_entries.append(
+                (tel_idx, 0 if active else 1, area_valid, area_key, size_valid, size_key)
+            )
 
-        tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4]))
+        tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4], x[5]))
         sort_indices[evt_idx] = np.array([t[0] for t in tel_entries])
 
     return sort_indices
@@ -960,6 +983,8 @@ def load_training_data(model_configs, file_list, analysis_type):
         _logger.info(f"Adding zenith binning: {model_configs.get('zenith_bins_deg', [])}")
 
     input_files = utils.read_input_file_list(file_list)
+    if not input_files:
+        raise ValueError(f"Input file list is empty: {file_list}")
 
     tmva_style = model_configs.get("tmva_style", False)
     if tmva_style and analysis_type == "classification":
@@ -974,12 +999,16 @@ def load_training_data(model_configs, file_list, analysis_type):
         branch_list = features_module.features(analysis_type, training=True)
     _logger.info(f"Branch list: {branch_list}")
     if max_events is not None and max_events > 0:
-        max_events_per_file = max_events // len(input_files)
+        # Reserve a bounded quota per file, then perform one deterministic
+        # final sample below.  Integer floor division used to turn a small
+        # global cap into zero (which silently disabled sampling).
+        max_events_per_file = max(1, int(np.ceil(max_events / len(input_files))))
     else:
         max_events_per_file = None
     _logger.info(f"Max events per file: {max_events_per_file}")
 
-    tel_config = None  # Will be read from first file
+    # Reuse/validate across signal and background loads.
+    tel_config = model_configs.get("tel_config")
     dfs = []
     executor = ThreadPoolExecutor(max_workers=model_configs.get("max_cores", 1))
     total_files = len(input_files)
@@ -990,21 +1019,30 @@ def load_training_data(model_configs, file_list, analysis_type):
                     _logger.warning(f"File: {f} does not contain a 'data' tree.")
                     continue
 
+                current_tel_config = read_telescope_config(root_file)
                 if tel_config is None:
-                    tel_config = read_telescope_config(root_file)
+                    tel_config = current_tel_config
                     model_configs["tel_config"] = tel_config
                 else:
-                    # Check if current file has a larger max_tel_id and update if needed
-                    current_tel_config = read_telescope_config(root_file)
-                    if current_tel_config["max_tel_id"] > tel_config["max_tel_id"]:
-                        _logger.info(
-                            f"Updating telescope configuration: max_tel_id from "
-                            f"{tel_config['max_tel_id']} to {current_tel_config['max_tel_id']} "
-                            f"(file: {f})"
+                    # A model cannot have a stable feature schema if telescope
+                    # IDs/areas change between input files.  The old code
+                    # silently replaced the configuration when max_tel_id grew.
+                    def _config_signature(config):
+                        return (
+                            int(config["max_tel_id"]),
+                            tuple(int(v) for v in config.get("tel_ids", [])),
+                            tuple(str(v) for v in config.get("tel_types", [])),
+                            tuple(
+                                float(v)
+                                for v in config.get("mirror_area", config.get("mirror_areas", []))
+                            ),
+                        )
+
+                    if _config_signature(current_tel_config) != _config_signature(tel_config):
+                        raise ValueError(
+                            "Classification/training input files have incompatible telescope "
+                            f"configurations: {input_files[0]} versus {f}."
                         )
-                        # Replace the full telescope configuration to keep all fields consistent
-                        tel_config = current_tel_config
-                        model_configs["tel_config"] = tel_config
 
                 _logger.info(f"Processing file: {f} (file {file_idx}/{total_files})")
                 tree = root_file["data"]
@@ -1015,7 +1053,9 @@ def load_training_data(model_configs, file_list, analysis_type):
                 raw_reservoir_chunks = []
                 reservoir_priorities = None
                 file_dfs = []
-                rng = np.random.default_rng(random_state)
+                rng = np.random.default_rng(
+                    None if random_state is None else int(random_state) + file_idx - 1
+                )
                 chunk_iterator = tree.iterate(
                     resolved_branch_list,
                     cut=model_configs.get("pre_cuts", None),
@@ -1103,6 +1143,14 @@ def load_training_data(model_configs, file_list, analysis_type):
                 if file_df is None or file_df.empty:
                     continue
 
+                if analysis_type == "classification":
+                    # Provenance is retained only as routing metadata and is
+                    # excluded by the feature profile before fitting.  It
+                    # enables grouped validation without rereading ROOT data.
+                    file_df["__source_file_id"] = file_idx - 1
+                    file_df["__source_file"] = str(f)
+                    file_df["__source_row"] = np.arange(len(file_df), dtype=np.int64)
+
                 _logger.info(
                     f"Number of events before / after event cut: {n_before} / "
                     f"{n_after_event_cut} (fraction retained: {n_after_event_cut / n_before:.4f})"
@@ -1119,10 +1167,22 @@ def load_training_data(model_configs, file_list, analysis_type):
                     file_df,
                     enabled=memory_profile,
                 )
+        except (FileNotFoundError, KeyError, ValueError):
+            raise
         except Exception as e:
-            raise FileNotFoundError(f"Error opening or reading file {f}: {e}") from e
+            raise RuntimeError(f"Error opening or reading file {f}: {e}") from e
 
+    if not dfs:
+        raise ValueError("No data loaded from input files.")
     df_final = pd.concat(dfs, ignore_index=True)
+    if analysis_type == "classification" and max_events is not None and max_events > 0:
+        if len(df_final) > max_events:
+            df_final = df_final.sample(
+                n=max_events,
+                random_state=random_state,
+                ignore_index=True,
+            )
+            _logger.info("Applied exact global classification event cap: %d", max_events)
     del dfs
     utils.log_memory_checkpoint("after final pandas concat", df_final, enabled=memory_profile)
     all_nan_columns = [col for col in df_final.columns if df_final[col].isna().all()]
@@ -1399,6 +1459,19 @@ def extra_columns(df, analysis_type, training, index, tel_config=None, observato
             "EChi2S": _to_numpy_1d(df["EChi2S"], np.float32),
             "EmissionHeight": _to_numpy_1d(df["EmissionHeight"], np.float32),
             "EmissionHeightChi2": _to_numpy_1d(df["EmissionHeightChi2"], np.float32),
+            # Keep routing quantities in the flattened frame for energy-bin
+            # weighting/diagnostics; feature-profile selection removes them
+            # before fitting.
+            "Erec": (
+                _to_numpy_1d(df["Erec"], np.float32)
+                if _has_field(df, "Erec")
+                else np.full(n, DEFAULT_FILL_VALUE, dtype=np.float32)
+            ),
+            "DispNImages": (
+                _to_numpy_1d(df["DispNImages"], np.float32)
+                if _has_field(df, "DispNImages")
+                else np.full(n, DEFAULT_FILL_VALUE, dtype=np.float32)
+            ),
         }
         if _has_field(df, "SizeSecondMax"):
             data["SizeSecondMax"] = _to_numpy_1d(df["SizeSecondMax"], np.float32)
@@ -1442,9 +1515,17 @@ def extra_columns(df, analysis_type, training, index, tel_config=None, observato
 
 def zenith_in_bins(zenith_angles, bins):
     """Apply zenith binning based on zenith angles and given bin edges."""
+    if bins is None or len(bins) < 2:
+        raise ValueError("At least two zenith-bin edges are required.")
     if isinstance(bins[0], dict):
+        if any("Ze_min" not in b or "Ze_max" not in b for b in bins):
+            raise ValueError("Zenith-bin dictionaries require Ze_min and Ze_max.")
         bins = [b["Ze_min"] for b in bins] + [bins[-1]["Ze_max"]]
     bins = np.asarray(bins, dtype=float)
+    if bins.ndim != 1 or len(bins) < 2 or not np.all(np.isfinite(bins)):
+        raise ValueError("Zenith-bin edges must be a finite one-dimensional sequence.")
+    if np.any(np.diff(bins) <= 0):
+        raise ValueError("Zenith-bin edges must be strictly increasing.")
     idx = np.clip(np.digitize(zenith_angles, bins) - 1, 0, len(bins) - 2)
     return idx.astype(np.int32)
 
diff --git a/src/eventdisplay_ml/evaluate.py b/src/eventdisplay_ml/evaluate.py
index 2e1ac6b..116a914 100644
--- a/src/eventdisplay_ml/evaluate.py
+++ b/src/eventdisplay_ml/evaluate.py
@@ -5,6 +5,7 @@
 import numpy as np
 import pandas as pd
 import xgboost as xgb
+from scipy.stats import beta as beta_distribution
 from sklearn.metrics import (
     classification_report,
     confusion_matrix,
@@ -25,8 +26,10 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
 
     for t in thresholds:
         pred = y_pred_proba >= t
-        eff_signal.append(((pred) & (y_test == 1)).sum() / n_signal if n_signal else 0)
-        eff_background.append(((pred) & (y_test == 0)).sum() / n_background if n_background else 0)
+        eff_signal.append(((pred) & (y_test == 1)).sum() / n_signal if n_signal else np.nan)
+        eff_background.append(
+            ((pred) & (y_test == 0)).sum() / n_background if n_background else np.nan
+        )
         _logger.info(
             f"{name}{context_label} Threshold: {t:.2f} | "
             f"Signal Efficiency: {eff_signal[-1]:.4f} | "
@@ -35,6 +38,14 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
 
     eff_signal = np.asarray(eff_signal, dtype=float)
     eff_background = np.asarray(eff_background, dtype=float)
+    background_survivors = n_background * eff_background
+    background_upper_limit = np.full(len(thresholds), np.nan, dtype=float)
+    if n_background:
+        for i, survivors in enumerate(background_survivors):
+            k = round(survivors)
+            background_upper_limit[i] = (
+                1.0 if k >= n_background else beta_distribution.ppf(0.95, k + 1, n_background - k)
+            )
 
     return pd.DataFrame(
         {
@@ -43,6 +54,7 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
             "background_efficiency": eff_background,
             "n_signal": n_signal * eff_signal,
             "n_background": n_background * eff_background,
+            "background_efficiency_upper95": background_upper_limit,
         }
     )
 
@@ -85,6 +97,37 @@ def evaluation_efficiency(name, model, x_test, y_test, return_by_zenith=False, z
     return efficiency_all, efficiencies_by_zenith
 
 
+def classification_thresholds_from_signal(signal_scores, efficiencies=(0.5, 0.7, 0.8, 0.9, 0.95)):
+    """Calibrate score thresholds to measured held-out signal efficiency.
+
+    Quantiles avoid treating XGBoost's score as a globally calibrated
+    probability.  ``method='higher'`` ensures ties do not exceed the requested
+    operating point.
+    """
+    scores = np.asarray(signal_scores, dtype=float)
+    scores = scores[np.isfinite(scores)]
+    if scores.size == 0:
+        raise ValueError("Cannot calibrate classification thresholds without signal scores.")
+    targets = np.asarray(efficiencies, dtype=float)
+    if np.any((targets <= 0) | (targets >= 1)):
+        raise ValueError("Signal efficiencies must be strictly between zero and one.")
+    thresholds = []
+    for efficiency in targets:
+        quantile = 1.0 - efficiency
+        try:
+            threshold = np.quantile(scores, quantile, method="higher")
+        except TypeError:  # NumPy < 1.22 compatibility
+            threshold = np.quantile(scores, quantile, interpolation="higher")
+        thresholds.append(float(np.clip(threshold, 0.0, 1.0)))
+    return pd.DataFrame(
+        {
+            "signal_efficiency_target": targets,
+            "threshold": thresholds,
+            "n_signal": len(scores),
+        }
+    )
+
+
 def evaluate_classification_model(model, x_test, y_test, df, x_cols, name):
     """Evaluate the trained model on the test set and log performance metrics.
 
diff --git a/src/eventdisplay_ml/features.py b/src/eventdisplay_ml/features.py
index d4b277d..7757276 100644
--- a/src/eventdisplay_ml/features.py
+++ b/src/eventdisplay_ml/features.py
@@ -47,6 +47,84 @@ def target_features(analysis_type):
     raise ValueError(f"Unknown analysis type: {analysis_type}")
 
 
+def classification_feature_columns(columns, profile="robust", ignore_ze_bin=False):
+    """Return safe classification columns from an already flattened frame.
+
+    The profile is applied after flattening, so it works with both VERITAS fixed
+    telescope indexing and the CTAO variable-length compatibility path.  Source
+    provenance and routing columns are deliberately never exposed to XGBoost.
+    """
+    if profile not in {"robust", "extended"}:
+        raise ValueError("classification feature profile must be 'robust' or 'extended'")
+
+    reserved = {
+        "label",
+        "Erec",
+        "ErecS",
+        "__source_file_id",
+        "__source_row",
+        "__source_file",
+    }
+    available = list(columns)
+    if profile == "extended":
+        selected = [name for name in available if name not in reserved]
+    else:
+        # Array/stereo quantities plus per-telescope image morphology.  The
+        # latter are essential gamma/hadron information; detector activity,
+        # telescope geometry and pointing remain excluded below.
+        stable = {
+            "MSCW",
+            "MSCL",
+            "EChi2S",
+            "EmissionHeight",
+            "EmissionHeightChi2",
+            "Core_Distance",
+            # Coarse zenith conditioning is important because atmospheric
+            # depth/projection changes the image morphology.  It can be
+            # removed explicitly with --ignore_ze_bin for a nuisance test.
+            "ze_bin",
+        }
+        image_bases = {
+            "size",
+            "cosphi",
+            "sinphi",
+            "loss",
+            "dist",
+            "width",
+            "length",
+            "asym",
+            "tgrad_x",
+        }
+        selected = [
+            name
+            for name in available
+            if name in stable or any(name.startswith(f"{base}_") for base in image_bases)
+        ]
+        # Small synthetic/unit-test frames (and old files missing derived
+        # quantities) should still be usable in isolated unit tests.  Real
+        # flattened frames contain at least one physics/activity name; fail
+        # loudly instead of silently falling back to nuisance columns there.
+        if not selected:
+            looks_like_flattened_physics = any(
+                name.startswith(("tel_", "ArrayPointing", "Xcore", "Ycore"))
+                or name in {"DispNImages", "Erec", "size"}
+                for name in available
+            )
+            if looks_like_flattened_physics:
+                raise ValueError(
+                    "Robust classification profile has no available stable morphology features."
+                )
+            selected = [name for name in available if name not in reserved]
+            if not ignore_ze_bin and "ze_bin" in available and "ze_bin" not in selected:
+                selected.append("ze_bin")
+
+    if ignore_ze_bin and "ze_bin" in selected:
+        selected.remove("ze_bin")
+    if not selected:
+        raise ValueError(f"No usable classification features for profile '{profile}'.")
+    return selected
+
+
 def excluded_features(analysis_type, ntel):
     """
     Features not to be used for training/prediction.
diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py
index 1125df6..edb303a 100644
--- a/src/eventdisplay_ml/models.py
+++ b/src/eventdisplay_ml/models.py
@@ -14,6 +14,7 @@
 import pandas as pd
 import uproot
 import xgboost as xgb
+from sklearn.metrics import roc_auc_score
 from sklearn.model_selection import train_test_split
 
 from eventdisplay_ml import data_processing, diagnostic_utils, features, utils
@@ -23,6 +24,7 @@
     zenith_in_bins,
 )
 from eventdisplay_ml.evaluate import (
+    classification_thresholds_from_signal,
     evaluate_classification_model,
     evaluate_regression_model,
     evaluation_efficiency,
@@ -196,9 +198,11 @@ def load_classification_models(model_prefix, model_name):
             raise KeyError(f"Model name '{model_name}' not found in file: {file}")
         models[e_bin]["features"] = model_data.get("features", [])
         models[e_bin]["efficiency"] = model_data["models"][model_name].get("efficiency")
+        calibration = model_data["models"][model_name].get("signal_threshold_calibration")
         models[e_bin]["thresholds"] = _calculate_classification_thresholds(
-            models[e_bin]["efficiency"]
+            models[e_bin]["efficiency"], calibration=calibration
         )
+        models[e_bin]["support"] = model_data["models"][model_name].get("support", {})
         energy_bin_metadata = _validate_energy_bin_metadata(
             model_data.get("energy_bins_log10_tev"),
             file,
@@ -223,7 +227,7 @@ def load_classification_models(model_prefix, model_name):
     return models, par
 
 
-def _calculate_classification_thresholds(efficiency, min_efficiency=0.2, steps=5):
+def _calculate_classification_thresholds(efficiency, min_efficiency=0.2, steps=5, calibration=None):
     """
     Calculate classification thresholds for given signal efficiencies.
 
@@ -243,7 +247,21 @@ def _calculate_classification_thresholds(efficiency, min_efficiency=0.2, steps=5
     dict[int, float]
         Mapping from efficiency (percent) to classification threshold.
     """
+    if efficiency is None or len(efficiency) == 0:
+        raise ValueError("Classification efficiency diagnostics are missing from the model file.")
     df = efficiency.copy()
+    if calibration is not None:
+        calibrated = pd.DataFrame(calibration)
+        if {"signal_efficiency_target", "threshold"}.issubset(calibrated.columns):
+            df = pd.concat(
+                [
+                    df[["signal_efficiency", "threshold"]],
+                    calibrated[["signal_efficiency_target", "threshold"]].rename(
+                        columns={"signal_efficiency_target": "signal_efficiency"}
+                    ),
+                ],
+                ignore_index=True,
+            ).drop_duplicates(subset=["signal_efficiency"], keep="last")
     df = df.sort_values("signal_efficiency")
     eff_targets = np.arange(min_efficiency * 100, 100, steps) / 100.0
     thresholds = np.interp(
@@ -524,10 +542,19 @@ def apply_classification_models(df, model_configs, threshold_keys):
             observatory=model_configs.get("observatory", "veritas"),
             preview_rows=model_configs.get("preview_rows", 20),
         )
-        model_lo = models[e_bin_lo]["model"]
-        model_hi = models[e_bin_hi]["model"]
-        flatten_lo = flatten_data.reindex(columns=models[e_bin_lo]["features"])
-        flatten_hi = flatten_data.reindex(columns=models[e_bin_hi]["features"])
+        resolved_lo = _resolve_classification_bin(models, e_bin_lo)
+        resolved_hi = _resolve_classification_bin(models, e_bin_hi)
+        model_lo = models[resolved_lo]["model"]
+        model_hi = models[resolved_hi]["model"]
+        missing_lo = sorted(set(models[resolved_lo]["features"]) - set(flatten_data.columns))
+        missing_hi = sorted(set(models[resolved_hi]["features"]) - set(flatten_data.columns))
+        if missing_lo or missing_hi:
+            raise ValueError(
+                "Classification model/input feature schema mismatch: "
+                f"low-bin missing={missing_lo}, high-bin missing={missing_hi}."
+            )
+        flatten_lo = flatten_data.loc[:, models[resolved_lo]["features"]]
+        flatten_hi = flatten_data.loc[:, models[resolved_hi]["features"]]
 
         class_probs_lo = model_lo.predict_proba(flatten_lo)[:, 1]
         if e_bin_lo == e_bin_hi:
@@ -538,8 +565,8 @@ def apply_classification_models(df, model_configs, threshold_keys):
             class_probs = (1.0 - alpha) * class_probs_lo + alpha * class_probs_hi
         class_probability[group_df.index] = class_probs
 
-        thresholds_lo = models[e_bin_lo].get("thresholds", {})
-        thresholds_hi = models[e_bin_hi].get("thresholds", {})
+        thresholds_lo = models[resolved_lo].get("thresholds", {})
+        thresholds_hi = models[resolved_hi].get("thresholds", {})
         for eff in threshold_keys:
             if eff in is_gamma:
                 thr_lo = thresholds_lo.get(eff)
@@ -558,6 +585,22 @@ def apply_classification_models(df, model_configs, threshold_keys):
     return class_probability, is_gamma
 
 
+def _resolve_classification_bin(models, requested_bin):
+    """Resolve a missing energy-bin model to the nearest available model."""
+    if requested_bin in models:
+        return requested_bin
+    available = sorted(models)
+    if not available:
+        raise ValueError("No classification models are available for application.")
+    nearest = min(available, key=lambda candidate: abs(candidate - requested_bin))
+    _logger.warning(
+        "No classification model for energy bin %d; borrowing nearest bin %d.",
+        requested_bin,
+        nearest,
+    )
+    return nearest
+
+
 def process_file_chunked(analysis_type, model_configs):
     """
     Stream events from an input file in chunks, apply XGBoost models, write events.
@@ -1024,47 +1067,123 @@ def train_classification(df, model_configs):
             f"signal_events={len(df[0])}, background_events={len(df[1])}."
         )
 
-    df[0]["label"] = 1
-    df[1]["label"] = 0
-    full_df = pd.concat([df[0], df[1]], ignore_index=True)
+    left_columns = set(df[0].columns)
+    right_columns = set(df[1].columns)
+    if left_columns != right_columns:
+        raise ValueError(
+            "Signal/background classification schemas differ. "
+            f"Only signal: {sorted(left_columns - right_columns)}; "
+            f"only background: {sorted(right_columns - left_columns)}"
+        )
+
+    signal = df[0].copy()
+    background = df[1].copy()
+    signal["label"] = 1
+    background["label"] = 0
+    full_df = pd.concat([signal, background], ignore_index=True)
     ze_data = full_df["ze_bin"] if "ze_bin" in full_df.columns else None
-    x_data = full_df.drop(columns=["label"])
-    if model_configs.get("ignore_ze_bin", False):
-        if model_configs.get("balance_class_zenith_weights", False):
-            raise ValueError("Cannot use ignore_ze_bin with balance_class_zenith_weights.")
-        if "ze_bin" in x_data.columns:
-            _logger.info("Removing ze_bin from classification training features.")
-            x_data = x_data.drop(columns=["ze_bin"])
+    if model_configs.get("balance_class_zenith_weights", False) and ze_data is None:
+        raise ValueError("Class/zenith balancing requires the derived ze_bin column.")
+
+    profile = model_configs.get("feature_profile", "robust")
+    feature_columns = features.classification_feature_columns(
+        full_df.columns,
+        profile=profile,
+        ignore_ze_bin=model_configs.get("ignore_ze_bin", False),
+    )
+    for column in feature_columns:
+        signal_all_nan = bool(signal[column].isna().all())
+        background_all_nan = bool(background[column].isna().all())
+        if signal_all_nan != background_all_nan:
+            raise ValueError(f"Classification feature '{column}' is all-NaN in only one class.")
+        if signal_all_nan:
+            raise ValueError(f"Classification feature '{column}' is all-NaN in both classes.")
+    x_data = full_df.loc[:, feature_columns]
     _logger.info(f"Features ({len(x_data.columns)}): {', '.join(x_data.columns)}")
     model_configs["features"] = list(x_data.columns)
     y_data = full_df["label"]
 
-    split_inputs = [x_data, y_data]
-    if ze_data is not None:
-        split_inputs.append(ze_data)
-
-    split_result = train_test_split(
-        *split_inputs,
-        train_size=model_configs.get("train_test_fraction", 0.5),
-        random_state=model_configs.get("random_state", None),
-        stratify=y_data,
+    train_idx, validation_idx, test_idx, split_metadata = _classification_split_indices(
+        y_data,
+        full_df.get("__source_file"),
+        train_fraction=model_configs.get("train_test_fraction", 0.5),
+        random_state=model_configs.get("random_state"),
+        grouped=model_configs.get("grouped_split", True),
     )
-    if ze_data is None:
-        x_train, x_test, y_train, y_test = split_result
-        ze_test = None
-    else:
-        x_train, x_test, y_train, y_test, _, ze_test = split_result
-
-    _logger.info(f"Training events: {len(x_train)}, Testing events: {len(x_test)}")
+    # Keep a small, explicitly reserved gamma subset for score-threshold
+    # calibration.  It is never used for fitting or assessment metrics.
+    test_signal_idx = test_idx[y_data.iloc[test_idx].to_numpy() == 1]
+    calibration_idx = np.asarray([], dtype=int)
+    test_signal_groups = (
+        full_df.iloc[test_signal_idx]["__source_file"]
+        if "__source_file" in full_df.columns
+        else None
+    )
+    if test_signal_groups is not None and test_signal_groups.nunique() >= 2:
+        calibration_groups, _assessment_groups = train_test_split(
+            test_signal_groups.unique(),
+            test_size=0.5,
+            random_state=model_configs.get("random_state"),
+        )
+        calibration_idx = test_signal_idx[test_signal_groups.isin(calibration_groups).to_numpy()]
+    elif len(test_signal_idx) >= 2:
+        calibration_idx, _assessment_signal_idx = train_test_split(
+            test_signal_idx, test_size=0.5, random_state=model_configs.get("random_state")
+        )
+    if len(calibration_idx):
+        test_idx = np.asarray(
+            sorted(set(test_idx) - set(calibration_idx)),
+            dtype=int,
+        )
+    x_train, x_validation, x_test = (
+        x_data.iloc[idx] for idx in (train_idx, validation_idx, test_idx)
+    )
+    y_train, y_validation, y_test = (
+        y_data.iloc[idx] for idx in (train_idx, validation_idx, test_idx)
+    )
+    ze_test = ze_data.iloc[test_idx] if ze_data is not None else None
+    _logger.info(
+        "Classification split: train=%d validation=%d test=%d (%s)",
+        len(x_train),
+        len(x_validation),
+        len(x_test),
+        split_metadata["method"],
+    )
+    model_configs["classification_split"] = split_metadata
+    model_configs["classification_split"]["n_signal_calibration"] = len(calibration_idx)
+    model_configs["classification_split"]["calibration_grouped"] = bool(
+        test_signal_groups is not None and test_signal_groups.nunique() >= 2
+    )
+    model_configs["classification_feature_profile"] = profile
+    model_configs["nuisance_diagnostics"] = _classification_nuisance_diagnostics(full_df, y_data)
     weights_train = None
+    weights_validation = None
     if model_configs.get("balance_class_zenith_weights", False):
-        weights_train = _class_zenith_balance_weights(x_train, y_train)
+        weights_train = _class_zenith_balance_weights(full_df.iloc[train_idx], y_train)
+        weights_validation = _class_zenith_balance_weights(
+            full_df.iloc[validation_idx], y_validation
+        )
         _logger.info(
             "Using class/zenith sample weights "
             f"(mean={weights_train.mean():.3f}, std={weights_train.std():.3f}, "
             f"min={weights_train.min():.3f}, max={weights_train.max():.3f})"
         )
-    eval_set = [(x_train, y_train), (x_test, y_test)]
+    eval_x, eval_y = x_validation, y_validation
+    eval_weights = weights_validation
+    eval_max_events = model_configs.get("eval_max_events", 0)
+    if eval_max_events and eval_max_events > 0 and len(eval_x) > eval_max_events:
+        eval_indices = eval_x.sample(
+            n=eval_max_events,
+            random_state=model_configs.get("random_state"),
+        ).index
+        eval_x = eval_x.loc[eval_indices]
+        eval_y = eval_y.loc[eval_indices]
+        if eval_weights is not None:
+            eval_weights = (
+                pd.Series(weights_validation, index=x_validation.index).loc[eval_indices].to_numpy()
+            )
+        _logger.info("Limited XGBoost validation set to %d events", eval_max_events)
+    eval_set = [(x_train, y_train), (eval_x, eval_y)]
 
     for name, cfg in model_configs.get("models", {}).items():
         _logger.info(f"Training {name}")
@@ -1072,6 +1191,7 @@ def train_classification(df, model_configs):
         fit_kwargs = {"eval_set": eval_set, "verbose": True}
         if weights_train is not None:
             fit_kwargs["sample_weight"] = weights_train
+            fit_kwargs["sample_weight_eval_set"] = [weights_train, eval_weights]
         model.fit(x_train, y_train, **fit_kwargs)
 
         shap_importance = evaluate_classification_model(
@@ -1091,12 +1211,138 @@ def train_classification(df, model_configs):
         for ze_bin, ze_efficiency in efficiencies_by_zenith.items():
             cfg[f"efficiency_ze{ze_bin}"] = ze_efficiency
         cfg["shap_importance"] = shap_importance
+        try:
+            calibration_frame = x_data.iloc[calibration_idx]
+            if calibration_frame.empty:
+                raise ValueError("no reserved gamma calibration events")
+            test_signal_scores = model.predict_proba(calibration_frame)[:, 1]
+            cfg["signal_threshold_calibration"] = classification_thresholds_from_signal(
+                test_signal_scores
+            )
+        except (TypeError, ValueError, IndexError) as exc:
+            # Lightweight mocks/legacy estimators may not expose probabilities;
+            # keep the model usable but make the missing calibration explicit.
+            _logger.warning("Could not compute held-out signal thresholds for %s: %s", name, exc)
+            cfg["signal_threshold_calibration"] = None
+        cfg["support"] = {
+            "n_train": len(x_train),
+            "n_validation": len(x_validation),
+            "n_test": len(x_test),
+            "n_signal_test": int((y_test == 1).sum()),
+            "n_background_test": int((y_test == 0).sum()),
+            "n_signal_calibration": len(calibration_idx),
+            "fallback_policy": "held_out_model_only; inspect support before applying",
+        }
 
     return model_configs
 
 
-def _class_zenith_balance_weights(x_train, y_train):
-    """Compute sample weights that equalize class distributions over ze_bin."""
+def _classification_split_indices(y_data, groups, train_fraction, random_state, grouped=True):
+    """Create class-stratified train/validation/test indices.
+
+    Grouping is attempted only when every class has at least six source
+    groups.  Sparse VERITAS lists commonly contain one file per class, so the
+    deterministic event-level fallback is intentional and recorded in model
+    metadata rather than pretending that grouping was achieved.
+    """
+    if not isinstance(y_data, pd.Series):
+        y_data = pd.Series(y_data)
+    if not 0.0 < train_fraction < 1.0:
+        raise ValueError("train_test_fraction must be between zero and one.")
+    rng = random_state
+    if groups is not None and not isinstance(groups, pd.Series):
+        groups = pd.Series(groups, index=y_data.index)
+    use_groups = grouped and groups is not None and groups.notna().all()
+    if use_groups:
+        use_groups = all(groups[y_data == label].nunique() >= 6 for label in y_data.unique())
+
+    train, validation, test = [], [], []
+    if use_groups:
+        for label in sorted(y_data.unique()):
+            label_mask = y_data.to_numpy() == label
+            label_groups = np.asarray(groups[label_mask].unique())
+            g_train, g_hold = train_test_split(
+                label_groups,
+                train_size=train_fraction,
+                random_state=rng,
+            )
+            g_validation, g_test = train_test_split(
+                g_hold,
+                train_size=0.5,
+                random_state=rng,
+            )
+            train.extend(np.flatnonzero(label_mask & groups.isin(g_train).to_numpy()))
+            validation.extend(np.flatnonzero(label_mask & groups.isin(g_validation).to_numpy()))
+            test.extend(np.flatnonzero(label_mask & groups.isin(g_test).to_numpy()))
+        method = "grouped_source_file"
+    else:
+        for label in sorted(y_data.unique()):
+            label_idx = np.flatnonzero(y_data.to_numpy() == label)
+            label_train, label_hold = train_test_split(
+                label_idx,
+                train_size=train_fraction,
+                random_state=rng,
+            )
+            label_validation, label_test = train_test_split(
+                label_hold,
+                train_size=0.5,
+                random_state=rng,
+            )
+            train.extend(label_train)
+            validation.extend(label_validation)
+            test.extend(label_test)
+        method = "stratified_event_fallback"
+
+    return (
+        np.asarray(sorted(train), dtype=int),
+        np.asarray(sorted(validation), dtype=int),
+        np.asarray(sorted(test), dtype=int),
+        {
+            "method": method,
+            "grouped_requested": bool(grouped),
+            "source_groups_available": bool(use_groups),
+        },
+    )
+
+
+def _classification_nuisance_diagnostics(df, labels):
+    """Measure separability of routing/activity proxies without serializing a model."""
+    candidates = {}
+    if "ze_bin" in df:
+        candidates["ze_bin"] = df["ze_bin"]
+    activity = [column for column in df.columns if column.startswith("tel_active_")]
+    if activity:
+        candidates["tel_active_count"] = df[activity].sum(axis=1, skipna=True)
+    telescope_columns = [
+        column
+        for column in df.columns
+        if column.endswith(tuple(f"_{index}" for index in range(64)))
+    ]
+    if telescope_columns:
+        candidates["feature_missing_fraction"] = df[telescope_columns].isna().mean(axis=1)
+    diagnostics = {}
+    for name, values in candidates.items():
+        numeric = pd.to_numeric(values, errors="coerce")
+        valid = numeric.notna() & labels.notna()
+        if valid.sum() < 4 or labels[valid].nunique() < 2 or numeric[valid].nunique() < 2:
+            diagnostics[name] = {"auc": np.nan, "n": int(valid.sum())}
+            continue
+        auc = float(roc_auc_score(labels[valid], numeric[valid]))
+        diagnostics[name] = {
+            "auc": auc,
+            "n": int(valid.sum()),
+            "shortcut_strength": max(auc, 1.0 - auc),
+        }
+    return diagnostics
+
+
+def _class_zenith_balance_weights(x_train, y_train, weight_cap=10.0, smoothing=1.0):
+    """Compute capped, smoothed weights equalizing class distributions over ``ze_bin``.
+
+    ``smoothing`` prevents a single sparse background bin from receiving an
+    arbitrarily large weight.  The cap is an explicit robustness guard for the
+    sparse-background regime common in VERITAS training lists.
+    """
     if "ze_bin" not in x_train.columns:
         raise ValueError(
             "Cannot apply class/zenith balancing because training features do not include ze_bin."
@@ -1118,7 +1364,9 @@ def _class_zenith_balance_weights(x_train, y_train):
     if total_valid == 0:
         raise ValueError("Cannot apply class/zenith balancing with no valid training events.")
 
-    target_fraction = ze_valid.value_counts(normalize=True).sort_index()
+    all_ze = np.sort(ze_valid.unique())
+    target_counts = ze_valid.value_counts().reindex(all_ze, fill_value=0).astype(float)
+    target_fraction = target_counts / target_counts.sum()
     weights = pd.Series(1.0, index=x_train.index, dtype=np.float64)
 
     _logger.info("Class/zenith balancing target distribution:")
@@ -1128,16 +1376,19 @@ def _class_zenith_balance_weights(x_train, y_train):
     for label in sorted(labels_valid.unique()):
         class_mask = labels_valid == label
         class_ze = ze_valid[class_mask]
-        observed_fraction = class_ze.value_counts(normalize=True).sort_index()
         _logger.info(f"Class/zenith balancing weights for label={label}:")
 
         for ze_bin, target_frac in target_fraction.items():
-            obs_frac = observed_fraction.get(ze_bin, 0.0)
-            if obs_frac <= 0:
-                _logger.info(f"  ze_bin={ze_bin}: no events for this class; no weight assigned")
-                continue
-
-            weight = target_frac / obs_frac
+            obs_count = float((class_ze == ze_bin).sum())
+            class_total = float(len(class_ze))
+            if obs_count == 0:
+                # There is no unbiased within-class estimate for an absent
+                # bin.  Give it a finite pseudo-count, then cap the resulting
+                # weight; events in absent bins remain at weight one.
+                obs_frac = smoothing / (class_total + smoothing * len(all_ze))
+            else:
+                obs_frac = obs_count / class_total
+            weight = min(float(target_frac / obs_frac), float(weight_cap))
             mask = valid & (labels == label) & (ze_bins == ze_bin)
             weights.loc[mask] = weight
             _logger.info(
@@ -1145,7 +1396,18 @@ def _class_zenith_balance_weights(x_train, y_train):
                 f"weight={weight:.6f}, events={int(mask.sum())}"
             )
 
-    weight_values = weights.to_numpy(dtype=np.float32)
+    # Equalize total influence of the two classes as well as their zenith
+    # shapes.  This prevents a large simulated signal sample from dominating
+    # a sparse background sample even when raw event counts differ.
+    class_labels = sorted(labels_valid.unique())
+    target_class_total = total_valid / len(class_labels)
+    for label in class_labels:
+        class_mask = valid & (labels == label)
+        class_sum = float(weights.loc[class_mask].sum())
+        if class_sum > 0:
+            weights.loc[class_mask] *= target_class_total / class_sum
+
+    weight_values = np.clip(weights.to_numpy(dtype=np.float32), 0.0, float(weight_cap))
     mean_weight = weight_values.mean()
     if mean_weight > 0:
         weight_values /= mean_weight
diff --git a/tests/test_classification_robustness.py b/tests/test_classification_robustness.py
new file mode 100644
index 0000000..5094c83
--- /dev/null
+++ b/tests/test_classification_robustness.py
@@ -0,0 +1,69 @@
+"""Focused tests for the code-only classification hardening contract."""
+
+import numpy as np
+import pandas as pd
+
+from eventdisplay_ml import features, models
+from eventdisplay_ml.evaluate import (
+    _efficiency_dataframe,
+    classification_thresholds_from_signal,
+)
+
+
+def test_robust_profile_excludes_routing_and_activity_columns():
+    columns = [
+        "MSCW",
+        "MSCL",
+        "EChi2S",
+        "EmissionHeight",
+        "EmissionHeightChi2",
+        "Core_Distance",
+        "size_0",
+        "width_0",
+        "length_0",
+        "tel_active_0",
+        "ze_bin",
+        "Erec",
+        "__source_file",
+    ]
+    assert features.classification_feature_columns(columns) == [
+        "MSCW",
+        "MSCL",
+        "EChi2S",
+        "EmissionHeight",
+        "EmissionHeightChi2",
+        "Core_Distance",
+        "size_0",
+        "width_0",
+        "length_0",
+        "ze_bin",
+    ]
+
+
+def test_extended_profile_retains_zenith_but_not_provenance():
+    columns = ["MSCW", "ze_bin", "Erec", "__source_file"]
+    assert features.classification_feature_columns(columns, profile="extended") == [
+        "MSCW",
+        "ze_bin",
+    ]
+
+
+def test_grouped_split_keeps_source_groups_disjoint_when_supported():
+    y = pd.Series(np.repeat([0, 1], 60))
+    groups = pd.Series(np.tile(np.arange(6), 20))
+    train, validation, test, metadata = models._classification_split_indices(
+        y, groups, train_fraction=0.5, random_state=7, grouped=True
+    )
+    assert metadata["method"] == "grouped_source_file"
+    assert set(groups.iloc[train]).isdisjoint(groups.iloc[validation])
+    assert set(groups.iloc[train]).isdisjoint(groups.iloc[test])
+    assert set(groups.iloc[validation]).isdisjoint(groups.iloc[test])
+
+
+def test_threshold_calibration_is_quantile_based_and_background_limit_nonzero():
+    calibration = classification_thresholds_from_signal(np.linspace(0.1, 0.9, 9))
+    assert calibration["threshold"].is_monotonic_decreasing
+    efficiency = _efficiency_dataframe(
+        "test", np.array([0.9, 0.1]), np.array([1, 0]), np.array([0.5])
+    )
+    assert efficiency.loc[0, "background_efficiency_upper95"] > 0

From 8b9481177cfce613ab6373463996390d42f494a4 Mon Sep 17 00:00:00 2001
From: GernotMaier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 14:51:00 +0200
Subject: [PATCH 02/10] unit tests

---
 src/eventdisplay_ml/data_processing.py | 144 ++++++++++++++++---------
 src/eventdisplay_ml/evaluate.py        |   4 +-
 src/eventdisplay_ml/features.py        |   3 +-
 src/eventdisplay_ml/models.py          |   7 ++
 4 files changed, 104 insertions(+), 54 deletions(-)

diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py
index a54f5fc..07ddb15 100644
--- a/src/eventdisplay_ml/data_processing.py
+++ b/src/eventdisplay_ml/data_processing.py
@@ -337,7 +337,7 @@ def _normalize_telescope_variable_to_tel_id_space(data, index_list, max_tel_id,
     row_indices, col_indices = np.where(~np.isnan(index_list))
     tel_ids = index_list[row_indices, col_indices].astype(int)
     # Filter for valid telescope IDs and valid column indices in data array
-    valid_mask = (tel_ids >= 0) & (tel_ids <= max_tel_id) & (col_indices < data.shape[1])
+    valid_mask = (tel_ids <= max_tel_id) & (col_indices < data.shape[1])
     full_matrix[row_indices[valid_mask], tel_ids[valid_mask]] = data[
         row_indices[valid_mask], col_indices[valid_mask]
     ]
@@ -468,6 +468,7 @@ def flatten_telescope_data_vectorized(
         Flattened DataFrame with per-telescope columns suffixed by ``_{i}``
     """
     flat_features = {}
+    classification_mode = analysis_type == "classification"
     tel_list_matrix = _to_dense_array(df["DispTelList_T"])
     n_evt = len(df)
     max_tel_id = tel_config["max_tel_id"] if tel_config else (n_tel - 1)
@@ -483,7 +484,7 @@ def flatten_telescope_data_vectorized(
     active_mask = np.zeros((n_evt, max_tel_id + 1), dtype=bool)
     row_indices, col_indices = np.where(~np.isnan(tel_list_matrix))
     tel_ids = tel_list_matrix[row_indices, col_indices].astype(int)
-    valid_tel_mask = (tel_ids >= 0) & (tel_ids <= max_tel_id)
+    valid_tel_mask = tel_ids <= max_tel_id
     active_mask[row_indices[valid_tel_mask], tel_ids[valid_tel_mask]] = True
 
     # Pre-load and normalize size to telescope-ID space for sorting
@@ -493,13 +494,27 @@ def flatten_telescope_data_vectorized(
     # A telescope absent from DispTelList_T is not a zero-sized image.  Keep it
     # explicitly missing so sorting and the XGBoost missing-value path cannot
     # learn a detector-slot/observing-condition proxy.
-    size_data = np.where(active_mask, size_data, np.nan)
+    if classification_mode:
+        size_data = np.where(active_mask, size_data, np.nan)
     size_data = _clip_size_array(size_data)
 
     core_x, core_y = _get_core_arrays(df)
 
     # Sorting by mirror area (desc; proxy for telescope type), then size (desc)
-    sort_indices = _compute_size_area_sort_indices(size_data, active_mask, tel_config, max_tel_id)
+    if classification_mode:
+        sort_indices = _compute_size_area_sort_indices(
+            size_data,
+            active_mask,
+            tel_config,
+            max_tel_id,
+            active_first=True,
+        )
+    else:
+        # Keep the historical call signature for regression callers and
+        # downstream monkeypatches.
+        sort_indices = _compute_size_area_sort_indices(
+            size_data, active_mask, tel_config, max_tel_id
+        )
 
     # Determine which telescope positions to keep (for feature reduction)
     if max_tel_per_type is not None and tel_config is not None:
@@ -540,14 +555,15 @@ def flatten_telescope_data_vectorized(
             )
             # Geometry is useful only for an active image; inactive slots must
             # not become a stable class/domain indicator.
-            for key in tuple(flat_features):
-                if key.startswith(f"{var}_"):
-                    sorted_tel = int(key.rsplit("_", 1)[1])
-                    flat_features[key] = np.where(
-                        active_mask[np.arange(n_evt), sort_indices[:, sorted_tel]],
-                        flat_features[key],
-                        np.nan,
-                    )
+            if classification_mode:
+                for key in tuple(flat_features):
+                    if key.startswith(f"{var}_"):
+                        sorted_tel = int(key.rsplit("_", 1)[1])
+                        flat_features[key] = np.where(
+                            active_mask[np.arange(n_evt), sort_indices[:, sorted_tel]],
+                            flat_features[key],
+                            np.nan,
+                        )
             continue
 
         data = _to_dense_array(df[var]) if _has_field(df, var) else np.full((n_evt, n_tel), np.nan)
@@ -564,7 +580,7 @@ def flatten_telescope_data_vectorized(
                 data, index_list_for_remapping, max_tel_id, n_evt
             )
 
-        if var != "tel_active":
+        if classification_mode and var != "tel_active":
             data_normalized = np.where(active_mask, data_normalized, np.nan)
 
         # All variables are now in telescope-ID space; apply sorting and flatten uniformly
@@ -651,7 +667,9 @@ def _get_core_arrays(df):
     return core_x, core_y
 
 
-def _compute_size_area_sort_indices(size_data, active_mask, tel_config, max_tel_id):
+def _compute_size_area_sort_indices(
+    size_data, active_mask, tel_config, max_tel_id, active_first=False
+):
     """Compute sorting indices: mirror area (desc) then size (desc).
 
     Missing telescopes (NaN size or no mirror area) are sorted to the end.
@@ -694,14 +712,20 @@ def _compute_size_area_sort_indices(size_data, active_mask, tel_config, max_tel_
             size_valid = 0 if not np.isnan(size_val) else 1
             area_key = -area if area_valid == 0 else 0.0
             size_key = -size_val if size_valid == 0 else 0.0
-            # Active images always precede inactive detector slots.  This keeps
-            # slot ordering deterministic without turning missing telescopes
-            # into a large/small-image classification feature.
-            tel_entries.append(
-                (tel_idx, 0 if active else 1, area_valid, area_key, size_valid, size_key)
-            )
+            if active_first:
+                # Classification-only safeguard: active images precede
+                # inactive detector slots without exposing slot activity.
+                tel_entries.append(
+                    (tel_idx, 0 if active else 1, area_valid, area_key, size_valid, size_key)
+                )
+            else:
+                # Preserve the historical regression ordering exactly.
+                tel_entries.append((tel_idx, area_valid, area_key, size_valid, size_key))
 
-        tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4], x[5]))
+        if active_first:
+            tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4], x[5]))
+        else:
+            tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4]))
         sort_indices[evt_idx] = np.array([t[0] for t in tel_entries])
 
     return sort_indices
@@ -968,6 +992,7 @@ def load_training_data(model_configs, file_list, analysis_type):
     pandas.DataFrame
         Flattened DataFrame ready for training.
     """
+    classification_mode = analysis_type == "classification"
     max_events = model_configs.get("max_events", None)
     random_state = model_configs.get("random_state", None)
     memory_profile = model_configs.get("memory_profile", False)
@@ -983,7 +1008,7 @@ def load_training_data(model_configs, file_list, analysis_type):
         _logger.info(f"Adding zenith binning: {model_configs.get('zenith_bins_deg', [])}")
 
     input_files = utils.read_input_file_list(file_list)
-    if not input_files:
+    if classification_mode and not input_files:
         raise ValueError(f"Input file list is empty: {file_list}")
 
     tmva_style = model_configs.get("tmva_style", False)
@@ -998,17 +1023,22 @@ def load_training_data(model_configs, file_list, analysis_type):
     else:
         branch_list = features_module.features(analysis_type, training=True)
     _logger.info(f"Branch list: {branch_list}")
-    if max_events is not None and max_events > 0:
+    if classification_mode and max_events is not None and max_events > 0:
         # Reserve a bounded quota per file, then perform one deterministic
         # final sample below.  Integer floor division used to turn a small
         # global cap into zero (which silently disabled sampling).
         max_events_per_file = max(1, int(np.ceil(max_events / len(input_files))))
     else:
-        max_events_per_file = None
+        if max_events is not None and max_events > 0:
+            # Preserve the historical regression quota behavior.
+            max_events_per_file = max_events // len(input_files)
+        else:
+            max_events_per_file = None
     _logger.info(f"Max events per file: {max_events_per_file}")
 
-    # Reuse/validate across signal and background loads.
-    tel_config = model_configs.get("tel_config")
+    # Classification reuses/validates configuration across signal/background;
+    # regression retains its historical first-file initialization.
+    tel_config = model_configs.get("tel_config") if classification_mode else None
     dfs = []
     executor = ThreadPoolExecutor(max_workers=model_configs.get("max_cores", 1))
     total_files = len(input_files)
@@ -1024,25 +1054,31 @@ def load_training_data(model_configs, file_list, analysis_type):
                     tel_config = current_tel_config
                     model_configs["tel_config"] = tel_config
                 else:
-                    # A model cannot have a stable feature schema if telescope
-                    # IDs/areas change between input files.  The old code
-                    # silently replaced the configuration when max_tel_id grew.
-                    def _config_signature(config):
-                        return (
-                            int(config["max_tel_id"]),
-                            tuple(int(v) for v in config.get("tel_ids", [])),
-                            tuple(str(v) for v in config.get("tel_types", [])),
-                            tuple(
-                                float(v)
-                                for v in config.get("mirror_area", config.get("mirror_areas", []))
-                            ),
-                        )
-
-                    if _config_signature(current_tel_config) != _config_signature(tel_config):
-                        raise ValueError(
-                            "Classification/training input files have incompatible telescope "
-                            f"configurations: {input_files[0]} versus {f}."
-                        )
+                    if classification_mode:
+                        # A classification model cannot have a stable feature
+                        # schema if telescope IDs/areas change between files.
+                        def _config_signature(config):
+                            return (
+                                int(config["max_tel_id"]),
+                                tuple(int(v) for v in config.get("tel_ids", [])),
+                                tuple(str(v) for v in config.get("tel_types", [])),
+                                tuple(
+                                    float(v)
+                                    for v in config.get(
+                                        "mirror_area", config.get("mirror_areas", [])
+                                    )
+                                ),
+                            )
+
+                        if _config_signature(current_tel_config) != _config_signature(tel_config):
+                            raise ValueError(
+                                "Classification/training input files have incompatible "
+                                f"telescope configurations: {input_files[0]} versus {f}."
+                            )
+                    elif current_tel_config["max_tel_id"] > tel_config["max_tel_id"]:
+                        # Preserve the historical regression behavior.
+                        tel_config = current_tel_config
+                        model_configs["tel_config"] = tel_config
 
                 _logger.info(f"Processing file: {f} (file {file_idx}/{total_files})")
                 tree = root_file["data"]
@@ -1053,8 +1089,12 @@ def _config_signature(config):
                 raw_reservoir_chunks = []
                 reservoir_priorities = None
                 file_dfs = []
-                rng = np.random.default_rng(
-                    None if random_state is None else int(random_state) + file_idx - 1
+                rng = (
+                    np.random.default_rng(
+                        None if random_state is None else int(random_state) + file_idx - 1
+                    )
+                    if classification_mode
+                    else np.random.default_rng(random_state)
                 )
                 chunk_iterator = tree.iterate(
                     resolved_branch_list,
@@ -1167,12 +1207,14 @@ def _config_signature(config):
                     file_df,
                     enabled=memory_profile,
                 )
-        except (FileNotFoundError, KeyError, ValueError):
-            raise
         except Exception as e:
-            raise RuntimeError(f"Error opening or reading file {f}: {e}") from e
+            if classification_mode:
+                if isinstance(e, (FileNotFoundError, KeyError, ValueError)):
+                    raise
+                raise RuntimeError(f"Error opening or reading file {f}: {e}") from e
+            raise FileNotFoundError(f"Error opening or reading file {f}: {e}") from e
 
-    if not dfs:
+    if classification_mode and not dfs:
         raise ValueError("No data loaded from input files.")
     df_final = pd.concat(dfs, ignore_index=True)
     if analysis_type == "classification" and max_events is not None and max_events > 0:
diff --git a/src/eventdisplay_ml/evaluate.py b/src/eventdisplay_ml/evaluate.py
index 116a914..554ab1a 100644
--- a/src/eventdisplay_ml/evaluate.py
+++ b/src/eventdisplay_ml/evaluate.py
@@ -26,9 +26,9 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
 
     for t in thresholds:
         pred = y_pred_proba >= t
-        eff_signal.append(((pred) & (y_test == 1)).sum() / n_signal if n_signal else np.nan)
+        eff_signal.append(((pred) & (y_test == 1)).sum() / n_signal if n_signal else 0.0)
         eff_background.append(
-            ((pred) & (y_test == 0)).sum() / n_background if n_background else np.nan
+            ((pred) & (y_test == 0)).sum() / n_background if n_background else 0.0
         )
         _logger.info(
             f"{name}{context_label} Threshold: {t:.2f} | "
diff --git a/src/eventdisplay_ml/features.py b/src/eventdisplay_ml/features.py
index 7757276..a23a437 100644
--- a/src/eventdisplay_ml/features.py
+++ b/src/eventdisplay_ml/features.py
@@ -104,7 +104,8 @@ def classification_feature_columns(columns, profile="robust", ignore_ze_bin=Fals
         # quantities) should still be usable in isolated unit tests.  Real
         # flattened frames contain at least one physics/activity name; fail
         # loudly instead of silently falling back to nuisance columns there.
-        if not selected:
+        morphology_selected = [name for name in selected if name != "ze_bin"]
+        if not morphology_selected:
             looks_like_flattened_physics = any(
                 name.startswith(("tel_", "ArrayPointing", "Xcore", "Ycore"))
                 or name in {"DispNImages", "Erec", "size"}
diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py
index edb303a..1496d69 100644
--- a/src/eventdisplay_ml/models.py
+++ b/src/eventdisplay_ml/models.py
@@ -1,6 +1,7 @@
 """Apply models for regression and classification tasks."""
 
 import logging
+import os
 import re
 import subprocess
 import sys
@@ -51,6 +52,11 @@ def _validate_saved_model(model_path):
         str(model_path),
         str(_MODEL_VALIDATION_MEMORY_BYTES),
     ]
+    validation_env = os.environ.copy()
+    # The validator applies a memory limit; unrestricted BLAS thread pools can
+    # allocate one workspace per host CPU and fail before the model is read.
+    for variable in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
+        validation_env[variable] = "1"
     try:
         result = subprocess.run(
             command,
@@ -58,6 +64,7 @@ def _validate_saved_model(model_path):
             check=False,
             text=True,
             timeout=_MODEL_VALIDATION_TIMEOUT_SECONDS,
+            env=validation_env,
         )
     except subprocess.TimeoutExpired as exc:
         raise RuntimeError(

From 3e78729a30bfc90a7bf12b52f9f3b62db44767c5 Mon Sep 17 00:00:00 2001
From: GernotMaier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 14:53:29 +0200
Subject: [PATCH 03/10] changelog

---
 docs/changes/79.bugfix.md | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 docs/changes/79.bugfix.md

diff --git a/docs/changes/79.bugfix.md b/docs/changes/79.bugfix.md
new file mode 100644
index 0000000..7d240fa
--- /dev/null
+++ b/docs/changes/79.bugfix.md
@@ -0,0 +1 @@
+Harden VERITAS gamma/hadron classification with robust image features, provenance-aware validation, sparse-background weighting, and reliable score calibration.

From fec6a2231b22e126a1a5f3185898dc3082ecc1ea Mon Sep 17 00:00:00 2001
From: GernotMaier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 15:02:50 +0200
Subject: [PATCH 04/10] copilot review

---
 src/eventdisplay_ml/evaluate.py         | 14 +++++++-------
 src/eventdisplay_ml/features.py         | 24 +++++++++++++++++++++++-
 tests/test_classification_robustness.py |  9 ++++++++-
 3 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/src/eventdisplay_ml/evaluate.py b/src/eventdisplay_ml/evaluate.py
index 554ab1a..fe4154d 100644
--- a/src/eventdisplay_ml/evaluate.py
+++ b/src/eventdisplay_ml/evaluate.py
@@ -23,13 +23,15 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
 
     eff_signal = []
     eff_background = []
+    background_survivor_counts = []
 
     for t in thresholds:
         pred = y_pred_proba >= t
-        eff_signal.append(((pred) & (y_test == 1)).sum() / n_signal if n_signal else 0.0)
-        eff_background.append(
-            ((pred) & (y_test == 0)).sum() / n_background if n_background else 0.0
-        )
+        n_signal_survivors = int(((pred) & (y_test == 1)).sum())
+        n_background_survivors = int(((pred) & (y_test == 0)).sum())
+        background_survivor_counts.append(n_background_survivors)
+        eff_signal.append(n_signal_survivors / n_signal if n_signal else 0.0)
+        eff_background.append(n_background_survivors / n_background if n_background else 0.0)
         _logger.info(
             f"{name}{context_label} Threshold: {t:.2f} | "
             f"Signal Efficiency: {eff_signal[-1]:.4f} | "
@@ -38,11 +40,9 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
 
     eff_signal = np.asarray(eff_signal, dtype=float)
     eff_background = np.asarray(eff_background, dtype=float)
-    background_survivors = n_background * eff_background
     background_upper_limit = np.full(len(thresholds), np.nan, dtype=float)
     if n_background:
-        for i, survivors in enumerate(background_survivors):
-            k = round(survivors)
+        for i, k in enumerate(background_survivor_counts):
             background_upper_limit[i] = (
                 1.0 if k >= n_background else beta_distribution.ppf(0.95, k + 1, n_background - k)
             )
diff --git a/src/eventdisplay_ml/features.py b/src/eventdisplay_ml/features.py
index a23a437..3c9dba4 100644
--- a/src/eventdisplay_ml/features.py
+++ b/src/eventdisplay_ml/features.py
@@ -66,8 +66,30 @@ def classification_feature_columns(columns, profile="robust", ignore_ze_bin=Fals
         "__source_file",
     }
     available = list(columns)
+    routing_or_activity = {
+        "label",
+        "Erec",
+        "ErecS",
+        "__source_file_id",
+        "__source_row",
+        "__source_file",
+    }
+
+    def is_excluded_routing_or_activity(name):
+        return name in routing_or_activity or name.startswith(
+            (
+                "__",
+                "tel_active_",
+                "mirror_area_",
+                "tel_rel_x_",
+                "tel_rel_y_",
+                "fpointing_dx_",
+                "fpointing_dy_",
+            )
+        )
+
     if profile == "extended":
-        selected = [name for name in available if name not in reserved]
+        selected = [name for name in available if not is_excluded_routing_or_activity(name)]
     else:
         # Array/stereo quantities plus per-telescope image morphology.  The
         # latter are essential gamma/hadron information; detector activity,
diff --git a/tests/test_classification_robustness.py b/tests/test_classification_robustness.py
index 5094c83..944633d 100644
--- a/tests/test_classification_robustness.py
+++ b/tests/test_classification_robustness.py
@@ -41,7 +41,14 @@ def test_robust_profile_excludes_routing_and_activity_columns():
 
 
 def test_extended_profile_retains_zenith_but_not_provenance():
-    columns = ["MSCW", "ze_bin", "Erec", "__source_file"]
+    columns = [
+        "MSCW",
+        "ze_bin",
+        "Erec",
+        "tel_active_0",
+        "mirror_area_0",
+        "__source_file",
+    ]
     assert features.classification_feature_columns(columns, profile="extended") == [
         "MSCW",
         "ze_bin",

From f022ccd9323d4dc31eaaf9f861914bcf9460c2ef Mon Sep 17 00:00:00 2001
From: GernotMaier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 15:17:49 +0200
Subject: [PATCH 05/10] models

---
 src/eventdisplay_ml/models.py           | 38 +++++++++++++++++++++----
 tests/test_classification_robustness.py | 29 +++++++++++++++++++
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py
index 1496d69..8a6b9cc 100644
--- a/src/eventdisplay_ml/models.py
+++ b/src/eventdisplay_ml/models.py
@@ -1268,11 +1268,27 @@ def _classification_split_indices(y_data, groups, train_fraction, random_state,
         for label in sorted(y_data.unique()):
             label_mask = y_data.to_numpy() == label
             label_groups = np.asarray(groups[label_mask].unique())
+            n_train_groups = int(np.ceil(len(label_groups) * train_fraction))
+            n_hold_groups = len(label_groups) - n_train_groups
+            if n_train_groups < 1 or n_hold_groups < 2:
+                raise ValueError(
+                    "Grouped classification split cannot create separate validation "
+                    "and test groups for label="
+                    f"{label}: train_test_fraction={train_fraction} leaves "
+                    f"{n_hold_groups} holdout groups. Reduce train_test_fraction or "
+                    "provide more source groups."
+                )
             g_train, g_hold = train_test_split(
                 label_groups,
                 train_size=train_fraction,
                 random_state=rng,
             )
+            if len(g_hold) < 2:
+                raise ValueError(
+                    "Grouped classification split cannot create separate validation "
+                    f"and test groups for label={label}: only {len(g_hold)} holdout "
+                    "groups remain after the training split."
+                )
             g_validation, g_test = train_test_split(
                 g_hold,
                 train_size=0.5,
@@ -1285,11 +1301,27 @@ def _classification_split_indices(y_data, groups, train_fraction, random_state,
     else:
         for label in sorted(y_data.unique()):
             label_idx = np.flatnonzero(y_data.to_numpy() == label)
+            n_train_events = int(np.ceil(len(label_idx) * train_fraction))
+            n_hold_events = len(label_idx) - n_train_events
+            if n_train_events < 1 or n_hold_events < 2:
+                raise ValueError(
+                    "Classification split cannot create separate validation and test "
+                    "events for label="
+                    f"{label}: train_test_fraction={train_fraction} leaves "
+                    f"{n_hold_events} holdout events. Reduce train_test_fraction or "
+                    "provide more events."
+                )
             label_train, label_hold = train_test_split(
                 label_idx,
                 train_size=train_fraction,
                 random_state=rng,
             )
+            if len(label_hold) < 2:
+                raise ValueError(
+                    "Classification split cannot create separate validation and test "
+                    f"events for label={label}: only {len(label_hold)} holdout events "
+                    "remain after the training split."
+                )
             label_validation, label_test = train_test_split(
                 label_hold,
                 train_size=0.5,
@@ -1320,11 +1352,7 @@ def _classification_nuisance_diagnostics(df, labels):
     activity = [column for column in df.columns if column.startswith("tel_active_")]
     if activity:
         candidates["tel_active_count"] = df[activity].sum(axis=1, skipna=True)
-    telescope_columns = [
-        column
-        for column in df.columns
-        if column.endswith(tuple(f"_{index}" for index in range(64)))
-    ]
+    telescope_columns = [column for column in df.columns if re.search(r"_\d+$", str(column))]
     if telescope_columns:
         candidates["feature_missing_fraction"] = df[telescope_columns].isna().mean(axis=1)
     diagnostics = {}
diff --git a/tests/test_classification_robustness.py b/tests/test_classification_robustness.py
index 944633d..93a0ce1 100644
--- a/tests/test_classification_robustness.py
+++ b/tests/test_classification_robustness.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import pandas as pd
+import pytest
 
 from eventdisplay_ml import features, models
 from eventdisplay_ml.evaluate import (
@@ -67,6 +68,34 @@ def test_grouped_split_keeps_source_groups_disjoint_when_supported():
     assert set(groups.iloc[validation]).isdisjoint(groups.iloc[test])
 
 
+def test_event_split_reports_insufficient_holdout_events():
+    y = pd.Series([0, 0, 1, 1])
+    with pytest.raises(ValueError, match="holdout events"):
+        models._classification_split_indices(
+            y, None, train_fraction=0.5, random_state=7, grouped=False
+        )
+
+
+def test_grouped_split_reports_insufficient_holdout_groups():
+    y = pd.Series(np.repeat([0, 1], 60))
+    groups = pd.Series(np.tile(np.arange(6), 20))
+    with pytest.raises(ValueError, match="holdout groups"):
+        models._classification_split_indices(
+            y, groups, train_fraction=0.9, random_state=7, grouped=True
+        )
+
+
+def test_nuisance_diagnostics_handles_telescope_ids_above_63():
+    frame = pd.DataFrame(
+        {
+            "size_64": [1.0, np.nan, 1.0, np.nan],
+            "width_128": [0.1, 0.2, np.nan, np.nan],
+        }
+    )
+    diagnostics = models._classification_nuisance_diagnostics(frame, pd.Series([0, 0, 1, 1]))
+    assert diagnostics["feature_missing_fraction"]["n"] == 4
+
+
 def test_threshold_calibration_is_quantile_based_and_background_limit_nonzero():
     calibration = classification_thresholds_from_signal(np.linspace(0.1, 0.9, 9))
     assert calibration["threshold"].is_monotonic_decreasing

From bd2204851561c94b2b992cd4ee6b85acdccffc89 Mon Sep 17 00:00:00 2001
From: GernotMaier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 15:39:26 +0200
Subject: [PATCH 06/10] copilot review

---
 src/eventdisplay_ml/data_processing.py        |  16 +-
 src/eventdisplay_ml/models.py                 | 199 ++++++++++++++++--
 ...test_classification_apply_interpolation.py |  65 ++++++
 tests/test_data_processing.py                 |   8 +-
 tests/test_train_classification_shap.py       |  42 ++++
 5 files changed, 306 insertions(+), 24 deletions(-)

diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py
index 07ddb15..f24f8be 100644
--- a/src/eventdisplay_ml/data_processing.py
+++ b/src/eventdisplay_ml/data_processing.py
@@ -1556,7 +1556,12 @@ def extra_columns(df, analysis_type, training, index, tel_config=None, observato
 
 
 def zenith_in_bins(zenith_angles, bins):
-    """Apply zenith binning based on zenith angles and given bin edges."""
+    """Apply zenith binning, marking out-of-range angles with ``-1``.
+
+    The final edge is included in the last bin.  Angles below the first edge,
+    above the final edge, or non-finite angles are invalid and receive ``-1``;
+    they are never silently assigned to an edge bin.
+    """
     if bins is None or len(bins) < 2:
         raise ValueError("At least two zenith-bin edges are required.")
     if isinstance(bins[0], dict):
@@ -1568,7 +1573,14 @@ def zenith_in_bins(zenith_angles, bins):
         raise ValueError("Zenith-bin edges must be a finite one-dimensional sequence.")
     if np.any(np.diff(bins) <= 0):
         raise ValueError("Zenith-bin edges must be strictly increasing.")
-    idx = np.clip(np.digitize(zenith_angles, bins) - 1, 0, len(bins) - 2)
+    zenith_angles = np.asarray(zenith_angles, dtype=float)
+    idx = np.full(zenith_angles.shape, -1, dtype=np.int32)
+    valid = np.isfinite(zenith_angles) & (zenith_angles >= bins[0]) & (zenith_angles <= bins[-1])
+    if np.any(valid):
+        idx[valid] = np.minimum(
+            np.digitize(zenith_angles[valid], bins) - 1,
+            len(bins) - 2,
+        )
     return idx.astype(np.int32)
 
 
diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py
index 8a6b9cc..c02916f 100644
--- a/src/eventdisplay_ml/models.py
+++ b/src/eventdisplay_ml/models.py
@@ -214,6 +214,9 @@ def load_classification_models(model_prefix, model_name):
             model_data.get("energy_bins_log10_tev"),
             file,
         )
+        models[e_bin]["energy_center"] = 0.5 * (
+            float(energy_bin_metadata["E_min"]) + float(energy_bin_metadata["E_max"])
+        )
         par = _update_parameters(
             par,
             model_data.get("zenith_bins_deg"),
@@ -323,7 +326,20 @@ def _validate_energy_bin_metadata(energy_bin, model_file):
             f"missing required key(s): {missing}."
         )
 
-    return energy_bin
+    try:
+        e_min = float(energy_bin["E_min"])
+        e_max = float(energy_bin["E_max"])
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            "Classification model file "
+            f"'{model_file}' has non-numeric energy-bin metadata for 'E_min'/'E_max'."
+        ) from exc
+    if not np.isfinite(e_min) or not np.isfinite(e_max) or e_min >= e_max:
+        raise ValueError(
+            "Classification model file "
+            f"'{model_file}' has invalid energy-bin metadata: require finite E_min < E_max."
+        )
+    return {"E_min": e_min, "E_max": e_max}
 
 
 def _update_parameters(full_params, zenith_bins, energy_bin, e_bin_number):
@@ -532,6 +548,18 @@ def apply_classification_models(df, model_configs, threshold_keys):
         if e_bin_lo == -1 or e_bin_hi == -1:
             _logger.warning("Skipping events with invalid energy interpolation bins")
             continue
+        if "ze_bin" in group_df:
+            zenith_values = pd.to_numeric(group_df["ze_bin"], errors="coerce")
+            valid_zenith = zenith_values.notna() & (zenith_values >= 0)
+            if not valid_zenith.all():
+                _logger.warning(
+                    "Skipping %d events with invalid/out-of-range zenith bins during "
+                    "classification apply",
+                    int((~valid_zenith).sum()),
+                )
+                group_df = group_df.loc[valid_zenith]
+                if group_df.empty:
+                    continue
 
         _logger.info(
             "Processing %d events with interpolation bins (%d, %d)",
@@ -564,11 +592,19 @@ def apply_classification_models(df, model_configs, threshold_keys):
         flatten_hi = flatten_data.loc[:, models[resolved_hi]["features"]]
 
         class_probs_lo = model_lo.predict_proba(flatten_lo)[:, 1]
-        if e_bin_lo == e_bin_hi:
+        interpolate_models = resolved_lo != resolved_hi
+        alpha = _classification_interpolation_alpha(
+            group_df,
+            e_bin_lo,
+            e_bin_hi,
+            resolved_lo,
+            resolved_hi,
+            models,
+        )
+        if not interpolate_models:
             class_probs = class_probs_lo
         else:
             class_probs_hi = model_hi.predict_proba(flatten_hi)[:, 1]
-            alpha = group_df["e_alpha"].to_numpy(dtype=np.float32)
             class_probs = (1.0 - alpha) * class_probs_lo + alpha * class_probs_hi
         class_probability[group_df.index] = class_probs
 
@@ -579,19 +615,70 @@ def apply_classification_models(df, model_configs, threshold_keys):
                 thr_lo = thresholds_lo.get(eff)
                 if thr_lo is None:
                     continue
-                if e_bin_lo == e_bin_hi:
+                if not interpolate_models:
                     threshold = thr_lo
                 else:
                     thr_hi = thresholds_hi.get(eff)
                     if thr_hi is None:
                         continue
-                    alpha = group_df["e_alpha"].to_numpy(dtype=np.float32)
                     threshold = (1.0 - alpha) * thr_lo + alpha * thr_hi
                 is_gamma[eff][group_df.index] = (class_probs >= threshold).astype(np.uint8)
 
     return class_probability, is_gamma
 
 
+def _classification_interpolation_alpha(
+    group_df,
+    requested_lo,
+    requested_hi,
+    resolved_lo,
+    resolved_hi,
+    models,
+):
+    """Return interpolation weights in the energy coordinate of resolved models.
+
+    For complete model grids, the precomputed ``e_alpha`` is already correct.
+    When a requested bin is borrowed, recompute the coordinate from the event
+    energy and the actual model centers so scores and calibrated thresholds use
+    the same models and energy geometry.
+    """
+    if resolved_lo == resolved_hi:
+        return np.zeros(len(group_df), dtype=np.float32)
+
+    fallback_alpha = group_df["e_alpha"].to_numpy(dtype=np.float32)
+    borrowed = (requested_lo, requested_hi) != (resolved_lo, resolved_hi)
+    if not borrowed:
+        return fallback_alpha
+
+    center_lo = models[resolved_lo].get("energy_center")
+    center_hi = models[resolved_hi].get("energy_center")
+    if (
+        center_lo is None
+        or center_hi is None
+        or not np.isfinite(center_lo)
+        or not np.isfinite(center_hi)
+        or center_hi <= center_lo
+    ):
+        raise ValueError(
+            "Cannot interpolate borrowed classification energy-bin models without "
+            "finite energy-center metadata."
+        )
+    if "Erec" not in group_df:
+        raise ValueError(
+            "Cannot interpolate borrowed classification energy-bin models without Erec."
+        )
+
+    erec = pd.to_numeric(group_df["Erec"], errors="coerce").to_numpy(dtype=np.float64)
+    valid_energy = np.isfinite(erec) & (erec > 0.0)
+    if not np.all(valid_energy):
+        raise ValueError(
+            "Cannot interpolate borrowed classification energy-bin models for events "
+            "with non-positive or non-finite Erec."
+        )
+    alpha = (np.log10(erec) - float(center_lo)) / float(center_hi - center_lo)
+    return np.clip(alpha, 0.0, 1.0).astype(np.float32)
+
+
 def _resolve_classification_bin(models, requested_bin):
     """Resolve a missing energy-bin model to the nearest available model."""
     if requested_bin in models:
@@ -1091,6 +1178,14 @@ def train_classification(df, model_configs):
     ze_data = full_df["ze_bin"] if "ze_bin" in full_df.columns else None
     if model_configs.get("balance_class_zenith_weights", False) and ze_data is None:
         raise ValueError("Class/zenith balancing requires the derived ze_bin column.")
+    if ze_data is not None:
+        zenith_values = pd.to_numeric(ze_data, errors="coerce")
+        invalid_zenith = zenith_values.isna() | (zenith_values < 0)
+        if invalid_zenith.any():
+            raise ValueError(
+                "Classification training contains out-of-range or invalid zenith bins: "
+                f"{int(invalid_zenith.sum())} events."
+            )
 
     profile = model_configs.get("feature_profile", "robust")
     feature_columns = features.classification_feature_columns(
@@ -1166,9 +1261,16 @@ def train_classification(df, model_configs):
     weights_train = None
     weights_validation = None
     if model_configs.get("balance_class_zenith_weights", False):
-        weights_train = _class_zenith_balance_weights(full_df.iloc[train_idx], y_train)
+        target_ze_fraction = _class_zenith_target_fraction(full_df.iloc[train_idx])
+        weights_train = _class_zenith_balance_weights(
+            full_df.iloc[train_idx],
+            y_train,
+            target_ze_fraction=target_ze_fraction,
+        )
         weights_validation = _class_zenith_balance_weights(
-            full_df.iloc[validation_idx], y_validation
+            full_df.iloc[validation_idx],
+            y_validation,
+            target_ze_fraction=target_ze_fraction,
         )
         _logger.info(
             "Using class/zenith sample weights "
@@ -1371,12 +1473,71 @@ def _classification_nuisance_diagnostics(df, labels):
     return diagnostics
 
 
-def _class_zenith_balance_weights(x_train, y_train, weight_cap=10.0, smoothing=1.0):
+def _class_zenith_target_fraction(x_train):
+    """Return the fixed zenith target distribution derived from training data."""
+    if "ze_bin" not in x_train.columns:
+        raise ValueError("Cannot derive a zenith target distribution without ze_bin.")
+    ze_bins = pd.to_numeric(x_train["ze_bin"], errors="coerce")
+    valid = ze_bins.notna() & (ze_bins >= 0)
+    if not valid.any():
+        raise ValueError("Cannot derive a zenith target distribution with no valid ze_bin.")
+    counts = ze_bins[valid].value_counts().sort_index().astype(float)
+    return counts / counts.sum()
+
+
+def _normalize_capped_weights(weights, weight_cap):
+    """Normalize positive weights to mean one while enforcing a hard upper bound."""
+    if not np.isfinite(weight_cap) or weight_cap <= 0:
+        raise ValueError("weight_cap must be a finite positive number.")
+    values = np.asarray(weights, dtype=np.float64)
+    if values.size == 0:
+        return values.astype(np.float32)
+    values = np.nan_to_num(values, nan=0.0, posinf=float(weight_cap), neginf=0.0)
+    values = np.clip(values, 0.0, float(weight_cap))
+    if not np.any(values):
+        return values.astype(np.float32)
+    if weight_cap < 1.0:
+        _logger.warning(
+            "weight_cap=%s is below one; returning capped weights without mean-one normalization.",
+            weight_cap,
+        )
+        return values.astype(np.float32)
+    if np.count_nonzero(values) * float(weight_cap) < values.size:
+        raise ValueError(
+            "Cannot normalize weights to mean one while enforcing weight_cap: "
+            "too many zero-weight events."
+        )
+
+    target_sum = float(values.size)
+    lower, upper = 0.0, 1.0
+    while np.minimum(values * upper, weight_cap).sum() < target_sum:
+        upper *= 2.0
+    for _ in range(64):
+        scale = 0.5 * (lower + upper)
+        if np.minimum(values * scale, weight_cap).sum() < target_sum:
+            lower = scale
+        else:
+            upper = scale
+    return np.minimum(values * upper, float(weight_cap)).astype(np.float32)
+
+
+def _class_zenith_balance_weights(
+    x_train,
+    y_train,
+    weight_cap=10.0,
+    smoothing=1.0,
+    target_ze_fraction=None,
+):
     """Compute capped, smoothed weights equalizing class distributions over ``ze_bin``.
 
     ``smoothing`` prevents a single sparse background bin from receiving an
     arbitrarily large weight.  The cap is an explicit robustness guard for the
     sparse-background regime common in VERITAS training lists.
+
+    ``target_ze_fraction`` optionally supplies the target distribution derived
+    from the training split.  Passing it when weighting validation data keeps
+    evaluation on the same target population rather than recalculating a
+    distribution from validation composition.
     """
     if "ze_bin" not in x_train.columns:
         raise ValueError(
@@ -1385,7 +1546,7 @@ def _class_zenith_balance_weights(x_train, y_train, weight_cap=10.0, smoothing=1
 
     labels = pd.Series(y_train, index=x_train.index, name="label")
     ze_bins = pd.Series(x_train["ze_bin"], index=x_train.index, name="ze_bin")
-    valid = labels.notna() & ze_bins.notna()
+    valid = labels.notna() & ze_bins.notna() & (ze_bins >= 0)
     n_invalid = int((~valid).sum())
     if n_invalid:
         _logger.warning(
@@ -1399,9 +1560,16 @@ def _class_zenith_balance_weights(x_train, y_train, weight_cap=10.0, smoothing=1
     if total_valid == 0:
         raise ValueError("Cannot apply class/zenith balancing with no valid training events.")
 
-    all_ze = np.sort(ze_valid.unique())
-    target_counts = ze_valid.value_counts().reindex(all_ze, fill_value=0).astype(float)
-    target_fraction = target_counts / target_counts.sum()
+    if target_ze_fraction is None:
+        target_fraction = _class_zenith_target_fraction(x_train)
+    else:
+        target_fraction = pd.Series(target_ze_fraction, dtype=float)
+        target_fraction = target_fraction.replace([np.inf, -np.inf], np.nan).dropna()
+        target_fraction = target_fraction[target_fraction > 0]
+        if target_fraction.empty:
+            raise ValueError("The zenith target distribution must contain positive mass.")
+        target_fraction = target_fraction / target_fraction.sum()
+    all_ze = np.asarray(target_fraction.index)
     weights = pd.Series(1.0, index=x_train.index, dtype=np.float64)
 
     _logger.info("Class/zenith balancing target distribution:")
@@ -1442,12 +1610,7 @@ def _class_zenith_balance_weights(x_train, y_train, weight_cap=10.0, smoothing=1
         if class_sum > 0:
             weights.loc[class_mask] *= target_class_total / class_sum
 
-    weight_values = np.clip(weights.to_numpy(dtype=np.float32), 0.0, float(weight_cap))
-    mean_weight = weight_values.mean()
-    if mean_weight > 0:
-        weight_values /= mean_weight
-
-    return weight_values
+    return _normalize_capped_weights(weights.to_numpy(dtype=np.float64), weight_cap)
 
 
 def _log_energy_bin_counts(df):
diff --git a/tests/test_classification_apply_interpolation.py b/tests/test_classification_apply_interpolation.py
index 198d71e..9787ed2 100644
--- a/tests/test_classification_apply_interpolation.py
+++ b/tests/test_classification_apply_interpolation.py
@@ -72,6 +72,71 @@ def test_apply_classification_models_interpolates_probabilities_and_thresholds(m
     np.testing.assert_array_equal(is_gamma[50], np.array([0, 1], dtype=np.uint8))
 
 
+def test_apply_borrowed_energy_bins_recomputes_alpha_from_resolved_centers(monkeypatch):
+    """Borrowed models must use their actual energy centers for score calibration."""
+    df = pd.DataFrame(
+        {
+            "Erec": [10.0],
+            "e_bin_lo": [1],
+            "e_bin_hi": [2],
+            "e_alpha": [0.1],
+            "dummy": [1.0],
+        }
+    )
+    model_configs = {
+        "models": {
+            0: {
+                "model": DummyXGBClassifier(0.0),
+                "features": ["dummy"],
+                "thresholds": {50: 0.2},
+                "energy_center": 0.0,
+            },
+            2: {
+                "model": DummyXGBClassifier(1.0),
+                "features": ["dummy"],
+                "thresholds": {50: 0.8},
+                "energy_center": 2.0,
+            },
+        }
+    }
+
+    monkeypatch.setattr(models, "flatten_feature_data", lambda *args, **kwargs: df[["dummy"]])
+
+    class_probability, is_gamma = models.apply_classification_models(df, model_configs, [50])
+
+    np.testing.assert_allclose(class_probability, np.array([0.5], dtype=np.float32))
+    np.testing.assert_array_equal(is_gamma[50], np.array([1], dtype=np.uint8))
+
+
+def test_apply_leaves_out_of_range_zenith_events_invalid(monkeypatch):
+    """Events outside the trained zenith range must not be scored by an edge model."""
+    df = pd.DataFrame(
+        {
+            "Erec": [1.0],
+            "e_bin_lo": [0],
+            "e_bin_hi": [0],
+            "e_alpha": [0.0],
+            "ze_bin": [-1],
+            "dummy": [1.0],
+        }
+    )
+    model_configs = {
+        "models": {
+            0: {
+                "model": DummyXGBClassifier(1.0),
+                "features": ["dummy"],
+                "thresholds": {50: 0.5},
+            }
+        }
+    }
+    monkeypatch.setattr(models, "flatten_feature_data", lambda *args, **kwargs: df[["dummy"]])
+
+    class_probability, is_gamma = models.apply_classification_models(df, model_configs, [50])
+
+    assert np.isnan(class_probability[0])
+    assert is_gamma[50][0] == 0
+
+
 def test_extra_columns_skip_tmva_only_size_second_max_when_branch_missing():
     """Standard classification should not synthesize the TMVA-only SizeSecondMax column."""
     df = pd.DataFrame(
diff --git a/tests/test_data_processing.py b/tests/test_data_processing.py
index e565113..6a071b7 100644
--- a/tests/test_data_processing.py
+++ b/tests/test_data_processing.py
@@ -6,14 +6,14 @@
 from eventdisplay_ml.data_processing import energy_interpolation_bins, zenith_in_bins
 
 
-def test_zenith_in_bins_numeric_edges_clips_and_handles_boundaries():
-    """Numeric bin edges should clip out-of-range values and place edge values consistently."""
+def test_zenith_in_bins_numeric_edges_marks_invalid_and_handles_boundaries():
+    """Numeric bin edges should mark out-of-range values instead of clipping them."""
     zenith_angles = np.array([-5.0, 0.0, 9.9, 10.0, 19.9, 20.0, 42.0], dtype=float)
     bins = [0.0, 10.0, 20.0, 30.0]
 
     result = zenith_in_bins(zenith_angles, bins)
 
-    np.testing.assert_array_equal(result, np.array([0, 0, 0, 1, 1, 2, 2], dtype=np.int32))
+    np.testing.assert_array_equal(result, np.array([-1, 0, 0, 1, 1, 2, -1], dtype=np.int32))
     assert result.dtype == np.int32
 
 
@@ -28,7 +28,7 @@ def test_zenith_in_bins_dict_bins_matches_numeric_definition():
 
     result = zenith_in_bins(zenith_angles, dict_bins)
 
-    np.testing.assert_array_equal(result, np.array([0, 0, 0, 1, 1, 2, 2], dtype=np.int32))
+    np.testing.assert_array_equal(result, np.array([-1, 0, 0, 1, 1, 2, -1], dtype=np.int32))
     assert result.dtype == np.int32
 
 
diff --git a/tests/test_train_classification_shap.py b/tests/test_train_classification_shap.py
index 0bfe0d0..a85300f 100644
--- a/tests/test_train_classification_shap.py
+++ b/tests/test_train_classification_shap.py
@@ -122,6 +122,48 @@ def test_class_zenith_balance_weights_equalize_class_zenith_distributions():
         assert ze1 / (ze0 + ze1) == pytest.approx(0.5)
 
 
+def test_class_zenith_balance_weights_enforce_hard_cap():
+    """Capped balancing weights must remain bounded after normalization."""
+    x_train = pd.DataFrame(
+        {
+            "f1": np.arange(20, dtype=np.float32),
+            "ze_bin": [0] * 18 + [1] * 2,
+        }
+    )
+    y_train = pd.Series([1] * 10 + [0] * 10, dtype=np.int32)
+
+    weights = models._class_zenith_balance_weights(x_train, y_train, weight_cap=1.5)
+
+    assert weights.max() <= 1.5 + 1e-7
+    assert weights.mean() == pytest.approx(1.0)
+
+
+def test_class_zenith_balance_weights_accept_training_target_for_validation():
+    """Validation weights use the training zenith target rather than validation priors."""
+    train = pd.DataFrame({"ze_bin": [0] * 8 + [1] * 2})
+    validation = pd.DataFrame({"ze_bin": [0] * 2 + [1] * 8})
+    labels = pd.Series([0] * 5 + [1] * 5)
+    target = models._class_zenith_target_fraction(train)
+
+    target_weights = models._class_zenith_balance_weights(
+        validation,
+        labels,
+        target_ze_fraction=target,
+    )
+    validation_weights = models._class_zenith_balance_weights(validation, labels)
+
+    assert not np.allclose(target_weights, validation_weights)
+
+
+def test_train_classification_rejects_invalid_zenith_bins():
+    """Invalid zenith routing states must fail before model fitting."""
+    signal = pd.DataFrame({"f1": [1.0, 2.0, 3.0], "ze_bin": [-1, 0, 0]})
+    background = pd.DataFrame({"f1": [-1.0, -2.0, -3.0], "ze_bin": [0, 0, 0]})
+
+    with pytest.raises(ValueError, match="out-of-range or invalid zenith bins"):
+        models.train_classification([signal, background], {"models": {}})
+
+
 def test_train_classification_applies_class_zenith_weights():
     """The optional class/zenith balance weights should be passed to XGBoost."""
     signal_df = pd.DataFrame(

From 86bc65505bad3cd1b4c7562e6d8cb6161dd4b9fb Mon Sep 17 00:00:00 2001
From: GernotMaier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 15:56:52 +0200
Subject: [PATCH 07/10] copilot review

---
 src/eventdisplay_ml/data_processing.py  |  36 +++++++-
 src/eventdisplay_ml/models.py           | 104 ++++++++++++++++++++----
 tests/test_classification_robustness.py |  15 ++++
 tests/test_data_processing.py           |  19 +++++
 4 files changed, 152 insertions(+), 22 deletions(-)

diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py
index f24f8be..de78313 100644
--- a/src/eventdisplay_ml/data_processing.py
+++ b/src/eventdisplay_ml/data_processing.py
@@ -1564,10 +1564,38 @@ def zenith_in_bins(zenith_angles, bins):
     """
     if bins is None or len(bins) < 2:
         raise ValueError("At least two zenith-bin edges are required.")
-    if isinstance(bins[0], dict):
-        if any("Ze_min" not in b or "Ze_max" not in b for b in bins):
-            raise ValueError("Zenith-bin dictionaries require Ze_min and Ze_max.")
-        bins = [b["Ze_min"] for b in bins] + [bins[-1]["Ze_max"]]
+    if any(isinstance(value, dict) for value in bins):
+        if not all(isinstance(value, dict) for value in bins):
+            raise ValueError("Zenith-bin definitions must be all numeric or all dictionaries.")
+        parsed_bins = []
+        for index, definition in enumerate(bins):
+            if "Ze_min" not in definition or "Ze_max" not in definition:
+                raise ValueError("Zenith-bin dictionaries require Ze_min and Ze_max.")
+            try:
+                ze_min = float(definition["Ze_min"])
+                ze_max = float(definition["Ze_max"])
+            except (TypeError, ValueError) as exc:
+                raise ValueError(
+                    f"Zenith-bin {index} has non-numeric Ze_min/Ze_max values."
+                ) from exc
+            if not np.isfinite(ze_min) or not np.isfinite(ze_max) or ze_min >= ze_max:
+                raise ValueError(
+                    f"Zenith-bin {index} must have finite Ze_min < Ze_max; "
+                    f"got ({ze_min}, {ze_max})."
+                )
+            if parsed_bins and not np.isclose(
+                ze_min,
+                parsed_bins[-1][1],
+                rtol=1e-9,
+                atol=1e-9,
+            ):
+                raise ValueError(
+                    "Zenith-bin dictionaries must be ordered and contiguous: "
+                    f"bin {index - 1} ends at {parsed_bins[-1][1]}, "
+                    f"but bin {index} starts at {ze_min}."
+                )
+            parsed_bins.append((ze_min, ze_max))
+        bins = [parsed_bins[0][0]] + [ze_max for _, ze_max in parsed_bins]
     bins = np.asarray(bins, dtype=float)
     if bins.ndim != 1 or len(bins) < 2 or not np.all(np.isfinite(bins)):
         raise ValueError("Zenith-bin edges must be a finite one-dimensional sequence.")
diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py
index c02916f..3a7ccd2 100644
--- a/src/eventdisplay_ml/models.py
+++ b/src/eventdisplay_ml/models.py
@@ -1362,43 +1362,110 @@ def _classification_split_indices(y_data, groups, train_fraction, random_state,
     if groups is not None and not isinstance(groups, pd.Series):
         groups = pd.Series(groups, index=y_data.index)
     use_groups = grouped and groups is not None and groups.notna().all()
+    groups_overlap_labels = False
     if use_groups:
         use_groups = all(groups[y_data == label].nunique() >= 6 for label in y_data.unique())
 
-    train, validation, test = [], [], []
+    def stable_group_key(value):
+        return (type(value).__name__, repr(value))
+
+    label_groups_by_label = {}
     if use_groups:
+        group_labels = {}
         for label in sorted(y_data.unique()):
             label_mask = y_data.to_numpy() == label
-            label_groups = np.asarray(groups[label_mask].unique())
-            n_train_groups = int(np.ceil(len(label_groups) * train_fraction))
-            n_hold_groups = len(label_groups) - n_train_groups
+            label_groups = np.asarray(
+                sorted(groups[label_mask].unique().tolist(), key=stable_group_key)
+            )
+            label_groups_by_label[label] = label_groups
+            for group in label_groups.tolist():
+                group_labels.setdefault(group, set()).add(label)
+        groups_overlap_labels = any(len(labels) > 1 for labels in group_labels.values())
+
+    train, validation, test = [], [], []
+    if use_groups:
+        if groups_overlap_labels:
+            # A group shared by signal and background must be assigned once
+            # globally; independent per-class splits could otherwise leak the
+            # same source into different partitions.
+            all_groups = np.asarray(
+                sorted(
+                    {
+                        group
+                        for label_groups in label_groups_by_label.values()
+                        for group in label_groups
+                    },
+                    key=stable_group_key,
+                )
+            )
+            n_train_groups = int(np.ceil(len(all_groups) * train_fraction))
+            n_hold_groups = len(all_groups) - n_train_groups
             if n_train_groups < 1 or n_hold_groups < 2:
                 raise ValueError(
                     "Grouped classification split cannot create separate validation "
-                    "and test groups for label="
-                    f"{label}: train_test_fraction={train_fraction} leaves "
-                    f"{n_hold_groups} holdout groups. Reduce train_test_fraction or "
-                    "provide more source groups."
+                    "and test groups: "
+                    f"train_test_fraction={train_fraction} leaves {n_hold_groups} "
+                    "holdout groups. Reduce train_test_fraction or provide more "
+                    "source groups."
                 )
             g_train, g_hold = train_test_split(
-                label_groups,
+                all_groups,
                 train_size=train_fraction,
                 random_state=rng,
             )
-            if len(g_hold) < 2:
-                raise ValueError(
-                    "Grouped classification split cannot create separate validation "
-                    f"and test groups for label={label}: only {len(g_hold)} holdout "
-                    "groups remain after the training split."
-                )
             g_validation, g_test = train_test_split(
                 g_hold,
                 train_size=0.5,
                 random_state=rng,
             )
-            train.extend(np.flatnonzero(label_mask & groups.isin(g_train).to_numpy()))
-            validation.extend(np.flatnonzero(label_mask & groups.isin(g_validation).to_numpy()))
-            test.extend(np.flatnonzero(label_mask & groups.isin(g_test).to_numpy()))
+            split_groups = (g_train, g_validation, g_test)
+            for label in sorted(y_data.unique()):
+                label_mask = y_data.to_numpy() == label
+                split_indices = [
+                    np.flatnonzero(label_mask & groups.isin(group_set).to_numpy())
+                    for group_set in split_groups
+                ]
+                if any(len(indices) == 0 for indices in split_indices):
+                    raise ValueError(
+                        "Grouped classification split cannot preserve all classes in "
+                        f"each partition for label={label} with overlapping group IDs."
+                    )
+                train.extend(split_indices[0])
+                validation.extend(split_indices[1])
+                test.extend(split_indices[2])
+        else:
+            for label in sorted(y_data.unique()):
+                label_mask = y_data.to_numpy() == label
+                label_groups = label_groups_by_label[label]
+                n_train_groups = int(np.ceil(len(label_groups) * train_fraction))
+                n_hold_groups = len(label_groups) - n_train_groups
+                if n_train_groups < 1 or n_hold_groups < 2:
+                    raise ValueError(
+                        "Grouped classification split cannot create separate validation "
+                        "and test groups for label="
+                        f"{label}: train_test_fraction={train_fraction} leaves "
+                        f"{n_hold_groups} holdout groups. Reduce train_test_fraction or "
+                        "provide more source groups."
+                    )
+                g_train, g_hold = train_test_split(
+                    label_groups,
+                    train_size=train_fraction,
+                    random_state=rng,
+                )
+                if len(g_hold) < 2:
+                    raise ValueError(
+                        "Grouped classification split cannot create separate validation "
+                        f"and test groups for label={label}: only {len(g_hold)} holdout "
+                        "groups remain after the training split."
+                    )
+                g_validation, g_test = train_test_split(
+                    g_hold,
+                    train_size=0.5,
+                    random_state=rng,
+                )
+                train.extend(np.flatnonzero(label_mask & groups.isin(g_train).to_numpy()))
+                validation.extend(np.flatnonzero(label_mask & groups.isin(g_validation).to_numpy()))
+                test.extend(np.flatnonzero(label_mask & groups.isin(g_test).to_numpy()))
         method = "grouped_source_file"
     else:
         for label in sorted(y_data.unique()):
@@ -1442,6 +1509,7 @@ def _classification_split_indices(y_data, groups, train_fraction, random_state,
             "method": method,
             "grouped_requested": bool(grouped),
             "source_groups_available": bool(use_groups),
+            "groups_overlap_labels": bool(groups_overlap_labels),
         },
     )
 
diff --git a/tests/test_classification_robustness.py b/tests/test_classification_robustness.py
index 93a0ce1..151b895 100644
--- a/tests/test_classification_robustness.py
+++ b/tests/test_classification_robustness.py
@@ -68,6 +68,21 @@ def test_grouped_split_keeps_source_groups_disjoint_when_supported():
     assert set(groups.iloc[validation]).isdisjoint(groups.iloc[test])
 
 
+def test_grouped_split_uses_global_assignment_for_overlapping_group_ids():
+    y = pd.Series(np.repeat([0, 1], 60))
+    groups = pd.Series(np.concatenate([np.tile(np.arange(6), 10), np.tile(np.arange(2, 8), 10)]))
+    train, validation, test, metadata = models._classification_split_indices(
+        y, groups, train_fraction=0.5, random_state=7, grouped=True
+    )
+    assert metadata["groups_overlap_labels"] is True
+    train_groups = set(groups.iloc[train])
+    validation_groups = set(groups.iloc[validation])
+    test_groups = set(groups.iloc[test])
+    assert train_groups.isdisjoint(validation_groups)
+    assert train_groups.isdisjoint(test_groups)
+    assert validation_groups.isdisjoint(test_groups)
+
+
 def test_event_split_reports_insufficient_holdout_events():
     y = pd.Series([0, 0, 1, 1])
     with pytest.raises(ValueError, match="holdout events"):
diff --git a/tests/test_data_processing.py b/tests/test_data_processing.py
index 6a071b7..589d3e6 100644
--- a/tests/test_data_processing.py
+++ b/tests/test_data_processing.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import pandas as pd
+import pytest
 
 from eventdisplay_ml.data_processing import energy_interpolation_bins, zenith_in_bins
 
@@ -32,6 +33,24 @@ def test_zenith_in_bins_dict_bins_matches_numeric_definition():
     assert result.dtype == np.int32
 
 
+def test_zenith_in_bins_rejects_noncontiguous_dict_bins():
+    bins = [
+        {"Ze_min": 0.0, "Ze_max": 10.0},
+        {"Ze_min": 12.0, "Ze_max": 20.0},
+    ]
+    with pytest.raises(ValueError, match="ordered and contiguous"):
+        zenith_in_bins([5.0], bins)
+
+
+def test_zenith_in_bins_rejects_invalid_dict_bin_bounds():
+    bins = [
+        {"Ze_min": 10.0, "Ze_max": 0.0},
+        {"Ze_min": 0.0, "Ze_max": 20.0},
+    ]
+    with pytest.raises(ValueError, match="finite Ze_min < Ze_max"):
+        zenith_in_bins([5.0], bins)
+
+
 def test_energy_interpolation_bins_interpolates_and_clamps_with_invalid_events():
     """Interpolation bins should handle interior, edge, and invalid energies robustly."""
     df_chunk = pd.DataFrame({"Erec": [0.0, 0.1, 1.0, 10.0, 100.0]})

From 6ae9cb22860e406c12075ea3f91ff2b51c442007 Mon Sep 17 00:00:00 2001
From: Gernot Maier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 21:21:20 +0200
Subject: [PATCH 08/10] code review

---
 docs/changes/79.bugfix.md                     |   2 +-
 src/eventdisplay_ml/config.py                 |   9 -
 src/eventdisplay_ml/data_processing.py        | 184 ++---
 src/eventdisplay_ml/evaluate.py               |  47 +-
 src/eventdisplay_ml/features.py               | 106 +--
 src/eventdisplay_ml/models.py                 | 659 ++++--------------
 ...test_classification_apply_interpolation.py |  36 -
 tests/test_classification_robustness.py       |  98 +--
 tests/test_data_processing.py                 |   8 +-
 tests/test_train_classification_shap.py       |  33 -
 10 files changed, 218 insertions(+), 964 deletions(-)

diff --git a/docs/changes/79.bugfix.md b/docs/changes/79.bugfix.md
index 7d240fa..bf44b91 100644
--- a/docs/changes/79.bugfix.md
+++ b/docs/changes/79.bugfix.md
@@ -1 +1 @@
-Harden VERITAS gamma/hadron classification with robust image features, provenance-aware validation, sparse-background weighting, and reliable score calibration.
+Harden gamma/hadron classification with robust features, inactive-telescope masking, source-aware validation splits, and stricter input validation.
diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py
index aff214d..ee8411e 100644
--- a/src/eventdisplay_ml/config.py
+++ b/src/eventdisplay_ml/config.py
@@ -113,14 +113,6 @@ def configure_training(analysis_type):
                 "available in existing data; 'extended' retains the historical feature set."
             ),
         )
-        parser.add_argument(
-            "--grouped_split",
-            action=argparse.BooleanOptionalAction,
-            default=True,
-            help=(
-                "Keep events from the same source file in one split when provenance is available."
-            ),
-        )
     parser.add_argument(
         "--max_cores",
         type=int,
@@ -212,7 +204,6 @@ def configure_training(analysis_type):
         )
         _logger.info(f"Ignore ze_bin feature: {model_configs.get('ignore_ze_bin')}")
         _logger.info(f"Classification feature profile: {model_configs.get('feature_profile')}")
-        _logger.info(f"Grouped classification split: {model_configs.get('grouped_split')}")
 
     model_configs["models"] = hyper_parameters(
         analysis_type, model_configs.get("hyperparameter_config")
diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py
index 082f1c9..1c6dce8 100644
--- a/src/eventdisplay_ml/data_processing.py
+++ b/src/eventdisplay_ml/data_processing.py
@@ -79,6 +79,14 @@ def read_telescope_config(root_file):
     }
 
 
+def _telescope_config_signature(config):
+    """Return fields that determine the flattened classification schema."""
+    return tuple(
+        tuple(np.asarray(config[key]).tolist())
+        for key in ("tel_ids", "mirror_area", "tel_x", "tel_y")
+    )
+
+
 def _resolve_branch_aliases(tree, branch_list):
     """
     Resolve branch name aliases (e.g. R_core vs R) and drop missing optional branches.
@@ -501,20 +509,7 @@ def flatten_telescope_data_vectorized(
     core_x, core_y = _get_core_arrays(df)
 
     # Sorting by mirror area (desc; proxy for telescope type), then size (desc)
-    if classification_mode:
-        sort_indices = _compute_size_area_sort_indices(
-            size_data,
-            active_mask,
-            tel_config,
-            max_tel_id,
-            active_first=True,
-        )
-    else:
-        # Keep the historical call signature for regression callers and
-        # downstream monkeypatches.
-        sort_indices = _compute_size_area_sort_indices(
-            size_data, active_mask, tel_config, max_tel_id
-        )
+    sort_indices = _compute_size_area_sort_indices(size_data, active_mask, tel_config, max_tel_id)
 
     # Determine which telescope positions to keep (for feature reduction)
     if max_tel_per_type is not None and tel_config is not None:
@@ -553,17 +548,6 @@ def flatten_telescope_data_vectorized(
                     sort_indices,
                 )
             )
-            # Geometry is useful only for an active image; inactive slots must
-            # not become a stable class/domain indicator.
-            if classification_mode:
-                for key in tuple(flat_features):
-                    if key.startswith(f"{var}_"):
-                        sorted_tel = int(key.rsplit("_", 1)[1])
-                        flat_features[key] = np.where(
-                            active_mask[np.arange(n_evt), sort_indices[:, sorted_tel]],
-                            flat_features[key],
-                            np.nan,
-                        )
             continue
 
         data = _to_dense_array(df[var]) if _has_field(df, var) else np.full((n_evt, n_tel), np.nan)
@@ -667,9 +651,7 @@ def _get_core_arrays(df):
     return core_x, core_y
 
 
-def _compute_size_area_sort_indices(
-    size_data, active_mask, tel_config, max_tel_id, active_first=False
-):
+def _compute_size_area_sort_indices(size_data, active_mask, tel_config, max_tel_id):
     """Compute sorting indices: mirror area (desc) then size (desc).
 
     Missing telescopes (NaN size or no mirror area) are sorted to the end.
@@ -702,7 +684,6 @@ def _compute_size_area_sort_indices(
         for tel_idx in range(max_tel_id + 1):
             area = mirror_lookup[tel_idx]
             size_val = sizes[evt_idx, tel_idx]
-            active = bool(active_mask[evt_idx, tel_idx])
             # Build sort key:
             #   1) valid area first (0), NaN area last (1)
             #   2) area descending via negative value
@@ -712,20 +693,9 @@ def _compute_size_area_sort_indices(
             size_valid = 0 if not np.isnan(size_val) else 1
             area_key = -area if area_valid == 0 else 0.0
             size_key = -size_val if size_valid == 0 else 0.0
-            if active_first:
-                # Classification-only safeguard: active images precede
-                # inactive detector slots without exposing slot activity.
-                tel_entries.append(
-                    (tel_idx, 0 if active else 1, area_valid, area_key, size_valid, size_key)
-                )
-            else:
-                # Preserve the historical regression ordering exactly.
-                tel_entries.append((tel_idx, area_valid, area_key, size_valid, size_key))
+            tel_entries.append((tel_idx, area_valid, area_key, size_valid, size_key))
 
-        if active_first:
-            tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4], x[5]))
-        else:
-            tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4]))
+        tel_entries.sort(key=lambda x: (x[1], x[2], x[3], x[4]))
         sort_indices[evt_idx] = np.array([t[0] for t in tel_entries])
 
     return sort_indices
@@ -1038,8 +1008,6 @@ def load_training_data(model_configs, file_list, analysis_type):
             max_events_per_file = None
     _logger.info(f"Max events per file: {max_events_per_file}")
 
-    # Classification reuses/validates configuration across signal/background;
-    # regression retains its historical first-file initialization.
     tel_config = model_configs.get("tel_config") if classification_mode else None
     dfs = []
     executor = ThreadPoolExecutor(max_workers=model_configs.get("max_cores", 1))
@@ -1057,28 +1025,14 @@ def load_training_data(model_configs, file_list, analysis_type):
                     model_configs["tel_config"] = tel_config
                 else:
                     if classification_mode:
-                        # A classification model cannot have a stable feature
-                        # schema if telescope IDs/areas change between files.
-                        def _config_signature(config):
-                            return (
-                                int(config["max_tel_id"]),
-                                tuple(int(v) for v in config.get("tel_ids", [])),
-                                tuple(str(v) for v in config.get("tel_types", [])),
-                                tuple(
-                                    float(v)
-                                    for v in config.get(
-                                        "mirror_area", config.get("mirror_areas", [])
-                                    )
-                                ),
-                            )
-
-                        if _config_signature(current_tel_config) != _config_signature(tel_config):
+                        if _telescope_config_signature(
+                            current_tel_config
+                        ) != _telescope_config_signature(tel_config):
                             raise ValueError(
                                 "Classification/training input files have incompatible "
-                                f"telescope configurations: {input_files[0]} versus {f}."
+                                f"telescope configurations: {f}."
                             )
                     elif current_tel_config["max_tel_id"] > tel_config["max_tel_id"]:
-                        # Preserve the historical regression behavior.
                         tel_config = current_tel_config
                         model_configs["tel_config"] = tel_config
 
@@ -1091,13 +1045,7 @@ def _config_signature(config):
                 raw_reservoir_chunks = []
                 reservoir_priorities = None
                 file_dfs = []
-                rng = (
-                    np.random.default_rng(
-                        None if random_state is None else int(random_state) + file_idx - 1
-                    )
-                    if classification_mode
-                    else np.random.default_rng(random_state)
-                )
+                rng = np.random.default_rng(random_state)
                 chunk_iterator = tree.iterate(
                     resolved_branch_list,
                     cut=model_configs.get("pre_cuts", None),
@@ -1186,12 +1134,7 @@ def _config_signature(config):
                     continue
 
                 if analysis_type == "classification":
-                    # Provenance is retained only as routing metadata and is
-                    # excluded by the feature profile before fitting.  It
-                    # enables grouped validation without rereading ROOT data.
-                    file_df["__source_file_id"] = file_idx - 1
                     file_df["__source_file"] = str(f)
-                    file_df["__source_row"] = np.arange(len(file_df), dtype=np.int64)
 
                 _logger.info(
                     f"Number of events before / after event cut: {n_before} / "
@@ -1210,23 +1153,25 @@ def _config_signature(config):
                     enabled=memory_profile,
                 )
         except Exception as e:
-            if classification_mode:
-                if isinstance(e, (FileNotFoundError, KeyError, ValueError)):
-                    raise
-                raise RuntimeError(f"Error opening or reading file {f}: {e}") from e
+            if classification_mode and isinstance(e, (KeyError, ValueError)):
+                raise
             raise FileNotFoundError(f"Error opening or reading file {f}: {e}") from e
 
     if classification_mode and not dfs:
         raise ValueError("No data loaded from input files.")
     df_final = pd.concat(dfs, ignore_index=True)
-    if analysis_type == "classification" and max_events is not None and max_events > 0:
-        if len(df_final) > max_events:
-            df_final = df_final.sample(
-                n=max_events,
-                random_state=random_state,
-                ignore_index=True,
-            )
-            _logger.info("Applied exact global classification event cap: %d", max_events)
+    if (
+        analysis_type == "classification"
+        and max_events is not None
+        and max_events > 0
+        and len(df_final) > max_events
+    ):
+        df_final = df_final.sample(
+            n=max_events,
+            random_state=random_state,
+            ignore_index=True,
+        )
+        _logger.info("Applied global classification event cap: %d", max_events)
     del dfs
     utils.log_memory_checkpoint("after final pandas concat", df_final, enabled=memory_profile)
     all_nan_columns = [col for col in df_final.columns if df_final[col].isna().all()]
@@ -1503,19 +1448,6 @@ def extra_columns(df, analysis_type, training, index, tel_config=None, observato
             "EChi2S": _to_numpy_1d(df["EChi2S"], np.float32),
             "EmissionHeight": _to_numpy_1d(df["EmissionHeight"], np.float32),
             "EmissionHeightChi2": _to_numpy_1d(df["EmissionHeightChi2"], np.float32),
-            # Keep routing quantities in the flattened frame for energy-bin
-            # weighting/diagnostics; feature-profile selection removes them
-            # before fitting.
-            "Erec": (
-                _to_numpy_1d(df["Erec"], np.float32)
-                if _has_field(df, "Erec")
-                else np.full(n, DEFAULT_FILL_VALUE, dtype=np.float32)
-            ),
-            "DispNImages": (
-                _to_numpy_1d(df["DispNImages"], np.float32)
-                if _has_field(df, "DispNImages")
-                else np.full(n, DEFAULT_FILL_VALUE, dtype=np.float32)
-            ),
         }
         if _has_field(df, "SizeSecondMax"):
             data["SizeSecondMax"] = _to_numpy_1d(df["SizeSecondMax"], np.float32)
@@ -1558,46 +1490,22 @@ def extra_columns(df, analysis_type, training, index, tel_config=None, observato
 
 
 def zenith_in_bins(zenith_angles, bins):
-    """Apply zenith binning, marking out-of-range angles with ``-1``.
-
-    The final edge is included in the last bin.  Angles below the first edge,
-    above the final edge, or non-finite angles are invalid and receive ``-1``;
-    they are never silently assigned to an edge bin.
-    """
-    if bins is None or len(bins) < 2:
-        raise ValueError("At least two zenith-bin edges are required.")
-    if any(isinstance(value, dict) for value in bins):
+    """Apply zenith binning, marking out-of-range angles with ``-1``."""
+    if bins is None or len(bins) == 0:
+        raise ValueError("Zenith-bin definitions must not be empty.")
+    if isinstance(bins[0], dict):
         if not all(isinstance(value, dict) for value in bins):
-            raise ValueError("Zenith-bin definitions must be all numeric or all dictionaries.")
-        parsed_bins = []
-        for index, definition in enumerate(bins):
-            if "Ze_min" not in definition or "Ze_max" not in definition:
-                raise ValueError("Zenith-bin dictionaries require Ze_min and Ze_max.")
-            try:
-                ze_min = float(definition["Ze_min"])
-                ze_max = float(definition["Ze_max"])
-            except (TypeError, ValueError) as exc:
-                raise ValueError(
-                    f"Zenith-bin {index} has non-numeric Ze_min/Ze_max values."
-                ) from exc
-            if not np.isfinite(ze_min) or not np.isfinite(ze_max) or ze_min >= ze_max:
-                raise ValueError(
-                    f"Zenith-bin {index} must have finite Ze_min < Ze_max; "
-                    f"got ({ze_min}, {ze_max})."
-                )
-            if parsed_bins and not np.isclose(
-                ze_min,
-                parsed_bins[-1][1],
-                rtol=1e-9,
-                atol=1e-9,
-            ):
-                raise ValueError(
-                    "Zenith-bin dictionaries must be ordered and contiguous: "
-                    f"bin {index - 1} ends at {parsed_bins[-1][1]}, "
-                    f"but bin {index} starts at {ze_min}."
-                )
-            parsed_bins.append((ze_min, ze_max))
-        bins = [parsed_bins[0][0]] + [ze_max for _, ze_max in parsed_bins]
+            raise ValueError("Zenith-bin definitions must use one format.")
+        try:
+            edges = [float(bins[0]["Ze_min"]), *(float(b["Ze_max"]) for b in bins)]
+            starts = np.asarray([float(b["Ze_min"]) for b in bins[1:]])
+        except (KeyError, TypeError, ValueError) as exc:
+            raise ValueError("Zenith-bin dictionaries require numeric Ze_min and Ze_max.") from exc
+        if not np.allclose(starts, edges[1:-1]):
+            raise ValueError("Zenith-bin dictionaries must be ordered and contiguous.")
+        bins = edges
+    elif len(bins) < 2:
+        raise ValueError("At least two zenith-bin edges are required.")
     bins = np.asarray(bins, dtype=float)
     if bins.ndim != 1 or len(bins) < 2 or not np.all(np.isfinite(bins)):
         raise ValueError("Zenith-bin edges must be a finite one-dimensional sequence.")
diff --git a/src/eventdisplay_ml/evaluate.py b/src/eventdisplay_ml/evaluate.py
index fe4154d..2e1ac6b 100644
--- a/src/eventdisplay_ml/evaluate.py
+++ b/src/eventdisplay_ml/evaluate.py
@@ -5,7 +5,6 @@
 import numpy as np
 import pandas as pd
 import xgboost as xgb
-from scipy.stats import beta as beta_distribution
 from sklearn.metrics import (
     classification_report,
     confusion_matrix,
@@ -23,15 +22,11 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
 
     eff_signal = []
     eff_background = []
-    background_survivor_counts = []
 
     for t in thresholds:
         pred = y_pred_proba >= t
-        n_signal_survivors = int(((pred) & (y_test == 1)).sum())
-        n_background_survivors = int(((pred) & (y_test == 0)).sum())
-        background_survivor_counts.append(n_background_survivors)
-        eff_signal.append(n_signal_survivors / n_signal if n_signal else 0.0)
-        eff_background.append(n_background_survivors / n_background if n_background else 0.0)
+        eff_signal.append(((pred) & (y_test == 1)).sum() / n_signal if n_signal else 0)
+        eff_background.append(((pred) & (y_test == 0)).sum() / n_background if n_background else 0)
         _logger.info(
             f"{name}{context_label} Threshold: {t:.2f} | "
             f"Signal Efficiency: {eff_signal[-1]:.4f} | "
@@ -40,12 +35,6 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
 
     eff_signal = np.asarray(eff_signal, dtype=float)
     eff_background = np.asarray(eff_background, dtype=float)
-    background_upper_limit = np.full(len(thresholds), np.nan, dtype=float)
-    if n_background:
-        for i, k in enumerate(background_survivor_counts):
-            background_upper_limit[i] = (
-                1.0 if k >= n_background else beta_distribution.ppf(0.95, k + 1, n_background - k)
-            )
 
     return pd.DataFrame(
         {
@@ -54,7 +43,6 @@ def _efficiency_dataframe(name, y_pred_proba, y_test, thresholds, context_label=
             "background_efficiency": eff_background,
             "n_signal": n_signal * eff_signal,
             "n_background": n_background * eff_background,
-            "background_efficiency_upper95": background_upper_limit,
         }
     )
 
@@ -97,37 +85,6 @@ def evaluation_efficiency(name, model, x_test, y_test, return_by_zenith=False, z
     return efficiency_all, efficiencies_by_zenith
 
 
-def classification_thresholds_from_signal(signal_scores, efficiencies=(0.5, 0.7, 0.8, 0.9, 0.95)):
-    """Calibrate score thresholds to measured held-out signal efficiency.
-
-    Quantiles avoid treating XGBoost's score as a globally calibrated
-    probability.  ``method='higher'`` ensures ties do not exceed the requested
-    operating point.
-    """
-    scores = np.asarray(signal_scores, dtype=float)
-    scores = scores[np.isfinite(scores)]
-    if scores.size == 0:
-        raise ValueError("Cannot calibrate classification thresholds without signal scores.")
-    targets = np.asarray(efficiencies, dtype=float)
-    if np.any((targets <= 0) | (targets >= 1)):
-        raise ValueError("Signal efficiencies must be strictly between zero and one.")
-    thresholds = []
-    for efficiency in targets:
-        quantile = 1.0 - efficiency
-        try:
-            threshold = np.quantile(scores, quantile, method="higher")
-        except TypeError:  # NumPy < 1.22 compatibility
-            threshold = np.quantile(scores, quantile, interpolation="higher")
-        thresholds.append(float(np.clip(threshold, 0.0, 1.0)))
-    return pd.DataFrame(
-        {
-            "signal_efficiency_target": targets,
-            "threshold": thresholds,
-            "n_signal": len(scores),
-        }
-    )
-
-
 def evaluate_classification_model(model, x_test, y_test, df, x_cols, name):
     """Evaluate the trained model on the test set and log performance metrics.
 
diff --git a/src/eventdisplay_ml/features.py b/src/eventdisplay_ml/features.py
index 3c9dba4..edbdc9e 100644
--- a/src/eventdisplay_ml/features.py
+++ b/src/eventdisplay_ml/features.py
@@ -47,102 +47,44 @@ def target_features(analysis_type):
     raise ValueError(f"Unknown analysis type: {analysis_type}")
 
 
-def classification_feature_columns(columns, profile="robust", ignore_ze_bin=False):
-    """Return safe classification columns from an already flattened frame.
-
-    The profile is applied after flattening, so it works with both VERITAS fixed
-    telescope indexing and the CTAO variable-length compatibility path.  Source
-    provenance and routing columns are deliberately never exposed to XGBoost.
-    """
+def classification_feature_columns(columns, profile="extended", ignore_ze_bin=False):
+    """Select classification features from a flattened data frame."""
     if profile not in {"robust", "extended"}:
         raise ValueError("classification feature profile must be 'robust' or 'extended'")
 
-    reserved = {
-        "label",
-        "Erec",
-        "ErecS",
-        "__source_file_id",
-        "__source_row",
-        "__source_file",
-    }
-    available = list(columns)
-    routing_or_activity = {
-        "label",
-        "Erec",
-        "ErecS",
-        "__source_file_id",
-        "__source_row",
-        "__source_file",
-    }
-
-    def is_excluded_routing_or_activity(name):
-        return name in routing_or_activity or name.startswith(
-            (
-                "__",
-                "tel_active_",
-                "mirror_area_",
-                "tel_rel_x_",
-                "tel_rel_y_",
-                "fpointing_dx_",
-                "fpointing_dy_",
-            )
-        )
-
-    if profile == "extended":
-        selected = [name for name in available if not is_excluded_routing_or_activity(name)]
-    else:
-        # Array/stereo quantities plus per-telescope image morphology.  The
-        # latter are essential gamma/hadron information; detector activity,
-        # telescope geometry and pointing remain excluded below.
-        stable = {
+    selected = [
+        name
+        for name in columns
+        if name not in {"label", "Erec", "ErecS"} and not name.startswith("__")
+    ]
+    if profile == "robust":
+        event_features = {
             "MSCW",
             "MSCL",
             "EChi2S",
             "EmissionHeight",
             "EmissionHeightChi2",
             "Core_Distance",
-            # Coarse zenith conditioning is important because atmospheric
-            # depth/projection changes the image morphology.  It can be
-            # removed explicitly with --ignore_ze_bin for a nuisance test.
             "ze_bin",
         }
-        image_bases = {
-            "size",
-            "cosphi",
-            "sinphi",
-            "loss",
-            "dist",
-            "width",
-            "length",
-            "asym",
-            "tgrad_x",
-        }
+        telescope_features = (
+            "cosphi_",
+            "sinphi_",
+            "loss_",
+            "dist_",
+            "width_",
+            "length_",
+            "asym_",
+            "tgrad_x_",
+        )
         selected = [
             name
-            for name in available
-            if name in stable or any(name.startswith(f"{base}_") for base in image_bases)
+            for name in selected
+            if name in event_features or name.startswith(telescope_features)
         ]
-        # Small synthetic/unit-test frames (and old files missing derived
-        # quantities) should still be usable in isolated unit tests.  Real
-        # flattened frames contain at least one physics/activity name; fail
-        # loudly instead of silently falling back to nuisance columns there.
-        morphology_selected = [name for name in selected if name != "ze_bin"]
-        if not morphology_selected:
-            looks_like_flattened_physics = any(
-                name.startswith(("tel_", "ArrayPointing", "Xcore", "Ycore"))
-                or name in {"DispNImages", "Erec", "size"}
-                for name in available
-            )
-            if looks_like_flattened_physics:
-                raise ValueError(
-                    "Robust classification profile has no available stable morphology features."
-                )
-            selected = [name for name in available if name not in reserved]
-            if not ignore_ze_bin and "ze_bin" in available and "ze_bin" not in selected:
-                selected.append("ze_bin")
-
-    if ignore_ze_bin and "ze_bin" in selected:
-        selected.remove("ze_bin")
+
+    if ignore_ze_bin:
+        selected = [name for name in selected if name != "ze_bin"]
     if not selected:
         raise ValueError(f"No usable classification features for profile '{profile}'.")
     return selected
diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py
index 68139d6..4b8f12d 100644
--- a/src/eventdisplay_ml/models.py
+++ b/src/eventdisplay_ml/models.py
@@ -1,7 +1,6 @@
 """Apply models for regression and classification tasks."""
 
 import logging
-import os
 import re
 import subprocess
 import sys
@@ -15,7 +14,6 @@
 import pandas as pd
 import uproot
 import xgboost as xgb
-from sklearn.metrics import roc_auc_score
 from sklearn.model_selection import train_test_split
 
 from eventdisplay_ml import data_processing, diagnostic_utils, features, utils
@@ -25,7 +23,6 @@
     zenith_in_bins,
 )
 from eventdisplay_ml.evaluate import (
-    classification_thresholds_from_signal,
     evaluate_classification_model,
     evaluate_regression_model,
     evaluation_efficiency,
@@ -52,11 +49,6 @@ def _validate_saved_model(model_path):
         str(model_path),
         str(_MODEL_VALIDATION_MEMORY_BYTES),
     ]
-    validation_env = os.environ.copy()
-    # The validator applies a memory limit; unrestricted BLAS thread pools can
-    # allocate one workspace per host CPU and fail before the model is read.
-    for variable in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS"):
-        validation_env[variable] = "1"
     try:
         result = subprocess.run(
             command,
@@ -64,7 +56,6 @@ def _validate_saved_model(model_path):
             check=False,
             text=True,
             timeout=_MODEL_VALIDATION_TIMEOUT_SECONDS,
-            env=validation_env,
         )
     except subprocess.TimeoutExpired as exc:
         raise RuntimeError(
@@ -205,18 +196,13 @@ def load_classification_models(model_prefix, model_name):
             raise KeyError(f"Model name '{model_name}' not found in file: {file}")
         models[e_bin]["features"] = model_data.get("features", [])
         models[e_bin]["efficiency"] = model_data["models"][model_name].get("efficiency")
-        calibration = model_data["models"][model_name].get("signal_threshold_calibration")
         models[e_bin]["thresholds"] = _calculate_classification_thresholds(
-            models[e_bin]["efficiency"], calibration=calibration
+            models[e_bin]["efficiency"]
         )
-        models[e_bin]["support"] = model_data["models"][model_name].get("support", {})
         energy_bin_metadata = _validate_energy_bin_metadata(
             model_data.get("energy_bins_log10_tev"),
             file,
         )
-        models[e_bin]["energy_center"] = 0.5 * (
-            float(energy_bin_metadata["E_min"]) + float(energy_bin_metadata["E_max"])
-        )
         par = _update_parameters(
             par,
             model_data.get("zenith_bins_deg"),
@@ -237,7 +223,7 @@ def load_classification_models(model_prefix, model_name):
     return models, par
 
 
-def _calculate_classification_thresholds(efficiency, min_efficiency=0.2, steps=5, calibration=None):
+def _calculate_classification_thresholds(efficiency, min_efficiency=0.2, steps=5):
     """
     Calculate classification thresholds for given signal efficiencies.
 
@@ -260,18 +246,6 @@ def _calculate_classification_thresholds(efficiency, min_efficiency=0.2, steps=5
     if efficiency is None or len(efficiency) == 0:
         raise ValueError("Classification efficiency diagnostics are missing from the model file.")
     df = efficiency.copy()
-    if calibration is not None:
-        calibrated = pd.DataFrame(calibration)
-        if {"signal_efficiency_target", "threshold"}.issubset(calibrated.columns):
-            df = pd.concat(
-                [
-                    df[["signal_efficiency", "threshold"]],
-                    calibrated[["signal_efficiency_target", "threshold"]].rename(
-                        columns={"signal_efficiency_target": "signal_efficiency"}
-                    ),
-                ],
-                ignore_index=True,
-            ).drop_duplicates(subset=["signal_efficiency"], keep="last")
     df = df.sort_values("signal_efficiency")
     eff_targets = np.arange(min_efficiency * 100, 100, steps) / 100.0
     thresholds = np.interp(
@@ -577,124 +551,47 @@ def apply_classification_models(df, model_configs, threshold_keys):
             observatory=model_configs.get("observatory", "veritas"),
             preview_rows=model_configs.get("preview_rows", 20),
         )
-        resolved_lo = _resolve_classification_bin(models, e_bin_lo)
-        resolved_hi = _resolve_classification_bin(models, e_bin_hi)
-        model_lo = models[resolved_lo]["model"]
-        model_hi = models[resolved_hi]["model"]
-        missing_lo = sorted(set(models[resolved_lo]["features"]) - set(flatten_data.columns))
-        missing_hi = sorted(set(models[resolved_hi]["features"]) - set(flatten_data.columns))
+        model_lo = models[e_bin_lo]["model"]
+        model_hi = models[e_bin_hi]["model"]
+        missing_lo = sorted(set(models[e_bin_lo]["features"]) - set(flatten_data.columns))
+        missing_hi = sorted(set(models[e_bin_hi]["features"]) - set(flatten_data.columns))
         if missing_lo or missing_hi:
             raise ValueError(
                 "Classification model/input feature schema mismatch: "
                 f"low-bin missing={missing_lo}, high-bin missing={missing_hi}."
             )
-        flatten_lo = flatten_data.loc[:, models[resolved_lo]["features"]]
-        flatten_hi = flatten_data.loc[:, models[resolved_hi]["features"]]
+        flatten_lo = flatten_data.loc[:, models[e_bin_lo]["features"]]
+        flatten_hi = flatten_data.loc[:, models[e_bin_hi]["features"]]
 
         class_probs_lo = model_lo.predict_proba(flatten_lo)[:, 1]
-        interpolate_models = resolved_lo != resolved_hi
-        alpha = _classification_interpolation_alpha(
-            group_df,
-            e_bin_lo,
-            e_bin_hi,
-            resolved_lo,
-            resolved_hi,
-            models,
-        )
-        if not interpolate_models:
+        if e_bin_lo == e_bin_hi:
             class_probs = class_probs_lo
         else:
             class_probs_hi = model_hi.predict_proba(flatten_hi)[:, 1]
+            alpha = group_df["e_alpha"].to_numpy(dtype=np.float32)
             class_probs = (1.0 - alpha) * class_probs_lo + alpha * class_probs_hi
         class_probability[group_df.index] = class_probs
 
-        thresholds_lo = models[resolved_lo].get("thresholds", {})
-        thresholds_hi = models[resolved_hi].get("thresholds", {})
+        thresholds_lo = models[e_bin_lo].get("thresholds", {})
+        thresholds_hi = models[e_bin_hi].get("thresholds", {})
         for eff in threshold_keys:
             if eff in is_gamma:
                 thr_lo = thresholds_lo.get(eff)
                 if thr_lo is None:
                     continue
-                if not interpolate_models:
+                if e_bin_lo == e_bin_hi:
                     threshold = thr_lo
                 else:
                     thr_hi = thresholds_hi.get(eff)
                     if thr_hi is None:
                         continue
+                    alpha = group_df["e_alpha"].to_numpy(dtype=np.float32)
                     threshold = (1.0 - alpha) * thr_lo + alpha * thr_hi
                 is_gamma[eff][group_df.index] = (class_probs >= threshold).astype(np.uint8)
 
     return class_probability, is_gamma
 
 
-def _classification_interpolation_alpha(
-    group_df,
-    requested_lo,
-    requested_hi,
-    resolved_lo,
-    resolved_hi,
-    models,
-):
-    """Return interpolation weights in the energy coordinate of resolved models.
-
-    For complete model grids, the precomputed ``e_alpha`` is already correct.
-    When a requested bin is borrowed, recompute the coordinate from the event
-    energy and the actual model centers so scores and calibrated thresholds use
-    the same models and energy geometry.
-    """
-    if resolved_lo == resolved_hi:
-        return np.zeros(len(group_df), dtype=np.float32)
-
-    fallback_alpha = group_df["e_alpha"].to_numpy(dtype=np.float32)
-    borrowed = (requested_lo, requested_hi) != (resolved_lo, resolved_hi)
-    if not borrowed:
-        return fallback_alpha
-
-    center_lo = models[resolved_lo].get("energy_center")
-    center_hi = models[resolved_hi].get("energy_center")
-    if (
-        center_lo is None
-        or center_hi is None
-        or not np.isfinite(center_lo)
-        or not np.isfinite(center_hi)
-        or center_hi <= center_lo
-    ):
-        raise ValueError(
-            "Cannot interpolate borrowed classification energy-bin models without "
-            "finite energy-center metadata."
-        )
-    if "Erec" not in group_df:
-        raise ValueError(
-            "Cannot interpolate borrowed classification energy-bin models without Erec."
-        )
-
-    erec = pd.to_numeric(group_df["Erec"], errors="coerce").to_numpy(dtype=np.float64)
-    valid_energy = np.isfinite(erec) & (erec > 0.0)
-    if not np.all(valid_energy):
-        raise ValueError(
-            "Cannot interpolate borrowed classification energy-bin models for events "
-            "with non-positive or non-finite Erec."
-        )
-    alpha = (np.log10(erec) - float(center_lo)) / float(center_hi - center_lo)
-    return np.clip(alpha, 0.0, 1.0).astype(np.float32)
-
-
-def _resolve_classification_bin(models, requested_bin):
-    """Resolve a missing energy-bin model to the nearest available model."""
-    if requested_bin in models:
-        return requested_bin
-    available = sorted(models)
-    if not available:
-        raise ValueError("No classification models are available for application.")
-    nearest = min(available, key=lambda candidate: abs(candidate - requested_bin))
-    _logger.warning(
-        "No classification model for energy bin %d; borrowing nearest bin %d.",
-        requested_bin,
-        nearest,
-    )
-    return nearest
-
-
 def process_file_chunked(analysis_type, model_configs):
     """
     Stream events from an input file in chunks, apply XGBoost models, write events.
@@ -1145,29 +1042,17 @@ def train_regression(df, model_configs):
 
 
 def train_classification(df, model_configs):
-    """
-    Train a single XGBoost model for gamma/hadron classification.
-
-    Parameters
-    ----------
-    df : list of pd.DataFrame
-        Training data.
-    model_configs : dict
-        Dictionary of model configurations.
-    """
+    """Train a single XGBoost model for gamma/hadron classification."""
     if df[0].empty or df[1].empty:
         raise ValueError(
             "Classification training requires non-empty signal and background data. "
             f"signal_events={len(df[0])}, background_events={len(df[1])}."
         )
-
-    left_columns = set(df[0].columns)
-    right_columns = set(df[1].columns)
-    if left_columns != right_columns:
+    if set(df[0].columns) != set(df[1].columns):
         raise ValueError(
             "Signal/background classification schemas differ. "
-            f"Only signal: {sorted(left_columns - right_columns)}; "
-            f"only background: {sorted(right_columns - left_columns)}"
+            f"Only signal: {sorted(set(df[0].columns) - set(df[1].columns))}; "
+            f"only background: {sorted(set(df[1].columns) - set(df[0].columns))}"
         )
 
     signal = df[0].copy()
@@ -1175,438 +1060,148 @@ def train_classification(df, model_configs):
     signal["label"] = 1
     background["label"] = 0
     full_df = pd.concat([signal, background], ignore_index=True)
-    ze_data = full_df["ze_bin"] if "ze_bin" in full_df.columns else None
-    if model_configs.get("balance_class_zenith_weights", False) and ze_data is None:
-        raise ValueError("Class/zenith balancing requires the derived ze_bin column.")
+    y_data = full_df["label"]
+    ze_data = full_df.get("ze_bin")
     if ze_data is not None:
-        zenith_values = pd.to_numeric(ze_data, errors="coerce")
-        invalid_zenith = zenith_values.isna() | (zenith_values < 0)
+        numeric_zenith = pd.to_numeric(ze_data, errors="coerce")
+        invalid_zenith = numeric_zenith.isna() | (numeric_zenith < 0)
         if invalid_zenith.any():
             raise ValueError(
                 "Classification training contains out-of-range or invalid zenith bins: "
                 f"{int(invalid_zenith.sum())} events."
             )
 
-    profile = model_configs.get("feature_profile", "robust")
+    profile = (
+        "extended"
+        if model_configs.get("tmva_style", False)
+        else model_configs.get("feature_profile", "extended")
+    )
     feature_columns = features.classification_feature_columns(
         full_df.columns,
         profile=profile,
         ignore_ze_bin=model_configs.get("ignore_ze_bin", False),
     )
-    for column in feature_columns:
-        signal_all_nan = bool(signal[column].isna().all())
-        background_all_nan = bool(background[column].isna().all())
-        if signal_all_nan != background_all_nan:
-            raise ValueError(f"Classification feature '{column}' is all-NaN in only one class.")
-        if signal_all_nan:
-            raise ValueError(f"Classification feature '{column}' is all-NaN in both classes.")
+    all_nan = [
+        column
+        for column in feature_columns
+        if signal[column].isna().all() or background[column].isna().all()
+    ]
+    if all_nan:
+        raise ValueError(
+            f"Classification features must contain values in both classes: {', '.join(all_nan)}"
+        )
+
     x_data = full_df.loc[:, feature_columns]
-    _logger.info(f"Features ({len(x_data.columns)}): {', '.join(x_data.columns)}")
-    model_configs["features"] = list(x_data.columns)
-    y_data = full_df["label"]
+    model_configs["features"] = feature_columns
+    _logger.info("Features (%d): %s", len(feature_columns), ", ".join(feature_columns))
 
-    train_idx, validation_idx, test_idx, split_metadata = _classification_split_indices(
+    train_idx, validation_idx, test_idx, split_method = _classification_split_indices(
         y_data,
         full_df.get("__source_file"),
-        train_fraction=model_configs.get("train_test_fraction", 0.5),
-        random_state=model_configs.get("random_state"),
-        grouped=model_configs.get("grouped_split", True),
-    )
-    # Keep a small, explicitly reserved gamma subset for score-threshold
-    # calibration.  It is never used for fitting or assessment metrics.
-    test_signal_idx = test_idx[y_data.iloc[test_idx].to_numpy() == 1]
-    calibration_idx = np.asarray([], dtype=int)
-    test_signal_groups = (
-        full_df.iloc[test_signal_idx]["__source_file"]
-        if "__source_file" in full_df.columns
-        else None
+        model_configs.get("train_test_fraction", 0.5),
+        model_configs.get("random_state"),
     )
-    if test_signal_groups is not None and test_signal_groups.nunique() >= 2:
-        calibration_groups, _assessment_groups = train_test_split(
-            test_signal_groups.unique(),
-            test_size=0.5,
-            random_state=model_configs.get("random_state"),
-        )
-        calibration_idx = test_signal_idx[test_signal_groups.isin(calibration_groups).to_numpy()]
-    elif len(test_signal_idx) >= 2:
-        calibration_idx, _assessment_signal_idx = train_test_split(
-            test_signal_idx, test_size=0.5, random_state=model_configs.get("random_state")
-        )
-    if len(calibration_idx):
-        test_idx = np.asarray(
-            sorted(set(test_idx) - set(calibration_idx)),
-            dtype=int,
-        )
     x_train, x_validation, x_test = (
-        x_data.iloc[idx] for idx in (train_idx, validation_idx, test_idx)
+        x_data.iloc[index] for index in (train_idx, validation_idx, test_idx)
     )
     y_train, y_validation, y_test = (
-        y_data.iloc[idx] for idx in (train_idx, validation_idx, test_idx)
+        y_data.iloc[index] for index in (train_idx, validation_idx, test_idx)
     )
     ze_test = ze_data.iloc[test_idx] if ze_data is not None else None
+    model_configs["classification_split"] = {
+        "method": split_method,
+        "n_train": len(train_idx),
+        "n_validation": len(validation_idx),
+        "n_test": len(test_idx),
+    }
     _logger.info(
         "Classification split: train=%d validation=%d test=%d (%s)",
         len(x_train),
         len(x_validation),
         len(x_test),
-        split_metadata["method"],
-    )
-    model_configs["classification_split"] = split_metadata
-    model_configs["classification_split"]["n_signal_calibration"] = len(calibration_idx)
-    model_configs["classification_split"]["calibration_grouped"] = bool(
-        test_signal_groups is not None and test_signal_groups.nunique() >= 2
+        split_method,
     )
-    model_configs["classification_feature_profile"] = profile
-    model_configs["nuisance_diagnostics"] = _classification_nuisance_diagnostics(full_df, y_data)
+
     weights_train = None
-    weights_validation = None
     if model_configs.get("balance_class_zenith_weights", False):
-        target_ze_fraction = _class_zenith_target_fraction(full_df.iloc[train_idx])
-        weights_train = _class_zenith_balance_weights(
-            full_df.iloc[train_idx],
-            y_train,
-            target_ze_fraction=target_ze_fraction,
-        )
-        weights_validation = _class_zenith_balance_weights(
-            full_df.iloc[validation_idx],
-            y_validation,
-            target_ze_fraction=target_ze_fraction,
-        )
-        _logger.info(
-            "Using class/zenith sample weights "
-            f"(mean={weights_train.mean():.3f}, std={weights_train.std():.3f}, "
-            f"min={weights_train.min():.3f}, max={weights_train.max():.3f})"
-        )
-    eval_x, eval_y = x_validation, y_validation
-    eval_weights = weights_validation
-    eval_max_events = model_configs.get("eval_max_events", 0)
-    if eval_max_events and eval_max_events > 0 and len(eval_x) > eval_max_events:
-        eval_indices = eval_x.sample(
-            n=eval_max_events,
-            random_state=model_configs.get("random_state"),
-        ).index
-        eval_x = eval_x.loc[eval_indices]
-        eval_y = eval_y.loc[eval_indices]
-        if eval_weights is not None:
-            eval_weights = (
-                pd.Series(weights_validation, index=x_validation.index).loc[eval_indices].to_numpy()
-            )
-        _logger.info("Limited XGBoost validation set to %d events", eval_max_events)
-    eval_set = [(x_train, y_train), (eval_x, eval_y)]
+        weights_train = _class_zenith_balance_weights(full_df.iloc[train_idx], y_train)
+    eval_set = [(x_train, y_train), (x_validation, y_validation)]
 
     for name, cfg in model_configs.get("models", {}).items():
-        _logger.info(f"Training {name}")
+        _logger.info("Training %s", name)
         model = xgb.XGBClassifier(**cfg.get("hyper_parameters", {}))
         fit_kwargs = {"eval_set": eval_set, "verbose": True}
         if weights_train is not None:
             fit_kwargs["sample_weight"] = weights_train
-            fit_kwargs["sample_weight_eval_set"] = [weights_train, eval_weights]
         model.fit(x_train, y_train, **fit_kwargs)
 
-        shap_importance = evaluate_classification_model(
-            model,
-            x_test,
-            y_test,
-            full_df,
-            x_data.columns.tolist(),
-            name,
-        )
         cfg["model"] = model
-        cfg["features"] = x_data.columns.tolist()  # Store feature names for diagnostics
-        efficiency_all, efficiencies_by_zenith = evaluation_efficiency(
+        cfg["features"] = feature_columns
+        cfg["shap_importance"] = evaluate_classification_model(
+            model, x_test, y_test, full_df, feature_columns, name
+        )
+        efficiency, efficiencies_by_zenith = evaluation_efficiency(
             name, model, x_test, y_test, return_by_zenith=True, ze_bins=ze_test
         )
-        cfg["efficiency"] = efficiency_all
+        cfg["efficiency"] = efficiency
         for ze_bin, ze_efficiency in efficiencies_by_zenith.items():
             cfg[f"efficiency_ze{ze_bin}"] = ze_efficiency
-        cfg["shap_importance"] = shap_importance
-        try:
-            calibration_frame = x_data.iloc[calibration_idx]
-            if calibration_frame.empty:
-                raise ValueError("no reserved gamma calibration events")
-            test_signal_scores = model.predict_proba(calibration_frame)[:, 1]
-            cfg["signal_threshold_calibration"] = classification_thresholds_from_signal(
-                test_signal_scores
-            )
-        except (TypeError, ValueError, IndexError) as exc:
-            # Lightweight mocks/legacy estimators may not expose probabilities;
-            # keep the model usable but make the missing calibration explicit.
-            _logger.warning("Could not compute held-out signal thresholds for %s: %s", name, exc)
-            cfg["signal_threshold_calibration"] = None
-        cfg["support"] = {
-            "n_train": len(x_train),
-            "n_validation": len(x_validation),
-            "n_test": len(x_test),
-            "n_signal_test": int((y_test == 1).sum()),
-            "n_background_test": int((y_test == 0).sum()),
-            "n_signal_calibration": len(calibration_idx),
-            "fallback_policy": "held_out_model_only; inspect support before applying",
-        }
 
     return model_configs
 
 
-def _classification_split_indices(y_data, groups, train_fraction, random_state, grouped=True):
-    """Create class-stratified train/validation/test indices.
-
-    Grouping is attempted only when every class has at least six source
-    groups.  Sparse VERITAS lists commonly contain one file per class, so the
-    deterministic event-level fallback is intentional and recorded in model
-    metadata rather than pretending that grouping was achieved.
-    """
-    if not isinstance(y_data, pd.Series):
-        y_data = pd.Series(y_data)
-    if not 0.0 < train_fraction < 1.0:
+def _classification_split_indices(labels, groups, train_fraction, random_state):
+    """Return source-grouped train, validation, and test row indices."""
+    if not 0 < train_fraction < 1:
         raise ValueError("train_test_fraction must be between zero and one.")
-    rng = random_state
-    if groups is not None and not isinstance(groups, pd.Series):
-        groups = pd.Series(groups, index=y_data.index)
-    use_groups = grouped and groups is not None and groups.notna().all()
-    groups_overlap_labels = False
-    if use_groups:
-        use_groups = all(groups[y_data == label].nunique() >= 6 for label in y_data.unique())
-
-    def stable_group_key(value):
-        return (type(value).__name__, repr(value))
-
-    label_groups_by_label = {}
-    if use_groups:
-        group_labels = {}
-        for label in sorted(y_data.unique()):
-            label_mask = y_data.to_numpy() == label
-            label_groups = np.asarray(
-                sorted(groups[label_mask].unique().tolist(), key=stable_group_key)
-            )
-            label_groups_by_label[label] = label_groups
-            for group in label_groups.tolist():
-                group_labels.setdefault(group, set()).add(label)
-        groups_overlap_labels = any(len(labels) > 1 for labels in group_labels.values())
-
-    train, validation, test = [], [], []
-    if use_groups:
-        if groups_overlap_labels:
-            # A group shared by signal and background must be assigned once
-            # globally; independent per-class splits could otherwise leak the
-            # same source into different partitions.
-            all_groups = np.asarray(
-                sorted(
-                    {
-                        group
-                        for label_groups in label_groups_by_label.values()
-                        for group in label_groups
-                    },
-                    key=stable_group_key,
-                )
-            )
-            n_train_groups = int(np.ceil(len(all_groups) * train_fraction))
-            n_hold_groups = len(all_groups) - n_train_groups
-            if n_train_groups < 1 or n_hold_groups < 2:
-                raise ValueError(
-                    "Grouped classification split cannot create separate validation "
-                    "and test groups: "
-                    f"train_test_fraction={train_fraction} leaves {n_hold_groups} "
-                    "holdout groups. Reduce train_test_fraction or provide more "
-                    "source groups."
-                )
-            g_train, g_hold = train_test_split(
-                all_groups,
-                train_size=train_fraction,
-                random_state=rng,
-            )
-            g_validation, g_test = train_test_split(
-                g_hold,
-                train_size=0.5,
-                random_state=rng,
-            )
-            split_groups = (g_train, g_validation, g_test)
-            for label in sorted(y_data.unique()):
-                label_mask = y_data.to_numpy() == label
-                split_indices = [
-                    np.flatnonzero(label_mask & groups.isin(group_set).to_numpy())
-                    for group_set in split_groups
-                ]
-                if any(len(indices) == 0 for indices in split_indices):
-                    raise ValueError(
-                        "Grouped classification split cannot preserve all classes in "
-                        f"each partition for label={label} with overlapping group IDs."
-                    )
-                train.extend(split_indices[0])
-                validation.extend(split_indices[1])
-                test.extend(split_indices[2])
-        else:
-            for label in sorted(y_data.unique()):
-                label_mask = y_data.to_numpy() == label
-                label_groups = label_groups_by_label[label]
-                n_train_groups = int(np.ceil(len(label_groups) * train_fraction))
-                n_hold_groups = len(label_groups) - n_train_groups
-                if n_train_groups < 1 or n_hold_groups < 2:
-                    raise ValueError(
-                        "Grouped classification split cannot create separate validation "
-                        "and test groups for label="
-                        f"{label}: train_test_fraction={train_fraction} leaves "
-                        f"{n_hold_groups} holdout groups. Reduce train_test_fraction or "
-                        "provide more source groups."
-                    )
-                g_train, g_hold = train_test_split(
-                    label_groups,
+
+    indices = np.arange(len(labels))
+    if groups is not None:
+        group_labels = pd.DataFrame({"group": groups, "label": labels}).drop_duplicates()
+        groups_are_class_specific = not group_labels["group"].duplicated().any()
+        if groups_are_class_specific:
+            try:
+                train_groups, holdout_groups = train_test_split(
+                    group_labels,
                     train_size=train_fraction,
-                    random_state=rng,
+                    random_state=random_state,
+                    stratify=group_labels["label"],
                 )
-                if len(g_hold) < 2:
-                    raise ValueError(
-                        "Grouped classification split cannot create separate validation "
-                        f"and test groups for label={label}: only {len(g_hold)} holdout "
-                        "groups remain after the training split."
-                    )
-                g_validation, g_test = train_test_split(
-                    g_hold,
+                validation_groups, test_groups = train_test_split(
+                    holdout_groups,
                     train_size=0.5,
-                    random_state=rng,
+                    random_state=random_state,
+                    stratify=holdout_groups["label"],
                 )
-                train.extend(np.flatnonzero(label_mask & groups.isin(g_train).to_numpy()))
-                validation.extend(np.flatnonzero(label_mask & groups.isin(g_validation).to_numpy()))
-                test.extend(np.flatnonzero(label_mask & groups.isin(g_test).to_numpy()))
-        method = "grouped_source_file"
-    else:
-        for label in sorted(y_data.unique()):
-            label_idx = np.flatnonzero(y_data.to_numpy() == label)
-            n_train_events = int(np.ceil(len(label_idx) * train_fraction))
-            n_hold_events = len(label_idx) - n_train_events
-            if n_train_events < 1 or n_hold_events < 2:
-                raise ValueError(
-                    "Classification split cannot create separate validation and test "
-                    "events for label="
-                    f"{label}: train_test_fraction={train_fraction} leaves "
-                    f"{n_hold_events} holdout events. Reduce train_test_fraction or "
-                    "provide more events."
+                return (
+                    indices[groups.isin(train_groups["group"])],
+                    indices[groups.isin(validation_groups["group"])],
+                    indices[groups.isin(test_groups["group"])],
+                    "grouped_source_file",
                 )
-            label_train, label_hold = train_test_split(
-                label_idx,
-                train_size=train_fraction,
-                random_state=rng,
-            )
-            if len(label_hold) < 2:
-                raise ValueError(
-                    "Classification split cannot create separate validation and test "
-                    f"events for label={label}: only {len(label_hold)} holdout events "
-                    "remain after the training split."
+            except ValueError:
+                _logger.warning(
+                    "Not enough source files for a grouped classification split; "
+                    "falling back to a stratified event split."
                 )
-            label_validation, label_test = train_test_split(
-                label_hold,
-                train_size=0.5,
-                random_state=rng,
-            )
-            train.extend(label_train)
-            validation.extend(label_validation)
-            test.extend(label_test)
-        method = "stratified_event_fallback"
-
-    return (
-        np.asarray(sorted(train), dtype=int),
-        np.asarray(sorted(validation), dtype=int),
-        np.asarray(sorted(test), dtype=int),
-        {
-            "method": method,
-            "grouped_requested": bool(grouped),
-            "source_groups_available": bool(use_groups),
-            "groups_overlap_labels": bool(groups_overlap_labels),
-        },
-    )
 
+    train_idx, holdout_idx = train_test_split(
+        indices,
+        train_size=train_fraction,
+        random_state=random_state,
+        stratify=labels,
+    )
+    validation_idx, test_idx = train_test_split(
+        holdout_idx,
+        train_size=0.5,
+        random_state=random_state,
+        stratify=labels.iloc[holdout_idx],
+    )
+    return train_idx, validation_idx, test_idx, "stratified_event"
 
-def _classification_nuisance_diagnostics(df, labels):
-    """Measure separability of routing/activity proxies without serializing a model."""
-    candidates = {}
-    if "ze_bin" in df:
-        candidates["ze_bin"] = df["ze_bin"]
-    activity = [column for column in df.columns if column.startswith("tel_active_")]
-    if activity:
-        candidates["tel_active_count"] = df[activity].sum(axis=1, skipna=True)
-    telescope_columns = [column for column in df.columns if re.search(r"_\d+$", str(column))]
-    if telescope_columns:
-        candidates["feature_missing_fraction"] = df[telescope_columns].isna().mean(axis=1)
-    diagnostics = {}
-    for name, values in candidates.items():
-        numeric = pd.to_numeric(values, errors="coerce")
-        valid = numeric.notna() & labels.notna()
-        if valid.sum() < 4 or labels[valid].nunique() < 2 or numeric[valid].nunique() < 2:
-            diagnostics[name] = {"auc": np.nan, "n": int(valid.sum())}
-            continue
-        auc = float(roc_auc_score(labels[valid], numeric[valid]))
-        diagnostics[name] = {
-            "auc": auc,
-            "n": int(valid.sum()),
-            "shortcut_strength": max(auc, 1.0 - auc),
-        }
-    return diagnostics
-
-
-def _class_zenith_target_fraction(x_train):
-    """Return the fixed zenith target distribution derived from training data."""
-    if "ze_bin" not in x_train.columns:
-        raise ValueError("Cannot derive a zenith target distribution without ze_bin.")
-    ze_bins = pd.to_numeric(x_train["ze_bin"], errors="coerce")
-    valid = ze_bins.notna() & (ze_bins >= 0)
-    if not valid.any():
-        raise ValueError("Cannot derive a zenith target distribution with no valid ze_bin.")
-    counts = ze_bins[valid].value_counts().sort_index().astype(float)
-    return counts / counts.sum()
-
-
-def _normalize_capped_weights(weights, weight_cap):
-    """Normalize positive weights to mean one while enforcing a hard upper bound."""
-    if not np.isfinite(weight_cap) or weight_cap <= 0:
-        raise ValueError("weight_cap must be a finite positive number.")
-    values = np.asarray(weights, dtype=np.float64)
-    if values.size == 0:
-        return values.astype(np.float32)
-    values = np.nan_to_num(values, nan=0.0, posinf=float(weight_cap), neginf=0.0)
-    values = np.clip(values, 0.0, float(weight_cap))
-    if not np.any(values):
-        return values.astype(np.float32)
-    if weight_cap < 1.0:
-        _logger.warning(
-            "weight_cap=%s is below one; returning capped weights without mean-one normalization.",
-            weight_cap,
-        )
-        return values.astype(np.float32)
-    if np.count_nonzero(values) * float(weight_cap) < values.size:
-        raise ValueError(
-            "Cannot normalize weights to mean one while enforcing weight_cap: "
-            "too many zero-weight events."
-        )
-
-    target_sum = float(values.size)
-    lower, upper = 0.0, 1.0
-    while np.minimum(values * upper, weight_cap).sum() < target_sum:
-        upper *= 2.0
-    for _ in range(64):
-        scale = 0.5 * (lower + upper)
-        if np.minimum(values * scale, weight_cap).sum() < target_sum:
-            lower = scale
-        else:
-            upper = scale
-    return np.minimum(values * upper, float(weight_cap)).astype(np.float32)
-
-
-def _class_zenith_balance_weights(
-    x_train,
-    y_train,
-    weight_cap=10.0,
-    smoothing=1.0,
-    target_ze_fraction=None,
-):
-    """Compute capped, smoothed weights equalizing class distributions over ``ze_bin``.
-
-    ``smoothing`` prevents a single sparse background bin from receiving an
-    arbitrarily large weight.  The cap is an explicit robustness guard for the
-    sparse-background regime common in VERITAS training lists.
 
-    ``target_ze_fraction`` optionally supplies the target distribution derived
-    from the training split.  Passing it when weighting validation data keeps
-    evaluation on the same target population rather than recalculating a
-    distribution from validation composition.
-    """
+def _class_zenith_balance_weights(x_train, y_train):
+    """Compute sample weights that equalize class distributions over ze_bin."""
     if "ze_bin" not in x_train.columns:
         raise ValueError(
             "Cannot apply class/zenith balancing because training features do not include ze_bin."
@@ -1614,7 +1209,7 @@ def _class_zenith_balance_weights(
 
     labels = pd.Series(y_train, index=x_train.index, name="label")
     ze_bins = pd.Series(x_train["ze_bin"], index=x_train.index, name="ze_bin")
-    valid = labels.notna() & ze_bins.notna() & (ze_bins >= 0)
+    valid = labels.notna() & ze_bins.notna()
     n_invalid = int((~valid).sum())
     if n_invalid:
         _logger.warning(
@@ -1628,16 +1223,7 @@ def _class_zenith_balance_weights(
     if total_valid == 0:
         raise ValueError("Cannot apply class/zenith balancing with no valid training events.")
 
-    if target_ze_fraction is None:
-        target_fraction = _class_zenith_target_fraction(x_train)
-    else:
-        target_fraction = pd.Series(target_ze_fraction, dtype=float)
-        target_fraction = target_fraction.replace([np.inf, -np.inf], np.nan).dropna()
-        target_fraction = target_fraction[target_fraction > 0]
-        if target_fraction.empty:
-            raise ValueError("The zenith target distribution must contain positive mass.")
-        target_fraction = target_fraction / target_fraction.sum()
-    all_ze = np.asarray(target_fraction.index)
+    target_fraction = ze_valid.value_counts(normalize=True).sort_index()
     weights = pd.Series(1.0, index=x_train.index, dtype=np.float64)
 
     _logger.info("Class/zenith balancing target distribution:")
@@ -1647,19 +1233,16 @@ def _class_zenith_balance_weights(
     for label in sorted(labels_valid.unique()):
         class_mask = labels_valid == label
         class_ze = ze_valid[class_mask]
+        observed_fraction = class_ze.value_counts(normalize=True).sort_index()
         _logger.info(f"Class/zenith balancing weights for label={label}:")
 
         for ze_bin, target_frac in target_fraction.items():
-            obs_count = float((class_ze == ze_bin).sum())
-            class_total = float(len(class_ze))
-            if obs_count == 0:
-                # There is no unbiased within-class estimate for an absent
-                # bin.  Give it a finite pseudo-count, then cap the resulting
-                # weight; events in absent bins remain at weight one.
-                obs_frac = smoothing / (class_total + smoothing * len(all_ze))
-            else:
-                obs_frac = obs_count / class_total
-            weight = min(float(target_frac / obs_frac), float(weight_cap))
+            obs_frac = observed_fraction.get(ze_bin, 0.0)
+            if obs_frac <= 0:
+                _logger.info(f"  ze_bin={ze_bin}: no events for this class; no weight assigned")
+                continue
+
+            weight = target_frac / obs_frac
             mask = valid & (labels == label) & (ze_bins == ze_bin)
             weights.loc[mask] = weight
             _logger.info(
@@ -1667,18 +1250,12 @@ def _class_zenith_balance_weights(
                 f"weight={weight:.6f}, events={int(mask.sum())}"
             )
 
-    # Equalize total influence of the two classes as well as their zenith
-    # shapes.  This prevents a large simulated signal sample from dominating
-    # a sparse background sample even when raw event counts differ.
-    class_labels = sorted(labels_valid.unique())
-    target_class_total = total_valid / len(class_labels)
-    for label in class_labels:
-        class_mask = valid & (labels == label)
-        class_sum = float(weights.loc[class_mask].sum())
-        if class_sum > 0:
-            weights.loc[class_mask] *= target_class_total / class_sum
-
-    return _normalize_capped_weights(weights.to_numpy(dtype=np.float64), weight_cap)
+    weight_values = weights.to_numpy(dtype=np.float32)
+    mean_weight = weight_values.mean()
+    if mean_weight > 0:
+        weight_values /= mean_weight
+
+    return weight_values
 
 
 def _log_energy_bin_counts(df):
diff --git a/tests/test_classification_apply_interpolation.py b/tests/test_classification_apply_interpolation.py
index 9787ed2..1b4196a 100644
--- a/tests/test_classification_apply_interpolation.py
+++ b/tests/test_classification_apply_interpolation.py
@@ -72,42 +72,6 @@ def test_apply_classification_models_interpolates_probabilities_and_thresholds(m
     np.testing.assert_array_equal(is_gamma[50], np.array([0, 1], dtype=np.uint8))
 
 
-def test_apply_borrowed_energy_bins_recomputes_alpha_from_resolved_centers(monkeypatch):
-    """Borrowed models must use their actual energy centers for score calibration."""
-    df = pd.DataFrame(
-        {
-            "Erec": [10.0],
-            "e_bin_lo": [1],
-            "e_bin_hi": [2],
-            "e_alpha": [0.1],
-            "dummy": [1.0],
-        }
-    )
-    model_configs = {
-        "models": {
-            0: {
-                "model": DummyXGBClassifier(0.0),
-                "features": ["dummy"],
-                "thresholds": {50: 0.2},
-                "energy_center": 0.0,
-            },
-            2: {
-                "model": DummyXGBClassifier(1.0),
-                "features": ["dummy"],
-                "thresholds": {50: 0.8},
-                "energy_center": 2.0,
-            },
-        }
-    }
-
-    monkeypatch.setattr(models, "flatten_feature_data", lambda *args, **kwargs: df[["dummy"]])
-
-    class_probability, is_gamma = models.apply_classification_models(df, model_configs, [50])
-
-    np.testing.assert_allclose(class_probability, np.array([0.5], dtype=np.float32))
-    np.testing.assert_array_equal(is_gamma[50], np.array([1], dtype=np.uint8))
-
-
 def test_apply_leaves_out_of_range_zenith_events_invalid(monkeypatch):
     """Events outside the trained zenith range must not be scored by an edge model."""
     df = pd.DataFrame(
diff --git a/tests/test_classification_robustness.py b/tests/test_classification_robustness.py
index 151b895..735b472 100644
--- a/tests/test_classification_robustness.py
+++ b/tests/test_classification_robustness.py
@@ -1,120 +1,62 @@
-"""Focused tests for the code-only classification hardening contract."""
+"""Focused tests for classification hardening."""
 
 import numpy as np
 import pandas as pd
-import pytest
 
 from eventdisplay_ml import features, models
-from eventdisplay_ml.evaluate import (
-    _efficiency_dataframe,
-    classification_thresholds_from_signal,
-)
 
 
 def test_robust_profile_excludes_routing_and_activity_columns():
     columns = [
         "MSCW",
         "MSCL",
-        "EChi2S",
-        "EmissionHeight",
-        "EmissionHeightChi2",
-        "Core_Distance",
-        "size_0",
         "width_0",
         "length_0",
         "tel_active_0",
+        "mirror_area_0",
         "ze_bin",
         "Erec",
         "__source_file",
     ]
-    assert features.classification_feature_columns(columns) == [
+
+    assert features.classification_feature_columns(columns, profile="robust") == [
         "MSCW",
         "MSCL",
-        "EChi2S",
-        "EmissionHeight",
-        "EmissionHeightChi2",
-        "Core_Distance",
-        "size_0",
         "width_0",
         "length_0",
         "ze_bin",
     ]
 
 
-def test_extended_profile_retains_zenith_but_not_provenance():
+def test_extended_profile_retains_historical_features_but_excludes_routing():
     columns = [
         "MSCW",
+        "DispNImages",
         "ze_bin",
         "Erec",
         "tel_active_0",
-        "mirror_area_0",
         "__source_file",
     ]
-    assert features.classification_feature_columns(columns, profile="extended") == [
+
+    assert features.classification_feature_columns(columns) == [
         "MSCW",
+        "DispNImages",
         "ze_bin",
+        "tel_active_0",
     ]
 
 
-def test_grouped_split_keeps_source_groups_disjoint_when_supported():
-    y = pd.Series(np.repeat([0, 1], 60))
-    groups = pd.Series(np.tile(np.arange(6), 20))
-    train, validation, test, metadata = models._classification_split_indices(
-        y, groups, train_fraction=0.5, random_state=7, grouped=True
+def test_grouped_split_keeps_source_files_disjoint():
+    labels = pd.Series(np.repeat([0, 1], 80))
+    groups = pd.Series(
+        np.concatenate([np.repeat(np.arange(8), 10), np.repeat(np.arange(8, 16), 10)])
     )
-    assert metadata["method"] == "grouped_source_file"
-    assert set(groups.iloc[train]).isdisjoint(groups.iloc[validation])
-    assert set(groups.iloc[train]).isdisjoint(groups.iloc[test])
-    assert set(groups.iloc[validation]).isdisjoint(groups.iloc[test])
-
 
-def test_grouped_split_uses_global_assignment_for_overlapping_group_ids():
-    y = pd.Series(np.repeat([0, 1], 60))
-    groups = pd.Series(np.concatenate([np.tile(np.arange(6), 10), np.tile(np.arange(2, 8), 10)]))
-    train, validation, test, metadata = models._classification_split_indices(
-        y, groups, train_fraction=0.5, random_state=7, grouped=True
+    train, validation, test, method = models._classification_split_indices(
+        labels, groups, train_fraction=0.5, random_state=7
     )
-    assert metadata["groups_overlap_labels"] is True
-    train_groups = set(groups.iloc[train])
-    validation_groups = set(groups.iloc[validation])
-    test_groups = set(groups.iloc[test])
-    assert train_groups.isdisjoint(validation_groups)
-    assert train_groups.isdisjoint(test_groups)
-    assert validation_groups.isdisjoint(test_groups)
-
 
-def test_event_split_reports_insufficient_holdout_events():
-    y = pd.Series([0, 0, 1, 1])
-    with pytest.raises(ValueError, match="holdout events"):
-        models._classification_split_indices(
-            y, None, train_fraction=0.5, random_state=7, grouped=False
-        )
-
-
-def test_grouped_split_reports_insufficient_holdout_groups():
-    y = pd.Series(np.repeat([0, 1], 60))
-    groups = pd.Series(np.tile(np.arange(6), 20))
-    with pytest.raises(ValueError, match="holdout groups"):
-        models._classification_split_indices(
-            y, groups, train_fraction=0.9, random_state=7, grouped=True
-        )
-
-
-def test_nuisance_diagnostics_handles_telescope_ids_above_63():
-    frame = pd.DataFrame(
-        {
-            "size_64": [1.0, np.nan, 1.0, np.nan],
-            "width_128": [0.1, 0.2, np.nan, np.nan],
-        }
-    )
-    diagnostics = models._classification_nuisance_diagnostics(frame, pd.Series([0, 0, 1, 1]))
-    assert diagnostics["feature_missing_fraction"]["n"] == 4
-
-
-def test_threshold_calibration_is_quantile_based_and_background_limit_nonzero():
-    calibration = classification_thresholds_from_signal(np.linspace(0.1, 0.9, 9))
-    assert calibration["threshold"].is_monotonic_decreasing
-    efficiency = _efficiency_dataframe(
-        "test", np.array([0.9, 0.1]), np.array([1, 0]), np.array([0.5])
-    )
-    assert efficiency.loc[0, "background_efficiency_upper95"] > 0
+    assert method == "grouped_source_file"
+    assert set(groups.iloc[train]).isdisjoint(groups.iloc[validation])
+    assert set(groups.iloc[train]).isdisjoint(groups.iloc[test])
+    assert set(groups.iloc[validation]).isdisjoint(groups.iloc[test])
diff --git a/tests/test_data_processing.py b/tests/test_data_processing.py
index 589d3e6..a0c6fe3 100644
--- a/tests/test_data_processing.py
+++ b/tests/test_data_processing.py
@@ -33,6 +33,12 @@ def test_zenith_in_bins_dict_bins_matches_numeric_definition():
     assert result.dtype == np.int32
 
 
+def test_zenith_in_bins_accepts_one_dictionary_bin():
+    result = zenith_in_bins([0.0, 10.0, 20.0], [{"Ze_min": 0.0, "Ze_max": 20.0}])
+
+    np.testing.assert_array_equal(result, np.array([0, 0, 0], dtype=np.int32))
+
+
 def test_zenith_in_bins_rejects_noncontiguous_dict_bins():
     bins = [
         {"Ze_min": 0.0, "Ze_max": 10.0},
@@ -47,7 +53,7 @@ def test_zenith_in_bins_rejects_invalid_dict_bin_bounds():
         {"Ze_min": 10.0, "Ze_max": 0.0},
         {"Ze_min": 0.0, "Ze_max": 20.0},
     ]
-    with pytest.raises(ValueError, match="finite Ze_min < Ze_max"):
+    with pytest.raises(ValueError, match="strictly increasing"):
         zenith_in_bins([5.0], bins)
 
 
diff --git a/tests/test_train_classification_shap.py b/tests/test_train_classification_shap.py
index a85300f..c0d6277 100644
--- a/tests/test_train_classification_shap.py
+++ b/tests/test_train_classification_shap.py
@@ -122,39 +122,6 @@ def test_class_zenith_balance_weights_equalize_class_zenith_distributions():
         assert ze1 / (ze0 + ze1) == pytest.approx(0.5)
 
 
-def test_class_zenith_balance_weights_enforce_hard_cap():
-    """Capped balancing weights must remain bounded after normalization."""
-    x_train = pd.DataFrame(
-        {
-            "f1": np.arange(20, dtype=np.float32),
-            "ze_bin": [0] * 18 + [1] * 2,
-        }
-    )
-    y_train = pd.Series([1] * 10 + [0] * 10, dtype=np.int32)
-
-    weights = models._class_zenith_balance_weights(x_train, y_train, weight_cap=1.5)
-
-    assert weights.max() <= 1.5 + 1e-7
-    assert weights.mean() == pytest.approx(1.0)
-
-
-def test_class_zenith_balance_weights_accept_training_target_for_validation():
-    """Validation weights use the training zenith target rather than validation priors."""
-    train = pd.DataFrame({"ze_bin": [0] * 8 + [1] * 2})
-    validation = pd.DataFrame({"ze_bin": [0] * 2 + [1] * 8})
-    labels = pd.Series([0] * 5 + [1] * 5)
-    target = models._class_zenith_target_fraction(train)
-
-    target_weights = models._class_zenith_balance_weights(
-        validation,
-        labels,
-        target_ze_fraction=target,
-    )
-    validation_weights = models._class_zenith_balance_weights(validation, labels)
-
-    assert not np.allclose(target_weights, validation_weights)
-
-
 def test_train_classification_rejects_invalid_zenith_bins():
     """Invalid zenith routing states must fail before model fitting."""
     signal = pd.DataFrame({"f1": [1.0, 2.0, 3.0], "ze_bin": [-1, 0, 0]})

From 66197293ce763ca7619f677390e5d91176ba6c4a Mon Sep 17 00:00:00 2001
From: Gernot Maier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 21:30:23 +0200
Subject: [PATCH 09/10] loader

---
 src/eventdisplay_ml/data_processing.py | 30 +++++++++++++++-----------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py
index 1c6dce8..5db36f9 100644
--- a/src/eventdisplay_ml/data_processing.py
+++ b/src/eventdisplay_ml/data_processing.py
@@ -995,20 +995,19 @@ def load_training_data(model_configs, file_list, analysis_type):
     else:
         branch_list = features_module.features(analysis_type, training=True)
     _logger.info(f"Branch list: {branch_list}")
+    if max_events is not None and max_events > 0:
+        max_events_per_file = max_events // len(input_files)
+    else:
+        max_events_per_file = None
     if classification_mode and max_events is not None and max_events > 0:
-        # Reserve a bounded quota per file, then perform one deterministic
-        # final sample below.  Integer floor division used to turn a small
-        # global cap into zero (which silently disabled sampling).
+        # Integer floor division can turn a small cap into zero, which means
+        # unlimited sampling. Classification applies an exact final cap below.
         max_events_per_file = max(1, int(np.ceil(max_events / len(input_files))))
-    else:
-        if max_events is not None and max_events > 0:
-            # Preserve the historical regression quota behavior.
-            max_events_per_file = max_events // len(input_files)
-        else:
-            max_events_per_file = None
     _logger.info(f"Max events per file: {max_events_per_file}")
 
-    tel_config = model_configs.get("tel_config") if classification_mode else None
+    tel_config = None  # Will be read from first file
+    if classification_mode:
+        tel_config = model_configs.get("tel_config")
     dfs = []
     executor = ThreadPoolExecutor(max_workers=model_configs.get("max_cores", 1))
     total_files = len(input_files)
@@ -1019,11 +1018,12 @@ def load_training_data(model_configs, file_list, analysis_type):
                     _logger.warning(f"File: {f} does not contain a 'data' tree.")
                     continue
 
-                current_tel_config = read_telescope_config(root_file)
                 if tel_config is None:
-                    tel_config = current_tel_config
+                    tel_config = read_telescope_config(root_file)
                     model_configs["tel_config"] = tel_config
                 else:
+                    # Check if current file has a larger max_tel_id and update if needed
+                    current_tel_config = read_telescope_config(root_file)
                     if classification_mode:
                         if _telescope_config_signature(
                             current_tel_config
@@ -1033,6 +1033,12 @@ def load_training_data(model_configs, file_list, analysis_type):
                                 f"telescope configurations: {f}."
                             )
                     elif current_tel_config["max_tel_id"] > tel_config["max_tel_id"]:
+                        _logger.info(
+                            f"Updating telescope configuration: max_tel_id from "
+                            f"{tel_config['max_tel_id']} to {current_tel_config['max_tel_id']} "
+                            f"(file: {f})"
+                        )
+                        # Replace the full telescope configuration to keep all fields consistent
                         tel_config = current_tel_config
                         model_configs["tel_config"] = tel_config
 

From d2c8a6db95dd85b17a18d1956ed944995db7a396 Mon Sep 17 00:00:00 2001
From: Gernot Maier <gernot.maier@desy.de>
Date: Tue, 4 Aug 2026 21:42:44 +0200
Subject: [PATCH 10/10] float comparison

---
 src/eventdisplay_ml/data_processing.py  | 23 ++++++++++++++---------
 tests/test_classification_robustness.py | 19 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/src/eventdisplay_ml/data_processing.py b/src/eventdisplay_ml/data_processing.py
index 5db36f9..789c529 100644
--- a/src/eventdisplay_ml/data_processing.py
+++ b/src/eventdisplay_ml/data_processing.py
@@ -79,12 +79,19 @@ def read_telescope_config(root_file):
     }
 
 
-def _telescope_config_signature(config):
-    """Return fields that determine the flattened classification schema."""
-    return tuple(
-        tuple(np.asarray(config[key]).tolist())
-        for key in ("tel_ids", "mirror_area", "tel_x", "tel_y")
-    )
+def _telescope_configs_match(first, second):
+    """Compare telescope configurations, tolerating float serialization noise."""
+    for key in ("tel_ids", "mirror_area", "tel_x", "tel_y"):
+        first_values = np.asarray(first[key])
+        second_values = np.asarray(second[key])
+        if first_values.shape != second_values.shape:
+            return False
+        if key == "tel_ids":
+            if not np.array_equal(first_values, second_values):
+                return False
+        elif not np.allclose(first_values, second_values, rtol=1e-7, atol=1e-7, equal_nan=True):
+            return False
+    return True
 
 
 def _resolve_branch_aliases(tree, branch_list):
@@ -1025,9 +1032,7 @@ def load_training_data(model_configs, file_list, analysis_type):
                     # Check if current file has a larger max_tel_id and update if needed
                     current_tel_config = read_telescope_config(root_file)
                     if classification_mode:
-                        if _telescope_config_signature(
-                            current_tel_config
-                        ) != _telescope_config_signature(tel_config):
+                        if not _telescope_configs_match(current_tel_config, tel_config):
                             raise ValueError(
                                 "Classification/training input files have incompatible "
                                 f"telescope configurations: {f}."
diff --git a/tests/test_classification_robustness.py b/tests/test_classification_robustness.py
index 735b472..8a4d4ad 100644
--- a/tests/test_classification_robustness.py
+++ b/tests/test_classification_robustness.py
@@ -3,7 +3,24 @@
 import numpy as np
 import pandas as pd
 
-from eventdisplay_ml import features, models
+from eventdisplay_ml import data_processing, features, models
+
+
+def test_telescope_config_comparison_tolerates_float_noise_and_nan():
+    first = {
+        "tel_ids": np.array([1, 2]),
+        "mirror_area": np.array([100.0, np.nan]),
+        "tel_x": np.array([0.0, 10.0]),
+        "tel_y": np.array([1.0, 2.0]),
+    }
+    second = {
+        "tel_ids": np.array([1, 2]),
+        "mirror_area": np.array([100.0 + 1e-8, np.nan]),
+        "tel_x": np.array([0.0, 10.0 + 1e-8]),
+        "tel_y": np.array([1.0, 2.0]),
+    }
+
+    assert data_processing._telescope_configs_match(first, second)
 
 
 def test_robust_profile_excludes_routing_and_activity_columns():