Coverage for pybmc/inference_utils.py: 97%
118 statements
« prev ^ index » next coverage.py v7.10.0, created at 2026-07-20 21:03 +0000
1import numpy as np
3from .error_models import VARIANCE_FLOOR
4from .rng import get_rng
7def gibbs_sampler(y, X, iterations, prior_info=None, burn=5000, seed=None):
8 """
9 Gibbs sampler for the homoscedastic (constant-variance) model.
11 The homoscedastic model is the constant-only special case of the
12 heteroscedastic likelihood (``sigma_i^2 = sigma^2`` for every point),
13 so this is a thin wrapper around `gibbs_sampler_heteroscedastic`
14 with a constant-only variance basis. There is no separate
15 homoscedastic likelihood implementation.
17 Args:
18 y (numpy.ndarray): Response vector (centered).
19 X (numpy.ndarray): Design matrix.
20 iterations (int): Number of retained posterior samples.
21 prior_info (dict, optional): Optional priors with keys
22 ``'b_mean_prior'``, ``'b_mean_cov'`` (Gaussian prior on the
23 coefficients) and ``'prior_spec'`` (Gamma ``(shape, scale)``
24 prior on ``sigma^2``). See `gibbs_sampler_heteroscedastic`.
25 burn (int, optional): Burn-in iterations (default: 5000).
26 seed (int | numpy.random.Generator | None, optional): Seed for
27 the sampler. None (default) draws from the shared
28 package-wide generator (see :mod:`pybmc.rng`).
30 Returns:
31 numpy.ndarray: Posterior samples ``[beta_1..beta_k, sigma^2]``.
32 """
33 prior_info = prior_info or {}
34 variance_basis = np.ones((len(y), 1))
35 samples, _ = gibbs_sampler_heteroscedastic(
36 y,
37 X,
38 variance_basis,
39 iterations,
40 burn=burn,
41 prior_spec=prior_info.get("prior_spec"),
42 b_mean_prior=prior_info.get("b_mean_prior"),
43 b_mean_cov=prior_info.get("b_mean_cov"),
44 seed=seed,
45 )
46 return samples
49def gibbs_sampler_simplex(
50 y, X, Vt_hat, S_hat, iterations, prior_info, burn=10000, stepsize=0.001,
51 seed=None,
52):
53 """
54 Performs Gibbs sampling with simplex constraints on model weights.
56 The likelihood is the same homoscedastic Gaussian model used by the
57 unconstrained sampler (``y_i ~ N(X_i . b, sigma^2)``); only the
58 weight constraint differs.
60 Args:
61 y (numpy.ndarray): Centered response vector.
62 X (numpy.ndarray): Design matrix of principal components.
63 Vt_hat (numpy.ndarray): Normalized right singular vectors.
64 S_hat (numpy.ndarray): Singular values.
65 iterations (int): Number of sampling iterations.
66 prior_info (list[float]): `[nu0, sigma20]` - prior parameters for variance.
67 burn (int, optional): Burn-in iterations (default: 10000).
68 stepsize (float, optional): Proposal step size (default: 0.001).
69 seed (int | numpy.random.Generator | None, optional): Seed for
70 the sampler. None (default) draws from the shared
71 package-wide generator (see :mod:`pybmc.rng`).
73 Returns:
74 numpy.ndarray: Posterior samples ``[beta_1..beta_k, sigma^2]``.
75 """
76 rng = get_rng(seed)
78 bias0 = np.full(len(Vt_hat.T), 1 / len(Vt_hat.T))
79 nu0, sigma20 = prior_info
80 cov_matrix_step = np.diag(S_hat**2 * stepsize**2)
81 n = len(y)
82 b_current = np.full(len(X.T), 0)
83 supermodel_current = X.dot(b_current)
84 residuals_current = y - supermodel_current
85 ssr_current = np.sum(residuals_current**2)
86 sigma2 = ssr_current / len(residuals_current)
87 samples = []
88 acceptance = 0
90 # Validate inputs
91 if burn < 0:
92 raise ValueError("Burn-in iterations must be non-negative.")
93 if stepsize <= 0:
94 raise ValueError("Stepsize must be positive.")
96 for i in range(burn + iterations):
97 b_proposed = rng.multivariate_normal(b_current, cov_matrix_step)
98 omegas_proposed = np.dot(b_proposed, Vt_hat) + bias0
100 # Skip proposals with negative weights
101 if not np.any(omegas_proposed < 0):
102 supermodel_proposed = X.dot(b_proposed)
103 residuals_proposed = y - supermodel_proposed
104 ssr_proposed = np.sum(residuals_proposed**2)
105 # Gaussian likelihood ratio at fixed sigma^2:
106 # exp(-(SSR' - SSR) / (2 sigma^2)).
107 acceptance_prob = min(
108 1,
109 np.exp((ssr_current - ssr_proposed) / (2 * sigma2)),
110 )
111 if rng.uniform() < acceptance_prob:
112 b_current = np.copy(b_proposed)
113 ssr_current = ssr_proposed
114 if i >= burn:
115 acceptance += 1
117 # Sample sigma^2 from its inverse-gamma full conditional
118 # (positive by construction).
119 shape_post = (nu0 + n) / 2.0
120 scale_post = (nu0 * sigma20 + ssr_current) / 2.0
121 sigma2 = 1 / rng.gamma(shape_post, 1 / scale_post)
122 if i >= burn:
123 samples.append(np.append(b_current, sigma2))
125 return np.array(samples)
128def gibbs_sampler_heteroscedastic(
129 y,
130 X,
131 variance_basis,
132 iterations,
133 burn=5000,
134 proposal_scales=None,
135 init_params=None,
136 prior_spec=None,
137 b_mean_prior=None,
138 b_mean_cov=None,
139 adapt_proposal=True,
140 target_acceptance=0.25,
141 seed=None,
142):
143 """
144 Gibbs-within-Metropolis sampler for the shared error-model likelihood.
146 The regression model is ``y_i ~ N(X_i . b, sigma_i^2)`` with a
147 per-point variance that is linear in basis functions of the
148 heteroscedasticity metrics:
150 ``sigma_i^2 = variance_basis[i] . theta``
152 The homoscedastic model is the special case where the basis is a
153 single column of ones, so ``theta = [sigma^2]`` is the constant
154 variance. All error models share this one likelihood.
156 The coefficients ``b`` are updated with a conjugate (weighted
157 least-squares) Gibbs step; the variance parameters ``theta`` with a
158 positivity-constrained Gaussian random-walk Metropolis-Hastings step
159 under Gamma priors. Because ``theta`` is kept strictly positive by
160 the MH step and ``variance_basis`` is validated to be non-negative
161 with a leading column of ones, every per-point variance satisfies
162 ``sigma_i^2 >= theta_1 > 0`` by construction — no variance flooring
163 is applied during sampling. The only floor is on the *initial*
164 constant term, which is estimated from the data residuals and could
165 otherwise be exactly zero for a perfectly fitting model.
167 Args:
168 y (numpy.ndarray): Centered response vector, shape ``(n,)``.
169 X (numpy.ndarray): Design matrix (principal components), shape
170 ``(n, k)``.
171 variance_basis (numpy.ndarray): Variance design matrix ``phi``
172 with a leading column of ones and no negative entries, shape
173 ``(n, p)``. See :func:`pybmc.error_models.variance_basis`.
174 iterations (int): Number of retained posterior samples.
175 burn (int, optional): Burn-in iterations discarded before
176 retention (default: 5000).
177 proposal_scales (list[float], optional): Diagonal of the Gaussian
178 random-walk proposal covariance for ``theta`` (length p).
179 Defaults to ``[1e-2, 1e-3, ..., 1e-3]``.
180 init_params (list[float], optional): Initial values for the
181 non-constant entries of ``theta`` (length p - 1, strictly
182 positive). The constant term starts at the OLS residual
183 variance. Defaults to 0.01 for every term.
184 prior_spec (list[tuple[float, float]], optional): Gamma prior
185 ``(shape, scale)`` for each entry of ``theta`` (length p).
186 Defaults to ``(2, 10)`` for every parameter.
187 b_mean_prior (numpy.ndarray, optional): Prior mean for ``b``
188 (default zeros).
189 b_mean_cov (numpy.ndarray, optional): Prior covariance for ``b``
190 (default ``1e6 * I``, i.e. weakly informative).
191 adapt_proposal (bool, optional): If True (default), rescale the
192 proposal covariance during burn-in toward
193 ``target_acceptance``, so the fixed defaults work across
194 data scales. Adaptation stops at the end of burn-in, which
195 preserves detailed balance for the retained samples.
196 target_acceptance (float, optional): Acceptance rate targeted by
197 the burn-in adaptation (default: 0.25).
198 seed (int | numpy.random.Generator | None, optional): Seed for
199 the sampler. None (default) draws from the shared
200 package-wide generator (see :mod:`pybmc.rng`).
202 Returns:
203 tuple[numpy.ndarray, float]:
204 - `samples` (numpy.ndarray): Posterior samples
205 ``[b_1..b_k, theta_1..theta_p]``, shape
206 ``(iterations, k + p)``.
207 - `acceptance_rate` (float): Post-burn-in MH acceptance rate.
208 """
209 if burn < 0:
210 raise ValueError("Burn-in iterations must be non-negative.")
212 rng = get_rng(seed)
214 n_points, n_betas = X.shape
215 n_params = variance_basis.shape[1]
217 if not np.allclose(variance_basis[:, 0], 1.0):
218 raise ValueError(
219 "The first column of 'variance_basis' must be ones "
220 "(the constant variance term)."
221 )
222 if np.any(variance_basis < 0):
223 raise ValueError(
224 "'variance_basis' must be non-negative so that positive "
225 "variance parameters guarantee positive variances."
226 )
228 if proposal_scales is None:
229 proposal_scales = [1e-2] + [1e-3] * (n_params - 1)
230 if len(proposal_scales) != n_params:
231 raise ValueError(
232 f"'proposal_scales' must have length {n_params} "
233 f"(got {len(proposal_scales)})."
234 )
235 proposal_cov = np.diag(np.asarray(proposal_scales, dtype=float))
237 if prior_spec is None:
238 prior_spec = [(2, 10)] * n_params
239 if len(prior_spec) != n_params:
240 raise ValueError(
241 f"'prior_spec' must have length {n_params} "
242 f"(got {len(prior_spec)})."
243 )
244 prior_shapes = np.array([shape for shape, _ in prior_spec], dtype=float)
245 prior_scales = np.array([scale for _, scale in prior_spec], dtype=float)
247 def log_likelihood(residuals, sigma2):
248 return -0.5 * float(
249 np.sum(np.log(2.0 * np.pi * sigma2) + residuals**2 / sigma2)
250 )
252 def log_prior(theta):
253 # Gamma log-density without the normalization constant, which
254 # cancels in the Metropolis-Hastings ratio.
255 return float(
256 np.sum((prior_shapes - 1.0) * np.log(theta) - theta / prior_scales)
257 )
259 if b_mean_prior is None:
260 b_mean_prior = np.zeros(n_betas)
261 if b_mean_cov is None:
262 b_mean_cov = np.eye(n_betas) * 1e6
263 b_mean_cov_inv = np.linalg.inv(b_mean_cov)
265 b_current = np.linalg.lstsq(X, y, rcond=None)[0]
266 if init_params is None:
267 init_params = [0.01] * (n_params - 1)
268 if len(init_params) != n_params - 1:
269 raise ValueError(
270 f"'init_params' must have length {n_params - 1} "
271 f"(got {len(init_params)})."
272 )
273 if np.any(np.asarray(init_params, dtype=float) <= 0):
274 raise ValueError("'init_params' must be strictly positive.")
275 theta_current = np.concatenate(
276 [[max(np.var(y - X @ b_current), VARIANCE_FLOOR)], init_params]
277 )
279 samples = []
280 accept_count = 0
282 # Burn-in adaptation of the proposal scale (Robbins-Monro style):
283 # every `adapt_interval` iterations the covariance is rescaled toward
284 # the target acceptance rate, then frozen for the sampling phase.
285 proposal_scale_factor = 1.0
286 adapt_interval = 100
287 accept_recent = 0
289 for i in range(burn + iterations):
290 # --- Gibbs step: b | theta (weighted least squares) ---
291 sigma2 = variance_basis @ theta_current
292 Xw = X / sigma2[:, None]
293 b_post_cov = np.linalg.inv(
294 X.T @ Xw + b_mean_cov_inv + np.eye(n_betas) * 1e-6
295 )
296 b_post_mean = b_post_cov @ (
297 Xw.T @ y + b_mean_cov_inv @ b_mean_prior
298 )
299 b_current = rng.multivariate_normal(b_post_mean, b_post_cov)
301 # --- MH step: theta | b (positivity-constrained random walk) ---
302 residuals = y - X @ b_current
303 theta_proposed = np.atleast_1d(
304 rng.multivariate_normal(
305 theta_current, proposal_cov * proposal_scale_factor**2
306 )
307 )
308 if np.all(theta_proposed > 0):
309 sigma2_proposed = variance_basis @ theta_proposed
310 log_ratio = (
311 log_likelihood(residuals, sigma2_proposed)
312 + log_prior(theta_proposed)
313 ) - (
314 log_likelihood(residuals, sigma2)
315 + log_prior(theta_current)
316 )
317 if np.log(rng.uniform()) < log_ratio:
318 theta_current = theta_proposed
319 if i >= burn:
320 accept_count += 1
321 else:
322 accept_recent += 1
324 if adapt_proposal and i < burn and (i + 1) % adapt_interval == 0:
325 recent_rate = accept_recent / adapt_interval
326 proposal_scale_factor *= np.exp(recent_rate - target_acceptance)
327 proposal_scale_factor = float(
328 np.clip(proposal_scale_factor, 1e-6, 1e6)
329 )
330 accept_recent = 0
332 if i >= burn:
333 samples.append(np.concatenate([b_current, theta_current]))
335 acceptance_rate = accept_count / max(iterations, 1)
336 return np.array(samples), acceptance_rate
339def USVt_hat_extraction(U, S, Vt, components_kept):
340 """
341 Extracts reduced-dimensionality matrices from SVD results.
343 Args:
344 U (numpy.ndarray): Left singular vectors.
345 S (numpy.ndarray): Singular values.
346 Vt (numpy.ndarray): Right singular vectors (transposed).
347 components_kept (int): Number of components to retain.
349 Returns:
350 tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
351 - `U_hat` (numpy.ndarray): Reduced left singular vectors.
352 - `S_hat` (numpy.ndarray): Retained singular values.
353 - `Vt_hat` (numpy.ndarray): Normalized right singular vectors.
354 - `Vt_hat_normalized` (numpy.ndarray): Original right singular vectors.
355 """
356 U_hat = np.array([U.T[i] for i in range(components_kept)]).T
357 S_hat = S[:components_kept]
358 Vt_hat = np.array([Vt[i] / S[i] for i in range(components_kept)])
359 Vt_hat_normalized = np.array([Vt[i] for i in range(components_kept)])
360 return U_hat, S_hat, Vt_hat, Vt_hat_normalized