An adaptive augmented Lagrangian method for training physics and
  equality constrained artificial neural networks
Physics and equality constrained artificial neural networks (PECANN) are
grounded in methods of constrained optimization to properly constrain the
solution of partial differential equations (PDEs) with their boundary and
initial conditions and any high-fidelity data that may be available. To this
end, adoption of the augmented Lagrangian method within the PECANN framework is
paramount for learning the solution of PDEs without manually balancing the
individual loss terms in the objective function used for determining the
parameters of the neural network. Generally speaking, ALM combines the merits
of the penalty and Lagrange multiplier methods while avoiding the ill
conditioning and convergence issues associated singly with these methods . In
the present work, we apply our PECANN framework to solve forward and inverse
problems that have an expanded and diverse set of constraints. We show that ALM
with its conventional formulation to update its penalty parameter and Lagrange
multipliers stalls for such challenging problems. To address this issue, we
propose an adaptive ALM in which each constraint is assigned a unique penalty
parameter that evolve adaptively according to a rule inspired by the adaptive
subgradient method. Additionally, we revise our PECANN formulation for improved
computational efficiency and savings which allows for mini-batch training. We
demonstrate the efficacy of our proposed approach by solving several forward
and PDE-constrained inverse problems with noisy data, including simulation of
incompressible fluid flows with a primitive-variables formulation of the
Navier-Stokes equations up to a Reynolds number of 1000.
Shamsulhaq Basir, Inanc Senocak
An adaptive augmented Lagrangian method for training physics and equality constrained artificial neural networks

###### Abstract

Physics and equality constrained artificial neural networks (PECANN) are grounded in methods of constrained optimization to properly constrain the solution of partial differential equations (PDEs) with their boundary and initial conditions and any high-fidelity data that may be available. To this end, adoption of the augmented Lagrangian method within the PECANN framework is paramount for learning the solution of PDEs without manually balancing the individual loss terms in the objective function used for determining the parameters of the neural network. Generally speaking, ALM combines the merits of the penalty and Lagrange multiplier methods while avoiding the ill conditioning and convergence issues associated singly with these methods. In the present work, we apply our PECANN framework to solve forward and inverse problems that have an expanded and diverse set of constraints. We show that ALM with its conventional formulation to update its penalty parameter and Lagrange multipliers stalls for such challenging problems. To address this issue, we propose an adaptive ALM in which each constraint is assigned a unique penalty parameter that evolve adaptively according to a rule inspired by the adaptive subgradient method. Additionally, we revise our PECANN formulation for improved computational efficiency and savings which allows for mini-batch training. We demonstrate the efficacy of our proposed approach by solving several forward and PDE-constrained inverse problems with noisy data, including simulation of incompressible fluid flows with a primitive-variables formulation of the Navier-Stokes equations up to a Reynolds number of 1000.

Augmented Lagrangian method constrained optimization incompressible flows inverse problems physics-informed neural networks

## 1 Introduction

Partial differential equations (PDEs) are used to describe a wide range of physical phenomena, including sound propagation, heat and mass transfer, fluid flow, and elasticity. The most common numerical methods for solving PDE problems involve domain discretization, such as finite difference, finite volume, finite element, and spectral element methods. However, the accuracy of the solution heavily depends on the quality of the mesh used for domain discretization. Additionally, mesh generation can be tedious and time-consuming, particularly for complex geometries or problems with moving boundaries. While numerical methods are efficient for solving forward problems, they are not well-suited for solving inverse problems, particularly data-driven modeling. In this regard, neural networks can be viewed as a promising alternative meshless approach to solving PDEs.

The use of neural networks for solving PDEs can be traced back the early 1990s [6, 32, 22, 17]. Recently, there has been a resurgence of interest in the use of neural networks to solve PDEs [7, 11, 29, 37], particularly after the introduction of
the term physics-informed neural networks (PINNs) [27]. In PINNs, the governing equations of a physical phenomenon of interest are embedded in an objective function along with any data to learn the solution to those governing equations. Although the performance of PINNs can be influenced by several factors such as the choice of activation function [14], sampling strategy [36], and architecture [33; 2], the formulation of the objective function that is used to determine the neural network parameters is paramount for satisfactory predictions.

In the baseline PINN formulation [27], several loss terms with different physical scales are aggregated into a single composite objective function with tunable weights, which is a rudimentary approach to a multi-objective optimization problem. Since neural networks are trained using gradient descent type optimizers, the model parameters can be influenced by the larger gradient of the loss function, irrespective of the physical scale. This can lead to unstable training, as the optimizer may prioritize one objective over another, sacrificing either the PDE loss or the boundary loss. Therefore, adjusting the interplay between the objective terms requires manual hyperparameter tuning, which can be a time-consuming and challenging task. Additionally, the absence of validation data or prior knowledge of the solution to the PDE for the purpose of tuning can render the baseline PINN approach impractical for the solution of PDEs [3].

Dynamic determination of the weights in the composite objective function of the baseline PINN approach has attracted the attention of several researchers. Wang et al. [33] proposed an empirical method that has several limitations that we discussed in a prior work [2]. van der Meer et al. [31] proposed an empirical method by considering a bi-objective loss function for the solution of linear PDEs. Liu and Wang [18] proposed a dual dimer method to adjust the interplay between the loss terms by searching for a saddle point. Wang et al. [34] studied PINNs using the Neural Tangent Kernel (NTK). NTK provides insights into the dynamics of fully-connected neural networks with infinite width during training via gradient descent. The authors proposed using the eigenvalues of the NTK to adapt the interplay between the objective terms. McClenny and Braga-Neto [21] proposed self-adaptive PINNs (SA-PINN) with a minimax formulation to adjust the interplay between the loss terms. Lagrange multipliers are sent through an empirical mask function such as sigmoid, which makes the dual unconstrained optimization formulation not equivalent to the constrained optimization problem. Because of that, equality constraints are not strictly enforced. The SA-PINN method produced results that are comparable or better than the NTK method [33]. In the present work, we compare our proposed method with both of these approaches for the solution of the wave equation.

By and larger the aforementioned works adopt an unconstrained optimization approach in the first place to formulate the objective function used to determine neural network parameters. In physics and equality constrained artificial neural networks (PECANNs) [2], we pursued a constrained optimization method to formulate the objective function. Specifically, we used the augmented Lagrangian method (ALM) [13; 26] to formally cast the constrained optimization problem into an unconstrained one in which PDE domain loss is constrained by boundary and initial conditions, and with any high-fidelity data that may be available. It is worth noting that ALM combines the merits of the penalty method and the Lagrange multiplier method. It balances feasibility and optimality by updating a penalty parameter to control the influence of constraint violations [24].

In what follows, we show that the conventional ALM with a single penalty parameter used in the PECANN model [2] struggles when applied to problems with multiple constraints of varying characteristics. To overcome this limitation, we propose an adaptive augmented Lagrangian method, which introduces multiple penalty parameters and independently updates them based on the characteristics of each constraint. Additionally, we propose a computationally efficient formulation of the objective function to handle a large number of constraints to enable mini-batch training while maintaining a small memory footprint during training. We solve several forward and inverse PDE problems, including the solution of incompressible fluid flow equations with a primitive-variables formulation up to a Reynolds number of 1000, to demonstrate the efficacy of our PECANN model with an adaptive augmented Lagrangian method. The codes used to produce the results in this paper are publicly available at [https://github.com/HiPerSimLab/PECANN](https://github.com/HiPerSimLab/PECANN)

## 2 Technical formulation

Let us consider a general constrained optimization problem with equality constraints

\[\min_{\theta}\mathcal{J}(\theta),\quad\text{ such that }\quad\mathcal{C}_{i}( \theta)=0,\quad\forall i\in\mathcal{E}, \tag{1}\]

where the objective function \(\mathcal{J}\) and the constraint functions \(\mathcal{C}_{i}\) are all smooth, real valued functions on a subset of \(R^{n}\) and \(\mathcal{E}\) is a finite set of equality constraints. We can write an equivalent _minimax_ unconstrained dual problem using the augmented Lagrangian method as follows

\[\min_{\theta}\max_{\lambda}\mathcal{L}(\theta;\lambda,\mu)=\mathcal{J}(\theta )+\sum_{i\in\mathcal{E}}\lambda_{i}\mathcal{C}_{i}+\frac{\mu}{2}\sum_{i\in \mathcal{E}}\mathcal{C}_{i}^{2}(\theta). \tag{2}\]
We can swap the order of the minimum and the maximum by using the following _minimax_ inequality concept or weak duality

\[\max_{\lambda}\min_{\theta}L(\theta;\lambda)\leq\min_{\theta}\max_{ \lambda}L(\theta;\lambda). \tag{3}\]

The minimization can be performed for a sequence of multipliers generated by

\[\lambda_{i}\leftarrow\lambda_{i}+\mu\:\mathcal{C}_{i}(\theta), \forall i\in\mathcal{E}. \tag{4a}\]

We should note that \(\mu\) can be viewed as a global learning rate for the Lagrange multipliers. ALM combines the merits of the penalty method and the Lagrange multiplier method by updating the penalty parameter in such a way that it balances the trade-off between feasibility and optimality, and ensures convergence. Traditionally, in ALM, textbook descriptions typically rely on a single penalty parameter to address all constraints as in (4a), and updating the penalty parameter is often done through empirical strategies. However, as we show in a later section, these update strategies become ineffective when there are multiple constraints with different characteristics. Next, we consider two existing strategies to update the penalty parameter in ALM.

First, in Algorithm 1, we monotonically increase the penalty parameter \(\mu\) at each training iteration until a maximum safeguarding value (i.e., \(\mu_{\max}\)) is reached. Safeguarding the penalty parameter is a common strategy and prevents it from reaching excessively large values that may lead to numerical instability or overflow. To establish clear distinction from other strategies and facilitate comparison, we will refer to Algorithm 1 as ALM with monotonically increasing penalty update (MPU).

```
1Input:\(\theta^{0},\mu_{\max},\beta\)
2\(\lambda_{i}^{0}=1\:\forall i\in\mathcal{E}\) /* Initializing Lagrange multipliers */
3\(\mu^{0}=1\) /* Initializing the penalty parameter */
4for\(t=1\)to...do
5\(\theta^{t}\leftarrow\underset{\theta}{\operatorname{argargmin}}\:\mathcal{L}( \theta^{t-1};\lambda^{t-1},\mu^{t-1})\) /* primal update */
6\(\lambda_{i}^{t}\leftarrow\lambda_{i}^{t-1}+\mu^{t-1}\mathcal{C}_{i}(\theta^{t})\) /* update Lagrange multipliers */
7\(\mu^{t}=\min(\beta\:\mu^{t-1},\mu_{\max})\) /* penalty parameter update */
8
9
10 end for Output:\(\theta^{t}\)
11Defaults:\(\beta=2,\:\mu_{\max}=1\times 10^{4}\)
```

**Algorithm 1**ALM with monotonically increasing penalty update (MPU) [20]

In the context of training the neural networks, inputs to Algorithm 1 are parameters for the neural network model \(\theta^{0}\), a maximum safeguarding penalty parameter \(\mu_{\max}\), and a multiplicative factor \(\beta\) for increasing the penalty parameter \(\mu\) at each iteration. We should note that in Algorithm 1, both the Lagrange multipliers and the penalty parameter are updated at each iteration [20]. A similar approach without the maximum safeguarding penalty parameter has been used in [19]. To prevent divergence, the maximum penalty parameter was set to \(10^{4}\). However, updating the penalty parameter at each epoch may aggressively increase the penalty parameter that could lead to divergence as we will demonstrate in a later section. Additionally, finding a suitable maximum penalty parameter can be challenging.

Another strategy to update the penalty parameter \(\mu\) is to update it only when the constraints have not decreased sufficiently at the current iteration [4, 35]. In Algorithm 2, we present an augmented Lagrangian method in which the penalty parameter is updated conditionally [35]. A similar strategy has been adopted in the work of Dener et al. [5] to train an encoder-decoder neural network for approximating the Fokker-Planck-Landau collision operator. We will refer to Algorithm 2 as ALM with conditional penalty update (MPU). Our inputs to Algorithm 2 are \(\eta\) which is a placeholder for the previous value of constraints, \(\mu_{\max}\) is a safeguarding penalty parameter, \(\beta\) is a multiplicative weight for increasing penalty parameter. Similar to the previous algorithm, the maximum penalty parameter is set to \(10^{4}\) to prevent numerical overflow and divergence. It should be noted that finding a suitable maximum penalty parameter can be challenging.

In our recent investigations, we have discovered that the rate of update for Lagrange multipliers can be a critical factor when dealing with problems that have different types of constraints (such as flux, data, or PDE constraints). Because the penalty parameter behaves like a learning rate for the Lagrange multiplier, employing the same update rate or penalty parameter for all constraints can lead to issues of instability during training. In some cases, the parameter may be too large or too small for certain constraints, which can adversely affect the optimization process. As such, we propose a
new approach to address this issue. Our method involves assigning a unique penalty parameter for each constraint type with its own update rate. By tailoring the penalty parameter to each specific constraint type, we can effectively manage the optimization process and ensure greater stability during training. We will show in the results section that we are able to learn the solution of challenging PDEs when other strategies fail or perform poorly.

### Adaptive augmented Lagrangian method

In this section, we propose an augmented Lagrangian method with a novel adaptive update strategy for multiple penalty parameters to handle diverse constraints in the optimization problem. We note that it is common to use a single penalty parameter that is monotonically increased in most implementations of the ALM. However, that strategy becomes insufficient to tackle challenging problems. As we discuss above, we need a unique, adaptive penalty parameter or learning rate for each Lagrange multiplier associated with a constraint. This tailored approach ensures that the penalty parameter conforms to the characteristics of individual Lagrange multipliers, enabling effective handling of diverse constraints. We formulate our unconstrained optimization problem for problem (1) as follows:

\[\max_{\lambda}\min_{\theta}\mathcal{L}(\theta,\lambda;\mu)=\mathcal{J}(\theta )+\sum_{i\in\mathcal{E}}\lambda_{i}\mathcal{C}_{i}(\theta)+\frac{1}{2}\sum_{i \in\mathcal{E}}\mu_{i}\mathcal{C}_{i}^{2}(\theta), \tag{5}\]

where \(\lambda_{i}\) is a vector of Lagrange multipliers and \(\mu_{i}\) is a vector of penalty parameters. The minimization can be performed using a variant of gradient descent type optimizer for a sequence of Lagrange multipliers generated by

\[\lambda_{i}\leftarrow\lambda_{i}+\mu_{i}\mathcal{C}_{i}(\theta), \forall i\in\mathcal{E}, \tag{6}\]

during training. Upon examining the dual update of the augmented Lagrangian method as shown in Eq. (6), we observe that it involves a gradient ascent step with learning rates denoted by \(\mu_{i}\) for each Lagrange multiplier \(\lambda_{i}\). Hence, an suitable approach is to adopt the strategy in RMSprop algorithm by G. Hinton [9] in finding an independent effective learning rate or penalty parameter for each Lagrange multiplier. The main reason behind our choice is that the update strategy in RMSprop is consistent with the dual update of ALM as in Eq. (6). Hence, we divide our global learning rates by the weighted moving average of the squared gradients of our Lagrange multipliers as follows

\[\bar{v}_{i} \leftarrow\alpha\bar{v}_{i}+(1-\alpha)\mathcal{C}_{i}(\theta)^{2}, \forall i\in\mathcal{E}, \tag{7}\] \[\mu_{i} \leftarrow\frac{\gamma}{\sqrt{\bar{v}_{i}}+\epsilon}, \forall i\in\mathcal{E},\] (8) \[\lambda_{i} \leftarrow\lambda_{i}+\mu_{i}\mathcal{C}_{i}(\theta), \forall i\in\mathcal{E}, \tag{9}\]

where \(\bar{v}_{i}\) are the weighted moving average of the squared gradient of our Lagrange multipliers, \(\gamma\) is a scheduled global learning rate, \(\epsilon\) is a term added to the denominator to avoid division by zero for numerical stability, and \(\alpha\) is a smoothing constant. In Algorithm 3, presents our training procedure. The input to the algorithm is an initialized set of parameters
for the network \(\theta^{0}\), a global learning rate \(\gamma\), a smoothing constant \(\alpha\), collocation points, noisy measurement data and high fidelity data to calculate our augmented Lagrangian. In Algorithm 3, the Lagrange multiplier vector is initialized to \(1.0\) with their respective averaged squared-gradients initialized to zero. In summary, our main approach involves reducing the penalty parameters for Lagrange multiplier updates with large and rapidly changing gradients, while increasing the penalty parameters for Lagrange multiplier updates with small and slowly changing gradients. Next, we delve into the subtle differences in learning the solution of forward and inverse problems and illustrate how they can be formulated as constrained optimization problems, which can then be cast as an unconstrained optimization problem using the augmented Lagrangian method.

```
1Defaults:\(\gamma=1\times 10^{-2},\;\alpha=0.99,\;\epsilon=1\times 10^{-8}\)
2Input:\(\theta^{0}\)
3\(\lambda_{i}^{0}=1\quad\forall i\in\mathcal{E}\) /* Initializing Lagrange multipliers */
4\(\mu_{i}^{0}=1\quad\forall i\in\mathcal{E}\) /* Initializing penalty parameters */
5\(\bar{v}_{i}^{0}=0\quad\forall i\in\mathcal{E}\) /* initializing averaged square-gradients */
6for\(t=1\)to...do
7\(\theta^{t}\leftarrow\underset{\theta}{\operatorname{argmin}}\;\mathcal{L}( \theta^{t-1};\lambda^{t-1},\mu^{t-1})\) /* primal update */
8\(\bar{v}_{i}^{t}\leftarrow\alpha\;\bar{v}_{i}^{t-1}+(1-\alpha)\;\mathcal{C}_{i} (\theta^{t})^{2},\quad\forall i\in\mathcal{E}\) /* square-gradient update */
9\(\mu_{i}^{t}\leftarrow\frac{1}{\sqrt{\bar{v}_{i}^{t}+\epsilon}},\quad\forall i \in\mathcal{E}\) /* penalty update */
10\(\lambda_{i}^{t}\leftarrow\lambda_{i}^{t-1}+\mu_{i}^{t}\;\mathcal{C}_{i}( \theta^{t}),\quad\forall i\in\mathcal{E}\) /* dual update */
11 end for Output:\(\theta^{t}\)
```

**Algorithm 3**Augmented Lagrangian method with adaptive penalty updates (APU)

### Constrained optimization formulation for solving forward and inverse problems

In this section, we explain our proposed formulation for solving a generic PDE problem with boundary and initial conditions. for demonstration purposes. We will also compare our new formulation to our previous approach and highlight the advantages and improvements of our new approach. We first formulate a constraint optimization problem by minimizing the loss on the governing PDE and constraining the individual points on the boundary and initial conditions assuming these conditions are noise-free and can be defined as equality constraints. Consider a scalar function \(u(\mathbf{x},t):\mathbb{R}^{d+1}\rightarrow\mathbb{R}\) on the domain \(\Omega\subset\mathbb{R}^{d}\) with its boundary \(\partial\Omega\) satisfying the following partial differential equation

\[\mathcal{F}(\mathbf{x},t;\frac{\partial u}{\partial t},\frac{\partial^ {2}u}{\partial t^{2}},\cdots,\frac{\partial u}{\partial\mathbf{x}},\frac{\partial^ {2}u}{\partial\mathbf{x}^{2}},\cdots,\mathbf{\nu})=0,\quad\forall(\mathbf{x},t)\in\mathcal{ U}, \tag{10}\] \[\mathcal{B}(\mathbf{x},t,g;u,\frac{\partial u}{\partial\mathbf{x}},\cdots) =0,\quad\forall(\mathbf{x},t)\in\partial\mathcal{U},\] (11) \[\mathcal{I}(\mathbf{x},t,h;u,\frac{\partial u}{\partial t},\cdots) =0,\quad\forall(\mathbf{x},t)\in\Gamma, \tag{12}\]

where \(\mathcal{F}\) is the residual form of the PDE containing differential operators, \(\mathbf{\nu}\) is a vector PDE parameters, \(\mathcal{B}\) is the residual form of the boundary condition containing a source function \(g(\mathbf{x},t)\) and \(\mathcal{I}\) is the residual form of the initial condition containing a source function \(h(\mathbf{x},t)\). \(\mathcal{U}=\{(\mathbf{x},t)\;|\;\mathbf{x}\in\Omega,t=[0,T]\}\), \(\partial\mathcal{U}=\{(\mathbf{x},t)\;|\;\mathbf{x}\in\partial\Omega,t=[0,T]\}\) and \(\Gamma=\{(\mathbf{x},t)\;|\;\mathbf{x}\in\partial\Omega,t=0\}\).

#### 2.2.1 Formulation for problems with point-wise constraints

In this section, we present a constrained optimization formulation, as represented in equation (1), within the context of solving partial differential equations (PDEs). Considering Eq.(10) with its boundary condition (11) and initial condition (12), we previously formulated the following constrained optimization problem [2]:

\[\min_{\theta}\sum_{i=1}^{N_{\mathcal{F}}}\|\mathcal{F}(\mathbf{x}^{(i)},t^{(i)}; \mathbf{\nu},\theta)\|_{2}^{2}, \tag{13}\]
subject to

\[\phi(\mathcal{B}(\mathbf{x}^{(i)},t^{(i)},g^{(i)};\theta)) =0,\;\forall(\mathbf{x}^{(i)},t^{(i)},g^{(i)})\in\partial\mathcal{U},\;i =1,\cdots,N_{\mathcal{B}} \tag{14}\] \[\phi(\mathcal{I}(\mathbf{x}^{(i)},t^{(i)},h^{(i)};\theta)) =0,\;\forall(\mathbf{x}^{(i)},t^{(i)},h^{(i)})\in\Gamma,\;i=1,\cdots,N_ {\mathcal{I}}, \tag{15}\]

where \(N_{\mathcal{F}}\), \(N_{\mathcal{B}}\), \(N_{\mathcal{I}}\) are the number of data points in \(\mathcal{U}\), \(\partial\mathcal{U}\) and \(\Gamma\) respectively. \(\phi\) is a distance function. We should note that number of Lagrange multipliers scale with the number of constraints \(N_{\mathcal{B}}\), \(N_{\mathcal{I}}\). When dealing with a substantial number of constraints (i.e.,\(N_{\mathcal{B}}\),\(N_{\mathcal{I}}\)), the efficiency of the optimization can be compromised due to the size of the computational graph and memory requirements. To mitigate that, we propose to minimize the expected loss on the PDE while incorporating an expected equality constraint on the boundary and initial conditions. Furthermore, we justify our proposed approach by examining the distribution of the Lagrange multipliers for a large number of constraints.

#### 2.2.2 Computationally efficient formulation based on expectation of constraints

In this section, we introduce an efficient constrained optimization formulation, represented in the form of equation (1), for solving partial differential equations (PDEs). In our original PECANN approach [2], we learned the solution of PDEs by constraining the optimization problem point-wise using the formulation described in the previous section. However, this strategy becomes computationally expensive for challenging problems, which benefit from increased number of collocation points. Constraining the learning problem pointwise is also not suitable for taking advantage of deep learning techniques designed to accelerate the training process. As we will demonstrate later in the present work, for large number of point-wise constraints, Lagrange multipliers assume a probabilistic distribution with a clear expected value. Based on this observation, we revise our formulation for point-wise constraints such that instead of constraining the loss on individual high fidelity data, we constrain the expected loss on a batch of our high fidelity data. Consequently, Lagrange multipliers for this particular approach will represent expected values, which has a direct impact on computational efficiency because Lagrange multipliers are not attached to any individual data point, we can do mini-batch training if our data does not fit in the memory for full batch training. Considering Eq.(10) with its boundary condition (11) and initial condition (12), we write the following constrained optimization problem:

\[\min_{\theta}\;\mathcal{J}(\theta)=\frac{1}{N_{\mathcal{F}}}\sum_{i=1}^{N_{ \mathcal{F}}}\|\mathcal{F}(\mathbf{x}^{(i)},t^{(i)};\mathbf{\nu},\theta)\|_{2}^{2}, \tag{16}\]

subject to

\[\frac{1}{N_{\mathcal{B}}}\sum_{i=1}^{N_{\mathcal{B}}}\phi(\mathcal{ B}(\mathbf{x}^{(i)},t^{(i)},g^{(i)};\theta)) :=0,\;\forall(\mathbf{x}^{(i)},t^{(i)},g^{(i)})\in\partial\mathcal{U}, \;i=1,\cdots,N_{\mathcal{B}}, \tag{17}\] \[\frac{1}{N_{\mathcal{I}}}\sum_{i=1}^{N_{\mathcal{I}}}\phi( \mathcal{I}(\mathbf{x}^{(i)},t^{(i)},h^{(i)};\theta)) :=0,\;\forall(\mathbf{x}^{(i)},t^{(i)},h^{(i)})\in\Gamma,\;i=1, \cdots,N_{\mathcal{I}}, \tag{18}\]

where \(N_{\mathcal{F}}\), \(N_{\mathcal{B}}\), \(N_{\mathcal{I}}\) are the number of data points in \(\mathcal{U}\), \(\partial\mathcal{U}\) and \(\Gamma\) respectively. \(\phi\) is a convex distance function. In all of the experiments \(\phi\) is a quadratic function unless specified otherwise. Next, we are present a typical PDE-constrained inverse problem.

#### 2.2.3 Formulation for PDE-constrained inverse problems

In our original formulation of the PECANN framework [2], we minimized the loss on the governing PDE and noisy data while constraining the loss with high fidelity data [2]. However, disparity in physical scales between the partial differential equations (PDEs) and noisy measurement data can exist, which couldcomplicate the inference problem. To address this challenge of solving inverse PDE problems using noisy and high fidelity data, we propose a PDE-constrained formulation for inverse problems by minimizing the loss on noisy data while considering the governing PDE and any available high fidelity data as constraints.

Les us consider a generic inverse PDE problem to demonstrate the new formulation. We assume that the initial condition is known exactly and can be treated as high fidelity data. Additionally, we assume that the underlying physical problem is governed by a typical PDE operator as shown in Eq. (10). Given a set of noisy measurement data \(\{(\mathbf{x}^{(i)},t^{(i)}),\hat{u}^{(i)}\}_{i=1}^{N_{\mathcal{M}}}\), we can minimize the following objective

\[\min_{\theta,\mathbf{\nu}}\;\frac{1}{N_{\mathcal{M}}}\sum_{i=1}^{N_{\mathcal{M}}} \|u(\mathbf{x}^{(i)},t^{(i)};\theta)-\hat{u}(\mathbf{x}^{(i)},t^{(i)})\|_{2}^{2}, \tag{19}\]
subject to

\[\frac{1}{N_{\mathcal{F}}}\sum_{i=1}^{N_{\mathcal{F}}}\phi(\mathcal{F} (\mathbf{x}^{(i)},t^{(i)};\mathbf{\nu},\theta)):=0,\;\forall(\mathbf{x}^{(i)},t^{(i)})\in \mathcal{U},\;i=1,\cdots,N_{\mathcal{F}}, \tag{20}\] \[\frac{1}{N_{\mathcal{I}}},\sum_{i=1}^{N_{\mathcal{I}}}\phi( \mathcal{I}(\mathbf{x}^{(i)},t^{(i)},h^{(i)};\theta)):=0,\;\forall(\mathbf{x}^{(i)},t ^{(i)},h^{(i)})\in\Gamma,\;i=1,\cdots,N_{\mathcal{I}}, \tag{21}\]

where \(N_{\mathcal{F}}\), \(N_{\mathcal{I}}\) are the number of data points in \(\mathcal{U}\) and \(\Gamma\) respectively. \(\phi\) is a convex distance function. In all of the experiments \(\phi\) is a quadratic function unless specified otherwise. It is worth noting that if additional high-fidelity data is available, we can incorporate it as an additional constraint into the above formulation along with other constraints.

## 3 Applications

In this section, we apply our proposed augmented Lagrangian method with an adaptive update strategy to learn the solution of forward and inverse problems. We provide additional examples in the Appendix as well. Given an \(n\)-dimensional vector of predictions \(\hat{\mathbf{u}}\in\mathbf{R}^{n}\) and an \(n\)-dimensional vector of exact values \(\mathbf{u}\in\mathbf{R}^{n}\), we define the relative Euclidean or \(l^{2}\)-norm, infinity norm \(l^{\infty}\)-norm, and the root-mean square (RMS), respectively, to assess the accuracy of our predictions

\[\mathcal{E}_{r}(\hat{u},u)=\frac{\|\hat{\mathbf{u}}-\mathbf{u}\|_{2}}{\|\mathbf{u}\|_{2}},\quad\mathcal{E}_{\infty}(\hat{u},u)=\|\hat{\mathbf{u}}-\mathbf{u}\|_{\infty},\quad \text{RMS}=\frac{1}{n}\sqrt{\sum_{i=1}^{n}(\hat{\mathbf{u}}^{(i)}-\mathbf{u}^{(i)})^{ 2}} \tag{22}\]

where \(\|\cdot\|_{2}\) denotes the Euclidean norm, and \(\|\cdot\|_{\infty}\) denotes the maximum norm. All the codes used in the following examples are available as open source at [https://github.com/HiPerSimLab/PECANN](https://github.com/HiPerSimLab/PECANN).

### Forward problem: unsteady heat transfer in a composite medium

In this section, we study a heat transfer problem in a composite material where temperature and heat fluxes are matched across the interface [1]. This canonical problem has initial condition, boundary condition and flux constraint which makes it ideal for highlighting the implementation intricacies of our method as well as demonstrating the significant improvement we achieve using our proposed formulation. Consider a time-dependent heat equation in a composite medium,

\[\frac{\partial u(x,t)}{\partial t}=\frac{\partial}{\partial x}[\kappa(x,t) \frac{\partial u(x,t)}{\partial x}]+s(x,t),\quad(x,t)\in\Omega\times[0,\tau] \tag{23}\]

along with Dirichlet boundary condition

\[u(x,t)=g(x,t),\quad(x,t)\in\partial\Omega\times(0,\tau], \tag{24}\]

and initial condition

\[u(x,0)=h(x),\quad x\in\Omega, \tag{25}\]

where \(u\) is the temperature, \(s\) is a heat source function, \(\kappa\) is thermal conductivity, \(g\) and \(h\) are source functions respectively. The composite medium consists of two non-overlapping sub-domains where \(\Omega=\Omega_{1}\cup\Omega_{2}\). We consider the thermal conductivity of the medium to vary as follows

\[\kappa(x,t)=\begin{cases}1,&(x,t)\in\Omega_{1}\times[0,2],\\ 3\pi,&(x,t)\in\Omega_{2}\times[0,2],\end{cases} \tag{26}\]

where \(\Omega_{1}=\{x|-1\leq x<0\}\) and \(\Omega_{2}=\{x|0<x\leq 1\}\). To accurately evaluate our model, we consider an exact solution of the form

\[u(x,t)=\begin{cases}\sin(3\pi x)t,&x\in\Omega_{1}\times[0,2],\\ tx,&x\in\Omega_{2}\times[0,2].\end{cases} \tag{27}\]

The corresponding source functions \(s(x,t)\), \(g(x,t)\) and \(h(x,t)\) can be calculated exactly using (27).
First, let us introduce an auxiliary flux parameter \(\sigma(x,t)=\kappa(x,t)\frac{\partial u}{\partial x}\) to obtain a system of first-order partial differential equation that reads

\[\mathcal{F}(x,t) :=\frac{\partial u(x,t)}{\partial t}-\frac{\partial\sigma(x,t)}{ \partial x}+q(x,t),\in\Omega\times[0,\tau], \tag{28a}\] \[\mathcal{Q}(x,t) :=\sigma(x,t)-\kappa(x,t)\frac{\partial u(x,t)}{\partial x},\in \Omega\times[0,\tau],\] (28b) \[\mathcal{B}(x,t) :=u(x,t)-g(x,t),\in\partial\Omega\times(0,\tau],\] (28c) \[\mathcal{I}(x,t) :=u(x,0)-h(x),\in\Omega,t=0, \tag{28d}\]

where \(\mathcal{F}\) is the residual form of our differential equation (used as our objective function), \(\mathcal{Q}\) is the residual form of our flux operator (equality constraint), \(\mathcal{B}\) is the residual form our boundary condition operator (equality constraint), and \(\mathcal{I}\) is the residual form of our initial condition operator (equality constraint).

We use a fully connected neural network architecture parameterized by \(\theta\) with two inputs for \((x,t)\) and two outputs for \(u\) and \(\sigma\) consisting of three hidden layers with 30 neurons per layer and tangent hyperbolic activation functions. We use L-BFGS optimizer [23] available in the Pytorch package [25] with _strong Wolfe_ line search function with its maximum iteration number set to five. For the purpose of this problem, we generate \(N_{\mathcal{F}}=N_{\mathcal{Q}}=10^{4}\) collocation points, \(N_{\mathcal{B}}=2\times 5000\) and \(N_{\mathcal{I}}=5000\) for approximating the boundary and initial conditions only once before training.

First, we solve this problem using the formulation discussed in section 2.2.1 using point-wise constraints. We present the results at \(t=1\) obtained using the point-wise constraint formulation in Fig.1 and compare the impact of three different strategies to update the penalty parameters in ALM. Fig.1(a) and (b) presents the temperature and heat flux distribution in the composite medium, respectively. We observe that adaptive penalty update strategy produces results that are in excellent agreement with the exact solution. However, results using monotonic and conditional penalty updates produce a solution that is markedly different than the exact solution.

With Fig.1 we show the importance of adaptively updating the penalty parameter when dealing with diverse constraint types (e.g. boundary condition, initial condition, and flux conservation) that exhibit different physical characteristics. However, we apply the constraints point-wise. As we mentioned previously, the computational cost of this particular implementation scales linearly with the number of points assigned to each constraint, especially when constraining the flux because it is applied to every collocation point in the domain. Furthermore, with the point-wise constraint formulation, we cannot train our model in mini-batch training fashion since each Lagrange multiplier is attached to every constraint data point.

In Fig. 2, we present the distribution of Lagrange multipliers resulting from the point-wise constraint formulation presented in section 2.2.1. Our goal is to find out if it is necessary to constrain the solution point-wise. We observe that the distribution of Lagrange multipliers assumes probabilistic distributions regardless of the penalty update strategy. This observation indicates that when there are a large number of point-wise constraints, we can constrain an expected value of the loss on each type of constraints using a single Lagrange multiplier as seen in Figs. 2(a)-(c). We should note that expected loss is also commonly used as an objective function in machine learning applications. Hence, we employ

Figure 1: Effects of augmented Lagrangian penalty update strategies on the heat transfer in a composite medium problem using the point-wise constraint formulation (i.e. section 2.2.1). (a) predicted temperature at t = 1, (b) predicted heat flux at t = 1
our proposed formulation to train our models using ALM with different penalty updating approaches to showcase the efficiency of our adaptive ALM and our proposed formulation.

Figure 3 presents the results of our PECANN model with the adaptive ALM and with the expected loss formulation given in section 2.2.2 applied to the heat transfer in a composite medium problem. Specifically, Fig.3(a)-(b) presents the space and time solution of the temperature and the absolute value of the point-wise error in temperature produced by our model, respectively. Space and time solution of the heat flux and its absolute value error are shown in Fig.3(c)-(d), respectively. Collectively, Fig. 3 demonstrates our predictions closely match the exact solutions. Our PECANN model using the monotonic (MPU) and conditional (CPU) penalty update strategies do not converge, and therefore we do not report temperature and heat flux predictions using those two strategies. However, we provide the error norms resulting from all three penalty update strategies for comparison purposes in Fig. 4 from which we observe that ALM with adaptive penalty update strategy outperforms the other two strategies. From Fig. 4, we also observe that adaptive penalty update (APU) strategy can provide acceptable error levels with even smaller number of epochs.

To further gain insight into how ALM with adaptive penalty updates (APU) works, we analyze the evolution of the penalty parameters during training. Fig.5(a) shows the evolution of penalty parameters with each epoch for enforcing boundary conditions \(\mu_{BC}\), initial conditions \(\mu_{IC}\), and flux constraints \(\mu_{Flux}\) using the adaptive penalty update (APU). The penalty parameters for each constraint type grow at different rates over the course of training, justifying the need for adaptive and independent penalty parameters for each constraint type. Fig.5(b) illustrates the evolution of the penalty parameter for enforcing boundary conditions, initial conditions, and flux constraints using conditional penalty update (CPU) and monotonically increasing penalty update (MPU) strategies. In both cases, the penalty parameter quickly reaches the maximum penalty value of \(\mu_{\max}=10^{4}\), causing the problem to become ill-conditioned and leading to divergence during optimization. This observation again highlights the importance of using an adaptive penalty update strategy to avoid ill-conditioning and ensure convergence. Additionally, in Fig.5(c), we present the evolution of the loss on our equality constraints during training, when using the adaptive penalty update strategy. We observe that all

Figure 2: Distribution of Lagrange multipliers arising from using the formulation with point-wise constraints (i.e. section 2.2.1). (a) ALM with monotonically increasing penalty update (Algorithm 1). (b) ALM with conditional penalty update (Algorithm 2). (c) ALM with adaptive penalty update (Algorithm 3, proposed).
our constraints decay at different rates, which provides further justification for the varying rates at which the penalty parameters increase.

Our investigation and findings from the heat transfer in a composite medium problem provide a strong support for our proposed augmented Lagrangian method with an adaptive penalty update strategy (i.e. Algorithm 3) and for the computationally efficient formulation of the objective function using expectation of constraints (i.e. formulation given in section 2.2.2). In the rest of the examples, we will adopt this combination to further showcase its efficacy and versatility.

### Forward problem: wave equation

Baseline PINN model [28] with a fully connected neural network architecture is known to struggle for certain type of PDEs [21, 34, 3]. Wang et al. [34] and McClenny and Braga-Neto [21] single out the one-dimensional (1D) wave equation as a challenge case for future PINN models because the baseline PINN model faces severe challenges when applied to this PDE problem. Here, we apply our PECANN model with the proposed adaptive ALM to learn the solution of the same 1D wave equation as studied in [34, 21] and compare the accuracy of our predictions and the size of our neural network model against the results and networks presented in those two studies. While comparing the accuracy levels is significant, we also believe it is imperative to take into perspective the size of the neural network architecture, and the number of the collocations points deployed to obtain the level of accuracy in predictions.

Consider the following 1D wave equation

\[\frac{\partial^{2}u}{\partial t^{2}}-4\frac{\partial^{2}u}{ \partial x^{2}}=0,\quad(x,t)\in\mathcal{U}, \tag{29}\]

with the boundary condition

\[u(x,t)=0,\quad(x,t)\in\partial\mathcal{U}, \tag{30}\]

Figure 3: Heat transfer in a composite medium using the adaptive penalty update strategy and the formulation based expectation of constraints (i.e. section 2.2.2). (a) predicted temperature (b) absolute point-wise error in predicted temperature, (c) predicted heat flux, (d) absolute point-wise error in heat flux prediction.
and the initial condition

\[\frac{\partial u}{\partial t}(x,t)=0,\quad(x,t)\in\partial\mathcal{U}. \tag{31}\]

The exact solution to the problem is

\[u(x,t)=\sin(\pi x)+\frac{1}{2}\sin(4\pi x),\quad(x,t)\in\mathcal{U}, \tag{32}\]

where \(\mathcal{U}=\{(x,t)\mid x\in\Omega,t=[0,1]\}\), \(\partial\mathcal{U}=\{(x,t)\mid x\in\partial\Omega,t=[0,1]\}\) and \(\Gamma=\{(x,t)\mid x\in\Omega,t=0\}\). \(\partial\Omega\) represents the boundary of the spatial domain \(\Omega=\{x|x=(0,1)\}\).

Wang et al. [34], McClenny and Braga-Neto [21] used a fully connected feed-forward neural network with five hidden layer \((H_{L}=5)\) and 500 neurons per hidden layer \((N_{H}=500)\) with \(N_{\mathcal{U}}=300\) collocation points, \(N_{\partial\mathcal{U}}=300\) boundary points and \(N_{\Gamma}=300\) initial points that are generated at every epoch to solve the same 1D wave equation. In the present work, we use a single hidden layer \((H_{L}=1)\) with 50 neurons \((N_{H}=50)\) and a tangent hyperbolic activation function. We sample \(N_{\mathcal{U}}=300\) collocation points, \(N_{\partial\mathcal{U}}=300\) boundary points and \(N_{\Gamma}=300\) only once

\begin{table}
\begin{tabular}{c c c c c c c} \hline \hline Models & Mean \(\mathcal{E}_{r}(u,\dot{u})\pm\) stdev; (trials) & No. trials & \(H_{L}\times N_{H}\) & \(N_{\mathcal{U}}\) & \(N_{\mathcal{U}}\) & \(N_{\mathcal{U}}\) \\ \hline Baseline PINN [34] & \(4.518\times 10^{-1}\) & \(1\) & \(5\times 500\) & \(24\times 10^{6}\) & \(24\times 10^{6}\) & \(24\times 10^{6}\) \\ PINN using NTR [34] & \(1.728\times 10^{-3}\) & \(1\) & \(5\times 500\) & \(24\times 10^{6}\) & \(24\times 10^{6}\) & \(24\times 10^{6}\) \\ Baseline PINN[21] & \(3.792\times 10^{-1}\pm 0.162\times 10^{-1}\) & \(10\) & \(5\times 500\) & \(24\times 10^{6}\) & \(8\times 10^{6}\) & \(8\times 10^{6}\) \\ SA-PINN [21] & \(8.105\times 10^{-1}\pm 1.591\times 10^{-1}\) & \(10\) & \(5\times 500\) & \(24\times 10^{6}\) & \(8\times 10^{6}\) & \(8\times 10^{6}\) \\ SA-PINN with SGD [21] & \(2.950\times 10^{-2}\pm 7.0\times 10^{-3}\) & \(10\) & \(5\times 500\) & \(24\times 10^{6}\) & \(8\times 10^{6}\) & \(8\times 10^{6}\) \\
**Current Method** & \(\mathbf{3.990\times 10^{-3}\pm 9.179\times 10^{-4}}\) & \(10\) & \(\mathbf{1\times 50}\) & \(\mathbf{300}\) & \(\mathbf{300}\) & \(\mathbf{300}\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Comparison of relative error \(\mathcal{E}_{r}\) levels and network size used in the solution of 1D wave equation.

Figure 4: Error levels during training using ALM with different penalty update strategy. Top row: Temperature (a) relative error (b) absolute error. Bottom row: Heat flux (c) relative error, (d) absolute error.
before training. Our objective function is the residual on the governing PDE given in (29) and our constraints are the boundary condition (30) and initial conditions (31). We assign three Lagrange multipliers and three penalty parameters for this problem. We train our model for \(10,000\) epochs even though we obtain an acceptable level of accuracy with \(2,000\) epochs. We use L-BFGS optimizer with its default parameters in the Pytorch package. Table 1 presents a summary of the prediction errors, architecture size and collocation points for this problem in comparison to other works.

We present the prediction of our model in Fig. 6. We can observe that the our model has successfully learned the underlying solution to a good level of accuracy. In Table 1 we compare our predictions against other works that tackled the same problem. We report the mean and standard deviation of the relative error calculated over 10 different trials. The results show that our method produces results with error levels that are one order of magnitude level than the results obtained with the self-adaptive PINNs [21], while using a much smaller neural network model.

### Forward problem: lid-driven cavity flow simulation up to a Reynolds number of 1000

In this section, we apply our adaptive augmented Lagrangian method to solve the steady-state flow in a two-dimensional lid-driven cavity benchmark problem [10] at Reynolds numbers of 100, 400, and 1000. Unlike other works where a stream function vorticity formulation was adopted, we use the primitive-variables formulation of the Navier-Stokes equations, which potentially makes our approach extensible to three-dimensional fluid flow problems. We should also mention that lid-driven cavity at Reynolds of 1000 is considered much more challenging than the same problem at lower Reynolds numbers, such as at \(Re=100\) or \(Re=400\).

Figure 5: Evolution of penalty parameters with different penalty update strategy: (a) evolution of the penalty parameter \(\mu_{\mathcal{B}}\) for enforcing boundary conditions, \(\mu_{\mathcal{I}}\) for enforcing initial conditions, and \(\mu_{\mathcal{Q}}\) for enforcing flux constraints using ALM with adaptive penalty update strategy. (b) evolution of penalty parameters for enforcing boundary condition, initial condition and flux equation using ALM with MPU and CPU. (c) evolution of constraints during training using ALM with APU
Steady-state Navier-Stokes equations in two dimensions and the continuity equation are written as follows:

\[u\frac{\partial u}{\partial x}+v\frac{\partial u}{\partial y}+ \frac{\partial p}{\partial x}-\frac{1}{Re}(\frac{\partial^{2}u}{\partial x^{2}} +\frac{\partial^{2}u}{\partial y^{2}}) =0,\quad(x,y)\in\Omega \tag{33a}\] \[u\frac{\partial v}{\partial x}+v\frac{\partial v}{\partial y}+ \frac{\partial p}{\partial y}-\frac{1}{Re}(\frac{\partial^{2}v}{\partial x^{2} }+\frac{\partial^{2}v}{\partial y^{2}}) =0,\quad(x,y)\in\Omega,\] (33b) \[\frac{\partial u}{\partial x}+\frac{\partial v}{\partial y} =0,\quad(x,y)\in\Omega \tag{33c}\]

where \(\mathbf{u}=(u,v)\) is the velocity vector, \(Re\) is the Reynolds number and \(p\) is the pressure. We aim to learn the solution of the momentum equations constrained by the continuity equation on \(\Omega=\{(x,y)\mid 0\leq x\leq 1,0\leq y\leq 1\}\) with its boundary \(\partial\Omega\). No-slip boundary conditions are imposed on all walls of the square cavity. The top wall of the square cavity moves at a velocity of \(u(x,1)=1\) in the \(x\) direction. Because we have a closed system with no inlets or outlets in which the pressure level is defined, we fix the pressure field at an arbitrary reference point as \(p(0.5,0)=0\). Fig.7 illustrates our simulation setup for the lid-driven cavity flow problem. We should note that our objective function in this problem is the loss on the total momentum \(\mathcal{F}\) field, and our constraints are the continuity equation (i.e. divergence-free velocity field) \(\mathcal{Q}\), boundary conditions and the pressure that is fixed at a reference point.

To formulate our constrained optimization problem, let us define \(\mathbf{\mathcal{F}}=(\mathcal{F}_{x},\mathcal{F}_{y})\), which represents the steady-state, two-dimensional momentum equations, and \(\mathcal{Q}\) as the divergence-free condition as follows

\[\mathcal{F}_{x}(x,y) :=u\frac{\partial u}{\partial x}+v\frac{\partial u}{\partial y}+ \frac{\partial p}{\partial x}-\frac{1}{Re}(\frac{\partial^{2}u}{\partial x^{2 }}+\frac{\partial^{2}u}{\partial y^{2}}),\quad(x,y)\in\Omega \tag{34a}\] \[\mathcal{F}_{y}(x,y) :=u\frac{\partial v}{\partial x}+v\frac{\partial v}{\partial y}+ \frac{\partial p}{\partial y}-\frac{1}{Re}(\frac{\partial^{2}v}{\partial x^{2 }}+\frac{\partial^{2}v}{\partial y^{2}}),\quad(x,y)\in\Omega,\] (34b) \[\mathcal{Q}(x,y) :=\frac{\partial u}{\partial x}+\frac{\partial v}{\partial y}, \quad(x,y)\in\Omega \tag{34c}\]

Figure 6: Solution of 1D wave equation obtained from a single hidden layer neural network with 50 neurons trained using the proposed augmented Lagrangian method: (a) predicted solution, (b) exact solution (c) absolute pointwise error

Figure 7: Boundary conditions for the lid-driven cavity problem
We sample \(N_{\mathcal{B}}=4\times 64\) points for approximating the boundary conditions, \(N_{\mathcal{Q}}=N_{\mathcal{F}}=10^{4}\) points for approximating the loss on conservation of mass and momentum only once before training, respectively. Additionally, we constrain a single point for pressure as shown in Fig.7. We use a fully connected neural network architecture consisting of four hidden layers with 30 neurons in each layer and the tangent hyperbolic activation function. We choose L-BFGS optimizer with its default parameters and _strong Wolfe_ line search function that are available in the PyTorch package. We train our network for 2000 epochs for each Reynolds number separately.

In Fig. 8, we investigate the effect of different penalty update strategies on the predictions. We observe that our model trained with the adaptive penalty update strategy in ALM learns the underlying solution for Re=100 well and matches the benchmark data closely. Figs. 8(a)-(b) also show that ALM with monotonic (MPU) or conditional (CPU) penalty update strategies are incapable of producing the expected solution. Next, we consider the case with Reynolds number 400 and train our model for 2000 epochs using the adaptive penalty update strategy only since the other two strategies did not produce acceptable results for \(Re=100\).

We observe from Fig. 9 and Fig. 10 that our model with the adaptive update strategy learns the underlying pattern for \(Re=400\) and \(Re=1000\), respectively. However, the predicted velocity fields do not closely match the benchmark numerical solution, despite the our model's superior performance for the \(Re=100\) case.

It is worth noting that our objective in these two simulation cases is to highlight the fact that despite implementing and constraining appropriate boundary conditions, divergence-free velocity field, and anchoring the pressure at a single point, the task of learning becomes increasingly more challenging as the Reynolds number increases. Learning the solution of lid-driven cavity problem at a high Reynolds numbers may be achievable with a larger neural network or a

Figure 8: Impact of three different ALM penalty update strategies on the solution of lid-driven cavity problem at Re = 100: (a) predicted horizontal velocity along the line at \(x=0.5\) compared against the benchmark data, (b) predicted vertical velocity along the line at \(y=0.5\) compared against the benchmark data.

Figure 9: Solution of lid-driven cavity problem at Re = 400 using ALM with APU: (a) predicted horizontal velocity along the line at \(x=0.5\) compared against the benchmark data, (b) predicted vertical velocity along the line at \(y=0.5\) compared against the benchmark data.
higher number of collocations. However, our aim here is not to demonstrate whether such learning is possible or not. Instead, we aim to illustrate the challenges involved for high Reynolds number flow problems, even when formulated properly as a constrained optimization problem. We hypothesize that for high Reynolds number flow problems, the underlying optimization problems become ill-conditioned and require additional regularization. Therefore, in order to further regularize the optimization process, we introduce a small set of exact, randomly sampled velocity data within the flow domain as shown in Fig. 11(a)-(b). We use these clean data points as equality constraints on the prediction of our model, similar to the way we handle boundary conditions.

We observe from Figs. 11(c)-(e) that our PECANN model with an adaptive ALM successfully predicts the fluid flow in the cavity at Reynolds number = 400 with a high degree of accuracy by leveraging a small sample of high fidelity data. Specifically, Figs. 11(c)-(e) show the predicted solution, x-component of the velocity over a line at \(x=0.5\)), and y-component of the velocity over a line at \(y=0.5\)) for \(Re=400\), while Figs. 11(f)-(h) show the corresponding results for \(Re=1000\). Compared to the results shown in Fig. 9 and Fig. 10, our new results suggest that a small set of randomly sampled labeled data is beneficial to regularize the optimization process for high Reynolds numbers. We aim to explore the feasibility of learning high Reynolds numbers in two and three dimensions without any labeled data using larger and different neural network architectures in the future.

### Inverse problem: estimating transient thermal boundary condition from noisy data

Here we aim to learn the transient thermal boundary condition and temperature based on measured data. Let us consider a parabolic partial differential of the form

\[\frac{\partial T}{\partial t}=\kappa\frac{\partial^{2}T}{\partial x^{2}}\quad (x,t)\in(0,1)\times(0,1], \tag{35}\]

where \(T\) is the temperature, \(\kappa\) is the thermal conductivity. The forward solution of Eq. 35 requires the following boundary information

\[T(0,t)\ \forall t\in(0,1],\quad T(1,t)\ \forall t\in(0,1],\quad T(x,0)\ \forall x \in(0,1). \tag{36}\]

However, in this problem, we do not know the boundary condition at \(T(0,t)\). Instead we have noisy data at the interior part of our domain. We constrain the prediction of our model on the governing physics and the noiseless boundary condition that we treat as high-fidelity data. Fig. 12(a)-(c) present the collocation points distributed over the space-time domain, and synthetically generated noisy data from the exact solution at two locations.

We use a fully connected artificial neural networks with three hidden layers and 30 neurons per layer. We adopt the tangent hyperbolic activation function. We randomly generate \(N_{\mathcal{F}}=512\) collocation points, \(N_{\mathcal{I}}=64\) number of training points for approximating the initial condition and \(N_{\mathcal{B}}=64\) number of training boundary data only once before training. For the noisy measurement data we generate \(N_{\mathcal{M}}=2\times 128\) with 10 percent noise (Gaussian noise with a standard deviation that is 10 \(\%\) of the standard deviation of the synthetically generated clean data). We use the L-BFGS optimizer with its default parameters and _strong Wolfe_ line search activation function. We train our model for \(5000\) epochs. The result of our numerical experiment is presented in Fig. 13.

Figure 13(a) presents the exact solution. From Fig. 13(b) we observe that our predicted solution is in great agreement with the exact solution. We provide the absolute point-wise error distribution in Fig.13(c). We have shown that by

Figure 10: Same caption as in Fig. 9, but for \(Re=1000\)
having two sensors at two different locations inside the domain that are measuring noisy data, we can sufficiently learn not only the boundary condition, but also the entire temperature distribution across the space-time domain.

Figure 11: Lid-driven cavity predictions with a small sample of high fidelity velocity data for different Reynolds numbers. ALM with APU is used in the computations. Locations of randomly sampled high fidelity data: (a) Re = 400, (b) Re = 1000. Results for Re=400: (c) predicted velocity streamlines, (d) predicted horizontal velocity along the line at \(x=0.5\) compared against the benchmark result, (e) predicted vertical velocity along the line at \(y=0.5\) compared against the benchmark result. Bottom row: Results for Re = 1000, same captions as in (c)-(e)
### Inverse problem: estimation of transient source term in diffusion problems

In this example, we aim to learn the temperature distribution as well as the heat source using noisy measurements and noiseless boundary or initial conditions, which is a more challenging inverse problem than the previous example. Hasanov [12] proposed a collocation algorithm, based on the piece-wise linear approximation of the unknown heat source. Several works have also been proposed to learn the unknown heat source with a single dependent variable [8, 15]. In our approach, we use a neural network model to represent the temperature and the heat source with time and space dependency. We then learn the heat distribution and the heat source simultaneously.

Let us consider the following initial-boundary value problem of the heat conduction equation as follows:

\[\frac{\partial u}{\partial t}=\kappa\frac{\partial^{2}u}{\partial x^{2}}+s, \quad(x,t)\in(0,1)\times(0,1], \tag{37}\]

where \(u\) is the temperature, \(\kappa\) is the thermal conductivity and \(s\) is the heat source. The objective is to learn the solution \(u\) and the heat source \(s\) from noiseless boundary condition and noisy measurements inside our domain. A schematic representation of data is given in Fig. 14.

Unlike typical inverse source problems that rely on a single input variable, this paper addresses a more challenging scenario where the problem is dependent not only on the spatial variable \(\mathbf{x}\) but also on the temporal variable t. We demonstrate that we obtain an accurate prediction without resorting to any stabilizing scheme or assuming any fundamental solution for the temperature or the heat source distribution. It is worth mentioning that our constrained optimization problem may also be called as PDE-constrained optimization problem. For this problem, we use three-hidden layer fully connected artificial neural networks with 30 neurons per layer with tangent hyperbolic activation function. We randomly generate \(N_{\mathcal{F}}=4096\) collocation points, \(N_{\mathcal{I}}=128\) number of training points for approximating the initial condition and \(N_{\mathcal{B}}=2\times 128\) number of training boundary data only once before training. For the noisy data we generate \(N_{\mathcal{M}}=512\) with 10 percent noise (Gaussian noise with a standard

Figure 12: Collocation points and measurement data for estimating transient thermal boundary condition from noisy data: (a) collocation points (magenta), boundary points (black) and collocation points for approximating the initial condition (blue), (b) noisy measured temperature \(\tilde{T}\) and exact temperature \(T\) at \(x=0.2\), (c) temperature \(\tilde{T}\) and exact temperature \(T\) at \(x=0.6\)

Figure 13: Estimating transient thermal boundary condition: (a) exact solution (b), predicted solution (c), absolute pointwise error
Figure 14: schematic representation of training data: (a) initial conditions (red points), boundary conditions (magenta points) and noisy measurement data (blue points), (b) collocation points.

Figure 15: Estimating transient heat source: (a) exact temperature distribution, (b) predicted temperature distribution, (c) absolute point wise error in temperature distribution, (d) exact heat source distribution, (e) predicted heat source distribution, (f) absolute point wise error in heat source distribution
deviation that is 10 \(\%\) of the standard deviation of the synthetically generated clean data). We use L-BFGS optimizer with its default parameters and _strong Wolfe_ line search activation function. We train our model for \(5000\) epochs.

We present the result from this numerical experiment is presented in Fig.15. The exact temperature distribution is plotted in Fig. 15(a). From Figs. 15(b)-(c) we observe that our predicted solution is in agreement with the exact solution to a great degree. The exact distribution of the heat source is plotted in Fig.15(d). Similar to the temperature distribution, the predicted heat source distribution agrees closely with the exact solution as shown in Figs.15(e)-(f).

## 4 Conclusion

Constrained optimization is central to the formulation of the PECANN framework [2] for learning the solution of forward and inverse problems and sets it apart from the popular PINN framework, which adopts an unconstrained optimization in which different loss terms are combined into a composite objective function through a linear scalarization. In the PECANN framework, the augmented Lagrangian method (ALM) enables us to constrain the solution of a PDE with its boundary and initial conditions and with any high-fidelity data that may be available in a principled fashion. However, for challenging forward and inverse PDE problems, the set of constraints can be large and diverse in their characteristics. In the present work, we have demonstrated that conventional approaches to update penalty and Lagrange multipliers within the ALM can fail for PDE problems with complex constraints. Consequently, we introduced an adaptive ALM in which a unique penalty parameter is assigned to each constraint and updated according to a proposed rule. This new method within the PECANN framework enables us to constrain the solution of PDEs flexibly toward convergence without the need to manually tune the hyperparameters in the objective function or resort to alternative strategies to balance the loss terms, which may not necessarily satisfy the optimality conditions.

When the problem size gets larger, constraining each collocation point associated with a constraint becomes taxing computationally. We have empirically demonstrated that Lagrange multipliers associated with a particular type of constraint exhibits a distribution with a clear mean value. Consequently, we have revised our constrained optimization formulation based on the expectation of each loss and constraint terms. Based on our experiments, the resulting formulation does not affect the efficacy of our PECANN framework, but makes it computational more efficient and enable mini-batch training.

We applied our proposed model to several forward and PDE-constrained inverse problems to demonstrate its efficacy and versatility. A noteworthy example is our simulation of the lid-driven cavity benchmark problem at a Reynolds number 1000 using the primitives-formulation of the Navier-Stokes equations. The \(Re=1000\) case is known to be a challenging case to simulate because of the dominant effect of advection in the flow physics and development of strong circulations in the bottom corners of the cavity. Unlike, streamfunction-vorticity formulation of the two-dimensional Navier-Stokes equation, primitive-variables formulation is more general, but at the same more challenging to solve computationally because the formulation retains pressure as one of the variables in addition to the velocity field and. In our flow simulations, we found that the key to obtaining results that are in close agreement with benchmark data [10] was to constrain the momentum equations with the divergence-free condition, boundary conditions, and with a small set of randomly sampled high-fidelity data while handling each type constraint with its own unique, adaptive penalty parameter. Another PDE problem that has been viewed as a challenge case for PINNs is the one-dimensional wave equation. Again, our proposed model performed better than the state-of-the-art PINN models for the wave equation problem as well. The codes used to produce the results in this paper are publicly available at [https://github.com/HiPerSimLab/PECANN](https://github.com/HiPerSimLab/PECANN). Two additional examples are provided in the Appendix.

## Acknowledgments

This material is based upon work supported by the National Science Foundation under Grant No. 1953204 and in part in part by the University of Pittsburgh Center for Research Computing through the resources provided.

## Appendix A Additional Examples

### Inviscid scalar transport in one dimension

In this section, we study the transport of a physical quantity dissolved or suspended in an inviscid fluid with a constant velocity. Numerical solutions of inviscid flows often exhibit a viscous or diffusive behavior owing to numerical dispersion. Let us consider a convection equation of the form

\[\frac{\partial\xi}{\partial t}+u\frac{\partial\xi}{\partial x}=0,\;\forall(x,t )\in\Omega\times[0,1], \tag{38}\]
satisfying the following boundary condition

\[\xi(0,t)=\xi(2\pi,t)\;\forall t\in[0,1], \tag{39}\]

and initial condition

\[\xi(x,0)=h(x),\;\forall x\in\partial\Omega, \tag{40}\]

where \(\xi\) is any physical quantity to be convected with the velocity \(u\), \(\Omega=\{x\mid 0<x<2\pi\}\) and \(\partial\Omega\) is its boundary. Eq.(38) is inviscid, so it lacks viscosity or diffusivity. For this problem, we consider \(u=40\) and \(h(x)=\sin(x)\). The analytical solution for the above problem is given as follows [16]

\[\xi(x,t)=F^{-1}(F(h(x))e^{-i\omega kt}), \tag{41}\]

where \(F\) is the Fourier transform, \(k\) is the frequency in the Fourier domain and \(i=\sqrt{-1}\). Since the PDE is already first-order, there is no need to introduce any auxiliary parameters. Following is the residual form of the PDE used to formulate our objective function,

\[\mathcal{F}(x,t)=\frac{\partial\xi(x,t)}{\partial t}+u\frac{ \partial\xi(x,t)}{\partial x}, \tag{42a}\] \[\mathcal{B}(t)=\xi(0,t)-\xi(2\pi,t),\] (42b) \[\mathcal{I}(x)=\xi(x,0)-\sin(x), \tag{42c}\]

where \(\mathcal{F}\) is the residual form of our differential equation, \(\mathcal{B}\) and \(\mathcal{I}\) are our boundary condition and initial condition constraints.

For this problem, we use a fully connected neural network architecture consisting of four hidden layers with 50 neurons and tangent hyperbolic activation functions. We generate \(N_{\mathcal{F}}=512\) collocation points from the interior part of the domain, \(N_{\mathcal{B}}=512\) from each boundary, and \(N_{\mathcal{I}}=512\) for approximating the initial conditions at each epoch. We use L-BFGS optimizer with its default parameters and _strong Wolfe_ line search function that is available in the PyTorch package. We train our network for 5000 epochs. We present the prediction of our neural network in Fig.16. We observe that our neural network model has successfully learned the underlying solution as shown in Figs. 16(b)-(c).

We also present a summary of the error norms from our approach and state-of-the-art results given in [16] in Table 2. We observe that our method achieves a relative \(\mathcal{E}_{r}=8.161\times 10^{-4}\), which is two orders of magnitude better than curriculum learning method presented in [16], and three orders of magnitude better than the baseline PINN model performance as reported in the same study.

\begin{table}
\begin{tabular}{l c c} \hline \hline Models & \(\mathcal{E}_{r}(\xi,\hat{\xi})\) & MAE \\ \hline Baseline PINN [16] & \(9.61\times 10^{-1}\) & \(5.82\times 10^{-1}\) \\ Curriculum learning [16] & \(5.33\times 10^{-2}\) & \(2.69\times 10^{-2}\) \\ Current method & \(\mathbf{8.161\times 10^{-4}}\) & \(\mathbf{6.810\times 10^{-4}}\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Inviscid scalar transport problem: summary of the relative error \(\mathcal{E}_{r}\) and the mean absolute error (MAE) metrics for training a fixed neural network architecture.

Figure 16: Inviscid scalar transport problem: (a) exact solution, (b) predicted solution, (c) absolute point-wise error
### Mixing of a hot front with a cold front

In this section, we explore formation of cold and warm fronts in a two-dimensional environment which can be described by the following convection equation:

\[\mathcal{F}:=\frac{\partial\xi}{\partial t}+u\frac{\partial\xi}{\partial x}+v \frac{\partial\xi}{\partial y}=0\ \in\Omega\times(0,T], \tag{43}\]

where \(\Omega=\{(x,y)\ |\ -4\leq x,y\leq 4\}\) and \(T=4\). The boundary conditions considered throughout this work are zero flux of the variable \(\xi\) along each boundary as shown in Fig. 17. The problem has the following analytical solution:

\[\xi(x,y,t)=-\tanh\Big{[}\frac{y}{2}\cos\omega t-\frac{x}{2}\sin\omega t\Big{]}, \tag{44}\]

where \(\omega\) is the frequency, which is defined as

\[\omega=\frac{\nu_{t}}{r\nu_{t,max}}. \tag{45}\]

The velocity field \(\nu_{t}=\mathrm{sech}\,r^{2}\tanh r\) is the tangential velocity around the center with a maximum value \(\nu_{t\max}=0.385\). The velocity components in the \(x\) and \(y\) directions can be obtained, respectively, as follows:

\[u(x,y)=-\omega y,\quad v(x,y)=\omega x. \tag{46}\]

We also present the velocity field in Fig.18(a). The initial scalar field varies gradually from positive values at the bottom to negative values at the top as can be seen Fig.18(b).

where the maximum and minimum values of the field is \(\xi_{\max}=0.964\) and \(\xi_{\min}=-0.964\), respectively.

Figure 17: Mixing hot and cold front: Schematic representation of the domain with its boundary condition

Figure 18: Velocity and initial state of mixing of hot and cold fronts: (a) velocity field, (b) initial state of our flow \(\xi_{\max}=0.964\) and \(\xi_{\min}=-0.964\), respectively.
We use a deep fully connected neural network with 4 hidden layers each with 30 neurons that we train for 5000 epochs total. We use a Sobol sequence to generate \(N_{\mathcal{F}}=10,000\) residual points from the interior part of the domain, \(N_{\mathcal{B}}=512\) points from each of the boundaries (faces) and \(N_{\mathcal{I}}=512\) points for enforcing the initial condition only once before training. Our optimizer is LBFGS with its default parameters and _strong Wolfe_ line search function that is available in the PyTorch package.

Figure 19 presents the temperature contours obtained from exact solution along with our predictions. We observe that our PECANN model is in good agreement with the exact solution. We also provide a summary of RMS errors obtained from our neural network model along with conventional numerical methods used in other works in Table. 3.

Table 3 compares the RMS error level of our predictions against the RMS error levels produced by five different finite-volume based advection schemes as presented in the work of Tamamidis and Assanis [30]. Among the finite-volume based methods, the most accurate numerical results were obtained with the QUICK scheme on all meshes consistently. We should note that our results on the Cartesian mesh sizes shown in Table 3 are post calculations from the trained neural network. For this reason, our error levels are nearly the same across all mesh sizes.
