---
title: "Solving multiphysics-based inverse problems with learned surrogates and constraints"
author:
  - name: Ziyi Yin*
    orcid: "0000-0002-5024-8771"
    email: ziyi.yin@gatech.edu
    url: https://ziyiyin97.github.io
    affiliations:
      - name: Georgia Institute of Technology
        url: https://slim.gatech.edu/
  - name: Rafael Orozco
    orcid: "0000-0003-0917-2442"
    email: rorozco@gatech.edu
    affiliations:
      - name: Georgia Institute of Technology
        url: https://slim.gatech.edu/
  - name: Mathias Louboutin
    orcid: "0000-0002-1255-2107"
    email: mathias.louboutin@gmail.com
    url: https://mloubout.github.io/
    affiliations:
      - name: Georgia Institute of Technology
        url: https://slim.gatech.edu/
  - name: Felix J. Herrmann
    orcid: "0000-0003-1180-2167"
    email: felix.herrmann@gatech.edu
    url: https://slim.gatech.edu/
    affiliations:
      - name: Georgia Institute of Technology
        url: https://slim.gatech.edu/
bibliography: paper.bib
format:
  html:
    link-external-newwindow: True
    css: ScholarlyMarkdown-forMarked-latest.css
execute:
  freeze: auto
---

```{=html}
<style type="text/css">
caption, .table-caption {
  text-align: left;
}
</style>
```

Affiliation: Georgia Institute of Technology

Address: 756 West Peachtree Street NW, Atlanta, Georgia, USA 30308

Emails:

- Ziyi Yin (corresponding author): [ziyi.yin@gatech.edu](mailto:ziyi.yin@gatech.edu)

- Rafael Orozco: [rorozco@gatech.edu](mailto:rorozco@gatech.edu)

- Mathias Louboutin: [mlouboutin3@gatech.edu](mailto:mlouboutin3@gatech.edu)

- Felix J. Herrmann: [felix.herrmann@gatech.edu](mailto:felix.herrmann@gatech.edu)

## Abstract

Solving multiphysics-based inverse problems for geological carbon storage monitoring can be challenging when multimodal time-lapse data are expensive to collect and costly to simulate numerically. We overcome these challenges by combining computationally cheap learned surrogates with learned constraints. Not only does this combination lead to vastly improved inversions for the important fluid-flow property, permeability, it also provides a natural platform for inverting multimodal data including well measurements and active-source time-lapse seismic data. By adding a learned constraint, we arrive at a computationally feasible inversion approach that remains accurate. This is accomplished by including a trained deep neural network, known as a normalizing flow, which forces the model iterates to remain in-distribution, thereby safeguarding the accuracy of trained Fourier neural operators that act as surrogates for the computationally expensive multiphase flow simulations involving partial differential equation solves. By means of carefully selected experiments, centered around the problem of geological carbon storage, we demonstrate the efficacy of the proposed constrained optimization method on two different data modalities, namely time-lapse well and time-lapse seismic data. While permeability inversions from both these two modalities have their pluses and minuses, their joint inversion benefits from either, yielding valuable superior permeability inversions and CO~2~ plume predictions near, and far away, from the monitoring wells. 

<!-- Utilizing cheap neural surrogates has the potential to significantly reduce these computational costs, making inversion more affordable. However, the accuracy of the surrogate cannot be guaranteed during inversion, as intermediate iterates might be outside the statistical distribution of the training samples for the surrogate. To address this issue, we propose using a learned constraint based on normalizing flows to capture the characteristics of the training samples and present a surrogate-assisted inversion algorithm with iteratively relaxed constraints. This ensures that the input to the surrogate during inversion remains in-distribution, leading to higher accuracy of the surrogate. Through three synthetic experiments for geological carbon storage monitoring, we demonstrate the efficacy of the proposed inversion technique using learned surrogate and constraints. -->

**Keywords**: Fourier neural operators, normalizing flows, multiphysics, deep learning, learned surrogates, learned constraints, inverse problems

## Introduction

In this paper, we introduce a novel learned inversion algorithm designed to address inverse problems based on partial differential equations (PDEs). These problems can be represented using the following general form:

$$
\mathbf{d} = \mathcal{H}\circ\mathcal{S}(\mathbf{K}) + \boldsymbol \epsilon.
$$ {#eq-inv}

In this expression, the nonlinear operator $\mathcal{S}$ represents the solution operator of a nonlinear parametric PDE mapping the coefficients $\mathbf{K}$ to the solution. Given numerical solutions of the PDE, partially observed data, collected in the vector $\mathbf{d}$, are modeled by compounding the solution operator with the measurement operator, $\mathcal{H}$, followed by adding the noise term $\boldsymbol{\epsilon}$ with noise level of $\sigma$---i.e., $\boldsymbol{\epsilon}\sim \mathcal{N}(\mathbf{0}, \sigma^2 \mathbf{I})$. This problem is quite general and pertinent to various physical applications, including geophysical exploration [@tarantola1984inversion; @tarantola2005inverse], medical imaging [@arridge1999optical], and experimental design [@alexanderian2021optimal].

Without loss of generality, we focus on time-lapse seismic monitoring of geological carbon storage (GCS), which involves underground storage of supercritical CO~2~ captured from the atmosphere or from industrial smoke stacks [@FURRE20173916]. We consider GCS in saline aquifers, which involves multiphase flow physics where CO~2~ replaces brine in the porous rocks [@nordbotten2011geological]. In this context, the PDE solution operator, $\mathcal{S}$, serves as the multiphase flow simulator, which takes the gridded spatially varying permeability in the reservoir, $\mathbf{K}$, as input and produces $n_t$ time-varying CO~2~ saturation snapshots, $\mathbf{c}=[\mathbf{c}_1, \mathbf{c}_2, \cdots,\mathbf{c}_{n_t}]$, as output. The governing equations for the multiphase flow involve Darcy's and the mass conservation law. Detailed information on the governing equations, initial and boundary conditions, and numerical solution schemes can be found in [@rasmussen2021open] and the references therein. To ensure safety, conformance, and containment of GCS projects, various kinds of time-lapse data are collected to monitor CO~2~ plumes. These different data modalities include measurements in wells [@freifeld2009recent; @nogues2011detecting], and the collection of gravity [@nooner2007constraints; @alnes2011results], electromagnetic [@carcione2012cross; @zhdanov2013electromagnetic], and seismic time-lapse data [@arts2004monitoring; @lumley20104d; @yin2023derisking] that can be used to follow the plume and invert for reservoir properties such as the permeability, $\mathbf{K}$. The latter is the property of interest in this exposition.

Overall, solving for the reservoir model parameter, $\mathbf{K}$, poses significant challenges for two primary reasons:

- the forward modeling operator, $\mathcal{H}\circ\mathcal{S}$, can be ill-posed, resulting in multiple model parameters that fit the observed data equally well. This necessitates the use of regularizers [@golub1999tikhonov;@tarantola2005inverse] in the form of penalties or constraints [@peters2019projection].

- The PDE modeling operator $\mathcal{S}$, and the sensitivity calculations with respect to the model parameters can be computationally expensive for large-scale problems, limiting the efficacy of iterative methods such as gradient-based [@liu1989limited] or Markov chain Monte Carlo [@cowles1996markov] methods.

To overcome the second challenge, numerous attempts have been made to replace computationally expensive PDE solves with more affordable approximate alternatives [@razavi2012review; @asher2015review], which include the use of radial basis functions to learn the complex models from few sample points [@powell1985radial] or reduced-order modeling where the dimension of the model space is reduced [@schilders2008model; @lu2019review]. More recently, various deep learning techniques have emerged as cheap alternatives to numerical PDE solves [@lu2019deeponet; @pestourie2020active; @qian2020lift; @karniadakis2021physics; @kovachki2021neural; @rahman2022u; @hijazi2023pod]. After incurring initial training costs, these neural operators lead to vastly improved computation of PDE solves. Data-driven methods have also been used successfully to learn coarse-to-fine grid mappings of PDEs solves. Because of their advertised performance on approximating solution operators of the multiphase flow in porous media [@wen2022u; @grady2023model; @witte2022industry; @witte2022sciai4industry; @wen2023real; @witte2023fast], we will consider Fourier neural operators [FNOs, @li2020fourier; @kovachki2021universal] in this work even though alternative choices can be made. Once trained, FNOs produce approximate PDE solutions orders of magnitude faster than traditional solvers [@li2020fourier; @li2021physics; @grady2023model; @de2022cost]. In addition, @yin2022SEGlci, @louboutin2022SEGais and @louboutin2023learned demonstrated that trained FNOs can replace PDE solution operators during inversion. This latest development is especially beneficial to applications such as GCS where trained FNOs can be used in lieu of numerically costly flow simulators [@lie2019introduction; @rasmussen2021open; @gross2021geosx]. However, despite their promising results, unconstrained inversion formulations offer little to no guarantees that the model iterates remain within the statistical distribution on which the FNO was trained initially during inversion. As a consequence, FNOs may no longer produce accurate fluid-flow simulations throughout the iterations, which can lead to erroneous inversion results when the errors become too large, possibly rendering surrogate modeling by FNOs ineffective. To avoid this situation, we propose a constrained formulation where a trained normalizing flow (NF, @rezende2015variational) is included as a learned constraint. This added constraint guarantees that the model iterates remain within the desired statistical distribution. Because our approach safeguards the FNO's accuracy, it allows FNOs to act as reliable low-cost neural surrogates replacing costly fluid-flow simulations and gradient calculations that rely on numerically expensive PDE solves during inversion.

The organization of this paper is as follows: First, we introduce FNOs and explore the possibility of replacing the forward modeling operator with a trained FNO surrogate. Next, NFs are introduced. By means of a motivational example, we demonstrate how these learned generative networks control the prediction error of FNOs by ensuring that the model iterates remain in distribution. Based on this motivational example, we propose our novel method for using trained NFs as a learned constraint to guarantee performance of FNO surrogates during inversion. Through four synthetic experiments related to GCS monitoring, the efficacy of our method will be demonstrated.

## Fourier neural operators

There is an extensive literature on training deep neural networks to serve as affordable alternatives to computationally expensive numerical simulators [@lu2019deeponet;@karniadakis2021physics;@kovachki2021neural;@kontolati2023learning; @benitez2023fine]. Without loss of generality, we limit ourselves in this exposition to the training of a special class of neural operators known as Fourier neural operators (FNOs). These FNOs are designed to approximate numerical solution operators of the PDE solution operator, $\mathcal{S}$, by minimizing the following objective:
$$
\underset{\boldsymbol{\theta}}{\operatorname{minimize}} \quad \, \frac{1}{N}\sum_{j=1}^{N}\|\mathcal{S}_{\boldsymbol{\theta}}(\mathbf{K}^{(j)})-\mathbf{c}^{(j)}\|_2^2\quad\text{where}\quad\mathbf{c}^{(j)}=\mathcal{S}(\mathbf{K}^{(j)}).
$$ {#eq-fno-train}

Here, $\mathcal{S}_{\boldsymbol{\theta}}$ denotes the FNO with network weights ${\boldsymbol{\theta}}$. The optimization aims to minimize the $\ell_2$ misfit between numerically simulated PDE solutions, $\mathbf{c}^{(j)}$, and solutions approximated by the FNO, across $N$ training samples (permeability models), $\{\mathbf{K}^{(j)}\}_{j=1}^N$ compiled by domain experts. Once trained, FNOs can generate approximate PDE solutions for unseen model parameters orders of magnitude faster than numerical simulations [@grady2023model;@de2022cost]. For model parameters that fall within the distribution used to train, approximation by FNOs are reasonably accurate---i.e., $\mathcal{S}_{\boldsymbol\theta^\ast}(\mathbf{K})\approx\mathcal{S}(\mathbf{K})$, with $\boldsymbol\theta^\ast$ being the minimizer of @eq-fno-train. We refer to the numerical examples section for details calculating these weights. Before studying the impact of applying these surrogates on samples for the permeability that are out of distribution, let us first consider an example where data is inverted using surrogate modeling.

## Inversion with learned surrogates

Replacing PDE solutions by approximate solutions yielded by trained FNO surrogates has two main advantages when solving inverse problems. First, as mentioned earlier, FNOs are orders of magnitude faster than numerical PDE solves, which allows for many simulations at negligible costs [@chandra2020surrogate;@lan2022scaling]. Second, existing softwate for multiphase flow simulations may not always support computationally efficient calculations of sensitivity, e.g. via adjoint-state calculations [@cao2003adjoint; @plessix2006review; @jansen2011adjoint] of the simulations with respect to their input. In such cases, FNO surrogates are favorable because automatic differentiation on the trained network [@griewank1989automatic; @louboutin2022SEGais; @yin2022SEGlci; @yang2023rapid; @louboutin2023learned] readily provides access to gradients with respect to model parameters. As a result, the PDE solver, $\mathcal{S}$, in @eq-inv can be replaced by trained surrogate, $\mathcal{S}_{\boldsymbol\theta^\ast}$---i.e., we have

$$
\underset{\mathbf{K}}{\operatorname{minimize}} \quad \|\mathbf{d} - \mathcal{H}\circ\mathcal{S}_{\boldsymbol\theta^\ast}(\mathbf{K})\|_2^2
$$ {#eq-inv-fno}

where $\boldsymbol\theta^\ast$ represent the optimized weights minimizing @eq-fno-train. While the above formulation in terms of trained surrogates has been applied successfully during permeability inversion from time-lapse seismic data [@li2020coupled; @yin2022SEGlci; @louboutin2023learned], this type of inversion is only valid as long as the (intermediate) permeabilities remain within distribution during the inversion. Practically, this means two things. First, the data need to be in the range of permeability models that are in distribution. This means that there can not be too much noise neither can the observed data be the result of an out-of-distribution permeability. Second, there are no guarantees that the permeability model iterates remain in distribution during inversion even though some bias of the gradients of the surrogate towards in-distribution permeabilities may be expected. To overcome this challenge, we propose to add a learned constraint to @eq-inv-fno that offers guarantees that the model iterates remain in distribution.

## Learned constraints with normalizing flows

As demonstrated by @peters2017constraints, @esser2018total, @peters2019projection regularization of non-linear inverse problems, such as full-waveform inversion, with constraints, e.g., total-variation [@esser2018total] or transform-domain sparsity with $\ell_1$-norms [@li2012fast], offers distinct advantages over regularizations based adding these norms as penalties. Even though constraint and penalty formulations are equivalent for linear inverse problems for the appropriate choice of the Lagrange multiplier, minimizing the constraint formulation leads to completely different solution paths compared to adding a penalty term to the data misfit objective [@hennenfent2008new]. In the constrained formulation, the model iterates remain at all times within the constraint set while model iterates yielded by the penalty formulation does not offer these guarantees. @peters2017constraints demonstrated this importance difference for the non-convex problem of full-waveform inversion. For this problem, it proved essential to work with a homotopy where the intersection of multiple handcrafted constraints (intersection of box and size of total-variation-norm ball constraints) are relaxed slowly during the inversion, so the model iterates remain physically feasible and local minima are avoided.

Motivated by these results, we propose a similar approach but now for "data-driven" learned constraints based on normalizing flows (NFs, @rezende2015variational). NFs are powerful deep generative neural networks capable of learning to generate samples from complex distributions [@dinh2016density;@doi:10.1190/geo2022-0472.1;@orozco2023MIDLanf;@louboutin2023learned]. Designed to be invertible, these NFs require the latent and model spaces to share identical dimensions, which confers several advantages:

- unlike variational autoencoders [@kingma2013auto] or generative adversarial networks [GANs, @goodfellow2014generative], which both have a lower-dimensional latent space, NFs do not impose any intrinsic dimensionality constraints. This flexibility lets NFs capture model space characteristics across high dimensions [@kobyzev2020normalizing]. Relevantly, concurrent literature has delved into the intrinsic dimensionality of NFs, indicating the potential to using NFs to generate models with inherently lower dimensions [@horvat2022intrinsic].

- NFs' inherent invertibility negates the need to store state variables during gradient calculations, enabling memory-efficient training and inversion in large-scale 3D applications, such as in geophysics [@zhang2020seismic; @siahkoohi2020ABIpto; @siahkoohi2021Seglbe; @zhang2021bayesian2; @zhao2020bayesian; @lensink2022fully; @doi:10.1190/geo2022-0472.1] and ultrasound imaging [@orozco2021photoacoustic; @adjoint; @orozco2022memory; @orozco2023MIDLanf; @orozco2023refining; @wen2023conditional].

- because of their invertibility NFs guarantee unique latent codes for all model space samples, including out-of-distribution ones. Therefore, they can still be used to invert for out-of-distribution model parameters, while other methods like GANs may introduce bias [@asim2020invertible].

Aside from being invertible, NFs are trained to map samples from a target distribution in the physical space to samples from the standard zero-mean Gaussian distribution noise in the latent space. After training is completed, samples from the target distribution are generated by running the NF in reverse on samples in the latent space from the standard Gaussian distribution. Below, we will demonstrate how NFs can be used to guarantee that the permeability remains in distribution during the inversion.

### Training normalizing flows

Given samples from the permeability distribution, $\{\mathbf{K}^{(j)}\}_{j=1}^N$, training NFs entails minimizing the Kullback-Leibler divergence between the base and target distributions [@ardizzone2018analyzing]. This involves solving the following variational problem:
$$
\underset{\mathbf{w}}{\operatorname{minimize}} \quad \frac{1}{N}\sum_{i=1}^{N} \left(\frac{1}{2}\|\mathcal{G}^{-1}_{\mathbf{w}}(\mathbf{K}_i)\|_2^2-\log\left|\det\mathrm{J}_{\mathcal{G}^{-1}_{\mathbf{w}}}(\mathbf{K}_i)\right|\right).
$$ {#eq-train-nf}

In this optimization problem, $\mathcal{G}^{-1}_{\mathbf{w}}$ represents the NF, which is parameterized by its network weights $\mathbf{w}$, while $\mathrm{J}_{\mathcal{G}^{-1}_{\mathbf{w}}}$ denotes its Jacobian. By minimizing the $\ell_2$-norm, the objective imposes a Gaussian distribution on the network's output and the second $\log\det$ term prevents trivial solutions, i.e., cases where $\mathcal{G}^{-1}_{\mathbf{w}}$ produces zeros. To ensure alignment between the permeability distributions, @eq-fno-train and @eq-train-nf are trained on the same dataset consisting of 2000 permeability models examples of which are included in @fig-trainingsample. Each $64\times64$ permeability model of consists of a randomly generated highly permeable channels ($120\,\mathrm{mD}$) in a low-permeable background of $20\,\mathrm{mD}$, where $\,\mathrm{mD}$ denotes millidarcy. Generative examples produced by the trained NF are included in the second row of @fig-trainingsample, which confirm the NF's ability to learn distributions from examples. Aside from generating samples from the learned distribution, trained NFs are also capable of carrying out density calculations, an ability we will exploit below.

::: {#fig-trainingsample layout="[[22, 22, 22], [22, 22, 22]]"}
![](./figs/K1.png)

![](./figs/K2.png)

![](./figs/K3.png)

![](./figs/gen_K1.png)

![](./figs/gen_K2.png)

![](./figs/gen_K3.png)

Permeability models. First row shows the realistic permeability samples for FNO and NF training. Second row shows the generative samples from the trained NF.
:::

### Trained normalizing flows as constraints

As we mentioned before, adding constraints to the solution of non-convex optimization problems offers guarantees that model iterates remain within constrained sets. When solving inverse problems with learned surrogates, it is important that model iterates remain "in distribution", which can be achieved by recasting the optimization problem in @eq-inv-fno into the following constrained form:

$$
\underset{\mathbf{z}}{\operatorname{minimize}} \quad \|\mathbf{d} - \mathcal{H}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2\quad\text{subject to}\quad\|\mathbf{z}\|_2\leq \tau.
$$ {#eq-inv-fno-nf}

To arrive at this constrained optimization problem, two important changes were made. First, the permeability $\mathbf{K}$ is replaced by the output of a trained NF with trained weights $\mathbf{w}^\ast$ obtained by minimizing @eq-train-nf [@asim2020invertible]. This reparameterization in terms of the latent variable, $\mathbf{z}$, produces permeabilities that are in distribution as long as $\mathbf{z}$ remains distributed according to the standard normal distribution. Second, we added a constraint on this latent space variable in @eq-inv-fno-nf, which ensures that the latent variable $\mathbf{z}$ remains within an $\ell_2$-norm ball of size $\tau$. 

<!-- To simplify notation, we drop the subscripts $\boldsymbol\theta^\ast$ and $\mathbf{w}^\ast$ from now on and consider the NF, $\mathcal{G}$, and FNO,, trained. -->

To better understand the behavior of a trained normalizing flow in conjunction with the $\ell_2$-norm constraint for in- and out-of-distribution examples, we include @fig-ood-samples and @fig-homotopy. In the latter Figure, nonlinear projections (via latent space shrinkage), 
$$
\widetilde{\mathbf{K}} = \mathcal{G}_{\mathbf{w}^\ast}\left(\alpha\mathbf{z}\right)\quad\text{where}\quad \mathbf{z}=\mathcal{G}_{\mathbf{w}^\ast}^{-1}(\mathbf{K})
$$ {#eq-shrinkage}

<!-- solutions of $\min_{\mathbf{z}}\|\mathbf{K}-\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2\,\text{subject to}\,\|\mathbf{z}\|_2\leq \tau$ are plotted for increasing $\tau$.  -->
are plotted as a function of increasing $\alpha$. We also plot in @fig-ood the NF's relative nonlinear approximation error, $\|\mathbf{\widetilde{K}}-\mathbf{K}\|_2/\|\mathbf{K}\|_2$, and the corresponding relative FNO prediction error,$\|\mathcal{S}_{\boldsymbol\theta^\ast}(\widetilde{\mathbf{K}})-\mathcal{S}(\widetilde{\mathbf{K}})\|_2/\|\mathcal{S}(\widetilde{\mathbf{K}})\|_2$ as a function of increasing $0\leq\alpha\leq 1$. From these plots, we can make the following observations. First, the latent representations (@fig-ind-latent and @fig-ood-latent) of the in- and out-of-distribution samples (@fig-ind-sample and @fig-ood-sample ) clearly show that NF applied to out-of-distribution samples produces a latent variable far from the standard normal distribution, while the latent variable corresponding to the in-distribution example is close to being white Gaussian noise. Quantitatively, the $\ell_2$ norm of the latent variables are $0.99\|\mathcal{N}(0,\mathbf{I})\|_2$ and $3.11\|\mathcal{N}(0,\mathbf{I})\|_2$, respectively, where $\|\mathcal{N}(0,\mathbf{I})\|_2$ corresponds to the $\ell_2$-norm of the standard normal distribution. Second, we observe from @fig-homotopy that for small $\ell_2$-norm balls ($\alpha\ll 1$) the projected solutions tend to be close to the most probable sample, which is a flat permeability channel in the middle. This is true for both the in- and out-of-distribution example. As $\alpha$ increases, the in-distribution example is reconstructed accurately when the $\ell_2$ norm of the scaled latent variable, $\|\alpha\mathbf{z}\|_2$, is close to the $\|\mathcal{N}(0,\mathbf{I})\|_2$. Clearly, this is not the case for the out-of-distribution example. When $\|\alpha\mathbf{z}\|_2\approx\|\mathcal{N}(0,\mathbf{I})\|_2$, the reconstruction still looks like an in-distribution permeability sample and is not close to the out-of-distribution sample. However, if $\alpha=1$, which makes $\|\alpha\mathbf{z}\|_2$ well beyond the norm of the standard normal distribution, the out-of-distribution example is recovered accurately by virtue of the invertibility of NFs, irrespective on their input and what they have been trained on. Third, the relative FNO prediction error for the in-distribution example (@fig-ind-alpha) remains flat while the error of the FNO surrogate increases as soon as $\alpha\approx 0.25$. At that value for $\alpha$, the projection, $\widetilde{\mathbf{K}}$, is gradually transitioning from being in-distribution to out-of-distribution, which occurs at a non-linear approximation error of about 45%. As expected the plots in @fig-ood also show a monotonous decay of the nonlinear approximation error as a function of increasing $\alpha$. To further analyze the effects of the nonlinear projections in  @eq-shrinkage, we draw 50 random realizations from the standard normal distribution, scale each of them by $0\leq\alpha\leq 2$, and calculate the FNO prediction errors on these samples. @fig-init-tau includes the results of this excercise where each column represents the FNO prediction error calculated for $0\leq \alpha \leq 2$. From these experiments, we make the following two observations. First, when $\alpha<0.8$, the FNO consistently makes accurate predictions for all projected samples. Second, as expected, the FNO starts to make less accurate predictions for $\alpha>1$ with errors that increase as the size of the $\ell_2$-norm ball of the latent space expands, demarcating the transition from being in distribution to being out-of-distribution.

:::{#fig-ood-samples layout="[[1,1], [1,1]]"}
![](./figs/in-dist-sample.png){#fig-ind-sample}

![](./figs/out-dist-sample.png){#fig-ood-sample}

![](./figs/normal_z.png){#fig-ind-latent}

![](./figs/ood_z.png){#fig-ood-latent}

Sample permeability models in the physical and latent space. *(a)* An in-distribution permeability model. *(b)* An out-of-distribution permeability model. *(c)* An in-distribution permeability model in the latent space. *(d)* An out-of-distribution permeability model in the latent space.
:::

:::{#fig-homotopy layout="[[15, 15, 15, 15, 15], [15, 15, 15, 15, 15]]"}

![](./figs/id-0.png){#fig-ind-sample-ho1}

![](./figs/id-0.1.png){#fig-ind-sample-ho2}

![](./figs/id-0.2.png){#fig-ind-sample-ho3}

![](./figs/id-0.4.png){#fig-ind-sample-ho4}

![](./figs/id-1.png){#fig-ind-sample-ho5}

![](./figs/ood-0.png){#fig-ood-sample-ho1}

![](./figs/ood-0.1.png){#fig-ood-sample-ho2}

![](./figs/ood-0.2.png){#fig-ood-sample-ho3}

![](./figs/ood-0.4.png){#fig-ood-sample-ho4}

![](./figs/ood-1.png){#fig-ood-sample-ho5}

Projections onto increasing $\ell_2$-norm balls for the in- and out-of-distribution examples of @fig-ood-samples. Top row: projections of in-distribution sample. Bottom row: projections of out-of-distribution sample. Each column corresponds to setting $\alpha=0,0.1,0.2,0.4,1$ in @eq-shrinkage\.
:::

::: {#fig-ood layout-ncol=2}
![](./figs/in-dist-alpha.png){#fig-ind-alpha}

![](./figs/out-dist-alpha.png){#fig-ood-alpha}

Latent space projection experiments.*(a)* Relative $\ell_2$ reconstruction error and FNO prediction error for in-distribution sample. *(b)* The same but for out-of-distribution sample. The blue curve shows the relative $\ell_2$ misfit between the permeability models before and after latent space shrinkage. The orange curve shows the FNO prediction error on the permeability model after shrinking the $\ell_2$-norm ball. The red dashed line denotes the amplitude of standard Gaussian noise.
:::

::: {#fig-init-tau}
![](./figs/init-tau.png){width=50%}

FNO prediction errors for the latent space shrinkage experiment in @eq-shrinkage for 50 random realizations of standard Gaussian noise.
:::

<!-- The first quantity signifies the discrepancy between the reconstructed model parameter and the initial model parameter, as illustrated by the blue curve in @fig-ood-alpha. This curve exhibits an almost monotonously decreasing trend with respect to the $\ell_2$ misfit. The second quantity represents the prediction error of the FNO. Since the model parameter is out-of-distribution, the FNO incurs a substantial prediction error when $\alpha$ approaches 1. As $\alpha$ diminishes, the reconstructed model parameter, $\widetilde{\mathbf{K}}$, progressively shifts towards the training sample distribution. This shift results in a reduced FNO prediction error, ultimately reaching approximately 7% -- the prediction error for the training samples. These findings corroborate that latent space shrinkage effectively projects out-of-distribution samples onto in-distribution samples, thereby enhancing the accuracy of FNO predictions. They also underscore the importance of constraining model parameters to remain within the training sample distribution during the inversion process. -->

In summary, the experiments of @fig-ood-samples to @fig-ood indicate that FNO errors remain small and relatively constant for the in-distribution example. Irrespective of the value of $\alpha$, the generated samples remain in distribution while moving from the most likely---i.e., a flat high-permeability channel in the middle, to the in-distribution sample as $\alpha$ increases. Conversely, the projection of the out-of-distribution example morphs from being in distribution to being out-of-distribution for $\alpha \geq 0.25$. The FNO prediction errors also increase during this transition from an in-distribution sample to an out-of-distribution sample. Therefore, shrinkage in the latent space by multiplying with a small $\alpha$ can serve as an effective projection that ensures relatively low FNO prediction errors. We will use this unique ability to control the distribution during inversion.

<!-- 
We also present an in-distribution model parameter, as depicted in @fig-ind-sample. For this in-distribution model parameter, the latent space shrinkage increases the likelihood of the model parameter. However, the FNO prediction error remains relatively constant, as the reconstructed model parameters in this sequence are all in-distribution, demonstrated in @fig-ind-alpha. Based on the motivational example provided in this section, we make the following observations:

- By training the FNO and the NF on the same samples, the NF is able to generate samples for which the FNO can produce relatively accurate predictions;
- The trained NF can gradually transform an out-of-distribution sample into an in-distribution one through shrinkage in the latent space, which consequently reduces the FNO prediction error.

These insightful findings motivate the use of both the trained FNO and NF to address the inverse problems shown in @eq-inv, thereby enhancing the reliability of the FNO during inversion with lower prediction errors. This can be achieved through the learned iterative algorithm introduced in the subsequent section, where we leverage the latent space shrinkage effects of the NF. -->


### Inversion with progressively relaxed learned constraints

Our main objective is to perform inversions where the multiphase flow equations are replaced with pretrained FNO surrogates. To make sure the learned surrogates remain accurate, we propose working with a continuation scheme where the learned constraint in @eq-inv-fno-nf is steadily relaxed by increasing the size of the $\ell_2$-norm ball constraint. Compared to the more common penalty formulation, where regularization entails adding a Lagrange-multiplier weighted $\ell_2$-norm squared, constrained formulations offer guarantees that the model iterates for the latent variable, $\mathbf{z}$, remain within the constraint set---i.e., within the $\ell_2$-norm ball of size $\tau$. Using the argument of the previous section, this implies that permeability distributions generated by the trained NF remain in distribution as long as the size of the initial $\ell_2$-norm ball, $\tau_{\mathrm{init}}$, is small enough (e.g., smaller than $0.6\|\mathcal{N}(0,\mathbf{I})\|_2$, following the observations from @fig-init-tau). Taking advantage of this trained NF in a homotopy, we propose the following algorithm:

#### Algorithm 1: Inversion with relaxed learned constraints
|    1. &emsp; **Input**: initial model parameter $\mathbf{K}_0\in \mathbb{R}^N$, observed data $\mathbf{d}$, noise level $\sigma$       
|    2. &emsp; **Input**: trained FNO $\mathcal{S}_{\boldsymbol{\theta}^\ast}$, trained NF $\mathcal{G}_{\mathbf{w}^\ast}$         
|    3. &emsp; **Input**: number of inner-loop iterations $maxiter$
|    4. &emsp; **Input**: initial $\ell_2$ ball size $\tau_{\mathrm{init}}$, multiplier $\beta>1$, final $\ell_2$ ball size $\tau_{\mathrm{final}}$       
|    5. &emsp; $\mathbf{z} = \mathcal{G}_{\mathbf{w}^\ast}^{-1}(\mathbf{K}_0)$
|    6. &emsp; $\tau = \tau_{\mathrm{init}}$
|    7. &emsp; **while** \ $\|\mathbf{d} - \mathcal{H}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2 > \sigma\|\mathcal{N}(0,\mathbf{I})\|_2$ \ **and** $\tau\leq\tau_{\mathrm{final}}$
|    8. &emsp; &emsp;**for** \ $iter = 1:maxiter$
|    9. &emsp; &emsp; &emsp;$\displaystyle\mathbf{g}=\nabla_{\mathbf{z}}\|\mathbf{d} - \mathcal{H}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2$
|  10. &emsp; &emsp; &emsp;$\mathbf{z} = \mathcal{P}_{\tau}(\mathbf{z} - \gamma\mathbf{g})$
|  11. &emsp; &emsp;**end**
|  12. &emsp; &emsp;$\tau = \beta\tau$
|  13. &emsp; **end**
|  14. &emsp; **Output**: inverted model parameter $\mathbf{K}=\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})$

Given observed data, $\mathbf{d}$, trained networks, $\mathcal{S}_{\boldsymbol{\theta}^\ast}$ and $\mathcal{G}_{\mathbf{w}^\ast}$, the initial guess for the permeability distribution, $\mathbf{K} _0$, the initial size of the $\ell_2$-norm ball, $\tau_{\mathrm{init}}$, and the final size of the $\ell_2$-norm ball, $\tau_{\mathrm{final}}$, Algorithm 1 proceeds by solving a series of constrained optimization problems where the size of the constraint set is increased by a factor of $\beta$ after each iteration (cf. line 12 in Algorithm 1). The constrained optimization subproblems themselves (cf. line 8 to 11 of Algorithm 1) are solved with projected gradient descent [@beck2014introduction]. Each iteration of the projected gradient descent method first calculates the gradient (cf. line 9 of Algorithm 1), followed by the much cheaper projection of the updated latent variable back onto the $\ell_2$-norm ball of size $\tau$ via the projection operator $\mathcal{P}_{\tau}$ (cf. line 10 in Algorithm 1). This projection is a trivial scaling operation if the updated latent variable $\ell_2$-norm exceeds the constraint --- i.e., 
$$
\mathcal{P}_{\tau}(\mathbf{z}) = 
\begin{cases} 
\mathbf{z} & \text{if } \|\mathbf{z}\|_2 \leq \tau \\
\tau\mathbf{z}/\|\mathbf{z}\|_2 & \text{if } \|\mathbf{z}\|_2 > \tau 
\end{cases}
$$ {#eq-proj}

A line search determines the steplength $\gamma$ [@stanimirovic2010accelerated] for each iteration shown in line 8 to 11. As is common in continuation methods, the relaxed gradient-descent iterations are warm-started with the optimization result from the previous iteration, which at the first iteration is initialized by the latent representation of the initial permeability model, $\mathbf{K}_0$ (cf. line 5 in Algorithm 1). Practically, each subproblem does not need to be fully solved, but only need a few iterations instead. The number of iterations to solve each subproblem is denoted by $maxiter$ in line 8 of Algorithm 1. This continuation strategy serves two purposes. First, for small $\tau$'s it makes sure the model iterates remain in distribution, so accuracy of the learned surrogate is preserved. Second, by relaxing the constraint slowly, the data residual is gradually allowed to decrease, bringing in more and more features derived from the data. By slowly relaxing the constraint, we find a careful balance between these two purposes as long as progress is made towards the solution when solving the subproblem (cf. line 8 to 11 in Algorithm 1). One notable distinction of the surrogate-assisted inversion, compared to the conventional inversion with relaxed constraints [@esser2018total], is that the size of the $\ell_2$-norm projection ball cannot increase far beyond the $\ell_2$-norm of the standard Gaussian white noise on which the NFs are trained. Otherwise, there is no guarantee the learned surrogate is accurate because the NF may generate samples that are out-of-distribution (cf. @fig-init-tau). This is explicitly incorporated into the stopping criteria, $\tau\leq\tau_{\mathrm{final}}$, in line 7 of Algorithm 1.

<!-- The inverse of the NF, $\mathcal{G}_{\mathbf{w}^\ast}$, initially generates model parameters within the training sample distribution from the latent variable $\mathbf{z}$. Subsequently, the FNO surrogate, $\mathcal{S}_{\boldsymbol\theta^\ast}$, predicts the PDE solution based on the model parameter. We aim to minimize the misfit between the predicted and observed PDE solutions using the measurement operator $\mathcal{H}$. To promote the input of the FNO surrogate to closely resemble the samples used in training, we explicitly impose a constraint on the $\ell_2$-norm of the latent variable $\mathbf{z}$, where the scalar $\tau$ dictates the size of the $\ell_2$-norm ball. A smaller $\ell_2$-norm corresponds to a more probable sample of model parameters, while a larger $\ell_2$-norm corresponds to a less probable one. We opt for the constraint formulation over a penalty term (such as the one proposed by @asim2020invertible) because a tight constraint at the start ensures higher accuracy of FNO predictions during the initial iterations. To devise an effective strategy for selecting $\tau$, we draw on existing literature on basis pursuit problems [@van2009probing] and total-variation regularized full-waveform inversion problems [@esser2018total]. We initiate iterations with a very small $\tau$ to achieve higher accuracy of the FNO surrogate at the beginning. In later iterations, we gradually relax the constraint to reduce the data misfit and explore potential model parameter choices that reside in the tail of the distribution. The optimization algorithm is outlined below. First, the algorithm converts the model parameter $\mathbf{K}$ to the latent variable $\mathbf{z}$ in line 4. Next, it initiates iterations with a very small projection ball size $\tau$ to find the optimal $\mathbf{z}$ that minimizes the data misfit, as shown in line 7. Note that this minimization problem for a fixed $\tau$ does not have to be fully solved. In practice, we can solve it for only a few iterations. The projection ball size $\tau$ is gradually increased, as displayed in line 8. Ultimately, the inverted model parameter can be obtained by converting from the latent variable $\mathbf{z}$, thanks to the invertibility of the NF. -->

<!-- ## Motivational example

Before introducing the surrogate-assisted inversion algorithm with learned constraint, we first showcase a basic example to demonstrate the relationship between shrinkage in the latent space and the accuracy of the PDE solutions predicted by the FNO. In this example, we consider training an FNO to replace the 2D multiphase flow modeling operator, $\mathcal{S}$, and train an NF to learn the statistical distribution of reservoir permeability.

We generate 2,000 synthetic permeability samples, each of which is 64 $\times$ 64 in size with a grid spacing of 15m in both vertical and horizontal directions. There is a high permeability channel (120 millidarcies) in the middle of each permeability model, and everywhere else in the model is low permeability (20 millidarcies). We use a Gaussian process to generate random curving boundaries for the high permeability channel in each model. Three of these permeability samples are shown in the first row of @fig-trainingsample. We set up a CO~2~ injection well on the left-hand side of the model and a brine production well on the right-hand side of the model. Supercritical CO~2~ is injected at a constant speed for 800 days. We use [Jutul.jl](https://github.com/sintefmath/Jutul.jl) [@Jutul; @JutulDarcy; @JutulDarcyRules] to simulate the CO~2~ dynamics for each permeability model by numerically solving the immiscible and compressible two-phase flow equations and obtain the CO~2~ saturation snapshot sover a period of 800 days with a time step of 100 days. Due to the buoyancy effect and the well controls, CO~2~ gradually moves from the left to the right and moves upward, as shown in the second row of @fig-fno-train-plume (where the corresponding permeability model for this simulation is shown in the first row).



Next, we use these 2,000 pairs of permeability and CO~2~ saturation snapshots to train an FNO. We split the pairs into a training set (with 1,900 pairs) and a validation set (with 100 pairs). After 200 epochs of training, we achieve an average of 7% prediction errors for permeability samples in the unseen test set. The prediction error that FNO makes is mostly at the edge of the CO2 plumes, potentially due to the inductive bias of FNO. We use the same 2,000 permeability samples (with the same training/validation split) to train an NF. We use the HINT network structure [@kruse2021hint] in the open-source package [InvertibleNetworks.jl](https://github.com/slimgroup/InvertibleNetworks.jl) [@InvertibleNetworks] and train it for 125 epochs. Three generative samples are shown in the second row of @fig-trainingsample. We can see that the trained NF is already capable of generating random permeability models that resemble the ones in the training samples, despite very minor artifacts.

::: {#fig-fno-train-plume}

![](./figs/FNO_train_plume.png)

FNO prediction result for CO~2~ saturation snapshots. First row shows a (time-invariant) permeability model in the unseen test set. Second row shows the CO~2~ saturation snapshots simulated by the PDE modeling operator. Third row shows the CO~2~ saturation snapshots predicted by FNO. Fourth row shows the $5\times$ difference between the second and the third row.
:::

### Quality of generative samples

With the trained FNO and NF, we first verify that the trained NF has already learned to capture the characteristics of the FNO training samples. To accomplish this, we draw 100 random model parameters from an unseen test set and calculate the FNO prediction error on each of them. We plot the histogram in @fig-err-hist-real-samples. Next, we draw 100 random realizations of white noise and pass them through the trained NF to obtain 100 generative samples. We then calculate the FNO prediction error on these 100 generative samples. The histogram is shown in @fig-err-hist-gen-samples.

From this comparison, we can see that the histogram of the FNO prediction error is roughly the same for the samples in the distribution and the generative samples from the NF. Despite a few outliers, most of the model parameters, either realistic test samples or samples generated by NF, lead to a 5%-8% FNO prediction error. This demonstrates not only that the NF has learned to generate samples that resemble the ones in the training samples but also is able to generate samples on which the FNO can yield a reasonably low prediction error, i.e., the statement, $\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\approx\mathcal{S}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})$, holds for white noise $\mathbf{z}$.

::: {#fig-qual-gen layout-ncol=2}

![](./figs/err-hist-real-samples.png){#fig-err-hist-real-samples}

![](./figs/err-hist-gen-samples.png){#fig-err-hist-gen-samples}

Histogram of FNO prediction errors on: *(a)* 100 test samples and *(b)* 100 generative samples from NF.
:::

### Effect of shrinkage in latent space

Next, we demonstrate the impact of shrinkage in the latent space on the model parameters and its influence on the prediction error of the FNO. We begin by introducing a permeability model that lies outside the statistical distribution of the training samples, as depicted in @fig-ood-sample. This particular permeability model lacks a high-permeability zone on the left-hand side, which could potentially hinder the CO2 plume's left-to-right propagation. Since it is out-of-distribution, The latent variable associated with the permeability is three times larger than standard white noise in the $\ell_2$-norm. We choose to examine a series of permeability models derived based on shrinkage in the latent space using various scaling factors $\alpha$. Specifically, we have:

$$
\widetilde{\mathbf{K}} = \mathcal{G}_{\mathbf{w}^\ast}\left(\alpha\mathbf{z}\right)\quad\text{where}\quad \mathbf{z}=\mathcal{G}_{\mathbf{w}^\ast}^{-1}(\mathbf{K}).
$$ {#eq-shrinkage}

In this context, $\alpha \in [0,1]$ represents the scaling factor, while $\mathbf{z}$ denotes the latent variable associated with the out-of-distribution permeability model. When $\alpha = 0$, $\widetilde{\mathbf{K}}$ corresponds to the most probable sample drawn from the distribution of training samples. As the value of $\alpha$ increases, the permeability model progressively deviates from this point and converges towards the out-of-distribution model when $\alpha = 1$. To investigate this behavior, we consider a range of $\alpha$ values within the interval $[0,1]$ and examine the following two quantities:

\begin{align}
\text{relative}\ \ell_2\ \text{misfit} &= \frac{\|\widetilde{\mathbf{K}}-\mathbf{K}\|}{\|\mathbf{K}\|} \\
\text{prediction error} &= \frac{\|\mathcal{S}_{\boldsymbol\theta^\ast}(\widetilde{\mathbf{K}})-\mathcal{S}(\widetilde{\mathbf{K}})\|}{\|\mathcal{S}(\widetilde{\mathbf{K}})\|}.
\end{align}



## Surrogate-assisted inversion with learned constraints

The motivational examples discussed earlier illustrate that shrinkage in the latent space can effectively guide an out-of-distribution sample towards an in-distribution sample, resulting in more accurate FNO predictions. To capitalize on this discovery, we propose the following optimization problem, which utilizes the FNO surrogate to minimize computational expenses and incorporates the NF-based constraint to ensure a more accurate and reliable inversion process:

$$
\underset{\mathbf{z}}{\operatorname{minimize}} \quad \|\mathbf{d} - \mathcal{H}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2\quad\text{subject to}\quad\|\mathbf{z}\|_2\leq \tau.
$$ {#eq-inv-fno-nf}

In this objective function, $\boldsymbol\theta^\ast$ and $\mathbf{w}^\ast$ denote the minimizers of @eq-fno-train and @eq-train-nf, respectively, both trained on identical samples of model parameters. The inverse of the NF, $\mathcal{G}_{\mathbf{w}^\ast}$, initially generates model parameters within the training sample distribution from the latent variable $\mathbf{z}$. Subsequently, the FNO surrogate, $\mathcal{S}_{\boldsymbol\theta^\ast}$, predicts the PDE solution based on the model parameter. We aim to minimize the misfit between the predicted and observed PDE solutions using the measurement operator $\mathcal{H}$. To promote the input of the FNO surrogate to closely resemble the samples used in training, we explicitly impose a constraint on the $\ell_2$-norm of the latent variable $\mathbf{z}$, where the scalar $\tau$ dictates the size of the $\ell_2$-norm ball. A smaller $\ell_2$-norm corresponds to a more probable sample of model parameters, while a larger $\ell_2$-norm corresponds to a less probable one. We opt for the constraint formulation over a penalty term (such as the one proposed by @asim2020invertible) because a tight constraint at the start ensures higher accuracy of FNO predictions during the initial iterations. To devise an effective strategy for selecting $\tau$, we draw on existing literature on basis pursuit problems [@van2009probing] and total-variation regularized full-waveform inversion problems [@esser2018total]. We initiate iterations with a very small $\tau$ to achieve higher accuracy of the FNO surrogate at the beginning. In later iterations, we gradually relax the constraint to reduce the data misfit and explore potential model parameter choices that reside in the tail of the distribution. The optimization algorithm is outlined below. First, the algorithm converts the model parameter $\mathbf{K}$ to the latent variable $\mathbf{z}$ in line 4. Next, it initiates iterations with a very small projection ball size $\tau$ to find the optimal $\mathbf{z}$ that minimizes the data misfit, as shown in line 7. Note that this minimization problem for a fixed $\tau$ does not have to be fully solved. In practice, we can solve it for only a few iterations. The projection ball size $\tau$ is gradually increased, as displayed in line 8. Ultimately, the inverted model parameter can be obtained by converting from the latent variable $\mathbf{z}$, thanks to the invertibility of the NF. -->

## Numerical Experiments

To showcase the advocasy of the proposed optimization method with relaxed learned constraints, a series of carefully chosen experiments of increasing complexity are conducted. These experiments are designed to be relevant to GCS, which in its ultimate form involves coupling of multiphase flow with the wave equation to perform end-to-end inversion for the permeability given multimodal data. To convince ourselves of the validity of our approach, at all times comparisons will be made between inversion results involving numerical solves of the multiphase equations and inversions yielded by approximations with our learned surrogate.

<!-- We showcase the effectiveness of our proposed method via a synthetic case study for reservoir property inversion during geological carbon sequestration (GCS). As GCS projects progress, increasing amounts of information about the CO~2~ plume become available, which can be utilized to infer reservoir permeability. In the following sections, we introduce three examples where we employ either relatively simple linear operators or a wave physics operator to measure CO~2~ saturation over time. For each example, we present 4 types of inversion results

1. Inversion using physics operators
2. Inversion using the FNO surrogate
3. Inversion using the FNO surrogate combined with the NF-based constraint
4. Inversion using physics operators combined with the NF-based constraint -->

For all numerical experiments, the "ground-truth" permeability model will be selected from the unseen test set and is shown in @fig-true-perm. The inversions will be initiated with the smooth permeability model depicted in @fig-init-perm. This initial model, $\mathbf{K}_0$, represents the arithmetic mean of all permeability samples in the training dataset. To ensure that the model iterates remain in distribution, we set the starting $\ell_2$-norm ball size to $\tau_{\mathrm{init}}=0.6 \|\mathcal{N}(0,\mathbf{I})\|_2$---i.e., $0.6\times$ the $\ell_2$-norm of standard white Gauss noise realizations for the discrete permeability model of 64 by 64 gridpoints. To gradually relax the learned constraint, the multiplier of the projection ball size is taken to be $\beta=1.2$, and we set the ultimate projection ball size $\tau_{\mathrm{final}}$ in Algorithm 1 to be $1.2$ times the norm of standard white noise. To limit computational costs of solving the subproblems, we allow each constrained subproblem (cf. line 8 to 11 in Algorithm 1) to perform 8 iterations of projected gradient descent to solve for the latent variable. From practical experience, we found that the proposed inversions are not very sensitive to the choice of these hyperparameters.

::: {#fig-perm-init-true layout="[[40, 40]]"}

![](./figs/true-perm.png){#fig-true-perm}

![](./figs/init-perm.png){#fig-init-perm}

Permeability models. *(a)* unknown "ground-truth" permeability model from unseen test set, where the symbols $\blacktriangleright$ and $\blacktriangleleft$ denote the CO~2~ injection and brine production location, respectively; *(b)* initial permeability model, $\mathbf{K}_0$.
:::

To simulate the evolution of injected CO~2~ plumes, we make use of the open-source software package [Jutul.jl](https://github.com/sintefmath/Jutul.jl) [@Jutul; @JutulDarcy; @JutulDarcyRules], which for each permeability model, $\mathbf{K}^{(j)}$, solves the immiscible and compressible two-phase flow equations for the CO~2~ and brine saturation. As shown in @fig-true-perm\, an injection well is set up on the left-hand side of the model, which injects supercritical CO~2~ with density 700 $\mathrm{kg}/\mathrm{m}^3$ at a constant rate of 0.005 $\mathrm{m}^3/\mathrm{s}$. To relieve pressure, a production well is included on the right-hand side of the model, which produces brine with density 1000 $\mathrm{kg}/\mathrm{m}^3$ with a constant rate of also 0.005 $\mathrm{m}^3/\mathrm{s}$. This finally results in approximately a $6\%$ storage capacity after 800 days of CO~2~ injection. From these simulations, we collect eight snapshots for the CO~2~ concentration, $\mathbf{c}=[\mathbf{c}_1, \mathbf{c}_2, \cdots,\mathbf{c}_{n_t}]$ with $n_t=8$ the number of snapshots that cover a total time period of 800 days. The last five snapshots of these simulations are included in the top row of @fig-true-perm. Due to buoyancy effects and well control, the CO~2~ plume gradually moves from the left to the right and upwards. 

::: {#fig-fno-train-plume layout="[[15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15]]"}

![](./figs/co2-true-1.png){#fig-co2-true-1}

![](./figs/co2-true-2.png){#fig-co2-true-2}

![](./figs/co2-true-3.png){#fig-co2-true-3}

![](./figs/co2-true-4.png){#fig-co2-true-4}

![](./figs/co2-true-5.png){#fig-co2-true-5}

![](./figs/co2-fno-1.png){#fig-co2-fno-1}

![](./figs/co2-fno-2.png){#fig-co2-fno-2}

![](./figs/co2-fno-3.png){#fig-co2-fno-3}

![](./figs/co2-fno-4.png){#fig-co2-fno-4}

![](./figs/co2-fno-5.png){#fig-co2-fno-5}

![](./figs/co2-diff-1.png){#fig-co2-diff-1}

![](./figs/co2-diff-2.png){#fig-co2-diff-2}

![](./figs/co2-diff-3.png){#fig-co2-diff-3}

![](./figs/co2-diff-4.png){#fig-co2-diff-4}

![](./figs/co2-diff-5.png){#fig-co2-diff-5}

Five CO~2~ saturation snapshots after 400, 500, 600, 700, and 800 days. First row shows the CO~2~ saturation simulated by the PDE. Second row shows the CO~2~ saturation predicted by the trained FNO. Third row shows the $5\times$ difference between the first row and the second row.

:::

Given these simulated CO~2~ concentrations, the optimized weights, $\mathbf{w}^\ast$, for the FNO surrogate are calculated by minimizing @eq-fno-train for $N=1900$ training pairs, $\{\mathbf{K}^{(j)},\, \mathbf{c}^{(j)}\}_{j=1}^N$. Another 100 training pairs are used for validation. After training with 350 epochs, an average of 7% prediction error is achieved for permeability samples from the unseen test set. As observed from @fig-fno-train-plume, the approximation errors of the FNO are mostly concentrated at the leading edge of the CO~2~ plumes. The same permeability models are used to train the NF by minimizing @eq-train-nf for 245 epochs using the open-source software package [InvertibleNetworks.jl](https://github.com/slimgroup/InvertibleNetworks.jl) [@InvertibleNetworks]. We use the HINT network structure [@kruse2021hint] for the NF. Three generative samples are shown in the second row of @fig-trainingsample. From these examples, we can see that the trained NF is capable of generating random permeability models that resemble the ones in the training samples closely, despite minor noisy artifacts.

### Unconstrained/constrained permeability inversion from CO~2~ saturation data

To demonstrate that permeability inversion with surrogates is indeed feasible, we first consider the idealized, impossible in practice, situation where we assume to have access to the time-lapse CO~2~ concentration, $\mathbf{c}=[\mathbf{c}_1, \mathbf{c}_2, \cdots,\mathbf{c}_{n_t}]$, everywhere, and for all $n_t=8$ timesteps. In that case, the measurement operator, $\mathcal{H}$ in @eq-inv, corresponds to the identity matrix. Given CO~2~ concentrations simulated from the "ground-truth" permeability distribution plotted in @fig-true-perm, we invert for the permeability by minimizing the unconstrained formulation (cf. @eq-inv-fno) for the correct, yielded by the PDE, and approximate fluid-flow physics, yielded by the trained FNO. The results of these inversions after 100 iterations of gradient descent with back-tracking linesearch [@stanimirovic2010accelerated] are plotted in @fig-flow-inv-jutul-all and @fig-flow-inv-fno-all. From these plots, we observe that the inversion results using PDE solvers delineates most of the upper boundary of the channel accurately. Because there is a null space in the fluid-flow modeling---i.e., this null space mostly corresponds to regions of the permeability model that are barely touched by the CO~2~ plume (e.g. bottom and right-hand side of the channel) --- artifacts are present in the high-permeability channel itself. As expected, the reconstruction of the permeability is also not perfect at the bottom and at the far right of the model. The inversion result with the FNO surrogate is similar but introduces unrealistic artifacts in the high-permeability channel and also outside the channel. These more severe artifacts can be explained by the behavior of the FNO approximation error plotted as the orange curve in @fig-flow-inv-fno-loss-all. The error value increases rapidly to 13%, and finally saturates at 10%. This behavior of the error is a manifestation of out-of-distribution model iterates that explain the erroneous behavior of the surrogate and its gradient with respect to the permeability.

Inversions yielded by the relaxed constrained formulation with the trained NF (see Algorithm 1), on the other hand, show virtually artifact free inversion results (see @fig-flow-inv-jutul-nf-all and @fig-flow-inv-fno-nf-all) that compare favorably with the "ground-truth" permeability plotted in @fig-perm-init-true. While adding the NF as a constraint obviously adds information, explaining the improved inversion for the accurate physics (@fig-flow-inv-jutul-nf-all), it also renders the approximate surrogates more accurate, as can be observed from the blue curve in @fig-flow-inv-fno-loss-all, where the FNO approximation error is controlled thanks to adding the constraint to the inversion. This behavior underlines the importance of ensuring model iterates to remain within distribution. It also demonstrates the benefits of a solution strategy where we start with a small $\tau$, followed by relaxing the constraint slowly by increasing the size of the constraint set gradually. 

<!-- With the assistance of the learned NF-based constraint, the inversion result shown in @fig-flow-inv-fno-nf-all (by solving @eq-inv-fno-nf) demonstrates a near-perfect recovery of the channel and the upper boundary of the channel, thanks to the generative power of the trained NF. More importantly, the FNO prediction error along these iterations is plotted in @fig-flow-inv-fno-nf-loss-all, which demonstrates that the prediction error directly decreases to about 5% due to the strict constraint at the beginning of the iterations. By gradually relaxing the constraint, our proposed method is able to maintain the FNO prediction error at roughly 5% throughout the process, while also recovering the permeability well enough to drive the objective down. As a reference, we present the inversion result using physics operators combined with the NF-based constraint in @fig-flow-inv-jutul-nf-all. It is evident that the difference between inversion results with and without the FNO surrogate is reduced when the NF-based constraint is employed. This further demonstrates that the NF-based constraint assists the FNO surrogate in more accurately approximating the correct physics operator during the inversion process. -->

::: {#fig-flow-inv-all layout="[[33, 33], [33, 33], [50]]"}

![](./figs/flow-inv-jutul-all.png){#fig-flow-inv-jutul-all}

![](./figs/flow-inv-fno-all.png){#fig-flow-inv-fno-all}

![](./figs/flow-inv-jutul-nf-all.png){#fig-flow-inv-jutul-nf-all}

![](./figs/flow-inv-fno-nf-all.png){#fig-flow-inv-fno-nf-all}

![](./figs/flow-inv-fno-all-error.png){#fig-flow-inv-fno-loss-all}

Permeability inversion from fully observed time-lapse CO~2~ saturations. *(a)* Inversion result with PDE solvers. *(b)* The same but via the approximate FNO surrogate. *(c)* Same as *(a)* but with NF constraint. *(d)* Same as *(b)* but with NF constraint. *(e)* The FNO approximation errors as a function of the number of iterations for the result plotted in *(b)* and *(d)*.
:::

### Unconstrained/constrained permeability inversion from well observations

While the example of the previous section established feasibility of constrained permeability inversion, it relied on having access to the CO~2~ saturation everywhere, which is unrealistic in practice. To address this issue, we first consider permeability inversion from CO~2~ saturations, collected at three equally spaced monitoring well locations, for only the first 6 timesteps over the period of 600 days [@mosser2019deepflow]. In this more realistic setting, the measurement operator, $\mathcal{H}$ in @eq-inv, corresponds to a restriction operator that extracts simulated CO~2~ saturations at each well location in first six snapshots. The objective function reads

$$
\underset{\mathbf{z}}{\operatorname{minimize}} \quad \|\mathbf{d}_{\mathrm{w}}-\mathbf{M}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2\quad\text{subject to}\quad\|\mathbf{z}\|_2\leq \tau,
$$ {#eq-inv-well}

where $\mathbf{d}_{\mathrm{w}}$ represents the well measurements collected at three well locations through the linear restriction operator $\mathbf{M}$. The goal is to invert for the permeability by minimizing the misfit of the well measurements of the CO~2~ saturation without and with constraints on the $\ell_2$-norm ball in the latent space. The results of these numerical experiments are included in the first row of @fig-flow-inv, where the differences with respect to the ground truth permeability shown in @fig-true-perm are plotted in the second row. Because the part of the permeability that is not touched by the CO~2~ plume lives in the null space, we highlight the CO~2~ plume in the difference plots by dark color and focus on analyzing errors within the plume region. As expected, the unconstrained inversions based on PDE solves (@fig-flow-inv-jutul) and surrogate approximations (@fig-flow-inv-fno) are both poorly resolved because of the limited spatial information on the saturation. Contrasting these unconstrained inversions with results for the constrained inversions for the PDE (@fig-flow-inv-jutul-nf) and surrogate (@fig-flow-inv-fno-nf) again shows the importance of adding constraints to the inversion. @fig-flow-inv-fno-loss-well clearly demonstrates that the FNO prediction errors remain relatively constant during constrained inversion while the error continues to grow during the unconstrained iterations eventually exceeding 14\%. Both constrained results improve significantly, even though they converge to different solutions in the end. This is because history matching is typically an ill-posed problem with many distinctive solutions [@canchumuni2019history]. This observation further motivates us to consider the experiment below, where time-lapse seismic data are jointly inverted for the subsurface permeability.

<!-- Next, we aim to estimate the permeability model from CO~2~ saturation snapshots at three equally spaced monitoring well locations at all 8 timesteps in the period of 800 days [@mosser2019deepflow]. In this case, the measurement operator, $\mathcal{H}$, is a linear restriction operator that extracts three columns from each CO~2~ saturation snapshot. The goal is to invert for the permeability model by minimizing the misfit of the well measurements of the CO~2~ saturation snapshots. The inversion result with the correct PDE modeling operator is shown in @fig-flow-inv-jutul. Due to the sparsity of the observational wells, the permeability model is poorly recovered between the wells. We also see artifacts in the permeability recovery that is near to the observational well in the center of the model. The inversion result using the FNO surrogate is presented in @fig-flow-inv-fno, where we also observe some artifacts. In @fig-flow-inv-fno-loss, we can see that the FNO prediction error during the surrogate-assisted inversion increases throughout the iterations because the intermediate iterates are pushed farther from the statistical distribution of the training samples. With the help of the NF-based constraint, we recover a reasonable permeability channel. Although part of the channel is not accurately recovered due to the ill-posedness of the measurement operator, the FNO prediction error remains low across the iterations, as shown in @fig-flow-inv-fno-nf-loss. -->

::: {#fig-flow-inv layout="[[24, 24, 24, 24], [24, 24, 24, 24], [50]]"}

![](./figs/flow-inv-jutul-well.png){#fig-flow-inv-jutul}

![](./figs/flow-inv-fno-well.png){#fig-flow-inv-fno}

![](./figs/flow-inv-jutul-nf-well.png){#fig-flow-inv-jutul-nf}

![](./figs/flow-inv-fno-nf-well.png){#fig-flow-inv-fno-nf}

![](./figs/diff-well-inv-jutul.png){#fig-diff-flow-inv-jutul}

![](./figs/diff-well-inv-fno.png){#fig-diff-flow-inv-fno}

![](./figs/diff-well-inv-jutul-nf.png){#fig-diff-flow-inv-jutul-nf}

![](./figs/diff-well-inv-fno-nf.png){#fig-diff-flow-inv-fno-nf}

![](./figs/flow-inv-fno-well-error.png){#fig-flow-inv-fno-loss-well}

Permeability inversions from CO~2~ saturations sampled at three well locations at 6 early snapshots. The well locations are denoted by the red vertical lines. *(a)* Unconstrained inversion result based on PDE solves. *(b)* Same as *(a)* but now with FNO surrogate approximation. *(c)* Constrained inversion result based on PDE  solves. *(d)* Same as *(c)* but now with FNO surrogate approximation. *(e)*-*(h)* The error of the permeability inversion results in *(a)*-*(d)* compared to the unseen ground truth shown in @fig-true-perm. *(i)* The FNO prediction errors as a function of the number of iterations for *(b)* and *(d)*.
:::

### Multiphysics end-to-end inversion

Next, we consider the alternative setting for seismic monitoring of geological carbon storage, where the dynamics of the CO~2~ plumes are indirectly observed from time-lapse seismic data. In this case, the measurement operator, $\mathcal{H}$, involves the composition of the rock physics modeling operator, $\mathcal{R}$, which converts CO~2~ saturations to decreases in the compressional wavespeeds for rocks within the reservoir [@avseth2010quantitative], and the seismic modeling operator, $\mathcal{F}$, which generates time-lapse seismic data recorded at the receiver locations and based on acoustic wave equation modeling [@sheriff1995exploration]. The multiphysics end-to-end inversion process estimates permeability from time-lapse seismic data via inversion of these nested physics operators for the flow, rock physics, and waves [@li2020coupled]. Following earlier work by @yin2022SEGlci and @louboutin2023learned, the fluid-flow PDE modeling is replaced by the trained FNO (cf. @eq-inv-fno-nf), resulting in the following optimization problem:

$$
\underset{\mathbf{z}}{\operatorname{minimize}} \quad \|\mathbf{d}_{\mathrm{s}} - \mathcal{F}\circ\mathcal{R}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2 \quad\text{subject to}\quad\|\mathbf{z}\|_2\leq \tau,
$$ {#eq-inv-end2end2end}

where $\mathbf{d}_{\mathrm{s}}$ represents the observed time-lapse seismic data. While this end-to-end inversion problem benefits from having remote access to changes in the compressional wavespeed, it may now suffer from null spaces associated with the flow, $\mathcal{S}_{\boldsymbol\theta^\ast}$, and the wave/rock physics, $\mathcal{F}\circ\mathcal{R}$. For instance, the latter suffers from bandwidth limitation of the source function and from limited aperture. Because important components are missing in the observed data, inversion based on the data objective alone in @eq-inv-end2end2end are likely to suffer from  artifacts that can easily drive the intermediate permeability model iterates out-of-distribution.

To demonstrate capabilities of the proposed relaxed inversion procedure with surrogates for the fluid flow, we assume the baseline to be known---i.e, we assume the brine-filled reservoir with $25\%$ porosity to be acoustically homogeneous prior to CO~2~ injection with a compressional wavespeed of $3500 \mathrm{m/s}$. We use the patchy saturation model [@avseth2010quantitative] to convert the time-dependent CO~2~ saturation resulting in $<300\mathrm{m/s}$ decreases in the wavespeed within the CO~2~ plumes. We collect six seismic surveys at the first six snapshots for the CO~2~ saturation from day 100 to day 600, which are the same snapshots as the ones used in the previous experiment. For each time-lapse seismic survey, 16 active-seismic sources are located within a well on the left-hand side of the model. We also position 16 sources on the top of the model. Each active source uses a Ricker wavelet with a central frequency of $50\mathrm{Hz}$. The transmitted and reflected wavefields are collected by 480 receivers on the top and 480 receivers on the right-hand side of the model. The seismic acquisition is shown in @fig-seismic-acquisition, where the plume at the last seismic vintage (at day 600) is plotted in the middle.

::: {#fig-seismic-acquisition}

![](./figs/acquisition.png){width=50%}

Seismic acquisition. The white $\boldsymbol{\times}$ represents the acoustic sources, and the red lines represent the dense receivers. The CO~2~ saturation snapshot at day 600 is plotted in the middle, which is the last snapshot that is monitored seismically.
:::

To avoid numerical dispersion, the velocity model is upsampled by a factor of two in both the horizontal and vertical directions, which results in a $7.5\mathrm{m}$ grid spacing. For the simulations, use is made of the open-source software package [JUDI.jl](https://github.com/slimgroup/JUDI.jl) [@witte2018alf;@JUDI] to generate the time-lapse seismic data at the first six snapshots. The fact that this software is based on [Devito](https://www.devitoproject.org/)'s wave propagators [@louboutin2018dae; @luporini2020architecture] allows us to do this quickly. For realism, we add 10 dB Gaussian noise to the time-lapse seismic data. Given these six time-lapse vintages, our goal is to invert for the permeability in the reservoir by minimizing the time-lapse seismic data misfit through the nested physics operators shown in @eq-inv-end2end2end. 

::: {#fig-end2end-inv layout="[[24, 24, 24, 24], [24, 24, 24, 24], [50]]"}

![](./figs/end2end-inv-jutul.png){#fig-end2end-inv-jutul}

![](./figs/end2end-inv-fno.png){#fig-end2end-inv-fno}

![](./figs/end2end-inv-jutul-nf.png){#fig-end2end-inv-jutul-nf}

![](./figs/end2end-inv-fno-nf.png){#fig-end2end-inv-fno-nf}

![](./figs/diff-end2end-inv-jutul.png){#fig-diff-end2end-flow-inv-jutul}

![](./figs/diff-end2end-inv-fno.png){#fig-diff-end2end-flow-inv-fno}

![](./figs/diff-end2end-inv-jutul-nf.png){#fig-diff-end2end-flow-inv-jutul-nf}

![](./figs/diff-end2end-inv-fno-nf.png){#fig-diff-end2end-flow-inv-fno-nf}

![](./figs/end2end-inv-fno-error.png){#fig-end2end-inv-fno-loss}

Permeability inversions from time-lapse seismic data. *(a)* Inversion result using PDE solvers. *(b)* The same as *(a)* but for the FNO surrogate. *(c)* The same as *(a)* but with the NF-based constraint. *(d)* The same as *(a)* but now for the FNO surrogate with the NF-based constraint. *(e)*-*(h)* The error of the permeability inversion results in *(a)*-*(d)* compared to the unseen ground truth shown in @fig-true-perm. *(i)* The FNO prediction errors as a function of the number of iterations for *(b)* and *(d)*.
:::

Inversion results obtained by solving the PDEs for the fluid flow during the inversion are shown in @fig-end2end-inv-jutul and @fig-end2end-inv-jutul-nf. As before, the inversions benefit majorly from adding the trained NF as a constraint. Remarkably, the end-to-end inversion results shown in @fig-end2end-inv-jutul, @fig-end2end-inv-jutul-nf, and @fig-end2end-inv-fno-nf are close to the results plotted in @fig-flow-inv-jutul-all, @fig-flow-inv-jutul-nf-all, and @fig-flow-inv-fno-nf-all, which was obtained with access to the CO~2~ saturation everywhere. This reaffirms the notion that time-lapse seismic can indeed provide useful spatial information away from the monitoring wells to estimate the reservoir permeability, which aligns with earlier observations by @li2020coupled, @yin2022SEGlci, and @louboutin2023learned. Juxtaposing the results for the FNO surrogate without (@fig-end2end-inv-fno) and with the constraint (@fig-end2end-inv-fno-nf) again underlines the importance of adding constraints especially in situations where the forward (wave) operator has a non-trivial nullspace. The presence of this nullspace has a detrimental affect on the unconstrained result obtained by the FNO. Contrary to solutions yielded by the PDE, trained FNOs offer little to no control on the feasibility of the solution, which explains the strong artifacts in @fig-end2end-inv-fno. As we can see from @fig-end2end-inv-fno-loss, these artifacts are mainly due to the FNO-approximation errors that dominate and grow after a few iterations. Conversely, the errors for the constrained case remain more or less flatlined between $7\%$ and $8\%$. In contrast, using the trained NF as a learned constraint yields better recovery where the errors are minor within the plume region and mostly live on the edges, shown in the second row of @fig-end2end-inv.

<!-- Using correct wave physics and fluid-flow physics modeling operators, the inversion result is shown in @fig-end2end-inv-jutul. The quality of this recovery result is similar to the ideal case where we assume to observe the full CO~2~ saturation snapshots, shown in @fig-flow-inv-jutul. This reaffirms that seismic monitoring is a viable and effective technology for CO~2~ plume monitoring, as long as inverting for the reservoir permeability [@ringrose2020store]. The inversion result using the FNO surrogate for the fluid-flow modeling is shown in @fig-end2end-inv-fno. Because the forward modeling operator is severely ill-posed in this case, there are enormous amount of artifacts in the permeability recovery. The curve for FNO prediction error along iterations is shown in @fig-end2end-inv-fno-loss, which demonstrates that the prediction error keeps increasing and ends up to be over 25%, which is 5 times more than normal. With the NF-based constraint, we recover a very good permeability model shown in @fig-end2end-inv-fno-nf, where the FNO prediction error goes down at the beginning to 5% and remains roughly constant, shown in @fig-end2end-inv-fno-loss. As a reference, the inversion result with full physics modeling operators and the NF-based constraint is shown in @fig-end2end-inv-jutul-nf, which is similar to the result when FNO replaces fluid-flow modeling operator. This again illustrates that the NF-based constraint is conducive for surrogate-assisted inversion, which reduces the prediction error of the surrogate and enables the surrogate to behave more like the correct physics modeling operator. -->

### Jointly inverting time-lapse seismic data and well measurements

Finally, we consider the most preferred scenario for GCS monitoring, where multiple modalities of data are jointly inverted for the reservoir permeability [@huang2022geophysical; @liu2023joint]. In our experiment, we consider to jointly invert time-lapse seismic data and well measurements by minimizing the following objective function:
$$
\underset{\mathbf{z}}{\operatorname{minimize}} \quad \|\mathbf{d}_{\mathrm{s}} - \mathcal{F}\circ\mathcal{R}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2 + \lambda \|\mathbf{d}_{\mathrm{w}}-\mathbf{M}\circ\mathcal{S}_{\boldsymbol\theta^\ast}\circ\mathcal{G}_{\mathbf{w}^\ast}(\mathbf{z})\|_2^2 \quad\text{subject to}\quad\|\mathbf{z}\|_2\leq \tau.
$$ {#eq-inv-combined}

This objective function includes both the time-lapse seismic data misfit from @eq-inv-end2end2end and the time-lapse well measurement misfit from @eq-inv-well with a balancing term $\lambda$. While better choices can be made, we select this $\lambda$ in our numerical experiment to be $10$, so that the magnitudes of the two terms are relatively the same. The inversion results and differences from the unseen ground truth permeability are shown in @fig-combine-inv, where we again observe large artifacts for the recovery when FNO surrogate is inverted without NF constraints. This behavior is confirmed by the plot for the FNO error curve as a function of the number of iterations. This error finally reaches a value over $15\%$.

::: {#fig-combine-inv layout="[[24, 24, 24, 24], [24, 24, 24, 24], [50]]"}

![](./figs/combine-inv-jutul.png){#fig-combine-inv-jutul}

![](./figs/combine-inv-fno.png){#fig-combine-inv-fno}

![](./figs/combine-inv-jutul-nf.png){#fig-combine-inv-jutul-nf}

![](./figs/combine-inv-fno-nf.png){#fig-combine-inv-fno-nf}

![](./figs/diff-combine-inv-jutul.png){#fig-combine-flow-inv-jutul}

![](./figs/diff-combine-inv-fno.png){#fig-diff-combine-flow-inv-fno}

![](./figs/diff-combine-inv-jutul-nf.png){#fig-diff-combine-flow-inv-jutul-nf}

![](./figs/diff-combine-inv-fno-nf.png){#fig-diff-combine-flow-inv-fno-nf}

![](./figs/combine-inv-fno-error.png){#fig-combine-inv-fno-loss}

Joint permeability inversions from both time-lapse seismic data and time-lapse well measurements. *(a)* Inversion result using PDE solvers. *(b)* The same as *(a)* but for the FNO surrogate. *(c)* The same as *(a)* but with the NF-based constraint. *(d)* The same as *(a)* but now for the FNO surrogate with the NF-based constraint. *(e)*-*(h)* The error of the permeability inversions in *(a)*-*(d)*, compared to the unseen ground truth shown in @fig-true-perm. *(i)* The FNO prediction errors as a function of the number of iterations for *(b)* and *(d)*.
:::

We report quantitative measures for the permeability inversions for all optimization methods and different types of observed data in @tbl-snrs-perm for the signal-to-noise ratios (S/Ns) and the structural similarity index measure [SSIM, @wang2004image]. To avoid undue influence of the null space for the permeability, we only calculate the S/N and SSIM values based on the parts of the models that are touched by CO~2~ plume. From these values, following observations can be made. First, the NF-constrained permeability inversion are superior in both S/Ns and SSIMs, which demonstrates the efficacy of the learned constraint. Second, by virtue of this NF constraint, the results yielded by either the PDE solver or by the FNO surrogate produce very similar S/Ns and SSIMs. This behavior reaffirms that the trained FNO behavior is similar to the behavior yielded by PDE solver when its inputs remain in-distribution, which is controlled by the NF constraints.

| Inversion method | Well measurement | Time-lapse seismic | Both |
|---------------------|----------|----------|--------|
| Unconstrained inversion with PDE solvers | 9.34 dB / 0.67 | 10.50 dB / 0.73 | 10.70 dB / 0.73 |
| Unconstrained inversion with FNO surrogates | 9.64 dB / 0.68 | 11.94 dB / 0.72 | 11.98 dB / 0.72 |
| Constrained inversion with PDE solvers | 12.2 dB / 0.77 | 14.18 dB / 0.80 | 15.20 dB / 0.85 |
| Constrained inversion with FNO surrogates | 11.06 dB / 0.74 | 14.16 dB / 0.81 | 14.92 dB / 0.83 |

: S/N (in dB) and SSIM values of permeability recovery. {#tbl-snrs-perm}

### CO~2~ plume estimation and forecast

While end-to-end permeability inversion from time-lapse data provides novel access to this important fluid-flow property, the real interest in monitoring GCS lies in determining where CO~2~ plumes are and will be in the foreseeable future, say of 100 and 200 days ahead. To demonstrate the value of the proposed surrogates and of the use of time-lapse seismic data, as opposed to time-lapse saturation data measured at the wells only, we in @fig-fno-train-plume juxtapose CO~2~ predictions obtained from fluid-flow simulations based on the inverted permeabilities in situations where either well data is available (first row), or where time-lapse seismic data is available (second row), or where both data modalities are available (third row). These results are achieved by first inverting for permeabilities using FNO surrogates and NF constraints, followed by running the fluid-flow simulations for additional time steps given the inverted permeabilities yielded by well-only (@fig-flow-inv-fno-nf), time-lapse data (@fig-end2end-inv-fno-nf), and both (@fig-combine-inv-fno-nf). From these plots, we draw the following two conclusion. First, the predicted CO~2~ plumes estimated from seismic data are significantly more accurate than those obtained by inverting time-lapse saturations measured at the wells only. As expected, there are large errors in the regions away from the wells for the CO~2~ plumes estimated from wells shown in the fourth row of @fig-co2-all. Second, thanks to the NF-constraint, the CO~2~ predictions obtained with the computationally beneficial surrogate approximation remain close to the ground truth CO~2~ plume plotted in the first row of @fig-fno-train-plume, with only minor artifacts at the edges. Third, using both seismic data and well measurements produces CO~2~ plume predictions with the smallest errors, while the uplift of well measurements on top of seismic observations is modest (comparing the second and the third rows of @fig-co2-all). Finally, while the CO~2~ plume estimates for the past (monitored) vintages (i.e. first three columns of the third row of @fig-co2-all) are accurate, the near-future forecasts without time-lapse well or seismic data (i.e. last two columns of the third row of @fig-co2-all) could be less accurate. This is because the right-hand side and the bottom of the permeability model are not touched yet by the CO~2~ plume during the first 600 days. As a result, the error on the permeability recovery on the right-hand side leads to the slightly larger errors on the CO~2~ plume forecast. Overall, these CO~2~ forecasts for the future 100 and 200 days match the general trend of the CO~2~ plume without any observed data despite minor errors. A continuous monitoring system, where multiple modalities of data are being acquired and inverted throughout the GCS project, could allow for updating the reservoir permeability and forecasting the CO~2~ plume consistently.

<!-- {++As a final attempt to conduct quality control, we also plot in figure the CO~2~ plume predictions for the four cases with well observations in @fig-co2-well, corresponding to the CO~2~ plumes simulated by @fig-flow-inv-jutul, @fig-flow-inv-fno, @fig-flow-inv-jutul-nf, and @fig-flow-inv-fno-nf, respectively. Following @yin2022SEGlci and @louboutin2023learned, we also use the inverted permeability in each case to forecase the CO~2~ plume in the near future, i.e., at day 700 and 800. This can be done by running the fluid-flow simulation for more time steps. While these recovery results all fit the CO~2~ plume at the observed well locations, the surrogate-assisted inversion without the learned constraint introduces larger error at the non-observed locations, shown in the second row of @fig-co2-well. On the other hand, the surrogate-assisted inversion with the learned constraint ends up with lower error at the non-observed locations thanks to the fact that the recovered permeability is in-distribution. The signal-to-noise ratios (S/Ns) of each row compared to the unseen ground truth CO~2~ plume plotted in the first row of @fig-fno-train-plume are 7.14 dB, 6.93dB, 8.13 dB, and 7.63 dB, respectively.++} -->

::: {#fig-co2-all layout="[[15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15]]"}

![](./figs/co2-fno-nf-well-400.png)

![](./figs/co2-fno-nf-well-500.png)

![](./figs/co2-fno-nf-well-600.png)

![](./figs/co2-fno-nf-well-700.png)

![](./figs/co2-fno-nf-well-800.png)

![](./figs/co2-fno-nf-end2end-400.png)

![](./figs/co2-fno-nf-end2end-500.png)

![](./figs/co2-fno-nf-end2end-600.png)

![](./figs/co2-fno-nf-end2end-700.png)

![](./figs/co2-fno-nf-end2end-800.png)

![](./figs/co2-fno-nf-combine-400.png)

![](./figs/co2-fno-nf-combine-500.png)

![](./figs/co2-fno-nf-combine-600.png)

![](./figs/co2-fno-nf-combine-700.png)

![](./figs/co2-fno-nf-combine-800.png)

![](./figs/diff-co2-fno-nf-well-400.png)

![](./figs/diff-co2-fno-nf-well-500.png)

![](./figs/diff-co2-fno-nf-well-600.png)

![](./figs/diff-co2-fno-nf-well-700.png)

![](./figs/diff-co2-fno-nf-well-800.png)

![](./figs/diff-co2-fno-nf-end2end-400.png)

![](./figs/diff-co2-fno-nf-end2end-500.png)

![](./figs/diff-co2-fno-nf-end2end-600.png)

![](./figs/diff-co2-fno-nf-end2end-700.png)

![](./figs/diff-co2-fno-nf-end2end-800.png)

![](./figs/diff-co2-fno-nf-combine-400.png)

![](./figs/diff-co2-fno-nf-combine-500.png)

![](./figs/diff-co2-fno-nf-combine-600.png)

![](./figs/diff-co2-fno-nf-combine-700.png)

![](./figs/diff-co2-fno-nf-combine-800.png)

CO~2~ plume estimation and forecast using FNO surrogates and NF constraints to invert different modalities of observed data. The first three columns represent past CO~2~ saturations at day 400, 500, and 600 of the first 600 days of CO~2~ saturation monitored either through the well measurements or time-lapse data. The last two columns include forecasts for the saturations at future days 700 and 800, where no observed data is available. The first row shows the past and future CO~2~ estimates yielded by inverting well measurements only. The second row is the same but now inverting time-lapse seismic data. The third row is the same but now jointly inverting well measurements and time-lapse seismic data. The fourth, fifth, and sixth rows show $5\times$ difference between the ground truth CO~2~ plume (first row of @fig-fno-train-plume) and the first, second, third row, respectively. The S/Ns for the first, the second, and the third rows are 15.26 dB, 20.14 dB, 20.46 dB, respectively.
:::

<!--
::: {#fig-co2-well layout="[[15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15]]"}

![](./figs/co2-true-1.png)

![](./figs/co2-true-2.png)

![](./figs/co2-true-3.png)

![](./figs/co2-true-4.png)

![](./figs/co2-true-5.png)

![](./figs/co2-jutul-nf-well-400.png)

![](./figs/co2-jutul-nf-well-500.png)

![](./figs/co2-jutul-nf-well-600.png)

![](./figs/co2-jutul-nf-well-700.png)

![](./figs/co2-jutul-nf-well-800.png)

![](./figs/co2-fno-nf-well-400.png)

![](./figs/co2-fno-nf-well-500.png)

![](./figs/co2-fno-nf-well-600.png)

![](./figs/co2-fno-nf-well-700.png)

![](./figs/co2-fno-nf-well-800.png)

![](./figs/diff-co2-jutul-nf-well-400.png)

![](./figs/diff-co2-jutul-nf-well-500.png)

![](./figs/diff-co2-jutul-nf-well-600.png)

![](./figs/diff-co2-jutul-nf-well-700.png)

![](./figs/diff-co2-jutul-nf-well-800.png)

![](./figs/diff-co2-fno-nf-well-400.png)

![](./figs/diff-co2-fno-nf-well-500.png)

![](./figs/diff-co2-fno-nf-well-600.png)

![](./figs/diff-co2-fno-nf-well-700.png)

![](./figs/diff-co2-fno-nf-well-800.png)

CO~2~ plume estimation and forecast from well measurements. The first three columns represent past CO~2~ saturations at day 400, 500, and 600 of the first 600 days of fluid-flow simulation monitored through the well measurements. The last two columns include estimates for the saturations at future days 700 and 800, where no observed data is available. The first row shows the ground truth simulated CO~2~ plume for reference. Second row contains past and future estimates yielded by inversion with the PDE. Third row, the same but now with the surrogate approximation. Fourth row shows 5X difference between the first row and the second row. Fifth row shows 5X difference between the first row and the third row.
:::

::: {#fig-co2-end2end layout="[[15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15]]"}

![](./figs/co2-true-1.png)

![](./figs/co2-true-2.png)

![](./figs/co2-true-3.png)

![](./figs/co2-true-4.png)

![](./figs/co2-true-5.png)

![](./figs/co2-jutul-nf-end2end-400.png)

![](./figs/co2-jutul-nf-end2end-500.png)

![](./figs/co2-jutul-nf-end2end-600.png)

![](./figs/co2-jutul-nf-end2end-700.png)

![](./figs/co2-jutul-nf-end2end-800.png)

![](./figs/co2-fno-nf-end2end-400.png)

![](./figs/co2-fno-nf-end2end-500.png)

![](./figs/co2-fno-nf-end2end-600.png)

![](./figs/co2-fno-nf-end2end-700.png)

![](./figs/co2-fno-nf-end2end-800.png)

![](./figs/diff-co2-jutul-nf-end2end-400.png)

![](./figs/diff-co2-jutul-nf-end2end-500.png)

![](./figs/diff-co2-jutul-nf-end2end-600.png)

![](./figs/diff-co2-jutul-nf-end2end-700.png)

![](./figs/diff-co2-jutul-nf-end2end-800.png)

![](./figs/diff-co2-fno-nf-end2end-400.png)

![](./figs/diff-co2-fno-nf-end2end-500.png)

![](./figs/diff-co2-fno-nf-end2end-600.png)

![](./figs/diff-co2-fno-nf-end2end-700.png)

![](./figs/diff-co2-fno-nf-end2end-800.png)

The same as @fig-co2-well but now from observed seismic data collected over the same time instances of the first 600 days.
:::

::: {#fig-co2-combine layout="[[15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15], [15, 15, 15, 15, 15]]"}

![](./figs/co2-true-1.png)

![](./figs/co2-true-2.png)

![](./figs/co2-true-3.png)

![](./figs/co2-true-4.png)

![](./figs/co2-true-5.png)

![](./figs/co2-jutul-nf-combine-400.png)

![](./figs/co2-jutul-nf-combine-500.png)

![](./figs/co2-jutul-nf-combine-600.png)

![](./figs/co2-jutul-nf-combine-700.png)

![](./figs/co2-jutul-nf-combine-800.png)

![](./figs/co2-fno-nf-combine-400.png)

![](./figs/co2-fno-nf-combine-500.png)

![](./figs/co2-fno-nf-combine-600.png)

![](./figs/co2-fno-nf-combine-700.png)

![](./figs/co2-fno-nf-combine-800.png)

![](./figs/diff-co2-jutul-nf-combine-400.png)

![](./figs/diff-co2-jutul-nf-combine-500.png)

![](./figs/diff-co2-jutul-nf-combine-600.png)

![](./figs/diff-co2-jutul-nf-combine-700.png)

![](./figs/diff-co2-jutul-nf-combine-800.png)

![](./figs/diff-co2-fno-nf-combine-400.png)

![](./figs/diff-co2-fno-nf-combine-500.png)

![](./figs/diff-co2-fno-nf-combine-600.png)

![](./figs/diff-co2-fno-nf-combine-700.png)

![](./figs/diff-co2-fno-nf-combine-800.png)

The same as @fig-co2-well but now observed seismic data and well measurements are jointly inverted.
:::

-->
<!--
| | well measurement | time-lapse seismic | both |
|-----------------------------|------------------|------------------|---|
| Unconstrained inversion with PDE solvers | 7.14 | 10.66 | 12.24 |
| Unconstrained inversion with FNO surrogates | 6.93 | 8.57 | 8.79 |
| Constrained inversion with PDE solvers | 8.13 | 10.47 | 10.75 |
| Constrained inversion with FNO surrogates | 7.63 | 10.07  | 10.23 |

: SNR values (in dB) of CO~2~ plume estimation. {#tbl-snrs}
-->

### Analysis of computational gains

FNOs, and deep neural surrogates in general, have the potential to be orders of magnitude faster than conventional PDE solvers [@li2020fourier], and this speed-up is generally problem-dependent. In our numerical experiments, the PDE solver from [Jutul.jl](https://github.com/sintefmath/Jutul.jl) [@Jutul; @JutulDarcy; @JutulDarcyRules] currently only supports CPUs and we find an average runtime for both the forward and gradient on the $64 \times 64$ model to be 10.6 seconds on average on an 8-core Intel(R) Xeon(R) W-2245 CPU. The trained FNO, implemented using modules from [Flux.jl](https://github.com/FluxML/Flux.jl) [@innes2018flux], takes 16.4 milliseconds on average for both the forward and gradient. This means that the trained FNO in our case provides $646\times$ speed up compared to conventional PDE solvers. The training of FNO takes about 4 hours on an NVIDIA T1000 8GB GPU. Given these numbers, we can calculate the break-even point --- i.e., the point where using FNO surrogate becomes cheaper in terms of the overall runtime, by the following formula:

$$
\textrm{breakeven} = \frac{\textrm{generating training set time} + \textrm{training time}}{\textrm{PDE solver runtime}-\textrm{FNO runtime}}
\approx 3364.
$$ {#eq-breakeven}

This means that after 3364 calls to the forward simulator and its gradients, the computational savings gained from using the FNO surrogate evaluations during the inversion process balances out the initial upfront costs. These upfront costs include the generation of the training dataset and the training of the FNO. Therefore, after this break-even point of 3364 calls, the use of the FNO surrogate becomes more cost-effective compared to the conventional PDE solver. Because the trained FNO has the potential to generalize to different kinds of inversion problems, and potentially also different GCS sites, 3364 calls is justifiable in practice. However, we acknowledge that a more detailed analysis on a more realistic 4D scale problem will be necessary to understand the potential computational gains and tradeoffs of the proposed methodology. For details on a high-performance computing parallel implementation of FNOs, we refer to @grady2023model who also conducted a realistic performance on large-scale 4D multiphase fluid-flow problems. Even in cases where the computational advances are perhaps challenging to justify, the use of FNOs has the additional benefit by providing access to the gradient with respect to model parameters (i.e. permeability) through automatic differentiation. This feature is important since it is an enabler for inversion problems that involve complex PDE solvers for which gradients are often not readily available, e.g. @gross2021geosx. By training FNOs on input-output pairs, "gradient-free" gradient-based inversion is made possible in situations where the simulator does not support gradients.

## Discussion and conclusions

Monitoring of geological carbon storage is challenging because of excessive computational needs and demands on data collection by drilling monitor wells or by collecting time-lapse seismic data. To offset the high computational costs of solving multiphase flow equations and to improve permeability inversions from possibly multimodal time-lapse data, we introduce the usage of trained Fourier neural operators (FNOs) that act as surrogates for the fluid-flow simulations. We propose to do this in combination with trained normalizing flows (NFs), which serve as regularizers to keep the inversion and the accuracy of the FNOs in check. Since the computational expense of FNO's online evaluation is negligible compared to numerically expensive partial differential equation solves, FNOs only incur upfront offline training costs. While this obviously presents a major advantage, the approximation accuracy of FNOs is, unfortunately, only guaranteed when its argument, the permeability, is in distribution---i.e., is drawn from the same distribution as the FNO was trained on. This creates a problem because there is, thanks to the non-trivial null space of permeability inversion, no guarantee the model iterates remain in-distribution. Quite the opposite, our numerical examples show that these iterates typically move out-of-distribution during the (early) iterations. This results in large errors in the FNO and in rather poor inversion results for the permeability.

To overcome this out-of-distribution dillema for the model iterates during permeability inversion with FNOs, we propose adding learned constraints, which ensure that model iterates remain in-distribution during the inversion. We accomplish this by training a NF on the same training set for the permeability used to train the FNO. After training, the NF is capable of generating in-distribution samples for the permeability from random realizations of standard Gaussian noise in the latent space. We employ this learned ability by parameterizing the unknown permeability in the latent space, which offers additional control on whether the model iterates remain in-distribution during the inversion. After establishing that out-of-distribution permeability models can be mapped to in-distribution models by restricting the $\ell_2$-norm of their latent representation, we introduce permeability inversion as a constrained optimization problem where the data misfit is minimized subject to a constraint on the $\ell_2$-norm of the latent space. Compared to adding pretrained NFs as priors via additative penalty terms, use of constraints ensures that model iterates remain at all times in-distribution. We show that this holds as long as the size of constraint set does not exceed the size of the $\ell_2$-norm ball of the standard normal distribution. As a result, we arrive at a computationally efficient continuation scheme, known as a homotopy, during which the $\ell_2$-norm constraint is relaxed slowly, so the data misfit objectives can be minimized while the model iterates remain in distribution. 

By means of a series of carefully designed numerical experiments, we were able to establish the advocasy of combining learned surrogates and constraints, yielding solutions to permeability inversion problems that are close to solutions yielded by costly PDE-based methods. The examples also clearly show the advantages of working with gradually relaxed constraints where model iterates remain at all times in distribution with the additional joint benefit of slowly building up the model while bringing down the data misfit, an approach known to mitgate the effects of local minima [@peters2017constraints; @esser2018total; @peters2019projection]. Consequently, the quality of all time-lapse inversions improved significantly without requiring information that goes beyond having access to the training set of permeability models. 

While we applied the proposed method to gradient-based iterative inversion, a similar approach can be used for other types of inversion methods, including inference with Markov chain Monte Carlo methods for uncertainty quantification [@lan2022scaling]. We also envisage extensions of the proposed method to other physics-based inverse problems [@heidenreich2014surrogate; @yang2023rapid] and simulation-based inference problems [@cranmer2020frontier], where numerical simulations often form the computational bottleneck.

Despite the encouraging results from the numerical experiments, the presented approach leaves room for improvements, which we will leave for future work. For instance, the gradient with respect to the model parameters (permeability) derived from the neural surrogate is not guaranteed to be accurate---e.g. close to the gradient yielded by the adjoint-state method. As recent work by @o2022derivate has shown, this potential source of error can be addressed by training neural surrogates on the simulator's gradient with respect to the model parameters, provided it is available. Unfortunately, deriving gradients of complex HPC implementatons of numerical PDE solvers is often extremely challenging, explaining why this information is often not available. Because our method solely relies on gradients of the surrogate, which are readily available through algorithmic differentiation, we only need access to numerical PDE solvers available in legacy codes. While this approach may go at the expense of some accuracy, this feature offers a distinct practical advantage. However, as with many other machine learning approaches, our learned methods may also suffer from time-lapse observations that are out-of-distribution---i.e., produced by a permeability model that is out-of-distribution. While this is a common problem in data-driven methods, recent developments [@siahkoohi2023reliable] may remedy this problem by applying latent space corrections, a solution that is amenable to our approach. On the other hand, expanding the latent space's $\ell_2$ norm ball during inversion would allow NFs to generate any out-of-distribution model parameter. However, in that case the accuracy of the learned surrogate is not guaranteed. For such cases, transitioning from the learned surrogate to the numerical solver during later iterations may be advantageous and merits further study. The choice for the size of the $\ell_2$-norm ball at the beginning and at the end can also be further investigated [@aster2018parameter].

While our paper primarily presents a proof of concept through a relatively small 2D experiment, our inversion strategy is designed to scale to large-scale 3D problems. NFs, with their inherent memory efficiency due to invertibility, are already primed for extension to 3D problems. For the learned surrogates, @grady2023model showcases model-parallel FNOs, demonstrating success in simulating 4D multiphase flow physics of over 2.6 billon variables. By combining these strengths, we are optimistic  scaling this inversion strategy to 3D.

To end on a positive note and forward looking note, we argue that the presented approach makes a strong case for the inversion of multimodal data, consisting of time-lapse well and seismic data. While inversions from time-lapse saturation data collected from wells are feasible and fall within the realm of reservoir engineering, their performance, as expected, degrades away from the well. We argue that adding active-source seismic provides essential fill-in away from the wells. As such, it did not come to our surprise that joint inversion of multimodal data resulted in the best permeability estimates. From our perspective, our successfull combination of these often disjoint data modalities holds future promise when addressing challenges that come with monitoring and control of geological carbon storage and enhanced geothermal systems.

<!-- Overall, we propose an inversion algorithm based on learned surrogates and constraints, which can in principle to be used to solve a variety of PDE-based inverse problems for simulation-based inference. The proposed method takes advantage of the affordable neural surrogate to invert for the model parameters, which could otherwise be extremely computationally expensive. To ensure a reliable inversion process, we explicitly use a learned constraint to force the input to the neural surrogate to be in the statistical distribution of the training samples, so that the surrogate is capable of making accurate predictions. -->

## Availability of data and materials

The scripts to reproduce the experiments are available on the SLIM GitHub page [https://github.com/slimgroup/FNO-NF.jl](https://github.com/slimgroup/FNO-NF.jl).

## Abbreviations
   
FNO: Fourier neural operator    
GCS: geological carbon storage    
NF: normalizing flow     
PDE: partial differential equations   

## Acknowledgements

The authors would like to thank Philipp A. Witte at Microsoft, Ali Siahkoohi at Rice University, and Olav Møyner at SINTEF for constructive discussions. The authors gratefully acknowledge the contribution of OpenAI's ChatGPT for refining sentence structure and enhancing the overall readability of this manuscript.

## Funding

This research was carried out with the support of Georgia Research Alliance and partners of the ML4Seismic Center. This research was also supported in part by the US National Science Foundation grant OAC 2203821 and the Department of Energy grant No. DE-SC0021515.

## Author information

### Authors and Affiliations

Georgia Institute of Technology, 756 West Peachtree Street NW, Atlanta, GA 30308

Ziyi Yin, Rafael Orozco, Mathias Louboutin, and Felix J. Herrmann

### Contributions

ZY concepted the general outline of the project, implemented the algorithm, designed and conducted the numerical experiments, and wrote the manuscript. RO helped concept the general outline of the project, discussed the approach and helped design the numerical experiments. ML helped concept the general outline of the project and discussed the approach. FH concepted the general outline of the project, discussed the approach, helped design the numerical experiments, edited the manuscript, and supervised the project. All authors reviewed the results and approved the final version of the manuscript.

### Corresponding author

Correspondence to [Ziyi Yin](mailto:ziyi.yin@gatech.edu).

## Ethics declarations

### Competing interests

The authors declare that they have no competing interests.
