From 7b04761f4447574b4c5ae1a6bd1bc929929a64e7 Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 14:49:36 +0200 Subject: [PATCH 01/11] High multiplicity --- src/eventdisplay_ml/config.py | 19 ++++++++ src/eventdisplay_ml/models.py | 82 +++++++++++++++++++--------------- tests/test_config.py | 30 +++++++++++++ tests/test_regression_apply.py | 42 +++++++++++++++++ 4 files changed, 138 insertions(+), 35 deletions(-) diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py index 595521d..3cc0d2a 100644 --- a/src/eventdisplay_ml/config.py +++ b/src/eventdisplay_ml/config.py @@ -232,6 +232,12 @@ def configure_apply(analysis_type): metavar="MODEL_PREFIX", help=("Path to directory containing XGBoost models (without energy bin suffix)."), ) + if analysis_type == "stereo_analysis": + parser.add_argument( + "--model_prefix_high_multiplicity", + metavar="MODEL_PREFIX", + help="Optional model prefix for events with DispNImages > 2.", + ) parser.add_argument( "--model_name", default="xgboost", @@ -293,6 +299,11 @@ def configure_apply(analysis_type): _logger.info(f"Observatory: {model_configs.get('observatory')}") _logger.info(f"Input file: {model_configs.get('input_file')}") _logger.info(f"Model prefix: {model_configs.get('model_prefix')}") + if model_configs.get("model_prefix_high_multiplicity"): + _logger.info( + "High-multiplicity model prefix: %s", + model_configs["model_prefix_high_multiplicity"], + ) _logger.info(f"Output file: {model_configs.get('output_file')}") _logger.info(f"Image selection: {model_configs.get('image_selection')}") _logger.info(f"Max events: {model_configs.get('max_events')}") @@ -307,5 +318,13 @@ def configure_apply(analysis_type): if analysis_type == "stereo_analysis": model_configs["target_mean"] = par.get("target_mean") model_configs["target_std"] = par.get("target_std") + if model_configs.get("model_prefix_high_multiplicity"): + model_configs["models_high_multiplicity"], high_par = load_models( + analysis_type, + model_configs["model_prefix_high_multiplicity"], + model_configs["model_name"], + ) + model_configs["target_mean_high_multiplicity"] = high_par.get("target_mean") + model_configs["target_std_high_multiplicity"] = high_par.get("target_std") return model_configs diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index b5a1491..a6e9829 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -302,8 +302,8 @@ def apply_regression_models(df, model_configs): """ Apply trained XGBoost model for stereo analysis to all events. - All events are processed with a single model trained on all multiplicities. - Features are created for all telescopes with DEFAULT_FILL_VALUE defaults for missing telescopes. + By default, all events are processed with a single model. If a high-multiplicity + model is configured, it is used for events with more than two images. Parameters ---------- @@ -336,41 +336,53 @@ def apply_regression_models(df, model_configs): preview_rows=model_configs.get("preview_rows", 20), ) - models = model_configs["models"] - model_data = next(iter(models.values())) - flatten_data = flatten_data.reindex(columns=model_data["features"]) - data_processing.print_variable_statistics(flatten_data) - - model = model_data["model"] - preds_scaled = model.predict(flatten_data) - - # Inverse transform predictions from standardized space back to original scale - # Model was trained on standardized targets (mean=0, std=1) - target_mean_cfg = model_configs.get("target_mean") - target_std_cfg = model_configs.get("target_std") - if not target_mean_cfg or not target_std_cfg: - raise ValueError( - "Missing target standardization parameters (target_mean/target_std). " - "Regenerate the regression model or load a model file that includes them." + def predict(models, target_mean_cfg, target_std_cfg, mask=None): + model_data = next(iter(models.values())) + model_input = flatten_data.reindex(columns=model_data["features"]) + if mask is not None: + model_input = model_input.loc[mask] + data_processing.print_variable_statistics(model_input) + + if not target_mean_cfg or not target_std_cfg: + raise ValueError( + "Missing target standardization parameters (target_mean/target_std). " + "Regenerate the regression model or load a model file that includes them." + ) + target_mean = np.array( + [target_mean_cfg[key] for key in ("Xoff_residual", "Yoff_residual", "E_residual")] ) + target_std = np.array( + [target_std_cfg[key] for key in ("Xoff_residual", "Yoff_residual", "E_residual")] + ) + return model_data["model"].predict(model_input) * target_std + target_mean + + high_models = model_configs.get("models_high_multiplicity") + if high_models is None: + preds = predict( + model_configs["models"], + model_configs.get("target_mean"), + model_configs.get("target_std"), + ) + else: + high_mask = df["DispNImages"].to_numpy() > 2 + preds = np.empty((len(df), 3), dtype=np.float64) + if np.any(~high_mask): + preds[~high_mask] = predict( + model_configs["models"], + model_configs.get("target_mean"), + model_configs.get("target_std"), + ~high_mask, + ) + if np.any(high_mask): + preds[high_mask] = predict( + high_models, + model_configs.get("target_mean_high_multiplicity"), + model_configs.get("target_std_high_multiplicity"), + high_mask, + ) - target_mean = np.array( - [ - target_mean_cfg["Xoff_residual"], - target_mean_cfg["Yoff_residual"], - target_mean_cfg["E_residual"], - ] - ) - target_std = np.array( - [ - target_std_cfg["Xoff_residual"], - target_std_cfg["Yoff_residual"], - target_std_cfg["E_residual"], - ] - ) - - # Inverse standardization: y = y_scaled * std + mean - preds = preds_scaled * target_std + target_mean + model_data = next(iter(model_configs["models"].values())) + flatten_data = flatten_data.reindex(columns=model_data["features"]) # Model predicts residuals, so add them to DispBDT baseline # Extract DispBDT predictions from the flattened data diff --git a/tests/test_config.py b/tests/test_config.py index e250738..9f6ff64 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -197,6 +197,36 @@ def test_configure_apply_stereo_loads_scalers(monkeypatch): assert result["target_std"]["Xoff_residual"] == pytest.approx(1.0) +def test_configure_apply_stereo_loads_high_multiplicity_model(monkeypatch): + monkeypatch.setattr( + sys, + "argv", + [ + "prog", + "--input_file", + "input.root", + "--model_prefix", + "model_two", + "--model_prefix_high_multiplicity", + "model_high", + "--output_file", + "output.root", + ], + ) + load_models = MagicMock( + side_effect=[ + ({"xgboost": {"model": "two"}}, {"target_mean": {}, "target_std": {}}), + ({"xgboost": {"model": "high"}}, {"target_mean": {}, "target_std": {}}), + ] + ) + monkeypatch.setattr(config, "load_models", load_models) + + result = config.configure_apply("stereo_analysis") + + assert load_models.call_args_list[1].args == ("stereo_analysis", "model_high", "xgboost") + assert result["models_high_multiplicity"]["xgboost"]["model"] == "high" + + def test_configure_apply_classification_omits_regression_scalers(monkeypatch): monkeypatch.setattr( sys, diff --git a/tests/test_regression_apply.py b/tests/test_regression_apply.py index 3d73743..b64b769 100644 --- a/tests/test_regression_apply.py +++ b/tests/test_regression_apply.py @@ -91,6 +91,48 @@ def _mock_flatten(*_args, **_kwargs): np.testing.assert_allclose(pred_yoff, expected_yoff, rtol=0, atol=1e-8) np.testing.assert_allclose(pred_erec_log, expected_erec_log, rtol=0, atol=1e-8) + def test_apply_regression_selects_model_by_multiplicity(self, monkeypatch): + """Test model selection by multiplicity.""" + df_flat = pd.DataFrame( + { + "Xoff_weighted_bdt": [10.0, 20.0, 30.0], + "Yoff_weighted_bdt": [10.0, 20.0, 30.0], + "ErecS": [1.0, 1.0, 1.0], + } + ) + unit_scaler = { + "Xoff_residual": 1.0, + "Yoff_residual": 1.0, + "E_residual": 1.0, + } + zero_scaler = dict.fromkeys(unit_scaler, 0.0) + model_configs = { + "models": { + "xgboost": { + "model": DummyXGBRegressor([[2.0, 2.0, 2.0]]), + "features": df_flat.columns, + } + }, + "models_high_multiplicity": { + "xgboost": { + "model": DummyXGBRegressor([[3.0, 3.0, 3.0], [4.0, 4.0, 4.0]]), + "features": df_flat.columns, + } + }, + "target_mean": zero_scaler, + "target_std": unit_scaler, + "target_mean_high_multiplicity": zero_scaler, + "target_std_high_multiplicity": unit_scaler, + } + monkeypatch.setattr(models, "flatten_feature_data", lambda *_args, **_kwargs: df_flat) + monkeypatch.setattr(models.data_processing, "print_variable_statistics", lambda *_: None) + + pred_xoff, _, _ = models.apply_regression_models( + pd.DataFrame({"DispNImages": [2, 3, 4]}), model_configs + ) + + np.testing.assert_allclose(pred_xoff, [12.0, 23.0, 34.0]) + def test_apply_regression_missing_standardization_params(self, monkeypatch): """Verify that missing target_mean/target_std raises clear error.""" df_flat = pd.DataFrame( From 062f15320c3ebee4240b211d9b7ea9dee0c70466 Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 14:51:22 +0200 Subject: [PATCH 02/11] changelog --- docs/changes/73.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 docs/changes/73.feature.md diff --git a/docs/changes/73.feature.md b/docs/changes/73.feature.md new file mode 100644 index 0000000..f1ba23c --- /dev/null +++ b/docs/changes/73.feature.md @@ -0,0 +1 @@ +Introduce `--model_prefix_high_multiplicity` parameter for stereo analyse to use different models for 2-tel and >2-tel multiplicity events. From c7de59b3b022265669786b9c5d5e440d2d1b5e99 Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 15:32:03 +0200 Subject: [PATCH 03/11] Correct low-multiplicity mask --- src/eventdisplay_ml/models.py | 12 +++++++----- tests/test_regression_apply.py | 10 +++++----- 2 files changed, 12 insertions(+), 10 deletions(-) diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index a6e9829..1d06fd0 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -364,14 +364,16 @@ def predict(models, target_mean_cfg, target_std_cfg, mask=None): model_configs.get("target_std"), ) else: - high_mask = df["DispNImages"].to_numpy() > 2 - preds = np.empty((len(df), 3), dtype=np.float64) - if np.any(~high_mask): - preds[~high_mask] = predict( + multiplicity = df["DispNImages"].to_numpy() + low_mask = multiplicity == 2 + high_mask = multiplicity > 2 + preds = np.full((len(df), 3), np.nan, dtype=np.float64) + if np.any(low_mask): + preds[low_mask] = predict( model_configs["models"], model_configs.get("target_mean"), model_configs.get("target_std"), - ~high_mask, + low_mask, ) if np.any(high_mask): preds[high_mask] = predict( diff --git a/tests/test_regression_apply.py b/tests/test_regression_apply.py index b64b769..b6dff61 100644 --- a/tests/test_regression_apply.py +++ b/tests/test_regression_apply.py @@ -95,9 +95,9 @@ def test_apply_regression_selects_model_by_multiplicity(self, monkeypatch): """Test model selection by multiplicity.""" df_flat = pd.DataFrame( { - "Xoff_weighted_bdt": [10.0, 20.0, 30.0], - "Yoff_weighted_bdt": [10.0, 20.0, 30.0], - "ErecS": [1.0, 1.0, 1.0], + "Xoff_weighted_bdt": [5.0, 10.0, 20.0, 30.0], + "Yoff_weighted_bdt": [5.0, 10.0, 20.0, 30.0], + "ErecS": [1.0, 1.0, 1.0, 1.0], } ) unit_scaler = { @@ -128,10 +128,10 @@ def test_apply_regression_selects_model_by_multiplicity(self, monkeypatch): monkeypatch.setattr(models.data_processing, "print_variable_statistics", lambda *_: None) pred_xoff, _, _ = models.apply_regression_models( - pd.DataFrame({"DispNImages": [2, 3, 4]}), model_configs + pd.DataFrame({"DispNImages": [1, 2, 3, 4]}), model_configs ) - np.testing.assert_allclose(pred_xoff, [12.0, 23.0, 34.0]) + np.testing.assert_allclose(pred_xoff, [np.nan, 12.0, 23.0, 34.0], equal_nan=True) def test_apply_regression_missing_standardization_params(self, monkeypatch): """Verify that missing target_mean/target_std raises clear error.""" From 5fbd7093b7688f62a82c7cf09d9196a05af623ab Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 15:38:24 +0200 Subject: [PATCH 04/11] model predicitons --- src/eventdisplay_ml/models.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index 1d06fd0..1703170 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -341,7 +341,6 @@ def predict(models, target_mean_cfg, target_std_cfg, mask=None): model_input = flatten_data.reindex(columns=model_data["features"]) if mask is not None: model_input = model_input.loc[mask] - data_processing.print_variable_statistics(model_input) if not target_mean_cfg or not target_std_cfg: raise ValueError( @@ -356,6 +355,10 @@ def predict(models, target_mean_cfg, target_std_cfg, mask=None): ) return model_data["model"].predict(model_input) * target_std + target_mean + model_data = next(iter(model_configs["models"].values())) + primary_model_input = flatten_data.reindex(columns=model_data["features"]) + data_processing.print_variable_statistics(primary_model_input) + high_models = model_configs.get("models_high_multiplicity") if high_models is None: preds = predict( @@ -383,8 +386,7 @@ def predict(models, target_mean_cfg, target_std_cfg, mask=None): high_mask, ) - model_data = next(iter(model_configs["models"].values())) - flatten_data = flatten_data.reindex(columns=model_data["features"]) + flatten_data = primary_model_input # Model predicts residuals, so add them to DispBDT baseline # Extract DispBDT predictions from the flattened data From ec88145b697330cf8572a23661fd78e7d8373ae3 Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 17:48:49 +0200 Subject: [PATCH 05/11] Training weights --- src/eventdisplay_ml/models.py | 144 ++++++++++++++---- .../test_train_regression_standardization.py | 71 ++++++--- 2 files changed, 166 insertions(+), 49 deletions(-) diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index 1703170..c2b457f 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -29,6 +29,8 @@ _EVAL_LOG_E_MIN = -2 _EVAL_LOG_E_MAX = 2.5 _EVAL_LOG_E_BINS = 9 +_MIN_WEIGHTED_ENERGY_BIN_EVENTS = 100 +_MAX_REGRESSION_SAMPLE_WEIGHT = 50.0 _logger = logging.getLogger(__name__) @@ -770,8 +772,10 @@ def train_regression(df, model_configs): train_erec_s, train_e_residual, train_disp_nimages, + return_weight_config=True, ) weights_train = bin_result[2] if bin_result else None + weight_config = bin_result[3] if bin_result else None del train_erec_s del train_e_residual del train_disp_nimages @@ -807,6 +811,21 @@ def train_regression(df, model_configs): model_configs.get("eval_max_events", 200000), model_configs.get("random_state", None), ) + weights_eval = None + if weight_config is not None: + weights_eval, _ = _regression_sample_weights( + df["ErecS"].to_numpy()[eval_idx], + df["E_residual"].to_numpy()[eval_idx], + df["DispNImages"].to_numpy()[eval_idx], + **weight_config, + ) + _logger.info( + "Validation sample weights: mean=%.3f, std=%.3f, min=%.3f, max=%.3f", + weights_eval.mean(), + weights_eval.std(), + weights_eval.min(), + weights_eval.max(), + ) _logger.info(f"XGBoost eval_set test events: {len(eval_idx)}") for name, cfg in model_configs.get("models", {}).items(): @@ -827,6 +846,7 @@ def train_regression(df, model_configs): x_train, y_train_scaled_array, sample_weight=weights_train, + sample_weight_eval_set=[weights_train, weights_eval], eval_set=eval_set, verbose=False, ) @@ -1065,8 +1085,8 @@ def _log_energy_bin_counts(df): (bin_edges, counts_dict, weights_array) where: - bin_edges: np.ndarray of bin boundaries - counts_dict: dict mapping intervals to event counts - - weights_array: np.ndarray of inverse-count weights for each event (normalized - for both energy and multiplicity) + - weights_array: np.ndarray of capped inverse-square-root energy weights combined + with multiplicity weights Returns None if E_residual not found. """ if "E_residual" not in df or "ErecS" not in df: @@ -1080,7 +1100,56 @@ def _log_energy_bin_counts(df): ) -def _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages): +def _regression_sample_weights( + erec_s, + e_residual, + disp_nimages, + *, + bins, + energy_bin_weights, + multiplicity_mean_square, + max_weight, + normalization_scale=None, +): + """Calculate capped regression weights using parameters derived from training data.""" + valid_erec = (erec_s > 0) & np.isfinite(erec_s) + mc_e0 = np.full(len(erec_s), np.nan, dtype=np.float64) + mc_e0[valid_erec] = e_residual[valid_erec] + np.log10(erec_s[valid_erec]) + bin_indices = pd.cut(mc_e0, bins=bins, include_lowest=True, labels=False) + + w_energy = np.zeros(len(erec_s), dtype=np.float64) + for i, weight in enumerate(energy_bin_weights): + w_energy[bin_indices == i] = weight + w_multiplicity = np.square(disp_nimages, dtype=np.float64) / multiplicity_mean_square + raw_weights = w_energy * w_multiplicity + + if normalization_scale is None: + if not np.any(raw_weights > 0): + raise ValueError("No regression events remain after energy-bin weighting.") + if np.count_nonzero(raw_weights) * max_weight < len(raw_weights): + raise ValueError("Too few weighted regression events to normalize within max_weight.") + + low, high = 0.0, 1.0 / raw_weights.mean() + while np.minimum(raw_weights * high, max_weight).mean() < 1.0: + high *= 2.0 + for _ in range(60): + middle = 0.5 * (low + high) + if np.minimum(raw_weights * middle, max_weight).mean() < 1.0: + low = middle + else: + high = middle + normalization_scale = high + + weights = np.minimum(raw_weights * normalization_scale, max_weight).astype(np.float32) + return weights, normalization_scale + + +def _log_energy_bin_counts_from_arrays( + erec_s, + e_residual, + disp_nimages, + return_weight_config=False, +): """Log counts and compute energy/multiplicity training weights from arrays.""" # Handle ErecS with proper checks for valid values (> 0) disp_erec_log = np.where(erec_s > 0, np.log10(erec_s), np.nan) @@ -1093,26 +1162,25 @@ def _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages): for interval, count in counts.items(): _logger.info(f" {interval.left:.2f} to {interval.right:.2f} : {int(count)}") - # Calculate inverse-count weights for balancing (events in low-count bins get higher weight) - # Bins with fewer than 10 events get zero weight (excluded from training) - bin_indices = pd.cut(mc_e0, bins=bins, include_lowest=True, labels=False) + # Use inverse-square-root balancing. Sparse bins are excluded and the final + # energy-times-multiplicity event weights are capped. count_per_bin = counts.values - # Only invert counts >= 10 to avoid divide-by-zero warning - inverse_counts = np.zeros_like(count_per_bin, dtype=np.float64) - mask = count_per_bin >= 10 - inverse_counts[mask] = 1.0 / count_per_bin[mask] - # Normalize by mean of non-zero weights only - valid_weights = inverse_counts[inverse_counts > 0] - if len(valid_weights) > 0: - inverse_counts = inverse_counts / valid_weights.mean() - - # Assign weight to each event based on its energy bin - w_energy = np.ones(len(erec_s), dtype=np.float32) - for i, inv_count in enumerate(inverse_counts): - mask = bin_indices == i - w_energy[mask] = inv_count - - _logger.info(f"Energy bin weights (inverse-count, normalized): {inverse_counts}") + eligible = count_per_bin >= _MIN_WEIGHTED_ENERGY_BIN_EVENTS + if not np.any(eligible): + _logger.warning( + "No energy bin has at least %d events; weighting all populated bins.", + _MIN_WEIGHTED_ENERGY_BIN_EVENTS, + ) + eligible = count_per_bin > 0 + reference_count = np.median(count_per_bin[eligible]) + energy_bin_weights = np.zeros_like(count_per_bin, dtype=np.float64) + energy_bin_weights[eligible] = np.sqrt(reference_count / count_per_bin[eligible]) + + _logger.info( + "Energy bin weights (inverse-sqrt count; bins with <%d events excluded): %s", + _MIN_WEIGHTED_ENERGY_BIN_EVENTS, + energy_bin_weights, + ) # Calculate multiplicity weights (prioritize higher-multiplicity events) mult_counts = pd.Series(disp_nimages).value_counts() @@ -1120,29 +1188,41 @@ def _log_energy_bin_counts_from_arrays(erec_s, e_residual, disp_nimages): for mult, count in mult_counts.items(): _logger.info(f" {int(mult)} telescopes: {int(count)}") - w_multiplicity = (disp_nimages**2).astype(np.float32) - w_multiplicity /= np.mean(w_multiplicity) + multiplicity_mean_square = np.mean(np.square(disp_nimages, dtype=np.float64)) + w_multiplicity = np.square(disp_nimages, dtype=np.float64) / multiplicity_mean_square _logger.info( - "Multiplicity weights (inverse-frequency, normalized): " + "Multiplicity weights (quadratic, normalized): " f"mean={w_multiplicity.mean():.3f}, " f"std={w_multiplicity.std():.3f}, " f"min={w_multiplicity.min():.3f}, " f"max={w_multiplicity.max():.3f}" ) - # Combine energy and multiplicity weights - combined_weights = w_energy * w_multiplicity - - # Normalize combined weights so mean is 1.0 to keep learning rate effective - combined_weights = combined_weights / np.mean(combined_weights) + weight_config = { + "bins": bins, + "energy_bin_weights": energy_bin_weights, + "multiplicity_mean_square": multiplicity_mean_square, + "max_weight": _MAX_REGRESSION_SAMPLE_WEIGHT, + } + combined_weights, normalization_scale = _regression_sample_weights( + erec_s, + e_residual, + disp_nimages, + **weight_config, + ) + weight_config["normalization_scale"] = normalization_scale _logger.info( f"Combined weights (energy x multiplicity): " f"mean={combined_weights.mean():.3f}, " f"std={combined_weights.std():.3f}, " f"min={combined_weights.min():.3f}, " - f"max={combined_weights.max():.3f}" + f"max={combined_weights.max():.3f} " + f"(cap={_MAX_REGRESSION_SAMPLE_WEIGHT:.1f})" ) - return bins, dict(counts.items()), combined_weights + result = (bins, dict(counts.items()), combined_weights) + if return_weight_config: + return (*result, weight_config) + return result diff --git a/tests/test_train_regression_standardization.py b/tests/test_train_regression_standardization.py index 3d8a349..37049a1 100644 --- a/tests/test_train_regression_standardization.py +++ b/tests/test_train_regression_standardization.py @@ -140,21 +140,18 @@ def test_log_energy_bin_counts_returns_correct_structure(self, regression_traini assert len(weights) == len(df), "weights array length should match dataframe rows" def test_log_energy_bin_counts_zeroes_low_count_bins(self): - """Verify bins with < 10 events get zero weight.""" - # Create minimal dataframe with specific energy distribution - rng = np.random.default_rng(42) - n_rows = 30 + """Verify bins below the minimum population get zero weight.""" + n_rows = 140 df = pd.DataFrame( { "ErecS": np.concatenate( [ - np.full(15, 100.0), - np.full(10, 10.0), - np.full(5, 1000.0), + np.full(120, 0.1), + np.full(20, 10.0), ] ), "E_residual": np.zeros(n_rows), - "DispNImages": rng.choice([2, 3], n_rows), + "DispNImages": np.full(n_rows, 3), } ) @@ -163,15 +160,39 @@ def test_log_energy_bin_counts_zeroes_low_count_bins(self): _, counts_dict, weights = result - # Find which bins have < 10 events - low_count_bins = {interval: count for interval, count in counts_dict.items() if count < 10} + assert any( + 0 < count < models._MIN_WEIGHTED_ENERGY_BIN_EVENTS for count in counts_dict.values() + ) + assert np.all(weights[-20:] == 0) + + def test_energy_bin_weights_use_inverse_square_root(self): + """A four-times smaller energy bin should get twice the per-event weight.""" + df = pd.DataFrame( + { + "ErecS": np.concatenate([np.full(400, 0.1), np.full(100, 10.0)]), + "E_residual": np.zeros(500), + "DispNImages": np.full(500, 3), + } + ) - # Events in low-count bins should have zero energy weight - # (multiplicity weight might still apply, but energy weight should be 0) - if low_count_bins: - zero_weights = weights[weights == 0] - if len(zero_weights) > 0: - assert len(zero_weights) > 0, "Expected some zero weights for low-count bins" + weights = models._log_energy_bin_counts(df)[2] + + assert weights[400] / weights[0] == pytest.approx(2.0) + + def test_regression_weights_are_capped_and_normalized(self): + """Combined weights must retain mean one without exceeding the configured cap.""" + weights, _ = models._regression_sample_weights( + np.concatenate([np.full(10, 0.1), np.full(90, 10.0)]), + np.zeros(100), + np.full(100, 3), + bins=np.array([-2.0, 0.0, 2.0]), + energy_bin_weights=np.array([1000.0, 1.0]), + multiplicity_mean_square=9.0, + max_weight=50.0, + ) + + assert weights.max() <= 50.0 + assert weights.mean() == pytest.approx(1.0) def test_log_energy_bin_counts_weight_normalization(self, regression_training_df): """Verify combined weights are normalized to mean ~1.0.""" @@ -251,8 +272,11 @@ def test_eval_max_events_limits_xgboost_eval_set( models.train_regression(df, cfg) eval_set = mock_model.fit.call_args.kwargs["eval_set"] + eval_weights = mock_model.fit.call_args.kwargs["sample_weight_eval_set"] assert len(eval_set[0][0]) == len(df) // 2 assert len(eval_set[1][0]) == 10 + assert len(eval_weights[0]) == len(df) // 2 + assert len(eval_weights[1]) == 10 class TestTrainRegressionIntegration: @@ -418,13 +442,26 @@ def test_memory_optimized_training_matches_reference_without_eval_cap( y_std = y_train.std() y_train_scaled = (y_train - y_mean) / y_std y_test_scaled = (y_test - y_mean) / y_std - reference_weights = models._log_energy_bin_counts(df.loc[y_train.index])[2] + train_weight_result = models._log_energy_bin_counts_from_arrays( + df.loc[y_train.index, "ErecS"].to_numpy(), + df.loc[y_train.index, "E_residual"].to_numpy(), + df.loc[y_train.index, "DispNImages"].to_numpy(), + return_weight_config=True, + ) + reference_weights = train_weight_result[2] + reference_eval_weights, _ = models._regression_sample_weights( + df.loc[y_test.index, "ErecS"].to_numpy(), + df.loc[y_test.index, "E_residual"].to_numpy(), + df.loc[y_test.index, "DispNImages"].to_numpy(), + **train_weight_result[3], + ) reference_model = xgb.XGBRegressor(**hyper_parameters) reference_model.fit( x_train, y_train_scaled, sample_weight=reference_weights, + sample_weight_eval_set=[reference_weights, reference_eval_weights], eval_set=[(x_train, y_train_scaled), (x_test, y_test_scaled)], verbose=False, ) From 5eab4198cf96e24d2943d5b04c2d69872ce2aeaa Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 17:59:36 +0200 Subject: [PATCH 06/11] changelog --- docs/changes/74.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 docs/changes/74.feature.md diff --git a/docs/changes/74.feature.md b/docs/changes/74.feature.md new file mode 100644 index 0000000..41a4fbf --- /dev/null +++ b/docs/changes/74.feature.md @@ -0,0 +1 @@ +Change training weights for stereo reconstruction, to avoid extreme weights leading to early training abortion. From 3af90d03e9c93456d126285264d8177cc97cb08c Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 18:23:16 +0200 Subject: [PATCH 07/11] reduced number of evaluation events --- src/eventdisplay_ml/config.py | 10 +++ src/eventdisplay_ml/models.py | 33 ++++++++-- .../test_train_regression_standardization.py | 65 +++++++++++++++++-- 3 files changed, 97 insertions(+), 11 deletions(-) diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py index 3cc0d2a..2a5e692 100644 --- a/src/eventdisplay_ml/config.py +++ b/src/eventdisplay_ml/config.py @@ -133,6 +133,15 @@ def configure_training(analysis_type): default=100000, help="Maximum number of test events used in XGBoost eval_set for early stopping.", ) + parser.add_argument( + "--diagnostic_max_events", + type=int, + default=100000, + help=( + "Maximum number of training events used for post-training generalization " + "diagnostics (set to 0 to use all training events)." + ), + ) if analysis_type == "stereo_analysis": parser.add_argument( "--min_images", @@ -154,6 +163,7 @@ def configure_training(analysis_type): _logger.info(f"Memory profiling: {model_configs['memory_profile']}") _logger.info(f"Read step size: {model_configs['read_step_size']}") _logger.info(f"Max eval events: {model_configs['eval_max_events']}") + _logger.info(f"Max diagnostic events: {model_configs['diagnostic_max_events']}") if model_configs.get("max_tel_per_type") is not None: _logger.info(f"Max telescopes per mirror area type: {model_configs['max_tel_per_type']}") if analysis_type == "stereo_analysis": diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index c2b457f..27bcf51 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -836,17 +836,26 @@ def train_regression(df, model_configs): y_eval_scaled_array = ((df.iloc[eval_idx, target_indices] - y_mean) / y_std).to_numpy( dtype=np.float32, copy=True ) - eval_set = [(x_train, y_train_scaled_array), (x_eval, y_eval_scaled_array)] + eval_set = [(x_eval, y_eval_scaled_array)] utils.log_memory_checkpoint("after building XGBoost fit arrays", enabled=memory_profile) utils.log_memory_checkpoint(f"{name}: before XGBRegressor init", enabled=memory_profile) - model = xgb.XGBRegressor(**cfg.get("hyper_parameters", {})) + hyper_parameters = dict(cfg.get("hyper_parameters", {})) + early_stopping_rounds = hyper_parameters.pop("early_stopping_rounds", None) + if early_stopping_rounds is not None: + hyper_parameters["callbacks"] = [ + xgb.callback.EarlyStopping( + rounds=early_stopping_rounds, + save_best=True, + ) + ] + model = xgb.XGBRegressor(**hyper_parameters) utils.log_memory_checkpoint(f"{name}: before model.fit", enabled=memory_profile) model.fit( x_train, y_train_scaled_array, sample_weight=weights_train, - sample_weight_eval_set=[weights_train, weights_eval], + sample_weight_eval_set=[weights_eval], eval_set=eval_set, verbose=False, ) @@ -864,10 +873,24 @@ def train_regression(df, model_configs): utils.log_memory_checkpoint(f"{name}: after releasing fit arrays", enabled=memory_profile) prediction_chunk_size = model_configs.get("prediction_chunk_size", 200000) + diagnostic_train_idx = _sample_eval_indices( + train_idx, + model_configs.get("diagnostic_max_events", 100000), + model_configs.get("random_state", None), + ) + y_train_diagnostic = ( + y_train + if diagnostic_train_idx is train_idx + else df.iloc[diagnostic_train_idx, target_indices] + ) + _logger.info( + "Post-training diagnostic training events: %d", + len(diagnostic_train_idx), + ) y_train_pred = _predict_unscaled_chunked( model, df, - train_idx, + diagnostic_train_idx, x_cols, y_mean, y_std, @@ -896,7 +919,7 @@ def train_regression(df, model_configs): ) generalization_metrics = diagnostic_utils.compute_generalization_metrics( - y_train, + y_train_diagnostic, y_train_pred, y_test, y_pred, diff --git a/tests/test_train_regression_standardization.py b/tests/test_train_regression_standardization.py index 37049a1..1e55e53 100644 --- a/tests/test_train_regression_standardization.py +++ b/tests/test_train_regression_standardization.py @@ -273,10 +273,63 @@ def test_eval_max_events_limits_xgboost_eval_set( eval_set = mock_model.fit.call_args.kwargs["eval_set"] eval_weights = mock_model.fit.call_args.kwargs["sample_weight_eval_set"] - assert len(eval_set[0][0]) == len(df) // 2 - assert len(eval_set[1][0]) == 10 - assert len(eval_weights[0]) == len(df) // 2 - assert len(eval_weights[1]) == 10 + assert len(eval_set) == 1 + assert len(eval_set[0][0]) == 10 + assert len(eval_weights) == 1 + assert len(eval_weights[0]) == 10 + + def test_early_stopping_callback_keeps_only_best_model( + self, regression_training_df, regression_model_config + ): + """Verify configured early stopping uses a save-best callback.""" + df = regression_training_df + cfg = regression_model_config + + with patch("xgboost.XGBRegressor") as mock_xgb: + mock_model = MagicMock() + mock_model.best_iteration = 5 + mock_model.best_score = 0.1 + mock_model.predict.side_effect = lambda x_values: np.zeros( + (len(x_values), len(cfg["targets"])) + ) + mock_xgb.return_value = mock_model + + with patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}): + models.train_regression(df, cfg) + + constructor_kwargs = mock_xgb.call_args.kwargs + assert "early_stopping_rounds" not in constructor_kwargs + assert len(constructor_kwargs["callbacks"]) == 1 + callback = constructor_kwargs["callbacks"][0] + assert callback.rounds == 2 + assert callback.save_best is True + + def test_diagnostic_max_events_limits_training_predictions( + self, regression_training_df, regression_model_config + ): + """Verify generalization diagnostics predict only a bounded training sample.""" + df = regression_training_df + cfg = regression_model_config.copy() + cfg["diagnostic_max_events"] = 10 + + with patch("xgboost.XGBRegressor") as mock_xgb: + mock_model = MagicMock() + mock_model.best_iteration = 5 + mock_model.best_score = 0.1 + prediction_lengths = [] + + def _predict(x_values): + prediction_lengths.append(len(x_values)) + return np.zeros((len(x_values), len(cfg["targets"]))) + + mock_model.predict.side_effect = _predict + mock_xgb.return_value = mock_model + + with patch("eventdisplay_ml.models.evaluate_regression_model", return_value={}): + models.train_regression(df, cfg) + + assert prediction_lengths[0] == 10 + assert sum(prediction_lengths[1:]) == len(df) // 2 class TestTrainRegressionIntegration: @@ -461,8 +514,8 @@ def test_memory_optimized_training_matches_reference_without_eval_cap( x_train, y_train_scaled, sample_weight=reference_weights, - sample_weight_eval_set=[reference_weights, reference_eval_weights], - eval_set=[(x_train, y_train_scaled), (x_test, y_test_scaled)], + sample_weight_eval_set=[reference_eval_weights], + eval_set=[(x_test, y_test_scaled)], verbose=False, ) From 5e5cd0e3f0e189f7e089fa9823586b156cf78f7c Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Sun, 28 Jun 2026 18:48:15 +0200 Subject: [PATCH 08/11] improved logging --- src/eventdisplay_ml/config.py | 22 ++++++++++++++++++++++ tests/test_config.py | 8 +++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py index 2a5e692..c22611b 100644 --- a/src/eventdisplay_ml/config.py +++ b/src/eventdisplay_ml/config.py @@ -2,9 +2,12 @@ import argparse import logging +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path import numpy as np +from eventdisplay_ml import models as models_module from eventdisplay_ml import utils from eventdisplay_ml.features import target_features from eventdisplay_ml.hyper_parameters import ( @@ -17,6 +20,16 @@ _logger = logging.getLogger(__name__) +def _log_runtime_provenance(): + """Log the installed version and imported model source path.""" + try: + package_version = version("eventdisplay-ml") + except PackageNotFoundError: + package_version = "not-installed" + _logger.info("eventdisplay-ml runtime version: %s", package_version) + _logger.info("eventdisplay-ml models source: %s", Path(models_module.__file__).resolve()) + + def configure_training(analysis_type): """Configure model training based on command-line arguments.""" parser = argparse.ArgumentParser(description=(f"Train XGBoost models for {analysis_type}.")) @@ -152,6 +165,7 @@ def configure_training(analysis_type): model_configs = vars(parser.parse_args()) + _log_runtime_provenance() _logger.info(f"--- XGBoost {analysis_type} training ---") _logger.info(f"Observatory: {model_configs.get('observatory')}") _logger.info(f"Model output prefix: {model_configs.get('model_prefix')}") @@ -168,6 +182,13 @@ def configure_training(analysis_type): _logger.info(f"Max telescopes per mirror area type: {model_configs['max_tel_per_type']}") if analysis_type == "stereo_analysis": _logger.info(f"Minimum images (DispNImages): {model_configs.get('min_images')}") + _logger.info( + "Regression weighting: energy=inverse-sqrt(count), min_bin_events=%d, " + "multiplicity=DispNImages**2, max_combined_weight=%.1f, " + "validation_weights=training-derived", + models_module._MIN_WEIGHTED_ENERGY_BIN_EVENTS, + models_module._MAX_REGRESSION_SAMPLE_WEIGHT, + ) if analysis_type == "classification": _logger.info( f"Balance class zenith weights: {model_configs.get('balance_class_zenith_weights')}" @@ -305,6 +326,7 @@ def configure_apply(analysis_type): model_configs = vars(parser.parse_args()) + _log_runtime_provenance() _logger.info(f"--- XGBoost {analysis_type} evaluation ---") _logger.info(f"Observatory: {model_configs.get('observatory')}") _logger.info(f"Input file: {model_configs.get('input_file')}") diff --git a/tests/test_config.py b/tests/test_config.py index 9f6ff64..520bbda 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -24,7 +24,8 @@ def model_parameters_file(tmp_path): return path -def test_configure_training_stereo_updates_models(monkeypatch): +def test_configure_training_stereo_updates_models(monkeypatch, caplog): + caplog.set_level("INFO") monkeypatch.setattr( sys, "argv", @@ -56,6 +57,11 @@ def test_configure_training_stereo_updates_models(monkeypatch): assert result["pre_cuts"] == "cut_4" assert result["targets"] == ["target_a"] assert result["memory_profile"] is False + assert "eventdisplay-ml runtime version:" in caplog.text + assert "eventdisplay-ml models source:" in caplog.text + assert "energy=inverse-sqrt(count), min_bin_events=100" in caplog.text + assert "max_combined_weight=50.0" in caplog.text + assert "validation_weights=training-derived" in caplog.text def test_configure_training_stereo_enables_memory_profile(monkeypatch): From 217d37806623c3357c6857f5ca04de644642aeae Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Wed, 1 Jul 2026 15:14:35 +0200 Subject: [PATCH 09/11] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/eventdisplay_ml/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index 27bcf51..b07adba 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -1191,7 +1191,7 @@ def _log_energy_bin_counts_from_arrays( eligible = count_per_bin >= _MIN_WEIGHTED_ENERGY_BIN_EVENTS if not np.any(eligible): _logger.warning( - "No energy bin has at least %d events; weighting all populated bins.", + "No energy bin has at least %d events; disabling min-bin exclusion and weighting all populated bins.", _MIN_WEIGHTED_ENERGY_BIN_EVENTS, ) eligible = count_per_bin > 0 From 1b65dc2e799f7d25c37550c9d8e634582ecb3d3a Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Wed, 1 Jul 2026 15:27:17 +0200 Subject: [PATCH 10/11] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- src/eventdisplay_ml/models.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index b07adba..24d9840 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -1140,10 +1140,11 @@ def _regression_sample_weights( mc_e0[valid_erec] = e_residual[valid_erec] + np.log10(erec_s[valid_erec]) bin_indices = pd.cut(mc_e0, bins=bins, include_lowest=True, labels=False) - w_energy = np.zeros(len(erec_s), dtype=np.float64) - for i, weight in enumerate(energy_bin_weights): - w_energy[bin_indices == i] = weight - w_multiplicity = np.square(disp_nimages, dtype=np.float64) / multiplicity_mean_square +w_energy = np.ones(len(erec_s), dtype=np.float64) +w_energy[~valid_erec] = 0.0 +for i, weight in enumerate(energy_bin_weights): + w_energy[bin_indices == i] = weight +w_multiplicity = np.square(disp_nimages, dtype=np.float64) / multiplicity_mean_square raw_weights = w_energy * w_multiplicity if normalization_scale is None: From b561ee60bb9adfeebcc216a9549ab4b7e471997b Mon Sep 17 00:00:00 2001 From: Gernot Maier Date: Wed, 1 Jul 2026 15:35:43 +0200 Subject: [PATCH 11/11] handle weights correctly --- src/eventdisplay_ml/models.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index 24d9840..79364bf 100644 --- a/src/eventdisplay_ml/models.py +++ b/src/eventdisplay_ml/models.py @@ -1140,11 +1140,11 @@ def _regression_sample_weights( mc_e0[valid_erec] = e_residual[valid_erec] + np.log10(erec_s[valid_erec]) bin_indices = pd.cut(mc_e0, bins=bins, include_lowest=True, labels=False) -w_energy = np.ones(len(erec_s), dtype=np.float64) -w_energy[~valid_erec] = 0.0 -for i, weight in enumerate(energy_bin_weights): - w_energy[bin_indices == i] = weight -w_multiplicity = np.square(disp_nimages, dtype=np.float64) / multiplicity_mean_square + w_energy = np.ones(len(erec_s), dtype=np.float64) + w_energy[~valid_erec] = 0.0 + for i, weight in enumerate(energy_bin_weights): + w_energy[bin_indices == i] = weight + w_multiplicity = np.square(disp_nimages, dtype=np.float64) / multiplicity_mean_square raw_weights = w_energy * w_multiplicity if normalization_scale is None: