Magnetohydrodynamics with physics informed neural operators

The goal of this work is to reproduce the incompressible MHD equations with PINOs. The incompressible MHD equations represent an incompressible fluid in the presence of a magnetic field B. These equations are given by

$$\begin{align} \partial_t \mathbf{u}+\mathbf{u} \cdot \nabla \mathbf{u} = -\nabla \left( p+\frac{B^2}{2} \right)/\rho_0 +\mathbf{B} \cdot \nabla \mathbf{B}+\nu \nabla^2 \mathbf{u}, \end{align} \tag{ 1 }$$

$$\begin{align} \partial_t \mathbf{B}+\mathbf{u} \cdot \nabla \mathbf{B} = \mathbf{B} \cdot \nabla \mathbf{u}+\eta \nabla^2 \mathbf{B}, \end{align} \tag{ 2 }$$

$$\begin{align} \nabla \cdot \mathbf{u} = 0, \end{align} \tag{ 3 }$$

$$\begin{align} \nabla \cdot \mathbf{B} = 0, \end{align} \tag{ 4 }$$

where u is the velocity field, p is the pressure, B is the magnitude of the magnetic field, $\rho_0 = 1$ is the density of the fluid, ν is the kinetic viscosity, and η is the magnetic resistivity. We have two equations for evolution and two constraint equations.

To ensure the zero velocity divergence condition of equation (3), the pressure is typically computed in such a way that this condition always holds true. This is generally done by solving a Poisson equation to calculate the pressure at each time step.

For the magnetic field divergence of equation (4), we lack any additional free parameters to ensure that the equation holds true at all times. There are several ways to ensure that this condition holds including hyperbolic divergence cleaning [28] and constrained transport [29]. For this work, we decide to preserve the magnetic field's zero divergence condition by instead evolving the magnetic vector potential A. This quantity is defined such that

$$\begin{align} \mathbf{B} = \nabla \times \mathbf{A}, \end{align} \tag{ 5 }$$

which ensures that the divergence of B is zero to numerical precision as the divergence of a curl of a vector field is zero. By evolving magnetic vector potential A instead of the magnetic field B, we have a new evolution equation for the vector potential A. This equation is given by

$$\begin{align} \partial_t \mathbf{A} + \mathbf{u} \cdot \nabla \mathbf{A} = \eta \nabla^2 \mathbf{A}. \end{align} \tag{ 6 }$$

To simulate the MHD equations, we employed the Dedalus code [30], an open source parallelized spectral python package that is designed to solve general PDEs. To obtain interesting results without additional computational difficulty, we elected to solve the incompressible MHD equations in 2D with periodic boundary conditions (BCs). This results in us solving a total of 3 evolution PDEs at each timestep-2 for the velocity evolution and 1 for the magnetic vector potential evolution. To visualize and more easily diagnose problems with the simulations, we include an additional PDE to evolve tracer particles denoted by s. The tracer particles had an evolution equation given by

$$\begin{align} \partial_t s + \mathbf{u} \cdot \nabla s = \nu \nabla^2 s. \end{align} \tag{ 7 }$$

Moreover, we employed a pressure gauge of $\int{p \, \mathrm{d}x^2} = 0$ for our equation. Numerically, most simulations were carried out on a unit length square grid at resolutions with a total number of grid points $N = 128^2$, with $N_x = N_y = 128$ points in the x and y directions, respectively. We also examined some lower resolution simulations with $N = 64^2$ and higher resolution simulations with $N = 256^2$ to explore how resolution may affect the results of the simulations.

As we have periodic BC, we employed a Fourier basis in each direction. To avoid aliasing error, we used a dealias factor of $3/2$ for these transformations. For timestepping, we used the RK4 integration method with a constant timestep of $\Delta t = 0.001$. Simulation outputs were stored every 10 timesteps resulting in us recording data at interval t = 0.01 time units. We ran these simulations until a time of t = 1. For each set of parameters, we produced 1000 simulations, each with different initial data.

An important consideration in this study was how the Reynolds number affects simulations and our AI model's ability to reproduce the results of the simulations. The Reynolds number is a dimensionless quantity that expresses the ratio of inertial to viscous forces. The higher the Reynolds number is, the more prevalent the effects of small scale phenomena are to the simulation. This often gives rise to turbulence at high Reynolds numbers. For systems described by MHD equations, there are actually two types of Reynolds numbers of interest. They are the kinetic Reynolds number, $\mathrm{Re}$, and the magnetic Reynolds number, $\mathrm{Re}_m$. We define the quantities in our simulation such that these Reynolds numbers are given by

$$\begin{align} \mathrm{Re} & = 1/\nu, \end{align} \tag{ 8 }$$

$$\begin{align} \mathrm{Re}_m & = 1/\eta. \end{align} \tag{ 9 }$$

The ratio between these two Reynolds numbers is called the magnetic Prandtl number $\mathrm{Pr}_m$ defined as

$$\begin{align} \mathrm{Pr}_m = \frac{\mathrm{Re}_m}{\mathrm{Re}} = \frac{\nu}{\eta}. \end{align} \tag{ 10 }$$

In this study, we sought to characterize how the Reynolds number affects our models and determine if there is a cutoff Reynolds number, after which point, the models' performance degrades considerably. We generally kept $\mathrm{Re} = \mathrm{Re}_m$ or, equivalently, $\mathrm{Pr}_m = 1$ unless otherwise noted and looked at Reynolds number values of 100, 250, 500, 750, 1000, and 10 000.

One difference between MHD and hydrodynamic simulations is that, compared to the pure hydrodynamic case, much of the magnetic field energy is stored at high wavenumbers which occur at smaller scales. Thus, the models must be able to characterize high frequency features if they are to successfully reproduce the results of the magnetic field evolution.

The goal of a NO is to reproduce the results of an operator given some input fields using neural networks. A common class of operator often studied by NOs are PDEs. To model PDEs, NO take the coordinates, initial conditions (ICs), BCs, and coefficient fields as inputs. As outputs, NOs provide the solution of the operator, in this case the PDE, at the desired spacetime coordinates.

There are a variety of different NOs that have been studied in recent works. These include DeepONets, physics informed DeepONets, low-rank NOs, graph NOs (GNO), multipole GNOs (MGNOs), Fourier NOs (FNO), factorized FNOs (fFNO), and PINO [22, 26, 27, 31–36]. In this study we use PINOs.

PINOs are a generic class of NOs that involve using an existing NO and employing physics informed methods in the training [26, 27]. Physics informed deep learning methods involves encoding known information about the physical system into the model during training [37–42]. This physics information may include governing PDEs, symmetries, constraint laws, ICs, and BCs. By including such physical knowledge, physics informed methods enable better generalization of the results of deep learning models [42]. In systems where we have a lot of knowledge about the physical system like PDEs, such methods are especially useful as theoretically we could learn from just this physics knowledge. In practice, we add data to help PINOs converge to the correct solution more quickly.

One of the most common ways of encoding this physics into a neural network model is to incorporate physics information into the loss function [37–40]. In other words, violations of these physics laws appear as terms in the loss function that are reduced over time during training. Thus, this technique treats physics laws as soft constraints learned by the neural network as it trains.

For the backbone NO, we selected a variant of fFNO [36] that employed the TensorLY [43] package to perform tensor factorization. We call this model a tensor FNO (tFNO). The base FNO model [35] applies the fast Fourier transform (FFT) to the data to separate it into its component frequencies and apply its weights before performing an inverse FFT to convert back to real space. One particular feature of the FNO is its ability to perform zero-shot super-resolution, in which the model predicts on higher resolution data than it was trained on [22, 35]. We hope such properties of the FNO model would allow it to learn the high frequency properties of the MHD equations.

To further augment the model, we utilized factorization with TensorLY within the spectral layers of the model, so that it becomes an tFNO. The factorization was found to significantly improve the generalization of the neural network in the initial tests. The decision to use TensorLY was inspired by experimental code found in the GitHub repository of [26]. We will describe the model architecture in more detail in section 4.2.

Now let us discuss in more detail how to apply physics informed methods to model the MHD equations with NOs. There are several aspects to modeling the MHD equations with such methods as denoted by each loss term. These are the data loss $\mathcal{L}_{\textrm{data}}$, the PDE loss $\mathcal{L}_{\textrm{PDE}}$, the constraint loss $\mathcal{L}_{\textrm{constr}}$, the IC loss $\mathcal{L}_{\textrm{IC}}$, and the BC loss $\mathcal{L}_{\textrm{BC}}$.

The data loss is modeled by obtaining simulation data and ensuring that the PINO output matches the simulation results. This requires us to produce a large quantity of simulations to provide sufficient training data like those described in section 2.2 to span the solution space. We will discuss in more detail how we produced a sufficient number of simulations in section 4.1. Then we use the relative mean squared error (MSE) between the simulation and the PINO predictions to define the value of the data loss $\mathcal{L}_{\textrm{data}}$.

The PDE loss describes the violations of the time evolution PDEs of the PINO outputs. Specifically, we encode these time evolution PDEs as part of the loss function. This requires us to represent the spatial and temporal derivatives of the output fields. Although some NOs are able to employ the automatic differentiation of the deep learning framework to calculate such derivatives [32], this method is too memory intensive for the tFNO architecture. Instead, we employ Fourier differentiation to represent the spacial derivatives, which computes highly accurate derivative values while conveniently assuming periodic BC that exists in our problem. For time derivatives, we use second order finite differencing. We note that such a time differencing technique cannot be used during the original simulation as this requires knowledge of future times to compute the time derivative of the current time.

Specifically, the PDEs we modeled with this technique were equation (1) for velocity evolution and equation (5) for magnetic potential evolution. In this case, the model has a total of 3 output fields in 2D, the velocity in the x direction u, the velocity in the y direction v, and the magnetic potential A. We also tested replacing the magnetic vector potential evolution with the magnetic field evolution of equation (2), which results in 4 output fields in 2D. The fields in this case are the velocity in the x direction u, the velocity in the y direction v, the magnetic field in the x direction B_x , and the magnetic field in the y direction B_y . We note that in testing, the former representation of the MHD equations produced better results. The PDE loss $\mathcal{L}_{\textrm{PDE}}$ is then defined as the MSE loss between zero and the PDE, after putting all the terms on the same side of the equation.

Due to the complexity of both representations of the MHD equations, we tested the PDE loss on the data produced by the simulations. We found that the loss was small, albeit nonzero. This nonzero PDE loss was expected since the numerical methods for computing the derivatives during the simulation differs from those of computing the derivatives of the PDE during the loss function. Some particular notable differences are having fewer time steps in the output data than was used during the simulation, using different time derivative methods (e.g. RK4 during simulation vs second order finite differencing during the loss computation), and lacking dealiasing in the spatial derivatives in the PDE loss function.

The constraint loss illustrates the deviations of the elliptic constraint equations of the MHD equations. Specifically, these refer to the velocity divergence free condition of equation (3) and the magnetic divergence free condition of equation (4). We implemented these constraints similarly to the time evolution equations in the PDE loss, but without any time derivative terms. The constraint loss $\mathcal{L}_{\textrm{constr}}$ is then the MSE between each of the constraint equations and zero. We note that we are representing the magnetic fields by the magnetic potential, the magnetic field divergence free condition is satisfied up to the numerical precision of the curl operator regardless of the prediction of the neural network. However, we still include the term for completeness and this condition is not guaranteed to be nonzero if we are trying to compute the magnetic fields directly.

The IC loss tells the model to associate the input field with the output at t = 0. Although this function can often be achieved with the data loss alone, the IC loss emphasizes the importance of predicting the correct IC and enables training in the absence of data. Both the significance of the IC and the ability to train without data stem from the PDE loss term. Theoretically, one can train by correctly predicting the IC, the BCs, and evolving the PDE correctly forward in time, although in practice, data helps the model converge more quickly. However, an incorrect IC results in the PDE evolving the wrong data forward in time. Thus, we add in the IC as its own term to encourage the model to first compute the IC correctly before learning the time evolution later in training. We calculate the IC loss $\mathcal{L}_{\textrm{IC}}$ by taking the input fields and computing the relative MSE between said input fields and the outputs fields at t = 0.

Finally, the BC loss $\mathcal{L}_{\textrm{BC}}$ describes the violations of the boundary terms. In our case, the tFNO model architecture ensures that we have the desired periodic BC. Therefore, we do not use such a term in our model. However, we mention this term for generality because not all BCs and the periodicity of the tFNO architecture can be removed by zero padding the inputs along the desired non-periodic axis.

In physics informed deep learning methods, one must be careful when combining fields of different magnitudes to ensure they all have an equal contribution to the loss. Therefore, our model has the ability to normalize the input fields and denormalize the output fields by multiplying by predetermined constants. Moreover, we assign a weight to each term when combining loss terms for different fields and equation. We add additional weights when combining different losses. Thus, our loss $\mathcal{L}$ is given as

$$\begin{align} \mathcal{L} & = w_{\textrm{data}} \mathcal{L}_{\textrm{data}} + w_{\textrm{PDE}} \mathcal{L}_{\textrm{PDE}} + w_{\textrm{constr}} \mathcal{L}_{\textrm{constr}} + w_{\textrm{IC}} \mathcal{L}_{\textrm{IC}}\,, \end{align} \tag{ 11 }$$

where w_data is the data weight, w_PDE is the PDE weight, w_constr is the constraint weight, and w_IC is the IC weight.

To generate the training data, we first needed to produce initial data before running it in the simulations described in section 2.2. These initial data fields are produced using Gaussian random field method similar to [27], in which the kernel was transformed into Fourier space to obey our desired periodic BCs. Specifically, we used the radial basis function kernel to produce smooth initial fields. This kernel k_l is defined as

$$\begin{align} k_{l}\left(x_{1}, x_{2}\right) = \exp\left(-\frac{\left\|x_{1}-x_{2}\right\|^{2}}{2 l^{2}}\right), \end{align} \tag{ 12 }$$

where l is the length scale of typical spatial deviations in the data. We used l = 0.1 for all fields in this works unless otherwise stated. We also needed to ensure that the velocity and magnetic fields are both divergence free. Therefore, we produced two initial data fields, the vorticity potential, ψ, and the magnetic potential A. We defined ψ such that

$$\begin{align} \mathbf{v} = \nabla \times \psi, \end{align} \tag{ 13 }$$

which guarantees the velocity fields are divergence free initially. The magnetic potential A is defined in equation (6) in a similar manner to prevent the presence of divergences in the initial magnetic fields.

We multiplied the resulting initial data fields by a constant to ensure that the resulting fields have appropriate magnitude and are numerically stable. For example, we need to prevent Courant–Friedrichs–Lewy condition violations, which can occur if the velocities are too high for a given resolution and timestep choice [44]. Numerical instabilities can also occur if the initial magnetic potential values are too high and may cause the simulation to fail. We chose these constants to be $c_{\psi} = 0.1$ for the vorticity potential ψ and $c_{A} = 0.005$ for the magnetic potential A.

A schematic of the tFNO models utilized in this study is presented in figure 1. The size and dimensions of the model can be described by three hyperparameters—the width, the number of Fourier modes, and the number of layers. We used four layers for all models in this work. The width and number of Fourier modes were the same across all layers of the model in this work and were both set to 8 unless stated otherwise. Although we suspect that we could get better results by increasing the these hyperparameters, especially the number of Fourier modes which may have helped the models reproduce small scale features in the flow, doing so proved too memory intensive for our GPUs.

Zoom In Zoom Out Reset image size

We factorized the weight tensors within the spectral layers with TensorLY to improve generalizability. Specifically, we used a canonical polyadic tensor decomposition [45] with a rank of 0.5 to perform this factorization. Prior to producing the outputs, the models employ a fully connected layer of width 128.

In addition, the model normalizes its data internally by dividing by a constant input normalization factor to ensure that magnitude varying inputs are treated the same way. In particular, the magnitude of the magnetic potential A was considerably less than that of the velocity fields u and v. Similarly, we multiplied the output by a constant output normalization factor to alleviate the tFNO model from having to produce results with significantly different magnitudes. For this work, we used the same value for the input and output normalization fields of 1 for the velocity fields and $0.000\,25$ for the magnetic potential A. Finally, we selected the Gaussian error linear unit for the nonlinearity of these models [46].

To accelerate the training, we would employ transfer learning across different Reynolds number Re. Specifically, we began by training a model at a resolution $N = 128^2$ at $\mathrm{Re} = 100$ from scratch. The checkpoint generated from this run was used as the starting point for all the other models. In turn, this transfer learning saved considerable time when training new models. We trained models initialized from the aforementioned checkpoint at resolution of $N = 128^2$ for $\mathrm{Re}$ values of $\mathrm{Re} = \{100, 250, 500, 750, 1000, 10\,000\}$. We then compared these models to additional ones trained at resolutions of $N = 64^2$ and $N = 256^2$. These models were for the same $\mathrm{Re}$ as the previous $N = 128^2$ resolutions, except for $\mathrm{Re} = 10\,000$. All models trained using 950 simulations that encompassed the training data and were evaluated using the remaining 50 simulations that served as the test data. In figure 2, we display some loss curves for the $\mathrm{Re} = 100$ model that was trained from scratch and for the $\mathrm{Re} = 250$ model both at resolutions of $N = 128^2$.

Zoom In Zoom Out Reset image size

For most models, we trained on all available timesteps. However, we wanted to see the effect of being asked to output fewer time steps had on the model. Therefore, we trained additional models that skipped several timesteps. These trained at $\mathrm{Re} = 250$ and $N = 128^2$ and are described further in section 5.

We trained these models using the PyTorch deep learning framework [47]. All models in this study trained for 100 epochs starting from the pretrained checkpoint. Most models used an initial learning rate of $5\times 10^{-4}$. However, the models at $N = 128^2$ with $\mathrm{Re}$ values of 750, 1000, and $10\,000$ had initial learning weights of 0.001. The different initial learning rate for these runs was unintentional. We suspect this higher may have improved the performance of these runs based on how they perform compared to the same $\mathrm{Re}$ runs at other resolutions, but no rigorous study of learning rate optimization was performed. We employed a scheduler to decrease the learning rate by a factor of 2 every 20 epoch to help finetune the models. We selected an AdamW optimizer [48] to optimize the models. This optimizer had a weight decay value of 0.1, $\beta_1 = 0.9$, $\beta_2 = 0.999$.

We set most of the weight hyperparameters of the various loss terms to 1 with a few notable exceptions. We set the weight of the magnetic potential evolution equation loss w_DA to 10⁶ as the magnitude of the term was small compared to that of the other evolution equations. In addition, we used a value of the constraint loss weight w_constr of 10. Finally, we selected $w_{\textrm{data}} = 5$ as our data loss weight. We logged specific hyperparameters used in training each PINO using WandB [49] for reproducibility. In addition, we ran an experiments where we modified the PDE loss weight w_PDE hyperparameter to explore its impact on our AI models. These were done at $\mathrm{Re} = 250$ and $\mathrm{Re} = 500$ both at $N = 128^2$ and are described further in section 5.

We initially computed a subset of the results on the Pittsburgh Supercomputing Center's Bridges-2 cluster. Specifically, we used the V100 GPUs on this cluster. We primarily used the 16 GB variant of these GPUs, though some cases utilized the 32 GB variation. To generate the training data, for the initial portion of the study, we utilized the CPUs on the Bridges-2 cluster's GPU nodes. We used a single GPU node for this process, which has two Intel Xeon Gold 6248 'Cascade Lake' CPUs, which have 20 cores, 2.50–3.90 GHz, 27.5 MB LLC, 6 memory channels each. After verifying that our models worked as expected, we scaled up our experiments using NCSA's Delta cluster. For training, we employed the cluster's NVIDIA A100 GPUs, training each model on a single GPU. To generate an expanded quantity of training data, we employed the Delta cluster's CPU nodes. These come equipped with 128 AMD EPYC 7763 'Milan' (PCIe Gen4) CPUs. We parallelized the generation of the training data such that multiple simulations were generated at one, but each only using one core. We considered low ($N = 64^2)$, standard ($N = 128^2)$, and high resolution ($N = 256^2)$ simulations.

To evaluate the performance of our PINO models, we look at the relative MSE. This allows us to compare errors of the various fields despite them differing in magnitude. We report the relative MSE for the predictions the PINOs on the test dataset for each field of interest. In addition, we compute the total relative MSE, $\textrm{MSE}_\mathrm{tot}$, for our PINOs on this data which we defined as

$$\begin{align} \textrm{MSE}_\mathrm{tot} = \textrm{MSE}_u + \textrm{MSE}_v + \textrm{MSE}_A, \end{align} \tag{ 14 }$$

where $\textrm{MSE}_u$, $\textrm{MSE}_v$, and $\textrm{MSE}_A$ are the MSE values of the u, v, and A fields respectively. In addition, we computed the kinetic energy spectra $E_\mathrm{kin}(k)$ and magnetic energy spectra, $E_\mathrm{mag}(k)$, of the PINO predictions and the ground truth simulations. This allowed us to compare how the models perform at various scales as specified by the wavenumber k.

To begin with, in the left panel of figure 3 we present how the data loss MSE varies with $\mathrm{Re}$ for $N = 128^2$. Therein we see the loss of each field $(u, v, A)$ as well as the total loss as a function of $\mathrm{Re}$. We observe that the two velocity fields u and v possess around the same MSE value as one would expect. In contrast, the MSE for the magnetic potential field A is larger than that of the velocity fields. Moreover, the MSE value of the A field increases faster compared to those of the velocity fields u and v. In the right panel of figure 3 we see the total MSE as a function of $\mathrm{Re}$ for three different resolutions, $N = \{64^2,128^2, 256^2\}$. If we combine these two sets of results, we realize that our AI surrogates can produce reliable MHD simulations for $\mathrm{Re} \leqslant 250$, and that the main factor that degrades the accuracy of our AI surrogates is the ability to capture detailed features of the vector potential A in turbulent flows. We also notice that the resolution of the training data does not have a major impact in the accuracy of our AI surrogates.

Zoom In Zoom Out Reset image size

Figure 4 illustrates how the MSE changes with PDE weight, w_PDE, for $\mathrm{Re} = 250$ and $\mathrm{Re} = 500$. We observe that the MSE improves as we increase w_PDE. The only exception of $\mathrm{Re} = 500$ and $w_{\textrm{PDE}} = 10$ which increases in MSE compared to $w_{\textrm{PDE}} = 5$, but still less than the other w_PDE values at $\mathrm{Re} = 500$. However, we should observe that while including violations of the PDE into the loss function improves the result, their contribution is only marginal.

Zoom In Zoom Out Reset image size

Figure 5 illustrates, for the $\mathrm{Re} = 250$ model, how the total MSE varies with the number of timesteps, t_step. We observe that the fewer timesteps the model is required to output data for (higher $t_\mathrm{step}$), the better the performance of the model—though this improvement is marginal.

Zoom In Zoom Out Reset image size

We now compare traditional MHD simulations with predictions from AI surrogates assuming a grid of size $N = 128^2$, and $t_\mathrm{step} = 1$. In figures 6–11 we present snapshots of the systems' evolution at time t = 1 to show the largest discrepancy between ground truth values (traditional MHD simulations) and AI predictions. These results illustrate ICs, targets, prediction, and MSEs between target and AI predictions. At a glance, we observe that AI models appear to possess the ability to resolve large scale features, but struggle to resolve cases with small scale structure. This small scale structure arise at higher $\mathrm{Re}$ and is observed most strongly in the A field. We summarize below the main findings of these studies using table 1, figures 3 and 6–11:

• $\mathrm{Re} = 100$. Figure 6 shows that PINO models provide an accurate description of these MHD simulations. Quantitatively and quantitatively PINOs can resolve the dynamics of the velocity field, (u, v), and the vector potential, A. We also notice that the largest discrepancy between AI predictions and traditional MHD simulations at t = 1 is ${\leqslant} 4\%$ for each of the fields.
• $\mathrm{Re} = 250$. Figure 7 shows that PINOs provide a reliable description of the dynamics of these simulations. The velocity field and vector potential potential are accurately described, with MSEs ${\leqslant} 7\%$ and ${\leqslant} 10\%$, respectively.
• $\mathrm{Re} = 500$. Figure 8 shows that PINOs capture well large scale features of these simulations. In particular, the velocity field can be recovered with MSE ${\leqslant} 14\%$. On the other hand, detailed features of the vector potential are not completely resolved, and the MSE is ${\leqslant} 26\%$.
• $\mathrm{Re} = 750$. Figure 9 shows that PINO MHD simulations can resolve large scale features of the velocity field, with MSE ${\leqslant} 16\%$, which is similar to simulations with $\mathrm{Re} = 500$. Similarly, large scale features of the vector potential are well described. However, as these systems become more turbulent, small scale features of the vector potential are not captured by PINOs, and report an increase in MSE, i.e. ${\leqslant} 31\%$.
• $\mathrm{Re} = 1000$. Figure 10 presents a similar story to the two previous cases. Large scale structure is well described, and the MSE for the velocity field is ${\leqslant} 18\%$. However, it is difficult to capture the small scale features of the magnetic field, which now evolve in a rather complex manner for these turbulent systems.
• $\mathrm{Re} = 10\,000$. Figure 11 indicates that, for very turbulent MHD systems, PINOs can only reproduce some of the large scale features of the flow, but struggle to reproduce many of the detailed features. One particular feature that the PINO misses is the small scale fluctuations in the velocity fields u and v that likely result from strong magnetic fields in those areas. In other words, the PINO cannot resolve the effect of the magnetic field on the fluid motion. For the magnetic potential A, the PINO model appears to miss most of the important features at first glance. If we look more closely, we observe that the PINO appears, to some extent, to reproduce the mean field value over some very large regions. However, this comes at the cost of missing any sort of interesting details in the magnetic field.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

In summary, figures 6–11 show that our PINO models can successfully reproduce large scale features of MHD simulations, but had difficulty resolving detailed features for large Reynolds numbers, especially for the magnetic potential, which is known to store its energy at higher wavenumbers. Thus, in the following section we explored the ability of PINOs to learn the right data features associated to low and high wavenumbers as the simulation evolves in time.

We analyzed the magnetic and kinetic energy spectra of the simulations in figure 12. Therein we observe that for both the kinetic and magnetic energy spectra, the PINO models performed well at low wavenumbers k. However, the PINOs were not able to accurately reproduce the energy spectra at high wavenumbers. Interestingly, at low $\mathrm{Re}$, the PINOs tended to overshoot the ground truth energies at high wavenumbers. While at high $\mathrm{Re}$, the PINOs tended to undershoot the ground truth energies at high wavenumbers.

Zoom In Zoom Out Reset image size

Moreover, large Reynolds number simulations usually store more energy at higher wavenumbers. This is especially true for the magnetic energy, which tends to peak at later wavenumbers in higher $\mathrm{Re}$ simulations. Therefore, the spectra plots in figure 12 indicate that the difficulties the PINOs had in reproducing the simulations at high $\mathrm{Re}$ and the magnetic potential field A may stem from their relatively poor performance at high wavenumbers. Therefore, future work should focus on developing methods that enable PINOs to capture data features contained at high wavenumbers.