Coverage for pybmc/bmc.py: 91%
144 statements
« prev ^ index » next coverage.py v7.10.0, created at 2026-07-20 21:03 +0000
1import numpy as np
2import pandas as pd
3import matplotlib.pyplot as plt
4from sklearn.model_selection import train_test_split
5import os
6from .inference_utils import (
7 gibbs_sampler_simplex,
8 gibbs_sampler_heteroscedastic,
9 USVt_hat_extraction,
10)
11from .sampling_utils import (
12 coverage,
13 rndm_m_heteroscedastic_calculator,
14 DEFAULT_PREDICTIVE_SEED,
15)
16from .error_models import (
17 VARIANCE_MODELS,
18 DEFAULT_SAMPLER_SETTINGS,
19 HeteroscedasticMetrics,
20 required_metrics,
21 variance_basis,
22)
25class BayesianModelCombination:
26 """
27 The main idea of this class is to perform Bayesian Model Combination (BMC) on the set of models that we choose
28 from the dataset class. What should this class contain:
29 + Orthogonalization step.
30 + Perform Bayesian inference on the training data that we extract from the Dataset class.
31 + Predictions for certain isotopes.
32 """
34 VALID_CONSTRAINTS = ("unconstrained", "simplex")
35 VALID_ERROR_MODELS = tuple(VARIANCE_MODELS)
37 def __init__(self, models_list, data_dict, truth_column_name, weights=None, constraint="unconstrained", error_model="homoscedastic"):
38 """
39 Initialize the BayesianModelCombination class.
41 :param models_list: List of model names
42 :param data_dict: Dictionary from `load_data()` where each key is a model name and each value is a DataFrame of properties
43 :param truth_column_name: Name of the column containing the truth values.
44 :param weights: Optional initial weights for the models.
45 :param constraint: Weight constraint mode. Options:
46 - ``"unconstrained"`` (default): No constraints on model weights.
47 - ``"simplex"``: Forces weights to lie on the probability simplex
48 (each weight between 0 and 1, weights sum to 1). Uses a
49 Metropolis-within-Gibbs sampler to enforce the constraint.
50 :param error_model: Noise structure of the combination. Options
51 (see :mod:`pybmc.error_models` for the variance forms):
52 - ``"homoscedastic"`` (default): A single constant variance.
53 - ``"hetero_pc_dist"`` / ``"hetero_pc_dist_quad"``: Variance
54 linear/quadratic in the distance from the training centroid
55 in principal-component space.
56 - ``"hetero_model_var"`` / ``"hetero_model_var_quad"``: Variance
57 linear/quadratic in the spread among model predictions.
58 - ``"hetero_combined_linear"`` / ``"hetero_combined_quadratic"``:
59 Variance depending on both metrics.
60 Heteroscedastic error models currently require the
61 ``"unconstrained"`` weight mode.
62 """
64 if not isinstance(models_list, list) or not all(isinstance(model, str) for model in models_list):
65 raise ValueError("The 'models' should be a list of model names (strings) for Bayesian Combination.")
66 if not isinstance(data_dict, dict) or not all(isinstance(df, pd.DataFrame) for df in data_dict.values()):
67 raise ValueError("The 'data_dict' should be a dictionary of pandas DataFrames, one per property.")
68 if constraint not in self.VALID_CONSTRAINTS:
69 raise ValueError(
70 f"Invalid constraint '{constraint}'. "
71 f"Must be one of {self.VALID_CONSTRAINTS}."
72 )
73 if error_model not in self.VALID_ERROR_MODELS:
74 raise ValueError(
75 f"Invalid error model '{error_model}'. "
76 f"Must be one of {self.VALID_ERROR_MODELS}."
77 )
78 if error_model != "homoscedastic" and constraint == "simplex":
79 raise ValueError(
80 "Heteroscedastic error models are not supported with the "
81 "'simplex' constraint; use constraint='unconstrained'."
82 )
84 self.data_dict = data_dict
85 self.models_list = models_list
86 self.models = [m for m in models_list if m != 'truth']
87 self.weights = weights if weights is not None else None
88 self.truth_column_name = truth_column_name
89 self.constraint = constraint
90 self.error_model = error_model
91 self.samples = None
92 self.Vt_hat = None
93 self.mh_acceptance_rate_ = None
94 self._trained_error_model = None
95 self._metrics_calculator = None
98 def orthogonalize(self, property, train_df, components_kept):
99 """
100 Perform orthogonalization for the specified property using training data.
102 :param property: The nuclear property to orthogonalize on (e.g., 'BE').
103 :param train_index: Training data from split_data
104 :param components_kept: Number of SVD components to retain.
105 """
106 # Store selected property
107 self.current_property = property
109 # Extract the relevant DataFrame for that property
110 df = self.data_dict[property].copy()
111 self.selected_models_dataset = df # Store for train() and predict()
113 # Extract model outputs (only the model columns)
114 models_output_train = train_df[self.models]
115 model_predictions_train = models_output_train.values
117 # Mean prediction across models (per nucleus)
118 predictions_mean_train = np.mean(model_predictions_train, axis=1)
120 # Experimental truth values for the property
121 centered_experiment_train = train_df[self.truth_column_name].values - predictions_mean_train
123 # Center model predictions
124 model_predictions_train_centered = model_predictions_train - predictions_mean_train[:, None]
126 # Perform SVD
127 U, S, Vt = np.linalg.svd(model_predictions_train_centered)
129 # Dimensionality reduction
130 U_hat, S_hat, Vt_hat, Vt_hat_normalized = USVt_hat_extraction(U, S, Vt, components_kept) #type: ignore
132 # Save for training
133 self.centered_experiment_train = centered_experiment_train
134 self.U_hat = U_hat
135 self.Vt_hat = Vt_hat
136 self.S_hat = S_hat
137 self.Vt_hat_normalized = Vt_hat_normalized
138 self._predictions_mean_train = predictions_mean_train
139 # Raw training predictions, needed to fit the heteroscedasticity
140 # metrics (PC-space centroid and normalization bounds).
141 self._train_model_predictions = model_predictions_train
144 def train(self, training_options=None):
145 """
146 Train the model combination using training data and optional training parameters.
148 All error models (the homoscedastic one included) share the
149 likelihood ``y_i ~ N(X_i . b, sigma_i^2)`` and are trained with
150 the same Gibbs-within-Metropolis sampler; the homoscedastic
151 model is simply the case where the variance basis is the
152 constant column alone, so ``sigma_i^2 = sigma^2``.
154 :param training_options: Dictionary of training options. Keys:
155 - 'iterations': (int) Number of retained Gibbs samples (default 50000)
156 - 'sampler': (str) Override the constraint mode for this training run.
157 ``"unconstrained"`` or ``"simplex"``. If not provided, uses the
158 instance-level ``self.constraint`` set at initialization.
159 - 'error_model': (str) Override the error model for this training
160 run (see ``VALID_ERROR_MODELS``). If not provided, uses the
161 instance-level ``self.error_model`` set at initialization.
162 - 'seed': (int) Seed for the sampler. If not provided, the
163 sampler draws from the shared package-wide generator
164 (see :mod:`pybmc.rng`), which is seeded at import.
165 - 'b_mean_prior': (np.ndarray) Prior mean vector (default zeros)
166 *(unconstrained sampler)*
167 - 'b_mean_cov': (np.ndarray) Prior covariance matrix (default diag(S_hat²))
168 *(unconstrained sampler)*
169 - 'nu0_chosen': (float) Degrees of freedom for variance prior (default 1.0)
170 *(simplex sampler only)*
171 - 'sigma20_chosen': (float) Prior variance (default 0.02)
172 *(simplex sampler only)*
173 - 'burn': (int) Burn-in iterations (default 10000 for simplex,
174 5000 otherwise)
175 - 'stepsize': (float) Proposal step size (default 0.001)
176 *(simplex sampler only)*
177 - 'proposal_scales': (list) Diagonal of the Metropolis-Hastings
178 proposal covariance for the variance parameters; sensible
179 per-model defaults are used if omitted
180 *(unconstrained sampler)*
181 - 'init_params': (list) Initial values for the non-constant
182 variance parameters *(unconstrained sampler)*
183 - 'prior_spec': (list of (shape, scale)) Gamma priors for the
184 variance parameters *(unconstrained sampler)*
185 - 'adapt_proposal': (bool) Rescale the proposal during burn-in
186 toward 'target_acceptance' (default True)
187 *(unconstrained sampler)*
188 - 'target_acceptance': (float) Acceptance rate targeted by the
189 burn-in adaptation (default 0.25)
190 *(unconstrained sampler)*
192 After training with the unconstrained sampler, the
193 Metropolis-Hastings acceptance rate of the variance parameters
194 is available in ``self.mh_acceptance_rate_``.
195 """
196 if training_options is None:
197 training_options = {}
199 # Determine which sampler to use: training_options overrides instance default
200 sampler_mode = training_options.get('sampler', self.constraint)
201 if sampler_mode not in self.VALID_CONSTRAINTS:
202 raise ValueError(
203 f"Invalid sampler '{sampler_mode}'. "
204 f"Must be one of {self.VALID_CONSTRAINTS}."
205 )
207 # Same override pattern for the error model.
208 error_model_mode = training_options.get('error_model', self.error_model)
209 if error_model_mode not in self.VALID_ERROR_MODELS:
210 raise ValueError(
211 f"Invalid error model '{error_model_mode}'. "
212 f"Must be one of {self.VALID_ERROR_MODELS}."
213 )
214 if error_model_mode != "homoscedastic" and sampler_mode == "simplex":
215 raise ValueError(
216 "Heteroscedastic error models are not supported with the "
217 "'simplex' sampler; use the unconstrained sampler."
218 )
220 iterations = training_options.get('iterations', 50000)
221 num_components = self.U_hat.shape[1]
222 S_hat = self.S_hat
223 seed = training_options.get('seed')
225 if sampler_mode == "simplex":
226 nu0_chosen = training_options.get('nu0_chosen', 1.0)
227 sigma20_chosen = training_options.get('sigma20_chosen', 0.02)
228 burn = training_options.get('burn', 10000)
229 stepsize = training_options.get('stepsize', 0.001)
230 self._metrics_calculator = None
231 self.mh_acceptance_rate_ = None
232 self.samples = gibbs_sampler_simplex(
233 self.centered_experiment_train,
234 self.U_hat,
235 self.Vt_hat,
236 self.S_hat,
237 iterations,
238 [nu0_chosen, sigma20_chosen],
239 burn=burn,
240 stepsize=stepsize,
241 seed=seed,
242 )
243 else:
244 settings = DEFAULT_SAMPLER_SETTINGS[error_model_mode]
245 terms = VARIANCE_MODELS[error_model_mode]
246 metric_names = required_metrics(error_model_mode)
248 if metric_names:
249 # Fit the metrics (PC centroid, normalization bounds) on the
250 # training predictions, then build the variance design matrix.
251 self._metrics_calculator = HeteroscedasticMetrics(
252 metric_names
253 ).fit(self._train_model_predictions, self.Vt_hat)
254 metrics_train = self._metrics_calculator.compute(
255 self._train_model_predictions
256 )
257 basis_train = variance_basis(metrics_train, terms)
258 else:
259 # Homoscedastic: the variance basis is the constant column
260 # alone, so theta = [sigma^2].
261 self._metrics_calculator = None
262 basis_train = np.ones(
263 (len(self.centered_experiment_train), 1)
264 )
266 b_mean_prior = training_options.get('b_mean_prior', np.zeros(num_components))
267 b_mean_cov = training_options.get('b_mean_cov', np.diag(S_hat**2))
268 self.samples, self.mh_acceptance_rate_ = gibbs_sampler_heteroscedastic(
269 self.centered_experiment_train,
270 self.U_hat,
271 basis_train,
272 iterations,
273 burn=training_options.get('burn', 5000),
274 proposal_scales=training_options.get(
275 'proposal_scales', settings['proposal_scales']
276 ),
277 init_params=training_options.get(
278 'init_params', settings['init_params']
279 ),
280 prior_spec=training_options.get(
281 'prior_spec', settings['prior_spec']
282 ),
283 b_mean_prior=b_mean_prior,
284 b_mean_cov=b_mean_cov,
285 adapt_proposal=training_options.get('adapt_proposal', True),
286 target_acceptance=training_options.get('target_acceptance', 0.25),
287 seed=seed,
288 )
290 # Remember which error model produced self.samples so that
291 # predict()/evaluate() use the matching predictive distribution.
292 self._trained_error_model = error_model_mode
296 def predict(self, property, seed=DEFAULT_PREDICTIVE_SEED):
297 """
298 Predict a specified property using the model weights learned during training.
300 :param property: The property name to predict (e.g., 'ChRad').
301 :param seed: Seed for the posterior predictive draws (subsampling
302 of the posterior samples and the noise added on top).
303 Defaults to a fixed constant so repeated calls are
304 reproducible; pass a different value for independent draws.
305 :return:
306 - rndm_m: array of shape (n_samples, n_points), full posterior draws
307 - lower_df: DataFrame with columns domain_keys + ['Predicted_Lower']
308 - median_df: DataFrame with columns domain_keys + ['Predicted_Median']
309 - upper_df: DataFrame with columns domain_keys + ['Predicted_Upper']
310 """
311 if self.samples is None or self.Vt_hat is None:
312 raise ValueError("Must call `orthogonalize()` and `train()` before predicting.")
314 if property not in self.data_dict:
315 raise KeyError(f"Property '{property}' not found in data_dict.")
317 df = self.data_dict[property].copy()
319 # Infer domain and model columns
320 full_model_cols = self.models
321 domain_keys = [col for col in df.columns if col not in full_model_cols and col != self.truth_column_name]
323 # Determine which models are present
324 available_models = [m for m in full_model_cols if m in df.columns]
326 if len(available_models) == 0:
327 raise ValueError("No available trained models are present in prediction DataFrame.")
329 # Filter predictions and model weights
330 model_preds = df[available_models].values
331 domain_df = df[domain_keys].reset_index(drop=True)
333 rndm_m, (lower, median, upper) = self._posterior_predictive(
334 model_preds, seed=seed
335 )
337 # Build output DataFrames
338 lower_df = domain_df.copy()
340 lower_df["Predicted_Lower"] = lower
342 median_df = domain_df.copy()
343 median_df["Predicted_Median"] = median
345 upper_df = domain_df.copy()
346 upper_df["Predicted_Upper"] = upper
348 return rndm_m, lower_df, median_df, upper_df
350 def evaluate(self, domain_filter=None, seed=DEFAULT_PREDICTIVE_SEED):
351 """
352 Evaluate the model combination using coverage calculation.
354 :param domain_filter: dict with optional domain key ranges, e.g., {"Z": (20, 30), "N": (20, 40)}
355 :param seed: Seed for the posterior predictive draws underlying
356 the coverage calculation. Defaults to a fixed constant so
357 repeated calls are reproducible.
358 :return: coverage list for each percentile
359 """
360 df = self.data_dict[self.current_property]
362 if domain_filter:
363 # Inline optimized filtering
364 for col, cond in domain_filter.items():
365 if col == 'multi' and callable(cond):
366 df = df[df.apply(cond, axis=1)]
367 elif callable(cond):
368 df = df[cond(df[col])]
369 elif isinstance(cond, tuple) and len(cond) == 2:
370 df = df[df[col].between(*cond)]
371 elif isinstance(cond, list):
372 df = df[df[col].isin(cond)]
373 else:
374 df = df[df[col] == cond]
376 # Coverage is only defined where truth values exist.
377 df = df.dropna(subset=[self.truth_column_name])
379 preds = df[self.models].to_numpy()
380 rndm_m, (lower, median, upper) = self._posterior_predictive(preds, seed=seed)
382 return coverage(np.arange(0, 101, 5), rndm_m, df, truth_column=self.truth_column_name)
384 def _posterior_predictive(self, model_preds, seed=DEFAULT_PREDICTIVE_SEED):
385 """
386 Posterior predictive draws for the given model predictions, using
387 the predictive distribution matching the trained error model.
389 :param model_preds: Array of shape (n_points, n_models) with one
390 column per model in ``self.models`` order.
391 :param seed: Seed for the posterior predictive draws.
392 :return: Tuple ``(rndm_m, (lower, median, upper))``.
393 """
394 error_model = self._trained_error_model or "homoscedastic"
395 terms = VARIANCE_MODELS[error_model]
396 if terms:
397 metrics = self._metrics_calculator.compute(model_preds)
398 basis = variance_basis(metrics, terms)
399 else:
400 # Homoscedastic: constant-only variance basis.
401 basis = np.ones((model_preds.shape[0], 1))
402 return rndm_m_heteroscedastic_calculator(
403 model_preds, self.samples, self.Vt_hat, basis, seed=seed
404 )
406 def get_weights(self, summary=True):
407 """
408 Compute model weights from posterior samples.
410 Converts the sampled coefficient vectors (beta) into model weights
411 using the transformation ``omega = beta @ Vt_hat + 1/M``, where M is
412 the number of models. In simplex-constrained mode, all weights are
413 guaranteed to be non-negative and sum to 1.
415 :param summary: If True (default), return a dictionary with
416 ``'mean'``, ``'std'``, ``'median'`` arrays keyed by statistic.
417 If False, return the full ``(n_samples, n_models)`` weight matrix.
418 :return: Weight summary dict or full weight matrix.
419 :raises ValueError: If ``train()`` has not been called.
420 """
421 if self.samples is None or self.Vt_hat is None:
422 raise ValueError("Must call `orthogonalize()` and `train()` before getting weights.")
424 # The first k columns are the PC coefficients; the remaining
425 # columns are variance parameters (a single sigma^2 for the
426 # homoscedastic model, several for heteroscedastic models).
427 betas = self.samples[:, : self.Vt_hat.shape[0]]
428 n_models = self.Vt_hat.shape[1]
429 default_weights = np.full(n_models, 1.0 / n_models)
430 weight_matrix = betas @ self.Vt_hat + default_weights
432 if summary:
433 return {
434 "mean": np.mean(weight_matrix, axis=0),
435 "std": np.std(weight_matrix, axis=0),
436 "median": np.median(weight_matrix, axis=0),
437 "models": self.models,
438 }
439 return weight_matrix