Convection-diffusion-reaction inverse problem#

This demo compares single-fidelity EKI with three multifidelity variants on a convection-diffusion-reaction (CDR) model with two inferred parameters:

  • EKI using only the full-order model (FOM),

  • MF-EKI using a tailored projection-based ROM,

  • MF-EKI using automatic on-the-fly Gaussian-process surrogate construction, and

  • MF-EKI using an automatically rebuilt Lipschitz-constrained neural-network surrogate.

The forward model solves a steady 2D CDR equation on a structured grid, and the QoI is the right-boundary flux functional used in the existing UQ examples. The example estimates the diffusion coefficient nu and reaction rate sigma from a synthetic observation.

Consistent solver configuration#

All four curves use the same EKI solver controls: an initial step size of 0.05, regularization parameter 1e-4, step-size growth and decay factors of 1.25 and 2.0, relaxation parameter 1.05, observation-error tolerance 1e-5, parameter-update tolerance 1e-6, and random seed 1. The three MF-EKI cases also use the same FOM and auxiliary ensemble sizes, ROM substep schedule, ROM error tolerance, training-history length, and policy for using the previous surrogate in an update when a rebuild occurs. This keeps the comparison focused on surrogate construction rather than differences in workflow defaults.

Neural-network surrogate#

The neural-network case uses the standard romtools neural surrogate with two hidden layers and a width of three times the parameter dimension. The published benchmark uses 5,000 Adam iterations for each surrogate rebuild. Parameters and targets are normalized before training.

The network is constrained to be Lipschitz using spectral normalization. The global Lipschitz constant is estimated from the maximum pairwise slope in the normalized training data and multiplied by a safety factor of 1.1. The resulting global bound is distributed evenly over the linear layers. This gives the neural surrogate an explicit smoothness constraint while retaining the same adaptive rebuild logic used by the other MF-EKI surrogates.

PyTorch is an optional dependency. Install neural-network support with

pip install "romtools[WithTorch]"

Run the demo#

The canonical implementation lives under examples/ and is validated against the current romtools checkout in CI.

python examples/eki_mf_eki_demo/example.py

A reduced configuration is available for quick validation. It uses the same four model paths but reduces the neural-network training to 100 Adam iterations:

python examples/eki_mf_eki_demo/example.py --smoke

Regenerate the published figure#

The documentation build does not execute this scientific demo implicitly. Regenerate the checked-in figure explicitly after changes that affect the benchmark:

python examples/eki_mf_eki_demo/example.py \
    --output docs/source/demos/notebooks/eki_mf_eki_demo.svg

Results#

EKI and MF-EKI error convergence

Mean observation error across iterations for single-fidelity EKI, multifidelity EKI with a tailored ROM, multifidelity EKI with automatic Gaussian-process surrogate construction, and multifidelity EKI with a Lipschitz-constrained neural-network surrogate.#

Implementation#

  1import argparse
  2import os
  3import shutil
  4import sys
  5
  6import matplotlib.pyplot as plt
  7import numpy as np
  8
  9from romtools.rom import LipschitzConfig, NeuralNetworkConfig
 10from romtools.workflows.inverse.eki_drivers import run_eki
 11from romtools.workflows.inverse.mf_eki_drivers import mf_eki_with_auto_rom, run_mf_eki
 12from romtools.workflows.parameter_spaces import HeterogeneousParameterSpace
 13from romtools.workflows.parameters import UniformParameter
 14
 15
 16EXAMPLE_DIR = os.path.abspath(os.path.dirname(__file__))
 17CDR_PATH = os.path.join(EXAMPLE_DIR, "convection_diffusion_reaction_system_code")
 18if CDR_PATH not in sys.path:
 19    sys.path.insert(0, CDR_PATH)
 20
 21import cdr  # noqa: E402
 22import cdr_rom  # noqa: E402
 23
 24
 25class CdrFomQoiModel:
 26    def __init__(self, system: cdr.AdvectionDiffusionSystem, b_vec: np.ndarray):
 27        self._system = system
 28        self._b_vec = b_vec
 29
 30    def populate_run_directory(self, run_directory: str, parameter_sample: dict) -> None:
 31        with open(os.path.join(run_directory, "params.txt"), "w", encoding="utf-8") as handle:
 32            for key, value in parameter_sample.items():
 33                handle.write(f"{key}: {value}\n")
 34
 35    def run_model(self, run_directory: str, parameter_sample: dict) -> int:
 36        nu = float(parameter_sample["nu"])
 37        sigma = float(parameter_sample["sigma"])
 38        u = cdr.solveFom(self._system, self._b_vec, nu, sigma)
 39        np.savez(os.path.join(run_directory, "solution.npz"), u=u)
 40        return 0
 41
 42    def compute_qoi(self, run_directory: str, parameter_sample: dict) -> np.ndarray:
 43        data = np.load(os.path.join(run_directory, "solution.npz"))
 44        u = data["u"]
 45        qoi = np.dot(self._system.C, u)
 46        return np.array([qoi])
 47
 48
 49class CdrRomQoiModel:
 50    def __init__(self, system: cdr.AdvectionDiffusionSystem, basis: np.ndarray, b_vec: np.ndarray):
 51        self._system = system
 52        self._basis = basis
 53        self._b_vec = b_vec
 54        self._rom = cdr_rom.primalGalerkinROM(system, basis)
 55
 56    def populate_run_directory(self, run_directory: str, parameter_sample: dict) -> None:
 57        with open(os.path.join(run_directory, "params.txt"), "w", encoding="utf-8") as handle:
 58            for key, value in parameter_sample.items():
 59                handle.write(f"{key}: {value}\n")
 60
 61    def run_model(self, run_directory: str, parameter_sample: dict) -> int:
 62        nu = float(parameter_sample["nu"])
 63        sigma = float(parameter_sample["sigma"])
 64        u_hat = cdr_rom.solveRom(self._rom, self._b_vec, nu, sigma)
 65        u = self._basis @ u_hat
 66        np.savez(os.path.join(run_directory, "solution.npz"), u=u, u_hat=u_hat)
 67        return 0
 68
 69    def compute_qoi(self, run_directory: str, parameter_sample: dict) -> np.ndarray:
 70        data = np.load(os.path.join(run_directory, "solution.npz"))
 71        u = data["u"]
 72        qoi = np.dot(self._system.C, u)
 73        return np.array([qoi])
 74
 75
 76class CdrRomBuilder:
 77    def __init__(self, system: cdr.AdvectionDiffusionSystem, b_vec: np.ndarray, rom_dim: int = 12):
 78        self._system = system
 79        self._b_vec = b_vec
 80        self._rom_dim = rom_dim
 81
 82    def build_from_training_dirs(
 83        self,
 84        offline_data_dir: str,
 85        training_data_dirs,
 86        training_parameters=None,
 87        training_qois=None,
 88    ):
 89        del offline_data_dir, training_parameters, training_qois
 90        snapshots = []
 91        for run_dir in training_data_dirs:
 92            solution_path = os.path.join(run_dir, "solution.npz")
 93            if os.path.exists(solution_path):
 94                data = np.load(solution_path)
 95                snapshots.append(data["u"])
 96        if not snapshots:
 97            raise RuntimeError("No training snapshots found for ROM construction.")
 98        snapshot_matrix = np.column_stack(snapshots)
 99        u, _, _ = np.linalg.svd(snapshot_matrix, full_matrices=False)
100        basis = u[:, : min(self._rom_dim, u.shape[1])]
101        return CdrRomQoiModel(self._system, basis, self._b_vec)
102
103
104def _collect_error_history(work_dir: str, mf: bool) -> list:
105    history = []
106    iteration = 0
107    while True:
108        restart_path = os.path.join(work_dir, f"iteration_{iteration}", "restart.npz")
109        if not os.path.exists(restart_path):
110            break
111        data = np.load(restart_path, allow_pickle=True)
112        if mf:
113            sample_one_fom_results = data["sample_one_fom_results"].item()
114            errors = sample_one_fom_results["errors"]
115        else:
116            errors = data["errors"]
117        history.append(float(np.mean(np.linalg.norm(errors, axis=0))))
118        iteration += 1
119    return history
120
121
122def main(smoke: bool = False, work_dir: str = None, output_path: str = None) -> None:
123    np.random.seed(1)
124
125    grid_size = 10 if smoke else 25
126    fom_ensemble_size = 3 if smoke else 4
127    rom_extra_ensemble_size = 4 if smoke else 12
128    max_iterations = 2 if smoke else 20
129    rom_dim = 4 if smoke else 12
130    rom_substep_start_iteration = 0 if smoke else 1
131    rom_substep_end_iteration = 1 if smoke else 15
132    num_rom_substeps = 1 if smoke else 4
133    max_rom_training_history = 2 if smoke else 3
134    nn_training_iterations = 100 if smoke else 5000
135
136    # Use the same EKI controls for every curve so differences are due to the
137    # surrogate strategy rather than hidden workflow defaults.
138    eki_solver_args = {
139        "initial_step_size": 0.05,
140        "regularization_parameter": 1.0e-4,
141        "step_size_growth_factor": 1.25,
142        "step_size_decay_factor": 2.0,
143        "max_step_size_decrease_trys": 5,
144        "relaxation_parameter": 1.05,
145        "error_norm_tolerance": 1.0e-5,
146        "delta_params_tolerance": 1.0e-6,
147        "random_seed": 1,
148    }
149    mf_eki_solver_args = {
150        **eki_solver_args,
151        "use_updated_rom_in_update_on_rebuild": False,
152    }
153
154    system = cdr.AdvectionDiffusionSystem(Nx=grid_size, Ny=grid_size)
155    b_vec = np.array([1.0, 1.0])
156
157    nu_true = 0.04
158    sigma_true = 0.3
159    u_true = cdr.solveFom(system, b_vec, nu_true, sigma_true)
160    observations = np.array([np.dot(system.C, u_true)])
161    observations_covariance = np.eye(1) * 1e-5
162
163    parameter_space = HeterogeneousParameterSpace(
164        [
165            UniformParameter("nu", 0.01, 0.08),
166            UniformParameter("sigma", 0.1, 0.6),
167        ]
168    )
169
170    base_dir = os.path.abspath(
171        work_dir if work_dir is not None else os.path.join(EXAMPLE_DIR, "eki_mf_eki_work")
172    )
173    eki_dir = os.path.join(base_dir, "eki")
174    mf_dir = os.path.join(base_dir, "mf_eki")
175    mf_auto_rom_dir = os.path.join(base_dir, "mf_eki_auto_rom")
176    mf_lipschitz_nn_dir = os.path.join(base_dir, "mf_eki_lipschitz_nn")
177    shutil.rmtree(base_dir, ignore_errors=True)
178    os.makedirs(base_dir, exist_ok=True)
179
180    fom_model = CdrFomQoiModel(system, b_vec)
181    rom_builder = CdrRomBuilder(system, b_vec, rom_dim=rom_dim)
182
183    run_eki(
184        model=fom_model,
185        parameter_space=parameter_space,
186        observations=observations,
187        observations_covariance=observations_covariance,
188        absolute_work_dir=eki_dir,
189        ensemble_size=fom_ensemble_size,
190        max_iterations=max_iterations,
191        evaluation_concurrency=1,
192        **eki_solver_args,
193    )
194
195    run_mf_eki(
196        model=fom_model,
197        rom_model_builder=rom_builder,
198        parameter_space=parameter_space,
199        observations=observations,
200        observations_covariance=observations_covariance,
201        absolute_work_dir=mf_dir,
202        rom_substep_start_iteration=rom_substep_start_iteration,
203        rom_substep_end_iteration=rom_substep_end_iteration,
204        num_rom_substeps=num_rom_substeps,
205        fom_ensemble_size=fom_ensemble_size,
206        rom_extra_ensemble_size=rom_extra_ensemble_size,
207        rom_tolerance=0.001,
208        max_iterations=max_iterations,
209        fom_evaluation_concurrency=1,
210        rom_evaluation_concurrency=1,
211        max_rom_training_history=max_rom_training_history,
212        **mf_eki_solver_args,
213    )
214
215    mf_eki_with_auto_rom(
216        model=fom_model,
217        parameter_space=parameter_space,
218        observations=observations,
219        observations_covariance=observations_covariance,
220        absolute_work_dir=mf_auto_rom_dir,
221        rom_substep_start_iteration=rom_substep_start_iteration,
222        rom_substep_end_iteration=rom_substep_end_iteration,
223        num_rom_substeps=num_rom_substeps,
224        fom_ensemble_size=fom_ensemble_size,
225        rom_extra_ensemble_size=rom_extra_ensemble_size,
226        rom_tolerance=0.001,
227        max_iterations=max_iterations,
228        fom_evaluation_concurrency=1,
229        rom_evaluation_concurrency=1,
230        rom_type="gp",
231        rom_args={
232            "normalize_parameters": True,
233            "normalize_targets": True,
234        },
235        max_rom_training_history=max_rom_training_history,
236        **mf_eki_solver_args,
237    )
238
239    mf_eki_with_auto_rom(
240        model=fom_model,
241        parameter_space=parameter_space,
242        observations=observations,
243        observations_covariance=observations_covariance,
244        absolute_work_dir=mf_lipschitz_nn_dir,
245        rom_substep_start_iteration=rom_substep_start_iteration,
246        rom_substep_end_iteration=rom_substep_end_iteration,
247        num_rom_substeps=num_rom_substeps,
248        fom_ensemble_size=fom_ensemble_size,
249        rom_extra_ensemble_size=rom_extra_ensemble_size,
250        rom_tolerance=0.001,
251        max_iterations=max_iterations,
252        fom_evaluation_concurrency=1,
253        rom_evaluation_concurrency=1,
254        rom_type="nn",
255        rom_args={
256            "network_config": NeuralNetworkConfig(
257                training_iterations=nn_training_iterations,
258            ),
259            "lipschitz_config": LipschitzConfig(
260                enabled=True,
261                safety_factor=1.1,
262            ),
263            "normalize_parameters": True,
264            "normalize_targets": True,
265        },
266        max_rom_training_history=max_rom_training_history,
267        **mf_eki_solver_args,
268    )
269
270    eki_history = _collect_error_history(eki_dir, mf=False)
271    mf_history = _collect_error_history(mf_dir, mf=True)
272    mf_auto_rom_history = _collect_error_history(mf_auto_rom_dir, mf=True)
273    mf_lipschitz_nn_history = _collect_error_history(mf_lipschitz_nn_dir, mf=True)
274
275    plt.figure(figsize=(7.2, 4.2))
276    plt.plot(eki_history, marker="o", label="EKI (FOM)")
277    plt.plot(mf_history, marker="s", label="MF-EKI (FOM+ROM)")
278    plt.plot(mf_auto_rom_history, marker="^", label="MF-EKI (FOM+GP auto-ROM)")
279    plt.plot(
280        mf_lipschitz_nn_history,
281        marker="D",
282        label="MF-EKI (FOM+Lipschitz NN auto-ROM)",
283    )
284    plt.yscale("log")
285    plt.xlabel("Iteration")
286    plt.ylabel("Mean observation error")
287    plt.title("EKI vs MF-EKI on a convection-diffusion-reaction model")
288    plt.grid(True, alpha=0.3)
289    plt.legend()
290
291    output_path = os.path.abspath(
292        output_path
293        if output_path is not None
294        else os.path.join(EXAMPLE_DIR, "eki_mf_eki_demo.png")
295    )
296    os.makedirs(os.path.dirname(output_path), exist_ok=True)
297    plt.tight_layout()
298    plt.savefig(output_path, dpi=180)
299    plt.close()
300    print(f"Wrote {output_path}")
301
302
303def _parse_args():
304    parser = argparse.ArgumentParser(description="Run the EKI/MF-EKI CDR demo.")
305    parser.add_argument(
306        "--smoke",
307        action="store_true",
308        help="Run a reduced configuration intended for CI validation.",
309    )
310    parser.add_argument(
311        "--work-dir",
312        default=None,
313        help="Override the directory used for EKI iteration data.",
314    )
315    parser.add_argument(
316        "--output",
317        default=None,
318        help="Override the output figure path.",
319    )
320    return parser.parse_args()
321
322
323if __name__ == "__main__":
324    args = _parse_args()
325    main(smoke=args.smoke, work_dir=args.work_dir, output_path=args.output)