Coverage for pybmc/error_models.py: 100%
60 statements
« prev ^ index » next coverage.py v7.10.0, created at 2026-07-20 21:03 +0000
1"""Heteroscedastic error models for Bayesian model combination.
3All error models share the mean structure of the orthogonalized BMC
4regression and differ only in how the noise variance depends on the
5data point:
7============================= ====================================================
8name sigma_i^2
9============================= ====================================================
10``homoscedastic`` sigma^2 (constant)
11``hetero_pc_dist`` alpha + beta * d_i
12``hetero_model_var`` alpha + beta * v_i
13``hetero_pc_dist_quad`` alpha + beta_1 d_i + beta_2 d_i^2
14``hetero_model_var_quad`` alpha + beta_1 v_i + beta_2 v_i^2
15``hetero_combined_linear`` alpha + beta_d d_i + beta_m v_i
16``hetero_combined_quadratic`` alpha + beta_d1 d_i + beta_d2 d_i^2
17 + beta_m1 v_i + beta_m2 v_i^2
18============================= ====================================================
20with two physics-informed, per-point metrics:
22- ``pc_dist`` (``d``): Euclidean distance from the training-data centroid
23 in principal-component space — grows away from the fitted region.
24- ``model_var`` (``v``): variance among the individual model predictions —
25 grows where the models disagree.
27Both metrics are min-max normalized using the *training* points only, so
28values on extrapolated points can exceed 1.
30Every model — the homoscedastic one included — is parametrized on the
31variance (sigma^2) scale and shares the single likelihood
32``y_i ~ N(X_i . b, sigma_i^2)`` sampled by
33:func:`pybmc.inference_utils.gibbs_sampler_heteroscedastic`; the
34homoscedastic model is simply the case where the variance basis is the
35constant column alone.
36"""
38import numpy as np
40#: Floor applied to a variance only where positivity is *not* guaranteed
41#: by construction: the sampler's initial constant term (estimated from
42#: data residuals, which can be exactly zero for a perfect fit) and
43#: prediction-time variances (the normalized metrics of extrapolated
44#: points can be negative, so ``phi . theta`` can dip below zero there).
45#: During sampling no floor is needed: the variance parameters are kept
46#: strictly positive by the Metropolis-Hastings step and the training
47#: variance basis is non-negative with a leading column of ones, so
48#: every per-point variance is positive by construction.
49VARIANCE_FLOOR = 1e-9
51#: Maps each error-model name to its variance-basis terms beyond the
52#: constant, as ``(metric_name, power)`` tuples. ``homoscedastic`` has no
53#: terms: its variance basis is the constant column alone, and it runs
54#: through the same likelihood and sampler as every other error model.
55VARIANCE_MODELS = {
56 "homoscedastic": [],
57 "hetero_pc_dist": [("pc_dist", 1)],
58 "hetero_model_var": [("model_var", 1)],
59 "hetero_pc_dist_quad": [("pc_dist", 1), ("pc_dist", 2)],
60 "hetero_model_var_quad": [("model_var", 1), ("model_var", 2)],
61 "hetero_combined_linear": [("pc_dist", 1), ("model_var", 1)],
62 "hetero_combined_quadratic": [
63 ("pc_dist", 1),
64 ("pc_dist", 2),
65 ("model_var", 1),
66 ("model_var", 2),
67 ],
68}
70#: Metropolis-Hastings tuning defaults per error model: random-walk
71#: proposal scales (variances of the diagonal Gaussian proposal), initial
72#: values for the non-constant variance parameters, and Gamma prior
73#: ``(shape, scale)`` pairs for every variance parameter (constant term
74#: first).
75DEFAULT_SAMPLER_SETTINGS = {
76 "homoscedastic": {
77 "proposal_scales": [0.05],
78 "init_params": [],
79 "prior_spec": [(2, 10)],
80 },
81 "hetero_pc_dist": {
82 "proposal_scales": [0.05, 0.005],
83 "init_params": [0.01],
84 "prior_spec": [(2, 10), (1, 1)],
85 },
86 "hetero_model_var": {
87 "proposal_scales": [0.05, 0.005],
88 "init_params": [0.01],
89 "prior_spec": [(2, 10), (1, 1)],
90 },
91 "hetero_pc_dist_quad": {
92 "proposal_scales": [5e-2, 5e-3, 5e-3],
93 "init_params": [0.01, 0.001],
94 "prior_spec": [(2, 10), (2, 10), (2, 10)],
95 },
96 "hetero_model_var_quad": {
97 "proposal_scales": [5e-2, 5e-3, 5e-3],
98 "init_params": [0.01, 0.001],
99 "prior_spec": [(2, 10), (2, 10), (2, 10)],
100 },
101 "hetero_combined_linear": {
102 "proposal_scales": [1e-2, 1e-3, 1e-3],
103 "init_params": [0.01, 0.01],
104 "prior_spec": [(2, 10), (2, 10), (2, 10)],
105 },
106 "hetero_combined_quadratic": {
107 "proposal_scales": [1e-2, 1e-3, 1e-3, 1e-3, 1e-3],
108 "init_params": [0.01, 0.001, 0.01, 0.001],
109 "prior_spec": [(2, 10)] * 5,
110 },
111}
114def required_metrics(error_model):
115 """
116 Returns the metric names an error model needs.
118 Args:
119 error_model (str): One of the keys of `VARIANCE_MODELS`.
121 Returns:
122 list[str]: Unique metric names (empty for ``homoscedastic``).
123 """
124 if error_model not in VARIANCE_MODELS:
125 raise ValueError(
126 f"Unknown error model '{error_model}'. "
127 f"Must be one of {tuple(VARIANCE_MODELS)}."
128 )
129 return sorted({metric for metric, _ in VARIANCE_MODELS[error_model]})
132def variance_parameter_names(error_model):
133 """
134 Returns human-readable names of the variance parameters of a model.
136 Args:
137 error_model (str): One of the keys of `VARIANCE_MODELS`.
139 Returns:
140 list[str]: Names such as ``['alpha', 'beta_pc_dist^1', ...]``;
141 ``['sigma^2']`` for the homoscedastic model (whose constant term
142 is the variance itself).
143 """
144 terms = VARIANCE_MODELS[error_model] if error_model in VARIANCE_MODELS else None
145 if terms is None:
146 raise ValueError(
147 f"Unknown error model '{error_model}'. "
148 f"Must be one of {tuple(VARIANCE_MODELS)}."
149 )
150 if not terms:
151 return ["sigma^2"]
152 return ["alpha"] + [f"beta_{metric}^{power}" for metric, power in terms]
155def pc_distance_metric(model_predictions, Vt_hat, pc_centroid):
156 """
157 Distance of each point from a reference centroid in PC space.
159 The principal-component coordinates of a point are obtained by
160 projecting its raw model predictions with ``Vt_hat`` (for a
161 row-centered SVD this is identical to projecting the centered
162 predictions, since the right singular vectors are orthogonal to the
163 all-ones direction).
165 Args:
166 model_predictions (numpy.ndarray): Model outputs, shape
167 ``(n_points, n_models)``.
168 Vt_hat (numpy.ndarray): Scaled right singular vectors, shape
169 ``(components_kept, n_models)``.
170 pc_centroid (numpy.ndarray): Training centroid in PC space,
171 shape ``(components_kept,)``.
173 Returns:
174 numpy.ndarray: Euclidean distances, shape ``(n_points,)``.
175 """
176 pc_coords = model_predictions @ Vt_hat.T
177 return np.linalg.norm(pc_coords - pc_centroid, axis=1)
180def model_variance_metric(model_predictions):
181 """
182 Variance among the model predictions for each point (NaN-aware).
184 Args:
185 model_predictions (numpy.ndarray): Model outputs, shape
186 ``(n_points, n_models)``.
188 Returns:
189 numpy.ndarray: Per-point variance across models, shape ``(n_points,)``.
190 """
191 return np.nanvar(model_predictions, axis=1)
194def variance_basis(metrics, terms):
195 """
196 Builds the variance design matrix ``phi`` for a heteroscedastic model.
198 The per-point variance is ``sigma_i^2 = phi[i] . theta`` where the
199 first column of ``phi`` is ones (the constant term ``alpha``).
201 Args:
202 metrics (dict[str, numpy.ndarray]): Metric arrays keyed by name.
203 terms (list[tuple[str, int]]): ``(metric_name, power)`` pairs.
205 Returns:
206 numpy.ndarray: Basis matrix, shape ``(n_points, 1 + len(terms))``.
207 """
208 n_points = len(next(iter(metrics.values())))
209 columns = [np.ones(n_points)]
210 for metric, power in terms:
211 columns.append(np.asarray(metrics[metric], dtype=float) ** power)
212 return np.column_stack(columns)
215class HeteroscedasticMetrics:
216 """
217 Computes and normalizes the per-point metrics of the error models.
219 The object is fit on the training model predictions: it stores the
220 training centroid in PC space and the training min/max of every
221 metric. Metrics for any other set of points are then computed
222 consistently and scaled with the *training* bounds, so extrapolated
223 points can legitimately exceed 1.
224 """
226 def __init__(self, metric_names):
227 """
228 :param metric_names: Metric names to compute
229 (subset of ``{'pc_dist', 'model_var'}``).
230 """
231 unknown = set(metric_names) - {"pc_dist", "model_var"}
232 if unknown:
233 raise ValueError(f"Unknown metric names: {sorted(unknown)}")
234 self.metric_names = list(metric_names)
235 self.Vt_hat = None
236 self.pc_centroid = None
237 self.bounds = {}
239 def fit(self, train_model_predictions, Vt_hat):
240 """
241 Fit the normalization on the training predictions.
243 Args:
244 train_model_predictions (numpy.ndarray): Training model
245 outputs, shape ``(n_train, n_models)``.
246 Vt_hat (numpy.ndarray): Scaled right singular vectors from
247 the orthogonalization step.
249 Returns:
250 HeteroscedasticMetrics: ``self``, for chaining.
251 """
252 self.Vt_hat = np.asarray(Vt_hat)
253 self.pc_centroid = np.mean(
254 train_model_predictions @ self.Vt_hat.T, axis=0
255 )
256 train_metrics = self._raw_metrics(train_model_predictions)
257 self.bounds = {}
258 for name, values in train_metrics.items():
259 lo, hi = float(np.min(values)), float(np.max(values))
260 self.bounds[name] = (lo, hi)
261 return self
263 def compute(self, model_predictions):
264 """
265 Computes normalized metrics for a set of points.
267 Args:
268 model_predictions (numpy.ndarray): Model outputs, shape
269 ``(n_points, n_models)``.
271 Returns:
272 dict[str, numpy.ndarray]: Normalized metric arrays keyed by name.
273 """
274 if self.Vt_hat is None:
275 raise ValueError("Call `fit()` before `compute()`.")
276 metrics = self._raw_metrics(model_predictions)
277 for name, values in metrics.items():
278 lo, hi = self.bounds[name]
279 if hi > lo:
280 metrics[name] = (values - lo) / (hi - lo)
281 # A metric that is constant on the training set is left
282 # unscaled; the sampler treats it like an extra constant.
283 return metrics
285 def _raw_metrics(self, model_predictions):
286 metrics = {}
287 if "pc_dist" in self.metric_names:
288 metrics["pc_dist"] = pc_distance_metric(
289 model_predictions, self.Vt_hat, self.pc_centroid
290 )
291 if "model_var" in self.metric_names:
292 metrics["model_var"] = model_variance_metric(model_predictions)
293 return metrics