This note explains the implementation corresponding to Sections 6 and 7.2 of Banerjee et al. (2025). The offline predictor for autoregressive design rows is an additional approximation used to remove the remaining PDE dependency.
The default sampling method is “marginal”, matching the smoothed-marginal procedure in equation (9) and Algorithm 4 of the paper. For each variance draw, \[ \Theta_e^{(m)}\mid\sigma^{2(m)},\mathcal D \sim MN(h_e,H_e,\sigma^{2(m)}R_{\rm episode}). \] The coefficients at distinct episodes are sampled independently conditional on the shared variance. The resulting sampling law is \[ q_{\rm marginal}(\Psi\mid\mathcal D) =p(\sigma^2\mid\mathcal D) \prod_e p(\Theta_e\mid\sigma^2,\mathcal D). \] It reproduces the smoothed marginals and does not retain the full posterior dependence across episode coefficients.
The option sampling = “joint” instead uses the backward conditional law in equation (7). With \(B_e=M_eG_{e+1}^{\mathsf T}A_{e+1}^{-1}\), it draws \[ \Theta_e^{(m)}\mid\Theta_{e+1}^{(m)},\sigma^{2(m)},\mathcal D \sim MN\left( m_e+B_e(\Theta_{e+1}^{(m)}-a_{e+1}), M_e-B_eA_{e+1}B_e^{\mathsf T}, \sigma^{2(m)}R_{\rm episode} \right). \] This preserves coefficient dependence and gives the joint posterior sampling law \(q_{\rm joint}(\Psi\mid\mathcal D)=p(\Psi\mid\mathcal D)\).
emulator_learn() defaults to marginal sampling. calibration_prepare() inherits the fitted method or resamples with an explicit sampling option. Saved objects record both the method and whether their samples are joint trajectories. The same two method choices are available for inverse-Wishart emulation, although calibration currently requires the scalar covariance model.
Let \(Y_e\) be an \(N\times C\) training response matrix for episode \(e\), and \(F_e\) its \(N\times P\) autoregressive design. For a fixed emulator posterior coefficient sample set \(m\), the coefficients are \(\Theta_e^{(m)}\) and the variance is \(\sigma^{2(m)}\).
The first GP predicts the design rows. Stack \(H=[F_1|\cdots|F_E]\), center its columns by \(c\), and scale active parameter coordinates to \([0,1]\). A shared squared-exponential kernel with fitted length scales and nugget gives
\[ (K_F)_{ij}=\exp\left[-\frac12\sum_d \frac{(x_{id}-x_{jd})^2}{\ell_d^2}\right] +\lambda_F\mathbf 1\{i=j\}. \]
With \(k_F(\eta)\) the vector of covariances to a new input,
\[ \widehat h(\eta)=c^\top+k_F(\eta)^\top K_F^{-1}(H-\mathbf 1c^\top). \]
Split this row into \(\widehat F_e(\eta)\). At runtime these rows require no PDE solution. Their GP interpolation variance is not currently propagated.
The second GP is the emulator’s original residual interpolation. Its exponential kernel uses the original mechanistic input coordinates:
\[ V_{ij}=a\exp(-\phi\|\eta_i-\eta_j\|)+\delta\mathbf 1\{i=j\}, \qquad J_i(\eta)=a\exp(-\phi\|\eta_i-\eta\|). \]
The nugget is independent new-run noise and is excluded from J even at a coincident input. Gaussian conditioning yields
\[ \mu_e^{(m)}(\eta)^\top =\widehat F_e(\eta)\Theta_e^{(m)} +J(\eta)^\top V^{-1}(Y_e-F_e\Theta_e^{(m)}), \]
\[ q(\eta)=a+\delta-J(\eta)^\top V^{-1}J(\eta). \]
The development name residual_weights means \(V^{-1}(Y_e-F_e\Theta_e^{(m)})\): coefficients over the training parameter inputs, not weights over spatial locations. Shared Cholesky factors are cached; one draw’s residual coefficients are cached at a time.
Each episode is a spatial block at an observation time. The recovery function places each block mean back into its original grid indices. This rearrangement does not estimate correlations. The full-field emulator covariance is supplied separately:
\[ \widetilde\Sigma^{(m)}(\eta) =\sigma^{2(m)}q(\eta)R_{\rm field}. \]
Observed locations select matching entries of the mean and covariance. The recorded kernel scale is preserved: R_field need not have unit diagonal. Its amplitude and nugget must not be silently normalized independently of the emulator variance scale.
The original exact FFBS prediction uses the smoothed coefficient mean. The extended function can instead condition on one fixed coefficient trajectory and return the conditional input covariance. Exact refers to Gaussian conditioning under the chosen surrogate, not an exact PDE solution.
For the observed sites, paper Equation (27) becomes
\[ z_t=y_t(\eta)+u_t+\epsilon_t,\qquad \epsilon_t\sim N(0,\tau_t^2 I), \]
\[ u_t=u_{t-1}+\omega_t,\qquad \omega_t\sim N(0,\tau_t^2 U). \]
The matrix U is a positive-definite discrepancy covariance scale; its diagonal need not equal one. Discrepancy increments have covariance tau2[t] * U, while observation errors retain covariance tau2[t] * I. It is distinct from the emulator covariance R_field. Initial states use \(u_0\mid\tau_0^2\sim N(m_0,\tau_0^2M_0)\) and \(\tau_0^2\sim IG(n_0,d_0)\). Subsequently \(\tau_t^2\sim IG(b^t n_0,b^t d_0)\), with the shape–scale inverse-gamma convention \(p(v)\propto v^{-a-1}\exp(-d/v)\). Stan samples reciprocal variances using the equivalent shape–rate gamma.
Integrating out the Gaussian emulator field gives
\[ z_t\mid\eta,u_t,\tau_t^2,\Psi^{(m)} \sim N\left(\mu_t^{(m)}(\eta)+u_t,\, \widetilde\Sigma^{(m)}(\eta)+\tau_t^2 I\right). \]
The observed input is z. It is not replaced by the GP prediction. Conditional emulator errors are independent across observation times given the fixed coefficient sample set. The joint sampling option preserves dependence between coefficients from different episodes. The marginal option retains the smoothed marginal distributions conditional on the shared variance.
The joint Stan program samples discrepancy innovations. The default marginal program integrates discrepancy states out exactly. Starting at \(a_0=m_0\), \(C_0=\tau_0^2M_0\), the Kalman recursion is
\[ P_t=C_{t-1}+\tau_t^2U,\quad Q_t=P_t+\widetilde\Sigma^{(m)}(\eta)+\tau_t^2 I,\quad e_t=z_t-\mu_t^{(m)}(\eta)-a_{t-1}, \]
\[ a_t=a_{t-1}+P_tQ_t^{-1}e_t,\qquad C_t=P_t-P_tQ_t^{-1}P_t. \]
Its log likelihood is
\[ -\frac12\sum_t\left[n\log(2\pi)+\log|Q_t|+e_t^\top Q_t^{-1}e_t\right]. \]
Gaussian backward sampling subsequently recovers coherent discrepancy paths. The joint and marginal programs target the same conditional calibration model.
For emulator parameters \(\Psi\) and selected sampling law \(q_s\), pooling uses
\[ p_{\rm mod}^{(s)}(\vartheta\mid z,\mathcal D) =\int p(\vartheta\mid z,\Psi,\mathcal D) q_s(\Psi\mid\mathcal D)\,d\Psi \approx \frac1M\sum_m p(\vartheta\mid z,\Psi^{(m)},\mathcal D). \]
Each conditional distribution is normalized separately. Equal numbers of retained samples are pooled. A single model using an averaged likelihood would generally reweight emulator draws according to their field-data evidence and target a different posterior.
Here s denotes the selected emulator sampling method. Joint sampling uses the full emulator posterior in equation (29); marginal sampling uses the product of smoothed conditional marginals defined above. Pooling follows the selected law in both cases.
There are two sampling dimensions: the number of emulator trajectories and the number of calibration transitions conditional on each trajectory. The package defaults to one calibration chain with iter = 1, warmup = 0, and one retained sample per emulator draw. Larger iteration counts are options; warmup defaults to floor(iter / 2) and the retained count defaults to at most 200 available samples per emulator draw. The controls disable adaptation for zero warmup and enable it for positive warmup.
Convergence must be assessed within each conditional fit. One transition from each of 500 initial states is not a converged approximation of this integral.
The implementation retains coefficient samples, emulator variance, conditional GP prediction error, dynamic discrepancy, and independent measurement error. For fixed eta before observing field data, emulator uncertainty decomposes as
\[ \operatorname{Var}(y_t\mid\mathcal D) =E_\Psi[\widetilde\Sigma_\Psi] +\operatorname{Var}_\Psi[\mu_{t,\Psi}]. \]
Uncertainty in the fitted design GP mean is excluded. Under an additional independence assumption, a design error covariance \(C_F(\eta)\) would contribute \(\Theta_e^\top C_F(\eta)\Theta_e\) to output covariance; a full extension would also need to specify dependence between design components and residual errors.
The initial public calibration interface supports the sigma2R emulator, fixed discrepancy covariance scale, finite complete observation matrices, and subsets of the original time/grid setup. It does not silently convert an inverse-Wishart fit, fill missing observations, or forecast new sites or times.