H2-air flame EKI and MF-EKI rejuvenation benchmark#

This example uses the four-parameter H2-air flame model model to study ensemble collapse in ensemble Kalman inversion and the effect of adaptive MF-EKI rejuvenation. The inferred parameters are diffusivity kappa, scaled activation energy, and the two advection components beta_x and beta_y.

Benchmark design#

The benchmark compares four cases using the same synthetic observation and random seed:

  • single-fidelity EKI with four FOM ensemble members,

  • the same EKI configuration with rejuvenation_strategy="adaptive",

  • GP auto-ROM MF-EKI with four FOM members and 32 additional ROM-only members, with rejuvenation disabled, and

  • the same MF-EKI configuration with rejuvenation_strategy="adaptive".

The full model uses a 64 x 32 mesh, dt = 1e-3, and t_end = 6e-2. The full benchmark permits up to 30 outer iterations. Adaptive rejuvenation is triggered when the componentwise normalized RMS parameter-update magnitude falls below 1e-3 while the observation residual is still above its convergence tolerance. At most three rejuvenations are allowed.

When the empirical ensemble covariance collapses, the default parameter-scaled reference term corresponds to a 5% standard-deviation perturbation around the current parameter mean. Parameter bounds supply the fallback scale for a near-zero mean.

Diagnostics#

The benchmark records three diagnostics at every saved outer iteration:

  • mean observation residual,

  • RMS componentwise relative error of the FOM-ensemble parameter mean, and

  • range-scaled ensemble spread. For MF-EKI, the spread uses the combined FOM and ROM-only parameter ensembles.

Run the example#

From the repository root:

python examples/h2_air_flame_eki_mf_eki/example.py

A reduced API-validation run is available with:

python examples/h2_air_flame_eki_mf_eki/example.py --smoke

The full benchmark writes an SVG/PNG figure and a JSON file containing the complete histories and final summary values.

Implementation#

  1"""Benchmark MF-EKI ensemble rejuvenation on the H2-air flame model."""
  2
  3from __future__ import annotations
  4
  5import argparse
  6import json
  7import os
  8import shutil
  9import sys
 10from pathlib import Path
 11
 12import matplotlib.pyplot as plt
 13import numpy as np
 14
 15from romtools.workflows.inverse.eki_drivers import run_eki
 16from romtools.workflows.inverse.mf_eki_drivers import mf_eki_with_auto_rom
 17from romtools.workflows.parameter_spaces import HeterogeneousParameterSpace
 18from romtools.workflows.parameters import UniformParameter
 19
 20
 21REPOSITORY_ROOT = Path(__file__).resolve().parents[2]
 22EXAMPLE_MODELS_PATH = REPOSITORY_ROOT / "examples" / "models"
 23if str(EXAMPLE_MODELS_PATH) not in sys.path:
 24    sys.path.insert(0, str(EXAMPLE_MODELS_PATH))
 25
 26from h2_air_flame_model import H2AirFlameQoiModel  # noqa: E402
 27
 28
 29TRUTH = np.array([2.0, 8.0, 40.0, 7.0])
 30PARAMETER_MINS = np.array([0.5, 4.0, 20.0, 1.0])
 31PARAMETER_MAXES = np.array([4.0, 12.0, 60.0, 20.0])
 32PARAMETER_RANGES = PARAMETER_MAXES - PARAMETER_MINS
 33
 34
 35def _parse_args() -> argparse.Namespace:
 36    parser = argparse.ArgumentParser(
 37        description=(
 38            "Compare EKI, baseline GP auto-ROM MF-EKI, and adaptive-rejuvenation "
 39            "MF-EKI on the H2-air flame inverse problem."
 40        )
 41    )
 42    parser.add_argument("--smoke", action="store_true")
 43    parser.add_argument(
 44        "--work-dir",
 45        type=Path,
 46        default=Path(__file__).resolve().parent / "work",
 47    )
 48    parser.add_argument(
 49        "--output-dir",
 50        type=Path,
 51        default=Path(__file__).resolve().parent / "results",
 52    )
 53    parser.add_argument("--max-iterations", type=int, default=None)
 54    return parser.parse_args()
 55
 56
 57def _parameter_space() -> HeterogeneousParameterSpace:
 58    return HeterogeneousParameterSpace(
 59        [
 60            UniformParameter("kappa", PARAMETER_MINS[0], PARAMETER_MAXES[0]),
 61            UniformParameter(
 62                "scaled_activation_energy", PARAMETER_MINS[1], PARAMETER_MAXES[1]
 63            ),
 64            UniformParameter("beta_x", PARAMETER_MINS[2], PARAMETER_MAXES[2]),
 65            UniformParameter("beta_y", PARAMETER_MINS[3], PARAMETER_MAXES[3]),
 66        ]
 67    )
 68
 69
 70def _load_history(work_dir: Path, mf: bool) -> dict[str, list]:
 71    residual = []
 72    parameter_relative_error = []
 73    scaled_ensemble_spread = []
 74    rejuvenation_count = []
 75    iteration = 0
 76    while True:
 77        restart_path = work_dir / f"iteration_{iteration}" / "restart.npz"
 78        if not restart_path.exists():
 79            break
 80        with np.load(restart_path, allow_pickle=True) as data:
 81            if mf:
 82                errors = data["sample_one_fom_results"].item()["errors"]
 83                primary_samples = data["parameter_samples_one"]
 84                all_samples = np.vstack(
 85                    [data["parameter_samples_one"], data["parameter_samples_two"]]
 86                )
 87            else:
 88                errors = data["errors"]
 89                primary_samples = data["parameter_samples"]
 90                all_samples = primary_samples
 91
 92            mean_parameters = np.mean(primary_samples, axis=0)
 93            relative_error = np.linalg.norm(
 94                (mean_parameters - TRUTH) / TRUTH
 95            ) / np.sqrt(TRUTH.size)
 96            spread = np.linalg.norm(
 97                np.std(all_samples, axis=0, ddof=0) / PARAMETER_RANGES
 98            ) / np.sqrt(TRUTH.size)
 99
100            residual.append(float(np.mean(np.linalg.norm(errors, axis=0))))
101            parameter_relative_error.append(float(relative_error))
102            scaled_ensemble_spread.append(float(spread))
103            rejuvenation_count.append(
104                int(data["rejuvenation_count"])
105                if "rejuvenation_count" in data
106                else 0
107            )
108        iteration += 1
109
110    return {
111        "residual": residual,
112        "parameter_relative_error": parameter_relative_error,
113        "scaled_ensemble_spread": scaled_ensemble_spread,
114        "rejuvenation_count": rejuvenation_count,
115    }
116
117
118def _summary(history: dict[str, list]) -> dict[str, float | int]:
119    if not history["residual"]:
120        raise RuntimeError("Benchmark produced no restart history.")
121    return {
122        "recorded_iterations": len(history["residual"]),
123        "final_residual": history["residual"][-1],
124        "minimum_residual": min(history["residual"]),
125        "final_parameter_relative_error": history["parameter_relative_error"][-1],
126        "final_scaled_ensemble_spread": history["scaled_ensemble_spread"][-1],
127        "rejuvenations": history["rejuvenation_count"][-1],
128    }
129
130
131def _plot_histories(histories: dict[str, dict[str, list]], output_dir: Path) -> None:
132    labels = {
133        "eki": "Single-fidelity EKI",
134        "eki_rejuvenated": "EKI + adaptive rejuvenation",
135        "mf_eki": "MF-EKI",
136        "mf_eki_rejuvenated": "MF-EKI + adaptive rejuvenation",
137    }
138    markers = {"eki": "o", "eki_rejuvenated": "D", "mf_eki": "s", "mf_eki_rejuvenated": "^"}
139
140    figure, axes = plt.subplots(1, 3, figsize=(13.0, 3.8))
141    for key, history in histories.items():
142        iterations = np.arange(len(history["residual"]))
143        axes[0].plot(
144            iterations,
145            history["residual"],
146            marker=markers[key],
147            label=labels[key],
148        )
149        axes[1].plot(
150            iterations,
151            history["parameter_relative_error"],
152            marker=markers[key],
153            label=labels[key],
154        )
155        axes[2].plot(
156            iterations,
157            history["scaled_ensemble_spread"],
158            marker=markers[key],
159            label=labels[key],
160        )
161
162    for axis in axes:
163        axis.set_yscale("log")
164        axis.set_xlabel("Outer iteration")
165        axis.grid(True, alpha=0.3)
166    axes[0].set_ylabel("Mean observation residual")
167    axes[1].set_ylabel("RMS relative parameter error")
168    axes[2].set_ylabel("Range-scaled ensemble spread")
169    axes[0].legend()
170    figure.suptitle("H2-air flame: effect of adaptive MF-EKI rejuvenation")
171    figure.tight_layout()
172    figure.savefig(output_dir / "h2_air_flame_eki_rejuvenation_benchmark.svg")
173    figure.savefig(
174        output_dir / "h2_air_flame_eki_rejuvenation_benchmark.png", dpi=180
175    )
176    plt.close(figure)
177
178
179def main() -> None:
180    args = _parse_args()
181    np.random.seed(1)
182
183    if args.smoke:
184        nx, ny = 20, 10
185        dt, t_end = 1.0e-3, 8.0e-3
186        fom_ensemble_size, rom_extra_ensemble_size = 4, 4
187        max_iterations = args.max_iterations or 2
188        concurrency = 1
189        rom_substep_start_iteration, rom_substep_end_iteration = 1, 2
190        num_rom_substeps = 1
191        max_rom_training_history = 2
192        adaptive_delta_tolerance = 1.0e12
193        max_rejuvenations = 1
194    else:
195        nx, ny = 64, 32
196        dt, t_end = 1.0e-3, 6.0e-2
197        fom_ensemble_size, rom_extra_ensemble_size = 4, 32
198        max_iterations = args.max_iterations or 30
199        concurrency = min(4, os.cpu_count() or 1)
200        rom_substep_start_iteration, rom_substep_end_iteration = 3, 15
201        num_rom_substeps = 3
202        max_rom_training_history = 3
203        adaptive_delta_tolerance = 1.0e-3
204        max_rejuvenations = 3
205
206    model = H2AirFlameQoiModel(
207        nx=nx,
208        ny=ny,
209        dt=dt,
210        t_end=t_end,
211        snapshot_stride=10,
212    )
213    truth = {
214        "kappa": TRUTH[0],
215        "scaled_activation_energy": TRUTH[1],
216        "beta_x": TRUTH[2],
217        "beta_y": TRUTH[3],
218    }
219    parameter_space = _parameter_space()
220
221    work_dir = args.work_dir.resolve()
222    output_dir = args.output_dir.resolve()
223    shutil.rmtree(work_dir, ignore_errors=True)
224    output_dir.mkdir(parents=True, exist_ok=True)
225
226    truth_dir = work_dir / "truth"
227    model.populate_run_directory(str(truth_dir), truth)
228    model.run_model(str(truth_dir), truth)
229    observations = model.compute_qoi(str(truth_dir), truth)
230    observations_covariance = np.eye(observations.size) * 1.0e-4
231
232    eki_dir = work_dir / "eki"
233    eki_rejuvenated_dir = work_dir / "eki_rejuvenated"
234    mf_dir = work_dir / "mf_eki"
235    mf_rejuvenated_dir = work_dir / "mf_eki_rejuvenated"
236
237    common_mf_args = dict(
238        model=model,
239        parameter_space=parameter_space,
240        observations=observations,
241        observations_covariance=observations_covariance,
242        parameter_mins=PARAMETER_MINS,
243        parameter_maxes=PARAMETER_MAXES,
244        fom_ensemble_size=fom_ensemble_size,
245        rom_extra_ensemble_size=rom_extra_ensemble_size,
246        max_rom_training_history=max_rom_training_history,
247        rom_substep_start_iteration=rom_substep_start_iteration,
248        rom_substep_end_iteration=rom_substep_end_iteration,
249        num_rom_substeps=num_rom_substeps,
250        max_iterations=max_iterations,
251        random_seed=1,
252        fom_evaluation_concurrency=concurrency,
253        rom_type="gp",
254        rom_args={
255            "normalize_parameters": True,
256            "normalize_targets": True,
257        },
258    )
259
260    print("\n=== Baseline MF-EKI ===", flush=True)
261    mf_eki_with_auto_rom(
262        **common_mf_args,
263        absolute_work_dir=str(mf_dir),
264        rejuvenation_strategy="none",
265    )
266
267    print("\n=== MF-EKI with adaptive rejuvenation ===", flush=True)
268    mf_eki_with_auto_rom(
269        **common_mf_args,
270        absolute_work_dir=str(mf_rejuvenated_dir),
271        rejuvenation_strategy="adaptive",
272        delta_params_tolerance=adaptive_delta_tolerance,
273        max_rejuvenations=max_rejuvenations,
274    )
275
276    print("\n=== Single-fidelity EKI with adaptive rejuvenation ===", flush=True)
277    run_eki(
278        model=model,
279        parameter_space=parameter_space,
280        observations=observations,
281        observations_covariance=observations_covariance,
282        parameter_mins=PARAMETER_MINS,
283        parameter_maxes=PARAMETER_MAXES,
284        absolute_work_dir=str(eki_rejuvenated_dir),
285        ensemble_size=fom_ensemble_size,
286        max_iterations=max_iterations,
287        random_seed=1,
288        evaluation_concurrency=concurrency,
289        delta_params_tolerance=adaptive_delta_tolerance,
290        rejuvenation_strategy="adaptive",
291        max_rejuvenations=max_rejuvenations,
292    )
293
294    print("\n=== Single-fidelity EKI reference ===", flush=True)
295    run_eki(
296        model=model,
297        parameter_space=parameter_space,
298        observations=observations,
299        observations_covariance=observations_covariance,
300        parameter_mins=PARAMETER_MINS,
301        parameter_maxes=PARAMETER_MAXES,
302        absolute_work_dir=str(eki_dir),
303        ensemble_size=fom_ensemble_size,
304        max_iterations=max_iterations,
305        random_seed=1,
306        evaluation_concurrency=concurrency,
307        delta_params_tolerance=adaptive_delta_tolerance,
308        rejuvenation_strategy="none",
309    )
310
311    histories = {
312        "eki": _load_history(eki_dir, mf=False),
313        "eki_rejuvenated": _load_history(eki_rejuvenated_dir, mf=False),
314        "mf_eki": _load_history(mf_dir, mf=True),
315        "mf_eki_rejuvenated": _load_history(mf_rejuvenated_dir, mf=True),
316    }
317    summaries = {key: _summary(history) for key, history in histories.items()}
318    results = {
319        "configuration": {
320            "smoke": args.smoke,
321            "nx": nx,
322            "ny": ny,
323            "dt": dt,
324            "t_end": t_end,
325            "max_iterations": max_iterations,
326            "fom_ensemble_size": fom_ensemble_size,
327            "rom_extra_ensemble_size": rom_extra_ensemble_size,
328            "adaptive_delta_params_tolerance": adaptive_delta_tolerance,
329            "rejuvenation_prior_weight": 0.0025,
330            "rejuvenation_fraction_when_collapsed": 0.05,
331            "random_seed": 1,
332        },
333        "summary": summaries,
334        "history": histories,
335    }
336
337    _plot_histories(histories, output_dir)
338    results_path = output_dir / "h2_air_flame_eki_rejuvenation_benchmark.json"
339    results_path.write_text(json.dumps(results, indent=2) + "\n", encoding="utf-8")
340
341    print("\nBenchmark summary:", flush=True)
342    for key, summary in summaries.items():
343        print(f"  {key}: {summary}", flush=True)
344    print(
345        "BENCHMARK_SUMMARY_JSON=" + json.dumps(summaries, sort_keys=True),
346        flush=True,
347    )
348    print(f"Wrote {results_path}", flush=True)
349
350
351if __name__ == "__main__":
352    main()