Analytic sine VI and MF-VI examples#

These examples use the analytically tractable sine-series inverse problem from the VI/MF-VI benchmark. The forward model is

\[u''(x)=f(x;\theta), \qquad u(0)=u(1)=0,\]

with

\[f(x;\theta)=\sum_{j=1}^{7}\theta_j\sin(j\pi x),\]

so the exact solution is

\[u(x;\theta) =-\sum_{j=1}^{7}\frac{\theta_j}{(j\pi)^2}\sin(j\pi x).\]

The prior is

\[\theta \sim \mathcal N(0,0.75^2 I),\]

with synthetic truth

\[\theta^\star =(1.00,-0.80,0.65,-0.50,0.35,-0.25,0.15)^T,\]

and observation-noise standard deviation \(\sigma_\eta=5\times10^{-3}\). Because the forward map is linear and the prior and likelihood are Gaussian, the exact posterior is available in closed form,

\[C_{\rm post}^{-1} = C_0^{-1}+G^T\Gamma_{\rm obs}^{-1}G,\]
\[m_{\rm post} = C_{\rm post} \left(C_0^{-1}m_0+G^T\Gamma_{\rm obs}^{-1}y\right).\]

This lets both examples measure convergence directly against the exact posterior.

VI-paper benchmark configuration#

The diagonal example uses the Newton and multifidelity settings from the VI-paper analytic-sine benchmark. The production run uses a full Hessian, lagged stochastic curvature with averaging factor \(\beta=0.25\), natural coordinates, and the paper line-search scales:

optimizer = VINewtonOptimizerConfig(
    newton_metric="natural",
    newton_hessian_type="full",
    newton_curvature_strategy="lagged",
    newton_hessian_averaging_factor=0.25,
    newton_regularization=5e-4,
    gradient_norm_tolerance=0.0,
    max_iterations=1000,
)

line_search = VIStochasticNonmonotoneLineSearchConfig(
    initial_step_size=0.25,
    max_step_size=1.0,
)

The shared VI/MF-VI arguments use the joint entropy estimator, a leave-one-out baseline, zero observation-covariance regularization, and inference seed 7. MF-VI uses a scalar multifidelity control-variate coefficient in both examples:

common_arguments = dict(
    ...,
    optimizer_method="newton",
    optimizer_config=optimizer,
    line_search_method="stochastic_nonmonotone",
    line_search_config=line_search,
    baseline_method="loo",
    score_function_entropy_strategy="joint",
    covariance_regularization=0.0,
    random_seed=7,
)

The full diagonal run uses 16 FOM samples per iteration. MF-VI uses 64 additional ROM samples, the automatic GP ROM, and four iterations of ROM training history:

workflows.mf_vi_with_auto_rom(
    ...,
    fom_sample_size=16,
    rom_extra_sample_size=64,
    max_rom_training_history=4,
    mfmc_control_variate_mode="scalar",
    rom_type="gp",
    rom_args={
        "normalize_parameters": True,
        "normalize_targets": True,
    },
)

prior_parameter_space defines the Bayesian prior, while the required initial_variational_parameter_space independently selects the variational family and supplies its initial moments.

Example 1: equispaced observations and a diagonal posterior#

The first example uses \(N=31\) equispaced interior observations,

\[x_i=\frac{i}{N+1}, \qquad i=1,\ldots,N.\]

For these points the discrete sine basis is orthogonal, so \(G^T\Gamma_{\rm obs}^{-1}G\) is diagonal. Since the prior covariance is also diagonal, the exact posterior covariance is diagonal to numerical precision.

The variational initializer is a GaussianParameterSpace. This selects the diagonal/mean-field VI implementation. This is the case that corresponds directly to the VI-paper analytic-sine benchmark and therefore uses the exact configuration listed above.

Run the example with

python examples/vi_mf_vi_demo/diagonal_example.py

Convergence#

The convergence plot reports the exact Gaussian \(D_{\rm KL}(q\|p_{\rm post})\) for the accepted VI and MF-VI states.

KL convergence for VI and MF-VI on the equispaced analytic sine problem

VI and MF-VI convergence to the exact diagonal posterior.#

Posterior mean#

Posterior mean comparison for the diagonal analytic sine problem

Final VI and MF-VI posterior means compared with the exact posterior mean and the synthetic truth.#

Posterior#

Posterior marginal densities for the diagonal analytic sine problem

Exact, VI, and MF-VI marginal posterior densities for the seven sine coefficients. The vertical dashed line in each panel is the synthetic truth.#

Implementation#

  1"""VI versus MF-VI for the diagonal analytic sine posterior."""
  2
  3from __future__ import annotations
  4
  5import argparse
  6from pathlib import Path
  7import shutil
  8import sys
  9from typing import Optional
 10
 11import numpy as np
 12
 13
 14REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
 15if str(REPOSITORY_ROOT) not in sys.path:
 16    sys.path.insert(0, str(REPOSITORY_ROOT))
 17
 18import romtools.workflows as workflows
 19from romtools.workflows.inverse.vi_optimization_methods import (
 20    VINewtonOptimizerConfig,
 21    VIStochasticNonmonotoneLineSearchConfig,
 22)
 23from romtools.workflows.parameter_spaces import GaussianParameterSpace, MonteCarloSampler
 24
 25from examples.vi_mf_vi_demo.analytic_sine_support import (
 26    AnalyticSineQoiModel,
 27    PRIOR_STD,
 28    build_problem,
 29    collect_history,
 30    write_convergence_plot,
 31    write_diagonal_posterior_plot,
 32    write_mean_plot,
 33)
 34
 35
 36def main(
 37    smoke: bool = False,
 38    work_dir: Optional[str] = None,
 39    output_dir: Optional[str] = None,
 40) -> None:
 41    """Run the equispaced-observation diagonal-posterior example."""
 42    problem = build_problem(equispaced=True)
 43    model = AnalyticSineQoiModel(problem)
 44
 45    sample_size = 8 if smoke else 16
 46    rom_extra_sample_size = 8 if smoke else 64
 47    max_iterations = 2 if smoke else 1000
 48
 49    prior = GaussianParameterSpace(
 50        parameter_names=list(problem.parameter_names),
 51        means=problem.prior_mean,
 52        stds=np.full(problem.prior_mean.size, PRIOR_STD),
 53        sampler=MonteCarloSampler,
 54    )
 55    initial_variational_parameter_space = GaussianParameterSpace(
 56        parameter_names=list(problem.parameter_names),
 57        means=problem.prior_mean,
 58        stds=np.full(problem.prior_mean.size, PRIOR_STD),
 59        sampler=MonteCarloSampler,
 60    )
 61    optimizer = VINewtonOptimizerConfig(
 62        newton_metric="natural",
 63        newton_hessian_type="full",
 64        newton_curvature_strategy="lagged",
 65        newton_hessian_averaging_factor=0.25,
 66        newton_regularization=5e-4,
 67        gradient_norm_tolerance=0.0,
 68        max_iterations=max_iterations,
 69    )
 70    line_search = VIStochasticNonmonotoneLineSearchConfig(
 71        initial_step_size=0.25,
 72        max_step_size=1.0,
 73    )
 74
 75    root = (
 76        Path(work_dir).resolve()
 77        if work_dir
 78        else Path(__file__).parent / "diagonal_work"
 79    )
 80    shutil.rmtree(root, ignore_errors=True)
 81    vi_dir = root / "vi"
 82    mf_vi_dir = root / "mf_vi"
 83
 84    common_arguments = dict(
 85        model=model,
 86        prior_parameter_space=prior,
 87        initial_variational_parameter_space=initial_variational_parameter_space,
 88        observations=problem.observations,
 89        observations_covariance=problem.observation_covariance,
 90        optimizer_method="newton",
 91        optimizer_config=optimizer,
 92        line_search_method="stochastic_nonmonotone",
 93        line_search_config=line_search,
 94        baseline_method="loo",
 95        score_function_entropy_strategy="joint",
 96        bounded_parameter_handling="clip",
 97        covariance_regularization=0.0,
 98        random_seed=7,
 99        restart_files_to_keep=max_iterations + 1,
100    )
101
102    workflows.run_vi(
103        **common_arguments,
104        absolute_work_dir=str(vi_dir),
105        sample_size=sample_size,
106        evaluation_concurrency=1,
107    )
108    workflows.mf_vi_with_auto_rom(
109        **common_arguments,
110        absolute_work_dir=str(mf_vi_dir),
111        fom_sample_size=sample_size,
112        rom_extra_sample_size=rom_extra_sample_size,
113        fom_evaluation_concurrency=1,
114        rom_evaluation_concurrency=1,
115        max_rom_training_history=4,
116        mfmc_control_variate_mode="scalar",
117        rom_type="gp",
118        rom_args={
119            "normalize_parameters": True,
120            "normalize_targets": True,
121        },
122    )
123
124    vi_history = collect_history(vi_dir, problem)
125    mf_history = collect_history(mf_vi_dir, problem)
126    output = (
127        Path(output_dir).resolve()
128        if output_dir
129        else Path(__file__).parent
130    )
131    write_convergence_plot(
132        output / "analytic_sine_diagonal_convergence.png",
133        vi_history,
134        mf_history,
135        "Equispaced observations: convergence to the exact posterior",
136    )
137    write_mean_plot(
138        output / "analytic_sine_diagonal_mean.png",
139        problem,
140        vi_history,
141        mf_history,
142        "Equispaced observations: posterior mean",
143    )
144    write_diagonal_posterior_plot(
145        output / "analytic_sine_diagonal_posterior.png",
146        problem,
147        vi_history,
148        mf_history,
149    )
150
151
152def _parse_args():
153    parser = argparse.ArgumentParser(description=__doc__)
154    parser.add_argument("--smoke", action="store_true")
155    parser.add_argument("--work-dir")
156    parser.add_argument("--output-dir")
157    return parser.parse_args()
158
159
160if __name__ == "__main__":
161    arguments = _parse_args()
162    main(
163        smoke=arguments.smoke,
164        work_dir=arguments.work_dir,
165        output_dir=arguments.output_dir,
166    )

Example 2: non-equispaced observations and a correlated posterior#

The second example keeps the same model, prior, truth, noise level, and number of observations, but changes only the observation locations. Starting from

\[t_i=\frac{i}{N+1},\]

the locations are warped according to

\[x_i =0.02+0.96\,\frac{t_i^3}{t_N^3}.\]

This breaks discrete sine orthogonality. The exact posterior therefore has nonzero cross-covariances (the largest absolute posterior correlation is about 0.30 for this setup).

The prior remains the same diagonal GaussianParameterSpace. The primary full-covariance runs use a MultivariateGaussianParameterSpace initialized with the same diagonal prior covariance. Its type selects true full-covariance VI, allowing off-diagonal covariance entries to develop during optimization. The same correlated target is also solved with a GaussianParameterSpace variational initializer for both VI and MF-VI. Those additional mean-field runs deliberately prohibit cross-covariances and expose the approximation error caused by the restricted variational family.

All four runs retain the same natural Newton metric, full Hessian, lagged curvature with \(\beta=0.25\), line-search scales, random seed, and sample allocation: 16 FOM samples per iteration and 64 additional ROM samples for MF-VI. MF-VI uses the same GP configuration, four iterations of ROM training history, and scalar multifidelity control variate for both variational families. The production optimization budget is 80 iterations. Regular VI uses Newton regularization \(5\times10^{-4}\), while MF-VI uses \(1\times10^{-4}\).

For the full-covariance family, Newton curvature is formed in \((\mu,\operatorname{svec}(\Sigma))\) coordinates and locally whitened with the exact Gaussian Fisher metric before the regularized Newton solve. The covariance update is then applied with the SPD-preserving exponential retraction described in Full-covariance variational inference.

Run the example with

python examples/vi_mf_vi_demo/full_covariance_example.py

Convergence#

KL convergence for full-covariance VI and MF-VI

Exact Gaussian KL convergence for full-covariance VI and MF-VI.#

Posterior mean#

Posterior mean comparison for the non-equispaced analytic sine problem

Final full-covariance VI and MF-VI posterior means compared with the exact posterior mean and truth.#

Posterior#

Posterior correlation and two-dimensional marginal for the non-equispaced sine problem

The exact posterior correlation matrix and the strongest correlated two-dimensional marginal. The ellipse panel compares one- and two-standard-deviation contours from the exact posterior, full-covariance VI, and full-covariance MF-VI.#

Variational-family error#

Because the target is Gaussian, the best possible reverse-KL mean-field approximation is known analytically. If \(P_{\rm post}=C_{\rm post}^{-1}\), the optimum over diagonal Gaussian covariances has

\[m_{\rm MF}^\star=m_{\rm post}, \qquad D_{ii}^\star=\frac{1}{(P_{\rm post})_{ii}}.\]

Its nonzero \(D_{\rm KL}(q\|p_{\rm post})\) therefore provides an exact variational-family error floor. The next plot overlays the full-covariance and mean-field VI/MF-VI histories and marks this floor. A mean-field method that converges to the horizontal line has exhausted its optimization error; the remaining KL discrepancy is caused by the diagonal covariance restriction.

Full-covariance and mean-field VI and MF-VI KL convergence on the correlated sine problem

Full-covariance and mean-field VI/MF-VI convergence on the same correlated target. The dashed horizontal reference is the exact optimal mean-field KL floor.#

The covariance restriction is also visible directly in posterior correlation. The exact target and full-covariance approximations can represent the induced cross-correlations, whereas the optimal mean-field approximation and both mean-field numerical runs are diagonal by construction.

Correlation matrices for exact, full-covariance, and mean-field VI and MF-VI approximations

Posterior correlation matrices for the exact target, full-covariance VI and MF-VI, the analytic optimal mean-field approximation, and mean-field VI and MF-VI.#

Implementation#

  1"""Full- and mean-field VI/MF-VI for non-equispaced sine observations."""
  2
  3from __future__ import annotations
  4
  5import argparse
  6from pathlib import Path
  7import shutil
  8import sys
  9from typing import Optional
 10
 11import matplotlib.pyplot as plt
 12import numpy as np
 13
 14
 15REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
 16if str(REPOSITORY_ROOT) not in sys.path:
 17    sys.path.insert(0, str(REPOSITORY_ROOT))
 18
 19import romtools.workflows as workflows
 20from romtools.workflows.inverse.vi_optimization_methods import (
 21    VINewtonOptimizerConfig,
 22    VIStochasticNonmonotoneLineSearchConfig,
 23)
 24from romtools.workflows.parameter_spaces import (
 25    GaussianParameterSpace,
 26    MonteCarloSampler,
 27    MultivariateGaussianParameterSpace,
 28)
 29
 30from examples.vi_mf_vi_demo.analytic_sine_support import (
 31    AnalyticSineQoiModel,
 32    PRIOR_STD,
 33    build_problem,
 34    collect_history,
 35    covariance_to_correlation,
 36    gaussian_kl,
 37    write_convergence_plot,
 38    write_correlated_posterior_plot,
 39    write_mean_plot,
 40)
 41
 42
 43def _mean_field_optimum(problem) -> tuple[np.ndarray, float]:
 44    """Return the reverse-KL optimal diagonal covariance and its KL floor."""
 45    posterior_precision = np.linalg.inv(problem.posterior_covariance)
 46    covariance = np.diag(1.0 / np.diag(posterior_precision))
 47    kl_floor = gaussian_kl(
 48        problem.posterior_mean,
 49        covariance,
 50        problem.posterior_mean,
 51        problem.posterior_covariance,
 52    )
 53    return covariance, kl_floor
 54
 55
 56def _write_family_convergence_plot(
 57    output_path: Path,
 58    problem,
 59    full_vi_history: dict[str, np.ndarray],
 60    full_mf_history: dict[str, np.ndarray],
 61    mean_field_vi_history: dict[str, np.ndarray],
 62    mean_field_mf_history: dict[str, np.ndarray],
 63) -> None:
 64    """Compare optimization error and the irreducible mean-field KL gap."""
 65    output_path.parent.mkdir(parents=True, exist_ok=True)
 66    _, mean_field_kl_floor = _mean_field_optimum(problem)
 67
 68    figure, axis = plt.subplots(figsize=(7.2, 4.4))
 69    histories = (
 70        ("Full-covariance VI", full_vi_history, "o"),
 71        ("Full-covariance MF-VI", full_mf_history, "^"),
 72        ("Mean-field VI", mean_field_vi_history, "s"),
 73        ("Mean-field MF-VI", mean_field_mf_history, "D"),
 74    )
 75    for label, history, marker in histories:
 76        axis.semilogy(
 77            np.arange(history["kl"].size),
 78            np.maximum(history["kl"], 1.0e-14),
 79            marker=marker,
 80            markevery=max(1, history["kl"].size // 12),
 81            label=label,
 82        )
 83    axis.axhline(
 84        max(mean_field_kl_floor, 1.0e-14),
 85        linestyle="--",
 86        linewidth=1.5,
 87        label="Optimal mean-field KL floor",
 88    )
 89    axis.set(
 90        xlabel="Accepted variational state",
 91        ylabel=r"$D_{\mathrm{KL}}(q\,\|\,p_{\mathrm{post}})$",
 92        title="Correlated posterior: variational-family comparison",
 93    )
 94    axis.grid(True, alpha=0.3)
 95    axis.legend()
 96    figure.tight_layout()
 97    figure.savefig(output_path, dpi=180)
 98    plt.close(figure)
 99
100
101def _write_family_correlation_plot(
102    output_path: Path,
103    problem,
104    full_vi_history: dict[str, np.ndarray],
105    full_mf_history: dict[str, np.ndarray],
106    mean_field_vi_history: dict[str, np.ndarray],
107    mean_field_mf_history: dict[str, np.ndarray],
108) -> None:
109    """Compare exact, full-covariance, and mean-field posterior correlations."""
110    output_path.parent.mkdir(parents=True, exist_ok=True)
111    mean_field_covariance, _ = _mean_field_optimum(problem)
112    panels = (
113        ("Exact posterior", problem.posterior_covariance),
114        ("Full-covariance VI", full_vi_history["covariance"][-1]),
115        ("Full-covariance MF-VI", full_mf_history["covariance"][-1]),
116        ("Optimal mean-field", mean_field_covariance),
117        ("Mean-field VI", mean_field_vi_history["covariance"][-1]),
118        ("Mean-field MF-VI", mean_field_mf_history["covariance"][-1]),
119    )
120
121    figure, axes = plt.subplots(2, 3, figsize=(11.2, 7.0))
122    image = None
123    for axis, (title, covariance) in zip(axes.flat, panels):
124        image = axis.imshow(
125            covariance_to_correlation(covariance),
126            vmin=-1.0,
127            vmax=1.0,
128            cmap="coolwarm",
129        )
130        axis.set(
131            title=title,
132            xticks=np.arange(problem.posterior_mean.size),
133            yticks=np.arange(problem.posterior_mean.size),
134            xticklabels=np.arange(1, problem.posterior_mean.size + 1),
135            yticklabels=np.arange(1, problem.posterior_mean.size + 1),
136        )
137    if image is not None:
138        figure.colorbar(image, ax=axes.ravel().tolist(), fraction=0.025, pad=0.02)
139    figure.suptitle("Posterior correlation: full covariance versus mean field")
140    figure.subplots_adjust(left=0.06, right=0.91, bottom=0.06, top=0.91, wspace=0.25, hspace=0.30)
141    figure.savefig(output_path, dpi=180)
142    plt.close(figure)
143
144
145def main(
146    smoke: bool = False,
147    work_dir: Optional[str] = None,
148    output_dir: Optional[str] = None,
149) -> None:
150    """Run the non-equispaced correlated-posterior example."""
151    problem = build_problem(equispaced=False)
152    model = AnalyticSineQoiModel(problem)
153
154    sample_size = 8 if smoke else 16
155    rom_extra_sample_size = 8 if smoke else 64
156    max_iterations = 2 if smoke else 80
157
158    prior = GaussianParameterSpace(
159        parameter_names=list(problem.parameter_names),
160        means=problem.prior_mean,
161        stds=np.full(problem.prior_mean.size, PRIOR_STD),
162        sampler=MonteCarloSampler,
163    )
164    full_covariance_initializer = MultivariateGaussianParameterSpace(
165        parameter_names=list(problem.parameter_names),
166        means=problem.prior_mean,
167        covariance=problem.prior_covariance,
168        sampler=MonteCarloSampler,
169    )
170    mean_field_initializer = GaussianParameterSpace(
171        parameter_names=list(problem.parameter_names),
172        means=problem.prior_mean,
173        stds=np.full(problem.prior_mean.size, PRIOR_STD),
174        sampler=MonteCarloSampler,
175    )
176    vi_optimizer = VINewtonOptimizerConfig(
177        newton_metric="natural",
178        newton_hessian_type="full",
179        newton_curvature_strategy="lagged",
180        newton_hessian_averaging_factor=0.25,
181        newton_regularization=5e-4,
182        gradient_norm_tolerance=0.0,
183        max_iterations=max_iterations,
184    )
185    mf_vi_optimizer = VINewtonOptimizerConfig(
186        newton_metric="natural",
187        newton_hessian_type="full",
188        newton_curvature_strategy="lagged",
189        newton_hessian_averaging_factor=0.25,
190        newton_regularization=1e-4,
191        gradient_norm_tolerance=0.0,
192        max_iterations=max_iterations,
193    )
194    line_search = VIStochasticNonmonotoneLineSearchConfig(
195        initial_step_size=0.25,
196        max_step_size=1.0,
197    )
198
199    root = (
200        Path(work_dir).resolve()
201        if work_dir
202        else Path(__file__).parent / "correlated_work"
203    )
204    shutil.rmtree(root, ignore_errors=True)
205    full_vi_dir = root / "full_vi"
206    full_mf_vi_dir = root / "full_mf_vi"
207    mean_field_vi_dir = root / "mean_field_vi"
208    mean_field_mf_vi_dir = root / "mean_field_mf_vi"
209
210    common_arguments = dict(
211        model=model,
212        prior_parameter_space=prior,
213        observations=problem.observations,
214        observations_covariance=problem.observation_covariance,
215        optimizer_method="newton",
216        line_search_method="stochastic_nonmonotone",
217        line_search_config=line_search,
218        baseline_method="loo",
219        score_function_entropy_strategy="joint",
220        bounded_parameter_handling="clip",
221        covariance_regularization=0.0,
222        random_seed=7,
223        restart_files_to_keep=max_iterations + 1,
224    )
225
226    workflows.run_vi(
227        **common_arguments,
228        initial_variational_parameter_space=full_covariance_initializer,
229        absolute_work_dir=str(full_vi_dir),
230        sample_size=sample_size,
231        evaluation_concurrency=1,
232        optimizer_config=vi_optimizer,
233    )
234    workflows.mf_vi_with_auto_rom(
235        **common_arguments,
236        initial_variational_parameter_space=full_covariance_initializer,
237        absolute_work_dir=str(full_mf_vi_dir),
238        fom_sample_size=sample_size,
239        rom_extra_sample_size=rom_extra_sample_size,
240        fom_evaluation_concurrency=1,
241        rom_evaluation_concurrency=1,
242        max_rom_training_history=4,
243        mfmc_control_variate_mode="scalar",
244        rom_type="gp",
245        rom_args={
246            "normalize_parameters": True,
247            "normalize_targets": True,
248        },
249        optimizer_config=mf_vi_optimizer,
250    )
251    workflows.run_vi(
252        **common_arguments,
253        initial_variational_parameter_space=mean_field_initializer,
254        absolute_work_dir=str(mean_field_vi_dir),
255        sample_size=sample_size,
256        evaluation_concurrency=1,
257        optimizer_config=vi_optimizer,
258    )
259    workflows.mf_vi_with_auto_rom(
260        **common_arguments,
261        initial_variational_parameter_space=mean_field_initializer,
262        absolute_work_dir=str(mean_field_mf_vi_dir),
263        fom_sample_size=sample_size,
264        rom_extra_sample_size=rom_extra_sample_size,
265        fom_evaluation_concurrency=1,
266        rom_evaluation_concurrency=1,
267        max_rom_training_history=4,
268        mfmc_control_variate_mode="scalar",
269        rom_type="gp",
270        rom_args={
271            "normalize_parameters": True,
272            "normalize_targets": True,
273        },
274        optimizer_config=mf_vi_optimizer,
275    )
276
277    full_vi_history = collect_history(full_vi_dir, problem)
278    full_mf_history = collect_history(full_mf_vi_dir, problem)
279    mean_field_vi_history = collect_history(mean_field_vi_dir, problem)
280    mean_field_mf_history = collect_history(mean_field_mf_vi_dir, problem)
281    output = (
282        Path(output_dir).resolve()
283        if output_dir
284        else Path(__file__).parent
285    )
286    write_convergence_plot(
287        output / "analytic_sine_correlated_convergence.png",
288        full_vi_history,
289        full_mf_history,
290        "Non-equispaced observations: full-covariance convergence",
291    )
292    write_mean_plot(
293        output / "analytic_sine_correlated_mean.png",
294        problem,
295        full_vi_history,
296        full_mf_history,
297        "Non-equispaced observations: posterior mean",
298    )
299    write_correlated_posterior_plot(
300        output / "analytic_sine_correlated_posterior.png",
301        problem,
302        full_vi_history,
303        full_mf_history,
304    )
305    _write_family_convergence_plot(
306        output / "analytic_sine_correlated_family_convergence.png",
307        problem,
308        full_vi_history,
309        full_mf_history,
310        mean_field_vi_history,
311        mean_field_mf_history,
312    )
313    _write_family_correlation_plot(
314        output / "analytic_sine_correlated_family_correlation.png",
315        problem,
316        full_vi_history,
317        full_mf_history,
318        mean_field_vi_history,
319        mean_field_mf_history,
320    )
321
322
323def _parse_args():
324    parser = argparse.ArgumentParser(description=__doc__)
325    parser.add_argument("--smoke", action="store_true")
326    parser.add_argument("--work-dir")
327    parser.add_argument("--output-dir")
328    return parser.parse_args()
329
330
331if __name__ == "__main__":
332    arguments = _parse_args()
333    main(
334        smoke=arguments.smoke,
335        work_dir=arguments.work_dir,
336        output_dir=arguments.output_dir,
337    )

Shared problem definition#

The forward model, exact posterior calculation, history processing, and plotting utilities are shared by both examples:

  1"""Shared analytic sine inverse problem and plotting utilities for VI examples."""
  2
  3from __future__ import annotations
  4
  5from dataclasses import dataclass
  6from pathlib import Path
  7
  8import matplotlib.pyplot as plt
  9from matplotlib.patches import Ellipse
 10import numpy as np
 11
 12
 13DIMENSION = 7
 14N_OBSERVATIONS = 31
 15OBSERVATION_NOISE_STD = 5.0e-3
 16PRIOR_STD = 0.75
 17RANDOM_SEED = 7
 18TRUTH = np.array([1.00, -0.80, 0.65, -0.50, 0.35, -0.25, 0.15], dtype=float)
 19
 20
 21@dataclass(frozen=True)
 22class AnalyticSineProblem:
 23    """Linear-Gaussian sine-series inverse problem with an exact posterior."""
 24
 25    parameter_names: tuple[str, ...]
 26    observation_locations: np.ndarray
 27    observation_matrix: np.ndarray
 28    truth: np.ndarray
 29    observations: np.ndarray
 30    observation_covariance: np.ndarray
 31    prior_mean: np.ndarray
 32    prior_covariance: np.ndarray
 33    posterior_mean: np.ndarray
 34    posterior_covariance: np.ndarray
 35
 36
 37def _observation_locations(equispaced: bool) -> np.ndarray:
 38    indices = np.arange(1, N_OBSERVATIONS + 1, dtype=float)
 39    if equispaced:
 40        return indices / (N_OBSERVATIONS + 1)
 41
 42    # Keep the same number of observations and nearly the same spatial extent,
 43    # but cluster points toward x=0. This breaks the discrete sine
 44    # orthogonality that makes the equispaced posterior diagonal.
 45    scaled = indices / (N_OBSERVATIONS + 1)
 46    warped = scaled**3 / scaled[-1]**3
 47    return 0.02 + 0.96 * warped
 48
 49
 50def build_problem(equispaced: bool) -> AnalyticSineProblem:
 51    """Build the sine-series inverse problem and its exact Gaussian posterior."""
 52    parameter_names = tuple(f"theta_{j}" for j in range(1, DIMENSION + 1))
 53    x = _observation_locations(equispaced)
 54    modes = np.arange(1, DIMENSION + 1, dtype=float)
 55    observation_matrix = -np.sin(np.pi * np.outer(x, modes)) / (
 56        np.pi * modes[None, :]
 57    ) ** 2
 58
 59    prior_mean = np.zeros(DIMENSION, dtype=float)
 60    prior_covariance = np.eye(DIMENSION) * PRIOR_STD**2
 61    observation_covariance = np.eye(N_OBSERVATIONS) * OBSERVATION_NOISE_STD**2
 62
 63    rng = np.random.default_rng(RANDOM_SEED)
 64    observations = (
 65        observation_matrix @ TRUTH
 66        + OBSERVATION_NOISE_STD * rng.standard_normal(N_OBSERVATIONS)
 67    )
 68
 69    prior_precision = np.linalg.inv(prior_covariance)
 70    observation_precision = np.linalg.inv(observation_covariance)
 71    posterior_precision = (
 72        prior_precision
 73        + observation_matrix.T @ observation_precision @ observation_matrix
 74    )
 75    posterior_covariance = np.linalg.inv(posterior_precision)
 76    posterior_mean = posterior_covariance @ (
 77        prior_precision @ prior_mean
 78        + observation_matrix.T @ observation_precision @ observations
 79    )
 80
 81    off_diagonal = posterior_covariance - np.diag(np.diag(posterior_covariance))
 82    if equispaced:
 83        scale = max(float(np.max(np.diag(posterior_covariance))), 1.0e-16)
 84        if np.max(np.abs(off_diagonal)) / scale > 1.0e-10:
 85            raise RuntimeError("Expected the equispaced posterior covariance to be diagonal.")
 86    else:
 87        correlation = covariance_to_correlation(posterior_covariance)
 88        correlation_off_diagonal = correlation - np.eye(DIMENSION)
 89        if np.max(np.abs(correlation_off_diagonal)) < 0.1:
 90            raise RuntimeError("Non-equispaced observations did not create meaningful correlation.")
 91
 92    return AnalyticSineProblem(
 93        parameter_names=parameter_names,
 94        observation_locations=x,
 95        observation_matrix=observation_matrix,
 96        truth=TRUTH.copy(),
 97        observations=observations,
 98        observation_covariance=observation_covariance,
 99        prior_mean=prior_mean,
100        prior_covariance=prior_covariance,
101        posterior_mean=posterior_mean,
102        posterior_covariance=posterior_covariance,
103    )
104
105
106class AnalyticSineQoiModel:
107    """ROMTools QoI model evaluating the exact sine-series solution."""
108
109    def __init__(self, problem: AnalyticSineProblem):
110        self.problem = problem
111
112    def _theta(self, parameter_sample: dict) -> np.ndarray:
113        return np.array(
114            [float(parameter_sample[name]) for name in self.problem.parameter_names],
115            dtype=float,
116        )
117
118    def populate_run_directory(self, run_directory: str, parameter_sample: dict) -> None:
119        del parameter_sample
120        Path(run_directory).mkdir(parents=True, exist_ok=True)
121
122    def run_model(self, run_directory: str, parameter_sample: dict) -> int:
123        qoi = self.problem.observation_matrix @ self._theta(parameter_sample)
124        np.save(Path(run_directory) / "qoi.npy", qoi)
125        return 0
126
127    def compute_qoi(self, run_directory: str, parameter_sample: dict) -> np.ndarray:
128        del parameter_sample
129        return np.load(Path(run_directory) / "qoi.npy")
130
131
132def gaussian_kl(
133    mean: np.ndarray,
134    covariance: np.ndarray,
135    reference_mean: np.ndarray,
136    reference_covariance: np.ndarray,
137) -> float:
138    """Return KL[N(mean,covariance) || N(reference_mean,reference_covariance)]."""
139    mean = np.asarray(mean, dtype=float)
140    covariance = np.asarray(covariance, dtype=float)
141    reference_mean = np.asarray(reference_mean, dtype=float)
142    reference_covariance = np.asarray(reference_covariance, dtype=float)
143
144    dimension = mean.size
145    delta = reference_mean - mean
146    sign_q, logdet_q = np.linalg.slogdet(covariance)
147    sign_p, logdet_p = np.linalg.slogdet(reference_covariance)
148    if sign_q <= 0 or sign_p <= 0:
149        return float("inf")
150    solve_covariance = np.linalg.solve(reference_covariance, covariance)
151    solve_delta = np.linalg.solve(reference_covariance, delta)
152    return float(
153        0.5
154        * (
155            np.trace(solve_covariance)
156            + delta @ solve_delta
157            - dimension
158            + logdet_p
159            - logdet_q
160        )
161    )
162
163
164def collect_history(work_dir: Path, problem: AnalyticSineProblem) -> dict[str, np.ndarray]:
165    """Read the accepted VI states and prepend the common initial state."""
166    with np.load(work_dir / "history.npz") as history:
167        means = np.asarray(history["vi_history_variational_mean"], dtype=float)
168        covariances = np.asarray(
169            history["vi_history_variational_covariance"], dtype=float
170        )
171
172    if (
173        means.shape[0] == 0
174        or not np.allclose(means[0], problem.prior_mean)
175        or not np.allclose(covariances[0], problem.prior_covariance)
176    ):
177        means = np.concatenate([problem.prior_mean[None, :], means], axis=0)
178        covariances = np.concatenate(
179            [problem.prior_covariance[None, :, :], covariances], axis=0
180        )
181
182    kl = np.asarray(
183        [
184            gaussian_kl(
185                mean,
186                covariance,
187                problem.posterior_mean,
188                problem.posterior_covariance,
189            )
190            for mean, covariance in zip(means, covariances)
191        ],
192        dtype=float,
193    )
194    return {"mean": means, "covariance": covariances, "kl": kl}
195
196
197def covariance_to_correlation(covariance: np.ndarray) -> np.ndarray:
198    standard_deviations = np.sqrt(np.maximum(np.diag(covariance), 0.0))
199    denominator = np.outer(standard_deviations, standard_deviations)
200    correlation = np.divide(
201        covariance,
202        denominator,
203        out=np.zeros_like(covariance, dtype=float),
204        where=denominator > 0.0,
205    )
206    np.fill_diagonal(correlation, 1.0)
207    return correlation
208
209
210def write_convergence_plot(
211    output_path: Path,
212    vi_history: dict[str, np.ndarray],
213    mf_history: dict[str, np.ndarray],
214    title: str,
215) -> None:
216    """Plot exact Gaussian KL convergence."""
217    output_path.parent.mkdir(parents=True, exist_ok=True)
218    figure, axis = plt.subplots(figsize=(6.5, 4.0))
219    axis.semilogy(
220        np.arange(vi_history["kl"].size),
221        np.maximum(vi_history["kl"], 1.0e-14),
222        marker="o",
223        label="VI",
224    )
225    axis.semilogy(
226        np.arange(mf_history["kl"].size),
227        np.maximum(mf_history["kl"], 1.0e-14),
228        marker="^",
229        label="MF-VI",
230    )
231    axis.set(
232        xlabel="Accepted variational state",
233        ylabel=r"$D_{\mathrm{KL}}(q\,\|\,p_{\mathrm{post}})$",
234        title=title,
235    )
236    axis.grid(True, alpha=0.3)
237    axis.legend()
238    figure.tight_layout()
239    figure.savefig(output_path, dpi=180)
240    plt.close(figure)
241
242
243def write_mean_plot(
244    output_path: Path,
245    problem: AnalyticSineProblem,
246    vi_history: dict[str, np.ndarray],
247    mf_history: dict[str, np.ndarray],
248    title: str,
249) -> None:
250    """Compare final posterior means against the exact posterior and truth."""
251    output_path.parent.mkdir(parents=True, exist_ok=True)
252    indices = np.arange(1, DIMENSION + 1)
253    figure, axis = plt.subplots(figsize=(7.2, 4.2))
254    axis.plot(indices, problem.truth, "k--", marker="x", label="Truth")
255    axis.plot(indices, problem.posterior_mean, marker="o", label="Exact posterior mean")
256    axis.plot(indices, vi_history["mean"][-1], marker="s", label="VI")
257    axis.plot(indices, mf_history["mean"][-1], marker="^", label="MF-VI")
258    axis.set(
259        xlabel="Sine coefficient index",
260        ylabel="Posterior mean",
261        title=title,
262        xticks=indices,
263    )
264    axis.grid(True, alpha=0.3)
265    axis.legend(ncol=2)
266    figure.tight_layout()
267    figure.savefig(output_path, dpi=180)
268    plt.close(figure)
269
270
271def _normal_density(grid: np.ndarray, mean: float, variance: float) -> np.ndarray:
272    variance = max(float(variance), 1.0e-16)
273    return np.exp(-0.5 * (grid - mean) ** 2 / variance) / np.sqrt(
274        2.0 * np.pi * variance
275    )
276
277
278def write_diagonal_posterior_plot(
279    output_path: Path,
280    problem: AnalyticSineProblem,
281    vi_history: dict[str, np.ndarray],
282    mf_history: dict[str, np.ndarray],
283) -> None:
284    """Plot one-dimensional marginals for the diagonal-posterior example."""
285    output_path.parent.mkdir(parents=True, exist_ok=True)
286    figure, axes = plt.subplots(2, 4, figsize=(12.0, 6.2))
287    final_states = (
288        ("VI", vi_history["mean"][-1], vi_history["covariance"][-1]),
289        ("MF-VI", mf_history["mean"][-1], mf_history["covariance"][-1]),
290    )
291    for index, axis in enumerate(axes.flat):
292        if index >= DIMENSION:
293            axis.axis("off")
294            continue
295        exact_mean = problem.posterior_mean[index]
296        exact_variance = problem.posterior_covariance[index, index]
297        width = 4.0 * np.sqrt(exact_variance)
298        grid = np.linspace(exact_mean - width, exact_mean + width, 300)
299        axis.plot(
300            grid,
301            _normal_density(grid, exact_mean, exact_variance),
302            linewidth=2.0,
303            label="Exact",
304        )
305        for label, mean, covariance in final_states:
306            axis.plot(
307                grid,
308                _normal_density(grid, mean[index], covariance[index, index]),
309                label=label,
310            )
311        axis.axvline(problem.truth[index], color="black", linestyle="--", linewidth=1.0)
312        axis.set_title(problem.parameter_names[index])
313        axis.grid(True, alpha=0.2)
314    axes.flat[0].legend()
315    figure.suptitle("Posterior marginal densities")
316    figure.tight_layout()
317    figure.savefig(output_path, dpi=180)
318    plt.close(figure)
319
320
321def _add_covariance_ellipse(
322    axis,
323    mean: np.ndarray,
324    covariance: np.ndarray,
325    label: str,
326    linestyle: str,
327    linewidth: float,
328) -> None:
329    eigenvalues, eigenvectors = np.linalg.eigh(covariance)
330    order = np.argsort(eigenvalues)[::-1]
331    eigenvalues = eigenvalues[order]
332    eigenvectors = eigenvectors[:, order]
333    angle = np.degrees(np.arctan2(eigenvectors[1, 0], eigenvectors[0, 0]))
334    for scale in (1.0, 2.0):
335        ellipse = Ellipse(
336            xy=mean,
337            width=2.0 * scale * np.sqrt(max(eigenvalues[0], 0.0)),
338            height=2.0 * scale * np.sqrt(max(eigenvalues[1], 0.0)),
339            angle=angle,
340            fill=False,
341            linestyle=linestyle,
342            linewidth=linewidth,
343            label=label if scale == 1.0 else None,
344        )
345        axis.add_patch(ellipse)
346
347
348def write_correlated_posterior_plot(
349    output_path: Path,
350    problem: AnalyticSineProblem,
351    vi_history: dict[str, np.ndarray],
352    mf_history: dict[str, np.ndarray],
353) -> None:
354    """Show posterior correlation and a two-dimensional Gaussian marginal."""
355    output_path.parent.mkdir(parents=True, exist_ok=True)
356    exact_correlation = covariance_to_correlation(problem.posterior_covariance)
357    off_diagonal = np.abs(exact_correlation - np.eye(DIMENSION))
358    first, second = np.unravel_index(np.argmax(off_diagonal), off_diagonal.shape)
359    pair = np.array([first, second], dtype=int)
360
361    figure, axes = plt.subplots(1, 2, figsize=(11.5, 4.6))
362    image = axes[0].imshow(exact_correlation, vmin=-1.0, vmax=1.0, cmap="coolwarm")
363    axes[0].set(
364        title="Exact posterior correlation",
365        xlabel="Parameter index",
366        ylabel="Parameter index",
367        xticks=np.arange(DIMENSION),
368        yticks=np.arange(DIMENSION),
369        xticklabels=np.arange(1, DIMENSION + 1),
370        yticklabels=np.arange(1, DIMENSION + 1),
371    )
372    figure.colorbar(image, ax=axes[0], fraction=0.046, pad=0.04)
373
374    exact_mean = problem.posterior_mean[pair]
375    exact_covariance = problem.posterior_covariance[np.ix_(pair, pair)]
376    _add_covariance_ellipse(
377        axes[1], exact_mean, exact_covariance, "Exact", "-", 2.2
378    )
379    _add_covariance_ellipse(
380        axes[1],
381        vi_history["mean"][-1, pair],
382        vi_history["covariance"][-1][np.ix_(pair, pair)],
383        "VI",
384        "--",
385        1.8,
386    )
387    _add_covariance_ellipse(
388        axes[1],
389        mf_history["mean"][-1, pair],
390        mf_history["covariance"][-1][np.ix_(pair, pair)],
391        "MF-VI",
392        ":",
393        2.0,
394    )
395    axes[1].plot(
396        problem.truth[first],
397        problem.truth[second],
398        "kx",
399        markersize=8,
400        label="Truth",
401    )
402
403    all_means = np.vstack(
404        [
405            exact_mean,
406            vi_history["mean"][-1, pair],
407            mf_history["mean"][-1, pair],
408        ]
409    )
410    max_std = max(
411        np.sqrt(np.max(np.linalg.eigvalsh(exact_covariance))),
412        np.sqrt(
413            np.max(
414                np.linalg.eigvalsh(
415                    vi_history["covariance"][-1][np.ix_(pair, pair)]
416                )
417            )
418        ),
419        np.sqrt(
420            np.max(
421                np.linalg.eigvalsh(
422                    mf_history["covariance"][-1][np.ix_(pair, pair)]
423                )
424            )
425        ),
426    )
427    lower = np.min(all_means, axis=0) - 3.0 * max_std
428    upper = np.max(all_means, axis=0) + 3.0 * max_std
429    axes[1].set_xlim(lower[0], upper[0])
430    axes[1].set_ylim(lower[1], upper[1])
431    axes[1].set(
432        title=(
433            f"Strongest correlated marginal "
434            f"(rho={exact_correlation[first, second]:.2f})"
435        ),
436        xlabel=problem.parameter_names[first],
437        ylabel=problem.parameter_names[second],
438    )
439    axes[1].grid(True, alpha=0.25)
440    axes[1].legend()
441
442    figure.suptitle("Full-covariance posterior structure")
443    figure.tight_layout()
444    figure.savefig(output_path, dpi=180)
445    plt.close(figure)

Regenerating the figures#

Run both production examples and regenerate all eight figures with

python examples/vi_mf_vi_demo/generate_docs_media.py

A reduced CI configuration exercises both examples with

python examples/vi_mf_vi_demo/example.py --smoke

ROM-only optimizer substeps#

MF-VI can optionally perform a fixed number of inexpensive ROM-only optimization updates between FOM-anchored iterations, mirroring the MF-EKI substep window. For example, add the following keyword arguments to run_mf_vi or mf_vi_with_auto_rom:

rom_substep_start_iteration=0,
rom_substep_end_iteration=20,
num_rom_substeps=2,

The defaults are 0, None, and 0, respectively; therefore ROM substeps are disabled unless requested. The start index is inclusive and the end index exclusive, counting only outer MF-VI iterations.

Each ROM-only substep uses the same optimization method and configuration as the outer MF-VI iteration. Gradient descent uses the configured standard or natural gradient; Adam uses the configured learning rate, Fisher damping, clipping, and moment parameters. Adam initializes a fresh optimizer at the start of every ROM-only sequence, reuses its moments within that sequence, and never changes the outer Adam moments. Newton substeps recompute both a fresh gradient and Hessian and use the configured metric, curvature type, and regularization. All inner updates respect the existing covariance and step-limiting safeguards.

The ROM remains frozen during substeps. The final proposed distribution is evaluated by the existing multifidelity FOM/ROM candidate check and rejected as a whole if its outer line search fails. ROM-only substeps do not perform FOM evaluations, refresh the ROM training set, or independently terminate the solve. They are distinct from the ABRIS sample-reuse batch-refresh loop.

Because the inner optimization direction uses an approximate ROM objective, this acceleration is not guaranteed to reduce FOM work on every problem. A finite end iteration can disable it near convergence if ROM directional errors cause oscillation.