diff --git a/CONTEXT.md b/CONTEXT.md index 6cd8a33..afedd35 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -142,7 +142,7 @@ An endogenous variable with no data, generated inside the model from its own VAR _Avoid_: "missing data" or "unobserved column": nothing is imputed, and a latent series is never passed in. Avoid "state" (state-space vocabulary for a model Impulso does not fit). **Latent block / stationarity constraint**: -The latent-on-latent part of the lag coefficients: the latent equations' coefficients on latent lags, across all lags. The observed series are data, so this block alone decides whether a generated path explodes. The stationarity constraint restricts it to the region where its companion matrix has spectral radius below 1. The posterior is the prior truncated to that region, and prior predictive sampling ignores the constraint. +The latent-on-latent part of the lag coefficients: the latent equations' coefficients on latent lags, across all lags. The observed series are data, so this block alone decides whether a generated path explodes. The stationarity constraint restricts it to the region where its companion matrix has spectral radius below 1, and the posterior is the prior truncated to that region. For one latent series with one lag the block is the latent own lag alone and the region is the interval (−1, 1), so the constraint is a truncated own-lag prior, which prior predictive sampling respects. In every other case it is a Potential, which prior predictive sampling ignores (ADR-0017). _Avoid_: "stationary VAR" for the constraint: the observed equations are unconstrained. Avoid "stationarity pretest", which is the unrelated frequentist test run before specification. **Model evidence (`ModelEvidence`)**: diff --git a/docs/adr/0012-exog-prior-scales-with-data.md b/docs/adr/0012-exog-prior-scales-with-data.md index f304cd7..6cdc08f 100644 --- a/docs/adr/0012-exog-prior-scales-with-data.md +++ b/docs/adr/0012-exog-prior-scales-with-data.md @@ -31,3 +31,7 @@ The default `exog_prior_scale = 100.0` is deliberately **loose, not tight**. Exo - **Not fixed here**: the prior scales but does not centre. An uncentred trend still gives NUTS a strong prior-mean/slope correlation, which costs sampling efficiency even though the inference is now correct. Internal centring for geometry is a separate change, and should be weighed against the same unstandardisation blast radius rejected above. - **Not fixed here**: `intercept ~ Normal(0, 1)` on the PyMC path has exactly the same units problem, on the term the conjugate engine is most diffuse about. Tracked separately. - Keying off `sigma_i` newly exposes the PyMC path to degenerate *endogenous* columns: a constant series gives `sigma_i ≈ 0`, which collapses that whole row of the prior towards zero. `VARData` does not reject constant endogenous columns, and `MinnesotaPrior` does not share the exposure precisely because it ignores the data. The conjugate path has always had it through `minnesota_dummies`. Not guarded here — the same guard belongs with the `MinnesotaPrior` scaling follow-up, where it can be applied once for both terms. + +## Amendment (0.1.4): a per-equation scale + +`exog_prior_scale` also accepts a tuple with one entry per endogenous equation, in `endog_names` order, mirroring `MinnesotaPrior.own_lag_mean`; the formula becomes `sd[i,j] = exog_prior_scale_i * sigma_i / s_j_eff`. The reason is the embedded path (ADR-0016). A latent endogenous series has no data of its own to pull a loose prior back, so the default of 100 lets one prior standard deviation of a regressor's coefficient move the latent series by 100 of its residual standard deviations: an MMM's latent sales baseline drew prior-predictive paths two orders of magnitude larger than the target. A single scalar could only tighten that equation by tightening every observed equation with it. The tuple's length is checked when the exogenous prior is built, because that is where the number of endogenous variables is known; every entry, like the scalar, must be positive and finite. What did not change: the belief is still stated in contribution space, the scalar form applies to every equation, its default stays at 100, and the scalar path gives the same numbers as before. diff --git a/docs/adr/0016-embeddable-var-path.md b/docs/adr/0016-embeddable-var-path.md index ff2ff9f..b98ce2c 100644 --- a/docs/adr/0016-embeddable-var-path.md +++ b/docs/adr/0016-embeddable-var-path.md @@ -35,9 +35,9 @@ The reason is the prototype's funnel. The first design passed a free latent colu **Stationarity of the latent block.** An explosive draw of the latent coefficients makes the generated path explode over the sample, and in the prototype that froze a chain (step size about 1e−278). Three things guard against it: -- **Initial values.** The model's initial point starts every latent equation's own first lag at its prior mean when that is inside the stationary region (|δᵢ| < 1), and at 0.5 otherwise, with every other coefficient in the latent rows at 0. PyMC's default start for a `Normal` is its mean. When every latent δᵢ is inside the region and the prior puts mean 0 on the rest of the latent rows, as the Minnesota prior does, the prior mean is already the stationary start, and `build_in_model` sets no initial value. Otherwise it sets one for `B` as a whole, so the observed rows start at their prior mean, which is PyMC's default start anyway. **PyMC cannot clone a model with a non-default initial value.** `Model.copy()` raises `NotImplementedError`, and so do `pm.do` and `pm.observe`, which clone the model. A latent own-lag prior mean outside the region, the default of 1 included, therefore leaves the host model unclonable. The sampler's jitter moves this start, and the Potential below is what rejects a jittered start that lands outside the region. -- **The `latent_stationarity` Potential.** It is 0 when the spectral radius of the companion matrix of the latent-on-latent block (latent rows, latent lags, all `n_lags`) is below 1, and −∞ otherwise. The observed series are data, so this block alone decides whether the path explodes. `B` stays a single `Normal`, so **the posterior is the prior truncated to the stationary region of the latent block.** A jittered start outside the region has −∞ log density, and PyMC redraws it. The Potential is piecewise constant, so its gradient is zero. PyTensor cannot differentiate the complex `eig`, and PyMC's rewrites strip `disconnected_grad`. The radius is therefore an `OpFromGraph` whose gradient is overridden with zeros (through `pullback`, or `lop_overrides` on older PyTensor). `eig` raises on a non-finite matrix, which would abort the sampler rather than record a divergence. A non-finite companion matrix is therefore swapped for `2·I` first, so the Potential is −∞. -- **What it does not do.** A Potential is not a random variable, so `sample_prior_predictive` ignores `latent_stationarity`. Prior-predictive latent paths come from the untruncated prior and can still explode. If the latent own-lag prior mean is near 1 and the data want persistence near 1, the posterior presses against the boundary and diverges there. On the 09b test setup with the default mean of 1, 212–290 of 400 draws diverged. That is one more reason to set δᵢ = 0 for a stationary deviation. +- **Initial values.** The model's initial point starts every latent equation's own first lag at its prior mean when that is inside the stationary region (|δᵢ| < 1), and at 0.5 otherwise, with every other coefficient in the latent rows at 0. PyMC's default start for a `Normal` is its mean. When every latent δᵢ is inside the region and the prior puts mean 0 on the rest of the latent rows, as the Minnesota prior does, the prior mean is already the stationary start, and `build_in_model` sets no initial value. Otherwise it sets one for `B` as a whole, so the observed rows start at their prior mean, which is PyMC's default start anyway. **PyMC cannot clone a model with a non-default initial value.** `Model.copy()` raises `NotImplementedError`, and so do `pm.do` and `pm.observe`, which clone the model. A latent own-lag prior mean outside the region, the default of 1 included, therefore leaves the host model unclonable. The sampler's jitter moves this start, and the Potential below is what rejects a jittered start that lands outside the region. For one latent series with `n_lags = 1` no initial value is set: the own lag is a truncated Normal that starts at 0 ([ADR-0017](0017-truncated-own-lag-for-one-latent-ar1.md)). +- **The `latent_stationarity` Potential.** It is 0 when the spectral radius of the companion matrix of the latent-on-latent block (latent rows, latent lags, all `n_lags`) is below 1, and −∞ otherwise. The observed series are data, so this block alone decides whether the path explodes. `B` stays a single `Normal`, so **the posterior is the prior truncated to the stationary region of the latent block.** A jittered start outside the region has −∞ log density, and PyMC redraws it. The Potential is piecewise constant, so its gradient is zero. PyTensor cannot differentiate the complex `eig`, and PyMC's rewrites strip `disconnected_grad`. The radius is therefore an `OpFromGraph` whose gradient is overridden with zeros (through `pullback`, or `lop_overrides` on older PyTensor). `eig` raises on a non-finite matrix, which would abort the sampler rather than record a divergence. A non-finite companion matrix is therefore swapped for `2·I` first, so the Potential is −∞. For one latent series with `n_lags = 1` there is no Potential: the own lag is a Normal truncated to (−1, 1) instead ([ADR-0017](0017-truncated-own-lag-for-one-latent-ar1.md)). +- **What it does not do.** A Potential is not a random variable, so `sample_prior_predictive` ignores `latent_stationarity`. Prior-predictive latent paths come from the untruncated prior and can still explode. If the latent own-lag prior mean is near 1 and the data want persistence near 1, the posterior presses against the boundary and diverges there. On the 09b test setup with the default mean of 1, 212–290 of 400 draws diverged. That is one more reason to set δᵢ = 0 for a stationary deviation. For one latent series with `n_lags = 1` neither applies: the truncated own lag is drawn by `sample_prior_predictive` and gives no boundary divergences ([ADR-0017](0017-truncated-own-lag-for-one-latent-ar1.md)). `latent_init`, `latent_innovations` and `latent` have no dims. A latent coordinate would be unprefixed and would collide between two VARs in one model, so the names travel on `handles.latent_names` instead, in the path's column order. @@ -55,7 +55,7 @@ The reason is the prototype's funnel. The first design passed a free latent colu - **Centred latent series with a `Potential` likelihood** — rejected. The caller passes a free latent column as symbolic `endog`, and the VAR logp constrains it. The prototype's evidence rules it out. It is a funnel in the latent innovation scale, with divergences on every seed. A zero-avoiding scale prior did not remove it. - **The conjugate route** — not needed. A closed-form marginal likelihood as a `Potential`, using `ConjugateVAR`'s machinery, would remove the coefficients from the sampler. The PRD set it as the remedy for *coupling*, meaning poor mixing along coefficient/latent ridges. The failure was a funnel. Divergent draws were displaced in log σ_b by −1.9 SD but in the loading by only +0.35 SD. corr(loading, own-lag) was −0.98 in every variant, and divergences did not sit along that ridge. The ridge does limit ESS (about 40–135 in the non-centred runs) and is the thing to watch if mixing becomes the problem. -- **A bounded own-lag variable** (e.g. `tanh` or a stationarity-respecting transform of the latent block) — rejected. For one latent series at VAR(1) a `tanh` bound on the own lag does bound the spectral radius, so there the rejection rests on cost alone. With cross-lags between latent series or `n_lags > 1`, bounding each own lag no longer bounds the radius, and only a full stationarity-respecting transform of the latent block would. Either way the costs are real. The latent rows' prior is no longer the Minnesota `Normal` that every other row gets, and the transform adds free variables. `B` could be reassembled as a `pm.Deterministic("B", dims=("var", "coeff"))`, so the posterior schema would survive. But `handles.B` would then be a Deterministic, and PyMC's initial point covers only free variables, so a `set_initval` on it would be silently ignored. The stationary start would have to move onto the underlying variables. A single `Normal` truncated by a Potential avoids all of this. It costs boundary divergences only when the posterior sits near the boundary, which the default own-lag mean of 1 makes likely. +- **A bounded own-lag variable** (e.g. `tanh` or a stationarity-respecting transform of the latent block) — rejected. For one latent series at VAR(1) a `tanh` bound on the own lag does bound the spectral radius, so there the rejection rests on cost alone. With cross-lags between latent series or `n_lags > 1`, bounding each own lag no longer bounds the radius, and only a full stationarity-respecting transform of the latent block would. Either way the costs are real. The latent rows' prior is no longer the Minnesota `Normal` that every other row gets, and the transform adds free variables. `B` could be reassembled as a `pm.Deterministic("B", dims=("var", "coeff"))`, so the posterior schema would survive. But `handles.B` would then be a Deterministic, and PyMC's initial point covers only free variables, so a `set_initval` on it would be silently ignored. The stationary start would have to move onto the underlying variables. A single `Normal` truncated by a Potential avoids all of this. It costs boundary divergences only when the posterior sits near the boundary, which the default own-lag mean of 1 makes likely. [ADR-0017](0017-truncated-own-lag-for-one-latent-ar1.md) adopts a bounded own lag for one latent series with `n_lags = 1`, as a truncated Normal that keeps the Minnesota prior on the region. - **Sampler-only guidance** — rejected. It means documenting that callers should initialise inside the stationary region, and doing nothing in the graph. Under nutpie's default initialisation the jitter started a chain outside the stationary region. The chain froze at an own-lag of 1.58 with 1000 divergences and R-hat 1.59, and the loading HDI excluded the truth. Guidance does not reach the caller who never reads it, and `build_in_model` knows exactly which rows need a stationary start. ## Consequences diff --git a/docs/adr/0017-truncated-own-lag-for-one-latent-ar1.md b/docs/adr/0017-truncated-own-lag-for-one-latent-ar1.md new file mode 100644 index 0000000..47c9562 --- /dev/null +++ b/docs/adr/0017-truncated-own-lag-for-one-latent-ar1.md @@ -0,0 +1,33 @@ +# One latent series at lag order 1 gets a truncated own lag, not the stationarity Potential + +ADR-0016 keeps a latent block stationary with the `latent_stationarity` Potential: 0 where the companion matrix of the latent-on-latent block has spectral radius below 1, −∞ elsewhere, and a zero gradient everywhere. NUTS cannot see that wall coming and can only diverge against it. What tells a latent series apart from white noise is its persistence, so a caller sets a persistent own-lag prior mean, 0.5 to 0.8, and the posterior then sits about two prior standard deviations from the wall. The adversarial review of pymc-labs/pymc-marketing#3293 (item P1.2) measured divergence rates of up to a third of draws on null data in that setup. Impulso's default own-lag mean of 1 is worse: it starts the latent equation outside the region, which forces an initial value on `B` and leaves the host model unclonable, and ADR-0016 recorded 212–290 divergent draws of 400 under it. The `eig` inside the Potential is also what keeps the JAX backend from differentiating the model (#378). + +For exactly one latent series and `n_lags = 1` the latent block is the 1×1 own lag `B[0, 0]`. Its companion matrix is that number, so the stationary region is the interval `|B[0, 0]| < 1`, and the constraint can be exact and smooth. `build_in_model` then registers three variables in place of the single `Normal` `B`: + +- `latent_own_lag`: the own lag's `Normal`, with its Minnesota mean and standard deviation, truncated to (−1, 1). PyMC samples it through its default interval transform. +- `B_free`: a one-dimensional `Normal` over every other lag coefficient, with the Minnesota means and standard deviations of those entries. +- `B`: the `Deterministic` scattered from the two, under its existing name and `(var, coeff)` dims. + +No Potential and no initial value are registered in this case. Every other case is unchanged. Without latent series `B` is the single `Normal`. With `n_lags > 1` or several latent series it is the single `Normal` with the Potential and the stationary initial value of ADR-0016. + +**The posterior does not change.** The Potential truncated `B`'s Normal prior to `|B[0, 0]| < 1` and left every other entry alone. A Normal truncated to that interval is the same density on the own lag, up to a normalising constant that depends only on the prior's fixed mean and standard deviation, and the other entries keep their Normals. What changes is the geometry NUTS sees: in the transformed coordinate the boundary is infinitely far away, and the log density is smooth all the way to it. + +`_register_coefficients` is written as a free block plus a constrained block, with the constrained block at the flat position of `B[0, 0]`. A general parametrisation of the latent block replaces the constrained block and widens its positions; it does not need a second rewrite. + +## Considered options + +- **A partial-autocorrelation parametrisation of the general latent block** — deferred to its own issue (#397). Mapping unconstrained variables onto the stationary region of the latent VAR(p) is the general fix for `n_lags > 1` and several latent series, where bounding each own lag does not bound the spectral radius. For one series at lag 1 the partial autocorrelation is the own lag, so that parametrisation reduces to a bounded own lag, which the truncated Normal already is. +- **A smooth penalty on the spectral radius** — rejected. A soft barrier in place of the wall changes the posterior. The truncation is the belief. A penalty moves prior mass inside the region, which nothing in the data asked for. +- **Sampler guidance alone** — rejected in ADR-0016 already. +- **ADR-0016's bounded own-lag variable, reconsidered.** ADR-0016 rejected a bounded own lag (`tanh` of a free variable) on cost. For this case its costs do not apply or are worth paying. The truncated Normal is the Minnesota Normal on the region, so the latent row's prior does not change, where a `tanh` of a Normal would change it. It adds no parameters: `B_free` and `latent_own_lag` together have as many entries as `B`. `B` survives as a `Deterministic` under its name and dims, so the posterior schema is unchanged. And it needs no initial value, so the ignored `set_initval` on a `Deterministic` that ADR-0016 worried about never arises. ADR-0016 accepted boundary divergences as the price of the Potential; under the default own-lag mean that price was most of the draws. + +## Consequences + +- The own lag starts at 0, the interval's midpoint, whatever its prior mean, and the free block starts at its prior mean. The latent equation therefore starts inside the region with no initial value, `_latent_b_initval` does not apply to this case, and the host model is clonable (`Model.copy()`, `pm.do`, `pm.observe`) for any own-lag prior mean, the default of 1 included. +- PyMC's start jitter acts in the transformed space, so every jittered start is inside the region and no retry is needed. +- `sample_prior_predictive` draws the own lag from the truncated prior, so prior-predictive latent paths are stationary in this case. The Potential cases still ignore the constraint in prior predictive sampling. +- There is no `eig` in the graph, so the JAX backend differentiates the model in this case. #378 stays open for the Potential cases. +- On the setup of the slow one-latent sampling test under the default own-lag mean of 1, the Potential gave 268–300 divergent draws of 400 on three seeds. The truncated own lag gives none, with R-hat at most 1.011 on the own lag. The slow test allows at most 1% divergent draws and requires R-hat below 1.05. +- The posterior carries two extra variables, `B_free` and `latent_own_lag`. `B` is a `Deterministic`, which PyMC records in the posterior like a free variable, so the posterior schema and `FittedVAR.from_posterior`, which reads `B` by name and ignores variables it does not know, are unchanged. `handles.B` is the `Deterministic`. +- The model's log density includes the truncation's normalising constant, `−log(Φ((1 − μ)/σ) − Φ((−1 − μ)/σ))` for the own lag's prior mean μ and standard deviation σ. It is a constant, so it changes no posterior, but a comparison of the joint density with a reference has to account for it. +- ADR-0016's "Initial values" and "The `latent_stationarity` Potential" bullets now apply only to the other latent cases, and so does its rejection of a bounded own-lag variable. diff --git a/src/impulso/spec.py b/src/impulso/spec.py index 4ea3f8f..e538dca 100644 --- a/src/impulso/spec.py +++ b/src/impulso/spec.py @@ -1,11 +1,12 @@ """VAR model specification.""" +import math from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, Self, cast import numpy as np -from pydantic import Field, model_validator +from pydantic import Field, field_validator, model_validator from impulso._arviz_compat import InferenceDataLike from impulso._base import ImpulsoBaseModel @@ -428,6 +429,59 @@ def _latent_companion(B: Any, n_latent: int, n_vars: int, n_lags: int) -> "pt.Te return pt.concatenate([top, shift], axis=0) +def _register_coefficients( + model: Any, b_mu: np.ndarray, b_sigma: np.ndarray, n_latent: int, n_lags: int +) -> "pt.TensorVariable": + """Register `B`, the VAR lag coefficients with `dims=("var", "coeff")`, and keep its latent block stationary. + + Without latent series `B` is one `Normal` with the prior's means and + standard deviations. With latent series an explosive latent block makes + the generated path explode and freezes the chain, so the block is + constrained to its stationary region, in one of two ways (ADR-0017): + + - One latent series with `n_lags == 1`: the block is the own lag + `B[0, 0]` alone, and the region is `|B[0, 0]| < 1`. `B` is a + `Deterministic` scattered from a free block, `B_free`, a `Normal` over + every other entry with that entry's prior mean and standard deviation, + and a constrained block, `latent_own_lag`, the own lag's `Normal` + truncated to (-1, 1). PyMC samples the own lag through an interval + transform, which starts it at 0 and keeps it inside the region, so no + initial value is set. + - Otherwise: `B` is one `Normal`, started inside the region by + `_set_latent_b_initval` and kept there by the `latent_stationarity` + Potential. + + Args: + model: The active `pymc.Model`. + b_mu: Prior means, shape `(n_vars, n_vars * n_lags)`. + b_sigma: Prior standard deviations, the same shape. + n_latent: Number of latent series (`0` for none). + n_lags: Lag order. + + Returns: + `B`: the `Deterministic` for one latent series with `n_lags == 1`, the + `Normal` otherwise. + """ + import pymc as pm + import pytensor.tensor as pt + + if n_latent == 1 and n_lags == 1: + b_mu = np.asarray(b_mu, dtype=float) + b_sigma = np.asarray(b_sigma, dtype=float) + # Flat, row-major positions in `B`; position 0 is the latent own lag `B[0, 0]`. + constrained = np.array([0]) + free = np.setdiff1d(np.arange(b_mu.size), constrained) + B_free = pm.Normal("B_free", mu=b_mu.ravel()[free], sigma=b_sigma.ravel()[free]) + own_lag = pm.TruncatedNormal("latent_own_lag", mu=b_mu[0, 0], sigma=b_sigma[0, 0], lower=-1.0, upper=1.0) + flat = pt.zeros(b_mu.size)[free].set(B_free)[constrained].set(own_lag[None]) + return pm.Deterministic(COEFFICIENTS, flat.reshape(b_mu.shape), dims=("var", "coeff")) + B = pm.Normal(COEFFICIENTS, mu=b_mu, sigma=b_sigma, dims=("var", "coeff")) + if n_latent: + _set_latent_b_initval(model, B, b_mu, n_latent) + _register_latent_stationarity(B, n_latent, np.shape(b_mu)[0], n_lags) + return B + + def _register_latent_stationarity(B: Any, n_latent: int, n_vars: int, n_lags: int) -> None: """Register the `latent_stationarity` Potential: 0 inside the stationary region, `-inf` outside. @@ -684,7 +738,7 @@ def _register_likelihood( def _exog_prior_sigma( sigma: np.ndarray, x_exog: np.ndarray, - scale: float, + scale: float | Sequence[float], exog_names: Sequence[str] | None = None, ) -> np.ndarray: """Prior standard deviations for the exogenous coefficients `B_exog`. @@ -696,15 +750,17 @@ def _exog_prior_sigma( on large-scale ones effectively unrestricted. This scales the prior so the belief lives in *contribution* space instead: - sd[i, j] = scale * sigma_i / s_j + sd[i, j] = scale_i * sigma_i / s_j where `sigma_i` is the AR(1) residual standard deviation of endogenous variable `i` (the same scale `MinnesotaPrior.build_priors` uses for its lag coefficients — see docs/adr/0015) and `s_j` is the sample standard deviation of exogenous column `j`. One prior standard deviation of - `B_exog[i, j]` then moves variable `i` by `scale` of its own residual - standard deviation when regressor `j` moves by one of its own. The default - `scale` is deliberately loose (see `VAR.exog_prior_scale`). + `B_exog[i, j]` then moves variable `i` by `scale_i` of its own residual + standard deviation when regressor `j` moves by one of its own. `scale_i` + is `scale` itself when `scale` is a scalar and its `i`-th entry when it is + a sequence. The default `scale` is deliberately loose (see + `VAR.exog_prior_scale`). Args: sigma: Per-endogenous-variable AR(1) residual standard deviation, @@ -715,7 +771,9 @@ def _exog_prior_sigma( x_exog: Exogenous regressor block of shape `(T_eff, n_exog)`, already trimmed to the rows the likelihood sees. scale: Multiplier in units of "residual standard deviations of the - dependent variable per standard deviation of the regressor". + dependent variable per standard deviation of the regressor". A + scalar applies to every equation; a sequence gives one entry per + equation, in the order of `sigma`. exog_names: Optional column names, used only to make the constant-column error message readable. @@ -723,6 +781,7 @@ def _exog_prior_sigma( Array of shape `(n_vars, n_exog)` of prior standard deviations. Raises: + ValueError: If `scale` is a sequence whose length is not `sigma`'s. ValueError: If a column of `x_exog` is exactly constant. `VARData` rejects columns that are constant over the whole sample, but trimming the first `n_lags` rows can flatten a column that did @@ -731,6 +790,12 @@ def _exog_prior_sigma( intercept, so the coefficient is not identified; the floor below would happily hand it a wide prior and hide that. """ + scale_arr = np.asarray(scale, dtype=float) + if scale_arr.ndim and scale_arr.shape != sigma.shape: + raise ValueError( + f"exog_prior_scale must be a scalar or have {sigma.size} entries to match n_vars, got {scale_arr.size}" + ) + s = x_exog.std(axis=0, ddof=1) # Checked before the floor is applied: the floor exists to tame columns with # tiny-but-real variation, not to manufacture a scale for columns with none. @@ -745,7 +810,7 @@ def _exog_prior_sigma( ) peak = np.abs(x_exog).max(axis=0) s_eff = np.maximum(s, _EXOG_SD_FLOOR_FRACTION * peak) - return scale * np.outer(sigma, 1.0 / s_eff) + return scale_arr.reshape(-1, 1) * np.outer(sigma, 1.0 / s_eff) @dataclass(frozen=True) @@ -763,7 +828,10 @@ class VARModelHandles: covering only the included equations, when `build_in_model` was given a strict subset via `intercept_equations`; `None` when `intercept_equations` excluded every equation. - B: VAR lag coefficients, `dims=("var", "coeff")`. + B: VAR lag coefficients, `dims=("var", "coeff")`. A `Normal`, except + for one latent series with `n_lags=1`, where it is a + `Deterministic` assembled from `B_free` and `latent_own_lag` (see + "Latent stationarity" in `VAR.build_in_model`). B_exog: Exogenous coefficients, `dims=("var", "exog")`, or `None` when no exogenous block was registered. L: Lower-triangular Cholesky factor of the structural-shock scale @@ -820,7 +888,15 @@ class VAR(ImpulsoBaseModel): engine uses `Vc = 10e6` on the intercept), and the prior's job here is to stop the scale of the regressor from silently setting the answer, not to shrink. Lower it to shrink `B_exog` towards zero. - Applies only to `VAR.fit`; `prior` governs the lag coefficients. + A scalar applies to every equation. A tuple gives one scale per + equation, in `endog_names` order, for a caller that needs one + equation tight (such as an embedded latent series, which has no + data to pull a loose prior back) and the rest loose; its length is + checked against the number of endogenous variables when the graph + is built. Every entry must be positive and finite. Governs the + exogenous block only; `prior` governs the lag coefficients. + `VAR.fit`, `VAR.prior_predictive` and `VAR.build_in_model` build + the same graph, so it applies to all three. error_dist: Observation error distribution — shorthand string (`"gaussian"`, the default, or `"student_t"`) or an `ErrorDistribution` protocol instance. The string form takes the @@ -828,9 +904,6 @@ class VAR(ImpulsoBaseModel): degrees of freedom; pass `StudentT(nu=5.0)` to fix them. Heavy- tailed errors are rejected in combination with time-varying volatility. - Governs the exogenous block only; `prior` governs the lag - coefficients. Both `VAR.fit` and `VAR.prior_predictive` build the - same graph, so it applies to either. """ lags: int | Literal["aic", "bic", "hq"] = Field(...) @@ -844,9 +917,17 @@ class VAR(ImpulsoBaseModel): # (StochasticVolatility) itself has a discriminated field (dynamics/prior). prior: Literal["minnesota"] | MinnesotaPrior | Prior = "minnesota" volatility: Literal["constant", "sv"] | Constant | StochasticVolatility | PyMCVolatilityProcess = "constant" - exog_prior_scale: float = Field(100.0, gt=0) + exog_prior_scale: float | tuple[float, ...] = 100.0 error_dist: Literal["gaussian", "student_t"] | Gaussian | StudentT | ErrorDistribution = "gaussian" + @field_validator("exog_prior_scale") + @classmethod + def _exog_prior_scale_is_positive(cls, value: float | tuple[float, ...]) -> float | tuple[float, ...]: + values = value if isinstance(value, tuple) else (value,) + if not all(math.isfinite(v) and v > 0 for v in values): + raise ValueError(f"exog_prior_scale must be positive and finite, got {value!r}") + return value + @model_validator(mode="after") def _validate_spec(self) -> Self: if self.max_lags is not None and isinstance(self.lags, int): @@ -1052,38 +1133,61 @@ def build_in_model( Latent stationarity: an explosive draw of the latent equations' coefficients makes the generated path explode over the - sample and can freeze a chain. Three things guard against it. The - `pm.Potential` `"latent_stationarity"` is 0 when the spectral radius - of the companion matrix of the latent-on-latent block (the latent - rows' coefficients on latent lags, over all `n_lags`) is below 1 and - `-inf` otherwise. The observed series are data, not generated, so - this block alone decides whether the path explodes. `B` stays a single - `Normal`, so the posterior is its prior truncated to the stationary - region of the latent block; with a latent own-lag prior mean near 1 - the posterior can press against that boundary and give divergences. - The prior sets that mean: `MinnesotaPrior(own_lag_mean=...)` takes - one entry per series, so a caller modelling a stationary deviation - passes 0 for the latent series and 1 for the random-walk observed - ones. And the model's initial point puts every latent equation - inside the stationary region: own first lag at its prior mean when - that is inside the region, otherwise 0.5, and every other - coefficient in the latent rows 0. PyMC's default start for a - `Normal` is its mean, so when every latent own-lag prior mean is - inside the region and the prior puts mean 0 on the rest of the - latent rows, as `MinnesotaPrior` does, no initial value is set. - Otherwise this method sets one for `B` as a whole, so the observed - rows start at their prior mean, which is PyMC's default start for - them anyway. PyMC cannot clone a model with a non-default initial - value: `model.copy()`, `pm.do` and `pm.observe` raise - `NotImplementedError`. A latent own-lag prior mean outside the - stationary region, the default of 1 included, therefore leaves the - host model unclonable. - PyMC's `jitter+adapt_diag` start moves this by up to +-1; a jittered - start outside the region has `-inf` log density, and PyMC redraws it - (`jitter_max_retries`). A Potential is not a random variable, so - `pm.sample_prior_predictive` ignores `"latent_stationarity"`: - prior-predictive latent paths are drawn from the untruncated prior - and can still explode. + sample and can freeze a chain. So the latent-on-latent block (the + latent rows' coefficients on latent lags, over all `n_lags`) is + constrained to its stationary region, where the spectral radius of + its companion matrix is below 1. The observed series are data, not + generated, so this block alone decides whether the path explodes. + Either way the posterior is the coefficients' prior truncated to that + region. How the constraint is enforced depends on the case + (ADR-0017). + + One latent series with `n_lags=1`: the block is the latent own lag + `B[0, 0]` alone, and the region is `|B[0, 0]| < 1`. The own lag is + registered as `latent_own_lag`, a `Normal` with its prior mean and + standard deviation truncated to (-1, 1); every other coefficient is + in `B_free`, a one-dimensional `Normal` with the prior means and standard + deviations of those entries; and `B` is a `Deterministic` assembled + from the two, with the same name and dims as in every other case. + PyMC samples the own lag through an interval transform, so the log + density is smooth up to the boundary and NUTS has no wall to diverge + against, whatever the own-lag prior mean. The own lag starts at 0, + the interval's midpoint, and the free block at its prior mean, so no + initial value is set and the host model can always be cloned + (`model.copy()`, `pm.do`, `pm.observe`); PyMC's start jitter acts in + the transformed space and stays inside the region. Prior predictive + draws of the own lag come from the truncated prior, so + prior-predictive latent paths are stationary. There is no `eig` in + the graph, so the JAX backend can differentiate the model. The + truncation's normalising constant is part of the model's log + density. + + Every other latent case (`n_lags > 1`, or several latent series): + `B` stays a single `Normal`, and the `pm.Potential` + `"latent_stationarity"` is 0 inside the region and `-inf` outside. + With a latent own-lag prior mean near 1 the posterior can press + against that boundary and give divergences. The prior sets that + mean: `MinnesotaPrior(own_lag_mean=...)` takes one entry per series, + so a caller modelling a stationary deviation passes 0 for the latent + series and 1 for the random-walk observed ones. The model's initial + point puts every latent equation inside the stationary region: own + first lag at its prior mean when that is inside the region, + otherwise 0.5, and every other coefficient in the latent rows 0. + PyMC's default start for a `Normal` is its mean, so when every + latent own-lag prior mean is inside the region and the prior puts + mean 0 on the rest of the latent rows, as `MinnesotaPrior` does, no + initial value is set. Otherwise this method sets one for `B` as a + whole, so the observed rows start at their prior mean, which is + PyMC's default start for them anyway. PyMC cannot clone a model with + a non-default initial value: `model.copy()`, `pm.do` and + `pm.observe` raise `NotImplementedError`. A latent own-lag prior mean + outside the stationary region, the default of 1 included, therefore + leaves the host model unclonable. PyMC's `jitter+adapt_diag` start + moves this by up to +-1; a jittered start outside the region has + `-inf` log density, and PyMC redraws it (`jitter_max_retries`). A + Potential is not a random variable, so `pm.sample_prior_predictive` + ignores `"latent_stationarity"`: prior-predictive latent paths are + drawn from the untruncated prior and can still explode. Nesting: open a `pm.Model(name=prefix)` before calling this method and every free random variable, `Deterministic` and the likelihood @@ -1296,13 +1400,9 @@ def build_in_model( # equations; excluding every equation drops the term entirely. intercept, intercept_term = _register_intercept(intercept_mask) - # VAR coefficients with Minnesota prior - B = pm.Normal( - COEFFICIENTS, - mu=prior_params["B_mu"], - sigma=prior_params["B_sigma"], - dims=("var", "coeff"), - ) + # VAR coefficients with Minnesota prior. With latent series this also keeps the latent block + # stationary (ADR-0016, ADR-0017). + B = _register_coefficients(model, prior_params["B_mu"], prior_params["B_sigma"], n_latent, n_lags) # Exogenous coefficients. The prior scales with the data so that it # encodes the same belief regardless of the units the regressors @@ -1337,10 +1437,6 @@ def build_in_model( # non-centred, so the observed rows are conditioned on their innovations `z` (ADR-0016). latent = z = None if n_latent: - # Start the latent equations inside the stationary region: an explosive own-lag - # makes the generated path explode and freezes the chain. - _set_latent_b_initval(model, B, prior_params["B_mu"], n_latent) - _register_latent_stationarity(B, n_latent, n_vars, n_lags) latent, z = _latent_path(endog, X_exog, n_lags, n_vars, n_latent, intercept_term, B, B_exog, L, init_sigma) full = pt.concatenate([latent, pt.as_tensor_variable(endog)], axis=1) Y, X_lag, _ = build_lag_design_matrix(full, n_lags) diff --git a/tests/test_build_in_model.py b/tests/test_build_in_model.py index 9d7491b..a04008e 100644 --- a/tests/test_build_in_model.py +++ b/tests/test_build_in_model.py @@ -78,6 +78,16 @@ def _pinned_intercept_prior_logp(full_model, point, pinned: list[int]) -> float: return float(np.sum(np.asarray(elementwise)[pinned])) +def _b_exog_prior_sigma(exog_prior_scale, **kwargs) -> np.ndarray: + """The prior standard deviations `build_in_model` gives `B_exog`, registered into a fresh model.""" + import pymc as pm + + with pm.Model(): + handles = VAR(lags=kwargs["n_lags"], exog_prior_scale=exog_prior_scale).build_in_model(**kwargs) + rv = handles.B_exog + return np.asarray(rv.owner.op.dist_params(rv.owner)[1].eval()) + + class TestWrapperLogpParity: """Pins `_build_pymc_model`'s numeric output across the 08a refactor. @@ -225,6 +235,38 @@ def test_returns_none_b_exog_when_no_exog(self, rng): assert handles.B_exog is None + def test_per_equation_exog_prior_scale_scales_each_row(self, rng): + """Row `i` of `B_exog`'s prior sd is `exog_prior_scale[i]` times the row a scale of 1 gives.""" + data = _make_data(rng, n_vars=3, exog_names=["z1", "z2"]) + kwargs = { + "endog": data.endog, + "exog": data.exog, + "n_lags": 1, + "endog_names": data.endog_names, + "exog_names": data.exog_names, + } + + unit = _b_exog_prior_sigma(1.0, **kwargs) + got = _b_exog_prior_sigma((1.0, 100.0, 5.0), **kwargs) + + np.testing.assert_allclose(got, np.array([[1.0], [100.0], [5.0]]) * unit) + + @pytest.mark.parametrize("exog_prior_scale", [(1.0, 2.0), (1.0, 2.0, 3.0, 4.0)]) + def test_exog_prior_scale_of_the_wrong_length_raises(self, rng, exog_prior_scale): + import pymc as pm + + data = _make_data(rng, n_vars=3, exog_names=["z"]) + spec = VAR(lags=1, exog_prior_scale=exog_prior_scale) + + with pm.Model(), pytest.raises(ValueError, match="exog_prior_scale must be a scalar or have 3 entries"): + spec.build_in_model( + endog=data.endog, + exog=data.exog, + n_lags=1, + endog_names=data.endog_names, + exog_names=data.exog_names, + ) + def test_registers_into_the_active_model_context(self, rng): """No prefix argument: `build_in_model` writes into `pymc.modelcontext(None)`.""" import pymc as pm @@ -1015,7 +1057,7 @@ def _stationary_point(model, seed: int = _STATIONARY_POINT_SEED) -> dict: point = model.initial_point(random_seed=0) for name, value in point.items(): shape = np.shape(value) - if name == "B": + if name in ("B", "B_free"): point[name] = 0.1 * local.standard_normal(shape) elif name == "sigma_sd_log__": point[name] = np.log(local.uniform(0.5, 1.5, size=shape)) @@ -1294,6 +1336,16 @@ def test_latent_init_sigma_sets_the_initial_value_prior(self, rng): }) np.testing.assert_allclose(np.ravel(logp), stats.norm(0, 3.0).logpdf([0.4, -1.2])) + def test_per_equation_exog_prior_scale_reaches_the_latent_row(self, rng): + """A tight scale on the latent equation leaves the observed equations on the default.""" + kwargs = _latent_setup(rng) + + default = _b_exog_prior_sigma(100.0, **kwargs) + got = _b_exog_prior_sigma((1.0, 100.0, 100.0), **kwargs) + + np.testing.assert_allclose(got[0], default[0] / 100.0) + np.testing.assert_allclose(got[1:], default[1:]) + @pytest.mark.slow @pytest.mark.parametrize("seed", [1, 2, 3]) def test_small_latent_var_samples(self, seed): @@ -1301,9 +1353,8 @@ def test_small_latent_var_samples(self, seed): The setup is gentle: a weak loading and a tight prior on the latent innovation scale (`_gentle_latent_volatility`). The latent own-lag - prior mean is 0, set through `MinnesotaPrior(own_lag_mean=...)`: with - the Minnesota mean of 1 this posterior presses against the - stationarity boundary and diverges heavily. + prior mean is 0, set through `MinnesotaPrior(own_lag_mean=...)`; the + default mean of 1 has its own test in `TestTruncatedOwnLag`. `build_in_model` starts the latent own-lag inside the stationary region and keeps it there, so no `initvals` are passed. The check is that sampling runs, no chain freezes and divergences stay a small @@ -1346,15 +1397,21 @@ def test_small_latent_var_samples(self, seed): def test_path_and_joint_logp_across_lag_orders_and_latent_counts(self, rng, n_lags, n_latent): """The recursion and the change-of-variables identity hold for any lag order and any number of latent series, including the `n_lags == 1` - and `n_latent == 1` edge cases whose static shapes scan must keep.""" + and `n_latent == 1` edge cases whose static shapes scan must keep. + Inside the stationary region the coefficients' density, constraint + included, is the Minnesota Normal prior, less the truncated own lag's + normalising constant for one latent series with `n_lags == 1`.""" import pymc as pm from scipy import stats + from impulso.priors import MinnesotaPrior + T, n_obs = 30, 2 n_vars = n_latent + n_obs obs = rng.standard_normal((T, n_obs)) exog = rng.standard_normal((T, 1)) latent_names = [f"b{i}" for i in range(n_latent)] + scales = np.linspace(0.5, 2.0, n_vars) with pm.Model() as model: VAR(lags=n_lags).build_in_model( endog=obs, @@ -1362,7 +1419,7 @@ def test_path_and_joint_logp_across_lag_orders_and_latent_counts(self, rng, n_la n_lags=n_lags, endog_names=[*latent_names, "y1", "y2"], exog_names=["x"], - endog_scales=np.linspace(0.5, 2.0, n_vars), + endog_scales=scales, latent_names=latent_names, ) @@ -1389,6 +1446,16 @@ def test_path_and_joint_logp_across_lag_orders_and_latent_counts(self, rng, n_la total = float(values["obs"]) + stats.norm.logpdf(z).sum() + jacobian assert total == pytest.approx(joint, rel=1e-10) + prior = MinnesotaPrior().build_priors(n_vars=n_vars, n_lags=n_lags, sigma=scales) + expected = stats.norm.logpdf(B, prior["B_mu"], prior["B_sigma"]).sum() + if (n_lags, n_latent) == (1, 1): + mu, sd = prior["B_mu"][0, 0], prior["B_sigma"][0, 0] + expected -= np.log(stats.norm.cdf((1.0 - mu) / sd) - stats.norm.cdf((-1.0 - mu) / sd)) + terms = [model["B_free"], model["latent_own_lag"]] + else: + terms = [model["B"], model["latent_stationarity"]] + assert float(model.compile_logp(vars=terms, jacobian=False)(point)) == pytest.approx(expected, rel=1e-10) + @pytest.mark.parametrize("symbolic", [False, True]) @pytest.mark.parametrize("n_lags", [1, 2]) def test_compiles_under_nutpie_with_a_single_latent_series(self, rng, n_lags, symbolic): @@ -1645,9 +1712,10 @@ def build_priors(self, n_vars, n_lags, *, sigma): class TestLatentModelIsClonable: """PyMC's `Model.copy()` and `pm.do` refuse a model with a non-default initial value. - `build_in_model` therefore sets one for `B` only when the stationary start - of the latent rows differs from the prior mean, which is PyMC's default - start for a `Normal`. + In the Potential cases `build_in_model` therefore sets one for `B` only when + the stationary start of the latent rows differs from the prior mean, which + is PyMC's default start for a `Normal`. One latent series with `n_lags=1` + never sets one (`TestTruncatedOwnLag`). """ @staticmethod @@ -1765,7 +1833,11 @@ def _numpy_companion(B: np.ndarray, n_latent: int, n_vars: int, n_lags: int) -> class TestLatentStationarity: - """The `latent_stationarity` Potential: -inf outside the stationary region of the latent block.""" + """The `latent_stationarity` Potential: -inf outside the stationary region of the latent block. + + One latent series with `n_lags == 1` has no Potential (`TestTruncatedOwnLag`), so these tests use the + other cases. + """ @staticmethod def _potential(model, B: np.ndarray) -> float: @@ -1773,8 +1845,7 @@ def _potential(model, B: np.ndarray) -> float: fn = model.compile_fn(potential, inputs=model.value_vars, on_unused_input="ignore") return float(fn({**model.initial_point(random_seed=0), "B": B})) - @pytest.mark.parametrize("n_lags", [1, 2]) - @pytest.mark.parametrize("n_latent", [1, 2]) + @pytest.mark.parametrize(("n_lags", "n_latent"), [(1, 2), (2, 1), (2, 2)]) def test_potential_is_zero_inside_and_minus_inf_outside(self, rng, n_lags, n_latent): import pymc as pm @@ -1800,8 +1871,7 @@ def test_potential_is_zero_inside_and_minus_inf_outside(self, rng, n_lags, n_lat combined[0, 0] = combined[0, n_vars] = 0.6 if n_latent == 2: combined[0, 1] = combined[1, 0] = 0.8 - if n_lags == 2 or n_latent == 2: - assert self._potential(model, combined) == -np.inf + assert self._potential(model, combined) == -np.inf @pytest.mark.parametrize("n_lags", [1, 2, 3]) @pytest.mark.parametrize("n_latent", [1, 2]) @@ -1816,22 +1886,22 @@ def test_companion_matches_numpy_reference(self, rng, n_lags, n_latent): ) def test_jittered_explosive_starts_are_rejected(self, rng): - """PyMC's jitter moves the own-lag start of 0.5 by up to 1, so without the - constraint about a quarter of starts are explosive; the init retry - (`jitter_max_retries`) must reject them all.""" + """PyMC's jitter moves each latent block entry (own lags 0.5, cross-lags 0) + by up to 1, so without the constraint many starts are explosive; the + init retry (`jitter_max_retries`) must reject them all.""" import pymc as pm from pymc.sampling.mcmc import _init_jitter with pm.Model() as model: - VAR(lags=1).build_in_model(**_latent_setup_n(rng, 1, 1)) + VAR(lags=1).build_in_model(**_latent_setup_n(rng, 2, 1)) points = _init_jitter(model, None, list(range(60)), jitter=True, jitter_max_retries=50) - own_lags = np.array([point["B"][0, 0] for point in points]) - assert np.abs(own_lags).max() < 1.0 + radii = [np.max(np.abs(np.linalg.eigvals(_numpy_companion(point["B"], 2, 4, 1)))) for point in points] + assert max(radii) < 1.0 - @pytest.mark.parametrize("n_lags", [1, 2]) + @pytest.mark.parametrize(("n_lags", "n_latent"), [(2, 1), (1, 2)]) @pytest.mark.parametrize("value", [np.nan, np.inf]) - def test_non_finite_latent_block_gives_minus_inf_not_an_error(self, rng, value, n_lags): + def test_non_finite_latent_block_gives_minus_inf_not_an_error(self, rng, value, n_lags, n_latent): """`eig` raises on a non-finite matrix; the Potential must turn that into -inf so NUTS records a divergence instead of aborting.""" import pymc as pm @@ -1839,7 +1909,7 @@ def test_non_finite_latent_block_gives_minus_inf_not_an_error(self, rng, value, from packaging.version import Version with pm.Model() as model: - VAR(lags=n_lags).build_in_model(**_latent_setup_n(rng, 1, n_lags)) + VAR(lags=n_lags).build_in_model(**_latent_setup_n(rng, n_latent, n_lags)) point = model.initial_point(random_seed=0) B = point["B"].copy() @@ -1863,6 +1933,146 @@ def test_logp_and_dlogp_compile_under_jax(self, rng): assert np.all(np.isfinite(model.compile_dlogp(mode="JAX")(point))) +def _one_latent_lag_one_model(rng: np.random.Generator, prior="minnesota"): + """A model with one latent series `b0` ahead of two observed series and `n_lags == 1`, and its setup.""" + import pymc as pm + + kwargs = _latent_setup_n(rng, 1, 1) + with pm.Model() as model: + VAR(lags=1, prior=prior).build_in_model(**kwargs) + return model, kwargs + + +def _point_with_own_lag(model, own_lag: float) -> dict: + """`model`'s initial point with the latent own lag at `own_lag`, set through its transform.""" + rv = model["latent_own_lag"] + transformed = model.rvs_to_transforms[rv].forward(np.float64(own_lag), *rv.owner.inputs).eval() + return {**model.initial_point(random_seed=0), model.rvs_to_values[rv].name: transformed} + + +class TestTruncatedOwnLag: + """One latent series with `n_lags == 1`: the latent block is the own lag alone, so the stationary region + is `|B[0, 0]| < 1`. The own lag is its Minnesota Normal truncated to (-1, 1), and `B` is a `Deterministic`.""" + + def test_b_is_a_deterministic_assembled_from_the_free_block_and_the_own_lag(self, rng): + from pymc.distributions.transforms import Interval + + model, kwargs = _one_latent_lag_one_model(rng) + + assert model["B"] not in model.free_RVs + assert model["B"] in model.deterministics + assert model.named_vars_to_dims["B"] == ("var", "coeff") + own_lag = model["latent_own_lag"] + assert own_lag in model.free_RVs + assert isinstance(model.rvs_to_transforms[own_lag], Interval) + assert model.initial_point(random_seed=0)["B_free"].size == 3 * 3 - 1 + assert "latent_stationarity" not in model.named_vars + + point = _point_with_own_lag(model, 0.7) + point["B_free"] = _prior_mu(model["B_free"]) + want = _minnesota_b_mu(kwargs) + want[0, 0] = 0.7 + np.testing.assert_allclose(_evaluate(model, point, ["B"])["B"], want) + + @pytest.mark.parametrize("own_lag", [0.95, 0.999]) + def test_the_gradient_near_the_boundary_is_finite_and_non_zero(self, rng, own_lag): + model, _ = _one_latent_lag_one_model(rng) + + grad = model.compile_dlogp(vars=[model["latent_own_lag"]])(_point_with_own_lag(model, own_lag)) + assert np.all(np.isfinite(grad)) + assert np.all(grad != 0.0) + + @pytest.mark.parametrize("transformed", [-50.0, 50.0]) + def test_logp_far_out_in_the_transformed_space_is_finite(self, rng, transformed): + model, _ = _one_latent_lag_one_model(rng) + + value = model.rvs_to_values[model["latent_own_lag"]] + point = {**model.initial_point(random_seed=0), value.name: np.float64(transformed)} + assert np.isfinite(model.compile_logp()(point)) + + def test_every_jittered_start_is_inside_the_region_without_a_retry(self, rng): + from pymc.sampling.mcmc import _init_jitter + + model, _ = _one_latent_lag_one_model(rng) + + # With no retry allowed, a start with -inf log density raises `SamplingError`. + points = _init_jitter(model, None, list(range(60)), jitter=True, jitter_max_retries=0) + own_lags = np.array([_evaluate(model, point, ["B"])["B"][0, 0] for point in points]) + assert np.abs(own_lags).max() < 1.0 + + @pytest.mark.parametrize("own_lag_mean", [0.0, 1.0, -0.9]) + def test_no_initial_value_and_the_own_lag_starts_at_zero_whatever_its_mean(self, rng, own_lag_mean): + from impulso.priors import MinnesotaPrior + + prior = MinnesotaPrior(own_lag_mean=(own_lag_mean, 1.0, 1.0)) + model, kwargs = _one_latent_lag_one_model(rng, prior) + + assert all(initval is None for initval in model.rvs_to_initial_values.values()) + want = prior.build_priors(n_vars=3, n_lags=1, sigma=kwargs["endog_scales"])["B_mu"] + want[0, 0] = 0.0 + np.testing.assert_array_equal(_evaluate(model, model.initial_point(random_seed=0), ["B"])["B"], want) + model.copy() + + def test_prior_predictive_own_lag_is_truncated(self, rng): + import pymc as pm + + from impulso.priors import MinnesotaPrior + + model, _ = _one_latent_lag_one_model(rng, MinnesotaPrior(own_lag_mean=(0.95, 1.0, 1.0))) + with model: + idata = pm.sample_prior_predictive(draws=2000, var_names=["B"], random_seed=1) + + own_lags = np.asarray(idata.prior["B"])[..., 0, 0] + assert own_lags.size == 2000 + assert np.abs(own_lags).max() < 1.0 + assert own_lags.mean() < 0.95 + + def test_logp_and_dlogp_compile_under_jax(self, rng): + pytest.importorskip("jax") + model, _ = _one_latent_lag_one_model(rng) + + point = _point_with_own_lag(model, 0.9) + assert np.isfinite(model.compile_logp(mode="JAX")(point)) + assert np.all(np.isfinite(model.compile_dlogp(mode="JAX")(point))) + assert np.all(model.compile_dlogp(vars=[model["latent_own_lag"]], mode="JAX")(point) != 0.0) + + @pytest.mark.slow + @pytest.mark.parametrize("seed", [1, 2, 3]) + def test_default_own_lag_mean_samples_without_a_divergence_storm(self, seed): + """`test_default_init_samples_with_own_lag_mean_zero`'s setup under Impulso's default own-lag + mean of 1, where the Potential made most draws divergent (ADR-0016).""" + import arviz as az + import pymc as pm + + full = _small_latent_var_data() + + draws, chains = 200, 2 + with pm.Model(): + VAR(lags=1, volatility=_gentle_latent_volatility()).build_in_model( + endog=full[:, 1:], + exog=None, + n_lags=1, + endog_names=["b", "y"], + endog_scales=[0.5, 0.3], + latent_names=["b"], + intercept_equations=["y"], + ) + idata = pm.sample( + draws=draws, + tune=300, + chains=chains, + cores=1, + random_seed=seed, + progressbar=False, + nuts_sampler="pymc", + ) + + divergences = int(np.asarray(idata.sample_stats["diverging"]).sum()) + assert divergences <= 0.01 * draws * chains, f"{divergences} divergences out of {draws * chains}" + rhat = float(az.rhat(np.asarray(idata.posterior["B"])[:, :, 0, 0])) + assert rhat < 1.05, f"R-hat of the latent own lag is {rhat:.3f}" + + class TestEmbeddedPathRejections: """`VAR.build_in_model` rejects options the embedded path cannot support. diff --git a/tests/test_fitted_from_posterior.py b/tests/test_fitted_from_posterior.py index 2b16147..23a3974 100644 --- a/tests/test_fitted_from_posterior.py +++ b/tests/test_fitted_from_posterior.py @@ -431,6 +431,32 @@ def test_accepts_nu_under_student_t(self, var_data): assert isinstance(fitted.error_dist, StudentT) +# --------------- Acceptance: truncated latent own lag --------------- + + +class TestFromPosteriorAcceptsTheTruncatedOwnLagLayout: + """`build_in_model` with one latent series and `n_lags=1` records `B` as a `Deterministic` assembled from + the free block `B_free` and the truncated own lag `latent_own_lag`, both also in the posterior.""" + + def test_coefficients_are_the_b_draws_and_the_extra_variables_are_ignored(self, var_data): + rng = np.random.default_rng(5) + own_lag = rng.uniform(-0.9, 0.9, size=(2, 5)) + free = 0.2 * rng.standard_normal((2, 5, 3)) + B = np.concatenate([own_lag[..., None], free], axis=-1).reshape(2, 5, 2, 2) + idata = _build_posterior( + extra={ + "B": xr.DataArray(B, dims=["chain", "draw", "var", "coeff"]), + "B_free": xr.DataArray(free, dims=["chain", "draw", "B_free_dim_0"]), + "latent_own_lag": xr.DataArray(own_lag, dims=["chain", "draw"]), + } + ) + + fitted = FittedVAR.from_posterior(idata, var_data, n_lags=1) + + np.testing.assert_array_equal(fitted.coefficients, B) + assert np.isfinite(fitted.forecast(steps=2).median().to_numpy()).all() + + # --------------- Acceptance: real estimator posteriors --------------- diff --git a/tests/test_spec.py b/tests/test_spec.py index 5a1f348..161a173 100644 --- a/tests/test_spec.py +++ b/tests/test_spec.py @@ -340,6 +340,18 @@ def test_frozen(self): with pytest.raises(ValidationError): spec.exog_prior_scale = 5.0 + @pytest.mark.parametrize("bad", [(1.0, 0.0), (1.0, -1.0), (1.0, np.inf), (np.nan,), np.inf]) + def test_rejects_non_positive_or_non_finite_entries(self, bad): + with pytest.raises(ValidationError, match="exog_prior_scale must be positive and finite"): + VAR(lags=1, exog_prior_scale=bad) + + def test_per_equation_scale_round_trips_through_json(self): + spec = VAR(lags=1, exog_prior_scale=(1.0, 100.0, 5.0)) + restored = VAR.model_validate(spec.model_dump(mode="json")) + assert restored == spec + assert restored.exog_prior_scale == (1.0, 100.0, 5.0) + assert VAR.model_validate_json(spec.model_dump_json()) == spec + class TestExogPriorSigma: """`_exog_prior_sigma` puts the B_exog prior in contribution space. @@ -489,6 +501,16 @@ def test_knob_reaches_the_graph(self, rng): _exog_prior_sigma(ar1_residual_sd(endog), exog[1:], 5.0), ) + def test_per_equation_knob_reaches_the_graph(self, rng): + endog = rng.standard_normal((150, 2)) + exog = rng.standard_normal((150, 2)) * np.array([0.01, 250.0]) + data = _exog_data(endog, exog, ["tiny", "huge"]) + + unit = _captured_sigma(self._capture(data, VAR(lags=1, exog_prior_scale=1.0)), "B_exog") + got = _captured_sigma(self._capture(data, VAR(lags=1, exog_prior_scale=(1.0, 100.0))), "B_exog") + + np.testing.assert_allclose(got, np.array([[1.0], [100.0]]) * unit) + def test_sigma_uses_lag_trimmed_rows(self, rng): """The prior keys off the rows the likelihood sees, not the raw column.""" endog = rng.standard_normal((150, 2))