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. diff --git a/src/eventdisplay_ml/config.py b/src/eventdisplay_ml/config.py index 3cc0d2a..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}.")) @@ -133,6 +146,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", @@ -143,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')}") @@ -154,10 +177,18 @@ 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": _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')}" @@ -295,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/src/eventdisplay_ml/models.py b/src/eventdisplay_ml/models.py index 1703170..79364bf 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(): @@ -817,16 +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_eval], eval_set=eval_set, verbose=False, ) @@ -844,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, @@ -876,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, @@ -1065,8 +1108,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 +1123,57 @@ 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.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: + 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 +1186,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; disabling min-bin exclusion and 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 +1212,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_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): diff --git a/tests/test_train_regression_standardization.py b/tests/test_train_regression_standardization.py index 3d8a349..1e55e53 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,64 @@ def test_eval_max_events_limits_xgboost_eval_set( models.train_regression(df, cfg) eval_set = mock_model.fit.call_args.kwargs["eval_set"] - assert len(eval_set[0][0]) == len(df) // 2 - assert len(eval_set[1][0]) == 10 + eval_weights = mock_model.fit.call_args.kwargs["sample_weight_eval_set"] + 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: @@ -418,14 +495,27 @@ 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, - 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, )