Neural Operators Learn the Local Physics of Magnetohydrodynamics
Taeyoung Kim Youngsoo Ha Myungjoo Kang






       
Neural Operators Learn the Local Physics of Magnetohydrodynamics
    Taeyoung Kim legend@snu.ac.kr 

       Department of Mathematical Science

       Seoul National University

       Seoul 08826, South Korea
       Youngsoo Ha youngamath@snu.ac.kr 

       Department of Mathematical Science

       Seoul National University

       Seoul 08826, South Korea
       Myungjoo Kang mkang@snu.ac.kr 

       Department of Mathematical Science

       Seoul National University

       Seoul 08826, South Korea
     
===========================================================================================================================================================================================================================================================================================================================================================================================================================================================





Magnetohydrodynamics (MHD) plays a pivotal role in describing the dynamics of plasma and conductive fluids, essential for understanding phenomena such as the structure and evolution of stars and galaxies, and in nuclear fusion for plasma motion through ideal MHD equations. Solving these hyperbolic PDEs requires sophisticated numerical methods, presenting computational challenges due to complex structures and high costs. Recent advances introduce neural operators like the Fourier Neural Operator (FNO) as surrogate models for traditional numerical analyses. This study explores a modified Flux Fourier neural operator model to approximate the numerical flux of ideal MHD, offering a novel approach that outperforms existing neural operator models by enabling continuous inference, generalization outside sampled distributions, and faster computation compared to classical numerical schemes.





§ INTRODUCTION




 §.§ Magnetohydrodynamics

In the mid-20th century, the advent of the Space Age and growing interest in nuclear fusion led to an intensified focus and understanding of plasma. During this period, significant theoretical studies on plasma were conducted, marking early research in kinetic equations by (<cit.>) and studies on magnetohydrodynamics (MHD) by (<cit.>). MHD describes the behavior of electrically conductive fluids and is primarily used for plasmas and liquid metals under the influence of large-scale, low-frequency magnetic fields (<cit.>). The applications of MHD extend across various fields, including astrophysics (<cit.>), solar physics (<cit.>), the study of Earth's magnetosphere (<cit.>), and research in nuclear fusion (<cit.>). Ideal MHD represents one of the simplest forms of these models, disregarding dissipative effects such as viscosity, thermal conductivity, and resistance. Despite its simplicity, ideal MHD proves to be a powerful tool capable of explaining a wide range of plasma phenomena (<cit.>). This includes its use in analyzing plasma within tokamaks and stellarators, understanding the dynamo process generating Earth's magnetic field (<cit.>), studying plasma phenomena inside and on the surface of the sun (<cit.>), and describing the formation of galactic structures and the universe (<cit.>).





 §.§ Numerical Schemes


Magnetohydrodynamics (MHD) can be viewed as a generalized model that combines hydrodynamics with electromagnetism. In particular, ideal MHD is considered as a hyperbolic conservation law, conserving mass, momentum, magnetic fields, and energy density. Numerical methods for solving hyperbolic conservation laws have a long history, with various approaches proposed to achieve stability, computational efficiency, and high-order accuracy. High-order accurate methodologies, such as essentially non-oscillatory (ENO) (<cit.>) and weighted ENO (WENO) (<cit.>) schemes, have been effective in solving hyperbolic conservation laws. Jiang and Shu introduced the WENO scheme with third and fifth-order accuracy (WENO-JS) (<cit.>), noted for its robust shock-capturing capability, albeit with a disadvantage of being dissipative in turbulent flows. To address this shortcoming, Henrick et al. developed the enhanced WENO-M scheme (<cit.>). However, applying WENO schemes to the system of conservation laws involves significant computational costs. Various attempts to reduce computational costs include Jiang and Shu's work (<cit.>), Pirozzoli's hybrid compact-WENO scheme (<cit.>), and the efforts by (<cit.>)(<cit.>) among others. For ideal MHD, maintaining the divergence-free condition for the magnetic field is essential, with methods like the Constrained Transport (CT) method by (<cit.>), and the projection scheme by (<cit.>) proposed. Comprehensive reviews on methods ensuring the divergence-free condition in ideal MHD are available in (<cit.>). Early studies on numerical solutions for ideal MHD include works by (<cit.>), (<cit.>), (<cit.>). And other various efforts to solve ideal MHD numerically have been develeoped (<cit.>)(<cit.>). In this research, we innovatively combine traditional numerical schemes with artificial neural network techniques to reduce computational costs and enforce the divergence-free condition for ideal MHD problems.





 §.§ Neural Operators and Flux Neural Operator

In recent times, several Neural Operators have been proposed as a method to replace conventional numerical analysis techniques with machine learning approaches. Notably, the Graph Kernel Network (<cit.>), Fourier Neural Operator (FNO) (<cit.>), DeepONet (<cit.>), and various adaptations thereof have been introduced (<cit.>)(<cit.>)(<cit.>). A common feature among these models is their ability to handle functional data, as they are not limited by the resolution of input data. This characteristic allows for the approximation of operators—mappings between functional spaces—and paves the way for these models to serve as surrogate models for numerical schemes. Unlike traditional numerical schemes that iteratively calculate solutions, most neural operator models learn global solvers that map initial data directly to the solution data at a specific time. One advantage of this method is its significantly faster computation speed compared to classical numerical schemes (<cit.>); however, it often lacks generalization capability (<cit.>), making it challenging to learn actual broad physical phenomena. The recently proposed Flux FNO model (<cit.>) combines the strengths of both numerical schemes and neural operators. In the respective study, it demonstrated a leap in generalization capability by learning numerical flux for hyperbolic conservation law problems, which enabled the neural operator model to learn local physics. While the study empirically showed the approximation of numerical flux by Flux FNO for one-dimensional scalar conservation laws, our paper takes this a step further by applying it to one of the most challenging problems, the ideal MHD problem.




 §.§ Our contribution

In this paper, we employ various techniques motivated from physical property of equation and numerical analysis to enhance the recently developed machine learning method, Flux FNO, for solving the ideal magnetohydrodynamic (MHD) problem. Firstly, as the problem transitions from scalar-valued to vector-valued, we redesign the Flux FNO model architecture to process each physical variable (density, velocity, magnetic field, energy) in each separated model. This is due to the experimental finding of limitations in the model's expressiveness when handling all variables simultaneously. Secondly, we design and apply a loss function that endows the approximated numerical flux with the Total Variation Diminishing (TVD) property to ensure stability. Unlike classical numerical schemes, if we only approximate the numerical flux without such measures, severe oscillations could arise as iterations accumulate. Furthermore, we enforce the divergence-free condition by taking loss over divergence of magnetic field variables. This loss allows us to impart an inductive bias suitable for the ideal MHD solver, which we verify through an ablation study. Utilizing these methodologies, we base our model to qualitatively assess generalization performance (inference over continuous time, inference on out-of-distribution samples) and to solve representative test problems of ideal MHD, comparing them with traditional numerical methodologies.



§ PRELIMINARIES




 §.§ Ideal MHD Equations

Ideal Magnetohydrodynamics (MHD) is described by a system of coupled partial differential equations that characterize conducting fluids. Among various formulations of MHD, the ideal MHD equations represent the simplest form, embodying a synthesis of fluid dynamics and Maxwell's equations, while excluding effects such as viscosity, resistance, and thermal conductivity. The conservative form of ideal MHD can be articulated as follows:


    ρ_t+∇·(ρu)   =0
    
        (ρu)_t+∇·[ρu⊗u+(p+1/2B^2)I-B⊗B]   =0 
    B_t+∇·(u⊗B-B⊗u)   =0 
    ℰ_t+∇·[(ℰ+p+1/2B^2)u-B(u·B)]   =0

Additionally, from Maxwell's equations, we require the divergence-free condition for the magnetic field, which is expressed as:


    ∇·B=0

In the scenario where the fluid is considered incompressible, an analogous divergence-free condition for velocity would be necessary. However, throughout this paper, we focus on compressible fluids, where such a condition for velocity is not explicitly required due to the fluid's capacity to vary in density. The pressure and energy is coupled through the following equation:

    p=(γ-1)(ℰ-1/2ρu^2-1/2B^2)

where γ is the ratio of specific heats.

One dimensional ideal MHD
The governing equation for ideal MHD in the one-dimensional case, expressed in conservative form, can be written as follows:


    ∂/∂ t[     ρ; ρ u_x; ρ u_y; ρ u_z;   B_y;   B_z;     ℰ ] +
        ∂/∂ x[                 ρ u_x; ρ u_x^2 + p^* - B_x^2;       ρ u_xu_y-B_xB_y;       ρ u_xu_z-B_xB_z;         B_yu_x-B_xu_y;         B_zu_x-B_xu_z;   (ℰ+p^*)u_x-B_x(u·B) ]
        =0

where p^*=p+B^2/2 represents the total pressure, incorporating magnetic pressure. Each variable with a subscript denotes components of velocity and magnetic field.


Two dimensional ideal MHD
For two-dimensional case, the equtions are expressed in conservative form, can be written as follows:


    ∂/∂ t[     ρ; ρ u_x; ρ u_y; ρ u_z;   B_x;   B_y;   B_z;     ℰ ] +
        ∂/∂ x[                 ρ u_x; ρ u_x^2 + p^* - B_x^2;       ρ u_xv_y-B_xB_y;       ρ u_xv_z-B_xB_z;                     0;         B_yu_x-B_xu_y;         B_zu_x-B_xu_z;   (ℰ+p^*)u_x-B_x(u·B) ] +
        ∂/∂ y[                 ρ u_y;       ρ u_yu_x-B_yB_x; ρ u_y^2 + p^* - B_y^2;       ρ u_yu_z-B_yB_z;         B_xu_y-B_yu_x;                     0;         B_zu_y-B_yu_z;   (ℰ+p^*)u_y-B_y(u·B) ]
        =0







 §.§ Numerical schemes

To apply our methodology, a dataset consisting of discretized functions over time is required. Although it is possible to generate the dataset using actual observations, we created both training and testing datasets using classical numerical analysis techniques. This section describes the numerical methods utilized to generate the datasets, along with concepts related to the stability of numerical solutions.

The WENO schemes
When using high-order numerical schemes, the phenomenon of oscillations at discontinuities, known as Gibbs phenomena, occurs. To address this issue, methods such as flux limiters, essentially non-oscillatory schemes (ENO), and slope limiters have been devised. Among these methods, the Weighted Essentially Non-Oscillatory (WENO) scheme stands out as it reconstructs the function values in a non-oscillatory manner by nonlinearly weighting each sub-stencil, utilizing the given function values. As an example, let's apply this to the one-dimensional conservation laws:


    u_t+f(u)_x=0,   x∈ℝ,   t≥0. 1

Let x_0<…<x_n be the uniform discretization of the computational domain. where x_j+1/2=x_j+x_j+1/2. The equation (<ref>) can be approximated with semi-discrete conservation schemes:


    du_j/dt=-∂ f/∂ x|_x=x_j2

where u_j(t) is numerical approximation of function u(x_j,t) on a grid. And by approximating right term of (<ref>) in a conservative manner, we get following formula:


    du_j/dt=-f̂_j+1/2-f̂_j-1/2/Δ x3

where f̂_j±1/2 is a numerical flux which satisfies lipschitz continuity and consistency with the physical flux f, namely, f̂(u,…,u)=f(u).
In fifth-order finite differnece WENO scheme (WENO-JS), the numerical flux is constructed using 5-point stencil which is subdivided into three sub-stencils. The computation of WENO-JS is as follows:


    f̂_j+1/2=∑_k=0^2ω_kf̂_k,j+1/2.

which is weighted sum of numerical fluxes f̂_k,j+1/2(k=0,1,2) which are as follows:


    f̂_0,j+1/2=1/3f_j-2-7/6f_j-1+11/6f_j,
    f̂_1,j+1/2=-1/6f_j-1+5/6f_j+1/3f_j+1,
    f̂_2,j+1/2=1/3f_j-5/6f_j+1-1/6f_j+2.

The nonlinear weights ω_k are contingent upon the variant WENO construction techniques, with WENO-JS and WENO-Z being predominantly utilized. The detailed calculation of these weights can be found in (<cit.>) and (<cit.>).


Total Variation Diminish Runge-Kutta method
The total variation of numerical solution is defined as follows:


    TV(u)=∑_j|u_j+1-u_j|

And, we say numerical scheme has total variation diminishing (TVD) property if it satisfies following condition:


    TV(u^n+1)≤ TV(u^n) 4



Let denote -L(u) is approximation of spatial derivative f(u)_x in (<ref>) then a general Runge-Kutta method for (<ref>) can be written as follows:


    u^0   = u^n, 
    
        u^i   = ∑_k=0^i-1(α_iku^(k)+Δ t β_ikL(u^(k))),   i=1,…,m, 5
    
        u^n+1   = u^m.

It is known that Ruge-Kutta method (<ref>) is TVD under the following conditon (Courant-Friedrichs-Lewy (CFL) condition):


    λ   ≤λ_0min_i,kα_ik/|β_ik|,
    λ   = Δ t/Δ x.

where λ_0 is a suitable CFL restriction.




 §.§ Flux Neural Operator

Fourier Neural Operator



Among the various types of neural operators for handling functional data, the Fourier Neural Operator (FNO) is a neural operator model that processes the global convolution operations of functions by leveraging the Fourier transform. Its most notable feature is that the convolution with the kernel function is expressed through the Fourier transform as multiplication with a high-order tensor. The structure of the FNO is depicted in Figure <ref>, where the values of discretized functional data are non-linearly lifted by the lifting layer, passed through the Fourier layers, and finally projected into the desired dimension vector by the projection layer composed of FNCs. As shown in diagram (b) of Figure <ref>, each Fourier layer consists of a function convolution operation that globally transforms the data and an auxiliary neural network, which must maintain resolution invariance, with CNN layers commonly adopted to handle local data processing. The mathematical formulation for a Fourier Neural Operator (FNO) with a CNN layers can be represented as follows:

Definition (FNO with CNN layers) 

v_0   :=𝒩_P(a|_X)=(𝒩_P(a_x·)_j)_x∈X, j=1,…,d_v_0 

v_t+1   :=𝒜_t+1(v_t)=σ(C_t+1(c_1,…,c_d)(ṽ_̃t̃)+ℱ^-1(R_t+1·(ℱ(v_t))))  

   =σ(∑_k=1^d_u∑_j_1=0^c_1-1⋯∑_j_d=0^c_d-1 K_t+1,jk,j_1,…,j_dṽ_t,x_1+j_1,…,x_d+j_d,k 


   + ∑_z,k∈K,k  D^†_xkR_t+1,k,jkD_kzv_t,zk)   (t=0,...,L-1) 

G(a;θ)   :=𝒩_Q(v_L)=(𝒩_Q(v_Lx·)_j)_x∈X, j=1,…,d_v_L

 
where 𝒩_P and 𝒩_Q are neural networks used for lifting and projection, respectively. a|_X represents the discretized functional data of a, and each C_i is a d-dimensional CNN layer with a kernel tensor K_i. The R_i tensors are weight tensors that parameterize the kernel function of the Fourier layers. D_kz represents the components of the discrete Fourier transform, also denoted as ℱ.



Flux Neural Operator




The Flux Fourier Neural Operator (Flux FNO) introduced in the paper (<cit.>), inspired by the numerical schemes of hyperbolic conservation laws, represents a method designed to learn the local physics of conservation laws directly through the approximation of fluxes, differing from original neural operators that predict target snapshots of solutions. The area where FNO is applied within Flux FNO can be substituted with other Neural Operators, leading to the generalized concept referred to as the Flux Neural Operator (Flux NO). The schematic of Flux NO for domain of one-dimensional case is shown in Figure <ref>, and unlike conventional neural operators, Flux NO aims to approximate the flux itself and calculates local residuals based on this approximation, offering a novel approach to handling local physical phenomena. For an N-dimensional problem, the Flux NO operates as described in formula (<ref>). Each G_i(·;θ_i) is approximated flux function along spatial direction x_i (θ_i is parameter of neural operator). And each U_j^l_i,k, U_j^r_i,k denote shifted U_j along spatial direction (which is corresponding spatial index).


    U_j+1 = U_j + ∑_i=1^kΔ t/Δ x_i[G_i(U_j^l_i,1,…,U_j^l_i,m;θ_i)-G_i(U_j^r_i,1,…,U_j^r_i,m;θ_i)] 6
 
To train the Flux NO, the required training dataset shape is [batch size, N_t, N_x_1, ⋯, N_x_k, N_u] where N_t represents the number of interations along time, N_x_i represent the number of grid points along each spatial dimensions across which the problem is defined, and N_u indicates the dimensionality of the problem (e.g., velocity components in fluid dynamics). Given this structured dataset, the loss function motivated from equation (<ref>) for training Flux NO can be constructed as follows:


    ℒ_tm({U_i}_i=1^B)   :=∑_i=1^B∑_j=1^N_t-1(U_i,j+1 - U_i,j) 
       -(∑_i=1^kΔ t/Δ x_i[G_i(U_i,j^l_i,1,…,U_i,j^l_i,m;θ_i)-G_i(U_i,j^r_i,1,…,U_i,j^r_i,m;θ_i)])^27
 
This loss function quantifies the difference between residual constructed from the predicted flux values by the Flux NO and the actual (or target) residual values, aiming to minimize this discrepancy during the training process. By focusing on flux approximation, the Flux NO offers a detailed and localized understanding of the underlying physical processes, making it especially suitable for problems governed by conservation laws where the flux plays a critical role. And we consider additional loss which guarantees the consistency of Flux function which is essential for convergence to weak solutions:


    ℒ_consi({U_i}_j=1^B):=∑_i=1^B∑_j=1^N^t∑_l=1^kG_l(U_i,j,…,U_i,j;θ_l)-F_l(U_i,j)^28
  
where F_i are actual physical fluxes. The actual training is implemented by optimizing weighted sum of these losses: ℒ_tm+λℒ_consi. These losses are presented also in (<cit.>), in this work, we consider more additional losses which make the convergnce of  training well and approximate propery with robustness and more generalization ability.




§ METHODS



 §.§ Adaptive FNO Architecture for Multidimensional Outputs in Ideal MHD

In the referenced study (<cit.>), the output dimension was one-dimensional since it dealt with one-dimensional scalar conservation laws. However, in the case of ideal Magnetohydrodynamics (MHD) discussed in this paper, the output dimension is seven-dimensional for one-dimensional problems and eight-dimensional for two-dimensional problems. We experimentally observed significant issues with expressiveness when a single FNO model was tasked with handling all outputs in cases of large output dimensions. To address this issue, we allocated separate FNO models to handle each physical quantity. Specifically, density, energy, velocity vectors, and magnetic field vectors were each processed by their own FNO models. For two-dimensional problems, given that the physical flux functions exist for both the x and y axes, we further segregated the considerations for each physical flux. This architectural approach is summarized in the schematic shown in the Figure <ref>. Modified version of (<ref>) for one-dimensional case can be written as follows:




    U_j+1   = U_j + Δ t/Δ x[G^l-G^r]  
    
        G^l   =(G_ρ(U_j^l_1,…,U_j^l_m;θ_ρ), 
        G_u(U_j^l_1,…,U_j^l_m;θ_u)^T,  
       G_B(U_j^l_1,…,U_j^l_m;θ_B)^T, 
        G_E(U_j^l_1,…,U_j^l_m;θ_E))^T9
    
        G^r   =(G_ρ(U_j^r_1,…,U_j^r_m;θ_ρ), 
        G_u(U_j^r_1,…,U_j^r_m;θ_u)^T,  
       G_B(U_j^r_1,…,U_j^r_m;θ_B)^T, 
        G_E(U_j^r_1,…,U_j^r_m;θ_E))^T

And the two-dimensional case can be written as follows:



    U_j+1   = U_j + Δ t/Δ x[G^l-G^r] + Δ t/Δ y[F^t-F^b]  
    
        G^l   =(G_ρ(U_j^l_1,…,U_j^l_m;θ_ρ^G), 
        G_u(U_j^l_1,…,U_j^l_m;θ_u^G)^T,  
       G_B(U_j^l_1,…,U_j^l_m;θ_B^G)^T, 
        G_E(U_j^l_1,…,U_j^l_m;θ_E^G))^T
    
        G^r   =(G_ρ(U_j^r_1,…,U_j^r_m;θ_ρ^G), 
        G_u(U_j^r_1,…,U_j^r_m;θ_u^G)^T,  
       G_B(U_j^r_1,…,U_j^r_m;θ_B^G)^T, 
        G_E(U_j^r_1,…,U_j^r_m;θ_E^G))^T10
    
        F^t   =(F_ρ(U_j^t_1,…,U_j^t_m;θ_ρ^F), 
        F_u(U_j^t_1,…,U_j^t_m;θ_u^F)^T,  
       F_B(U_j^t_1,…,U_j^t_m;θ_B^F)^T, 
        F_E(U_j^t_1,…,U_j^t_m;θ_E^F))^T
    
        F^b   =(G_ρ(U_j^b_1,…,U_j^b_m;θ_ρ^F), 
        F_u(U_j^b_1,…,U_j^b_m;θ_u^F)^T,  
       F_B(U_j^b_1,…,U_j^b_m;θ_B^F)^T, 
        F_E(U_j^b_1,…,U_j^b_m;θ_E^F))^T
 
Let the residual in (<ref>) and (<ref>) be ΔR(U_j;θ):=U_j+1-U_j where θ is collection of parameters. Then the loss (<ref>) can be now written as follows:



    ℒ_tm({U_i}_i=1^B)   :=∑_i=1^B∑_j=1^N_t-1(U_i,j+1 - U_i,j) - ΔR(U_i,j;θ)^211
 








 §.§ Enhanced Loss Function for Ideal MHD

We have developed additional loss functions to enhance the existing loss function introduced in (<cit.>), and to impart suitable inductive biases for specialized issues such as ideal MHD, which is the focus of this paper. Three further loss functions have been considered; one associated with the Total Variation Diminishing (TVD) property, another related to the divergence-free condition, and a third concerning local information about the flux.

TVD Loss
According to Godunov's theorem, the WENO method, which performs high-order polynomial approximations for all stencils, does not inherently possess the Total Variation Diminishing (TVD) property. Nevertheless, WENO exhibits a considerably non-oscillatory manner. The results we will present in Section 4 are based on models trained using training data generated by the high-order WENO method. Although the training data generated through the WENO method do not satisfy the TVD property, we have designed a loss function that imparts an inductive bias towards the TVD property to our model, inspired by equation (<ref>). Since the training data inherently lack the TVD property, we adopted an l^2 rather than l^1 function as a loss to apply the inductive bias in a soft manner. The specific formulation of this loss is as follows:


ℒ_TVD({U_i}_i=1^B):=∑_i=1^B∑_n=1^N_t-1⌊TV(Ũ_i,n+1)-TV(U_i,n)⌋_+^2   12 

where ⌊·⌋_+:=max(0,x) and Ũ is an output of Flux NO.


Divergence Free Loss Since the magnetic vector field in multidimensional ideal MHD is divergence-free, our approximated fluxes should also construct a divergence-free vector field. To achieve this, we have designed a loss function that imparts this condition. This approach can also be applied to other problems, such as incompressible fluids, which require a divergence-free condition for the velocity vector field.
Each of F and G represents the physical fluxes for the two-dimensional ideal MHD equations, as introduced in Section 2.1. The subscript B in (<ref>) indicates that the components of the vector are specifically related to the magnetic vector fields.


    ∇·∂U/∂ t=  ∂∇·U/∂ t = ∇·(-∂G/∂ x -∂F/∂ y) 
       ⇒ 0= ∂∇·B/∂ t  = ∇·(-∂G/∂ x -∂F/∂ y)_B13
 

According to equation (<ref>), the residual of the magnetic field must also satisfy the divergence-free condition. Therefore, we will design a loss function to ensure that the residuals constructed from our approximated flux satisfy this condition. However, the training datasets generated through numerical analysis may still exhibit non-zero divergence of the magnetic field, despite the application of divergence-free relaxation. Consequently, instead of implementing a hard constraint that forces the divergence-related loss to converge nearly to zero, or requiring the model architecture to output the residual of the magnetic field as the curl of a specific vector field, we empirically calculate the average or maximum values of the magnetic field's divergence in the training dataset. Using this information, we set a threshold such that the l^2 loss related to the divergence is applied only when the model's output exceeds this threshold. The specific mathematical formulation of this loss is as follows:


ℒ_div({U_i}_i=1^B):=∑_i=1^B∑_j=1^N_t-1 ⌊∇·ΔR(U_i,j;θ)_B_2^2-θ_div/| ∇·ΔR(U_i,j;θ)_B_2^2-θ_div|⌋_+∇·ΔR(U_i,j;θ)_B _2^2   14 

where θ_div is the threshold value.


l^∞ Loss
In addition to the l^2 norm of the residuals proposed in (<cit.>), we also consider the l^∞ norm, which is used to handle outlier values (<cit.>). Utilizing the l^∞ norm allows for the expectation of pointwise convergence, thereby enabling a more accurate approximation of the flux embedded in the training dataset. In cases like ideal MHD where the output is vector-valued, we apply the l^∞ norm to each vector component and then sum these values. The specific mathematical expression for this is as follows:

ℒ_∞({U_i}_i=1^B):=∑_i=1^B∑_n=1^N_t-1∑_l=1^N_usup_j,k ((U_i,n+1,j,k,l - U_i,n,j,k,l) - ΔR(U_i,n,j,k,l;θ))   15 








 §.§ Overall Algorithm
 
Now, by combining the defined loss functions in Equations (<ref>), (<ref>), (<ref>), (<ref>), and (<ref>), we demonstrate the training algorithm (Algorithm <ref>) for our modified Flux Neural Operator (Flux NO). Since inference is straightforwardly conducted using Equation (<ref>), we omit the details. We focus on the two-dimensional case, as the training process for the one-dimensional case is similar to the algorithm proposed in (<cit.>), with the exception of the enhanced loss function.










§ RESULTS
 


 §.§ Training Dataset and Architecture of Flux Neural Operator

One-dimensional case
For the one-dimensional problem, we constructed the initial conditions using Gaussian random fields. Specifically, we set the initial values for density and the B_x to 1, and specified the velocity and magnetic field components as vectors discretized to size 256, derived from Gaussian random fields with a covariance of k(x,y)=e^-100(x-y)^2. To ensure the well-posedness of the problem, each component of the velocity and magnetic field was scaled by 0.3. Given that the γ value was set to 2.0, the initial value of the energy is determined by the following equation: E=1.0+1/2ρu^2+1/2B^2. Based on these initial conditions, we solved the problem using the WENO-JS scheme for spatial order 5 and the RK method of order 4 for temporal discretization. The CFL number was 0.3. Based on this scheme, we divided the solution from t=0.0 to t=0.6 into 600 segments and then extracted snapshots only from t=0.1 to t=0.6 to form a dataset consisting of 500 snapshots for each function. We generated 500 tensors representing functions in C^∞([0.1,0.6]×[0,1]) to create the training dataset. In the model architecture, the Neural Operator component fundamentally employs the Fourier neural operator; however, as FNO performs poorly with non-periodic functions, all datasets were composed as periodic functions. We anticipate finding suitable models for non-periodic functions as well, since the Neural Operator component can be replaced with other models. For the one-dimensional problem discussed in this paper, the model’s architecture uses 5 modes, width 64, and depth 3, with output dimensions of 1, 3, 2, and 1 for density, velocity, magnetic field, and energy, respectively. The kernel size of the CNN layer in the Fourier layer is uniformly set to 1.




Two-dimensional case
For the two-dimensional problem, initial conditions were generated using two-dimensional Gaussian random fields. We independently sampled the components of velocity and magnetic field vectors from Gaussian random fields. And set the initial density value to γ^2. The Gaussian random fields we used have a power spectrum P(k)∝ k^-2.5. Unlike the one-dimensional case, in two dimensions, the divergence of the magnetic field must be zero for physically meaningful results and for the numerical scheme to be stable; therefore, the randomly sampled magnetic fields were relaxed using the Poisson equation (<cit.>). The grid size for the initial conditions was set to 64 along both the x and y axes, and the γ value was set to 5/3. The initial energy value will be determined by the following equation: E=5/2+1/2ρu^2+1/2B^2. Based on these initial conditions, we used the WENO-Z scheme for spatial order 5, temporal discretization of fourth order RK method, and set the CFL numbers to 0.4. Under this scheme, we divided the solution from  t=0.0 to t=0.75 into 150 segments, and then extracted snapshots only from t=0.25 to t=0.75 to form a dataset consisting of 100 snapshots per function. In this way, we created 100 tensors, each representing a function in C^∞([0.25,0.75]×[0,2π]^2) for training. The model architecture adopted a 2D FNO for the neural operator component, with architecture hyperparameters including 4 modes each for x and y axes, width 72, and depth 2. The output dimensions for models concerning density, velocity, magnetic field, and energy are 1, 3, 3, and 1, respectively. The kernel size of the CNN layer in the Fourier layer is uniformly set to 1. Both one-dimensional and two-dimensional problems use the GELU activation function. Each neural operator takes concatenated tensors composed of eight tensors, which are shifted circularly from -3 to 4 (where an i shift means j → j+i) for G^l and F^t, and from -4 to 3 for G^r and F^b for corresponding indices.









 §.§ Generalization Ability

The results in sections 4.2 to 4.4 were obtained using models trained under the following hyperparameters. For the one-dimensional case: the optimizer was Adam with a learning rate of 1e-3 and weight decay of 1e-3. The scheduler was CosineAnnealingWarmRestarts with T_0=100 and eta_min=1e-4. The batch size was 5, λ_tm=1.0, λ_TVD=5e-3, λ_∞=1.0, and λ_consi=1.0. For the two-dimensional case: the optimizer was Adam with a learning rate of 4e-4 and weight decay of 2e-4. The scheduler was CosineAnnealingWarmRestarts with T_0=100 and eta_min=1e-5. The batch size was 1, λ_tm=4.0, λ_TVD=2.5e-4, λ_div=1e-2, λ_∞=1.0, and θ_div=150. Where terms such as T_0 and eta_min follow the conventions used in PyTorch. In this section, we analyze the results of inferences made by our model on examples that are within the test distribution but were not utilized as training data, both qualitatively and quantitatively. Additionally, we experimentally demonstrate that our model can make inferences over longer time periods than those covered by the functions used in the training dataset for the test samples, and we also analyze results under conditions of higher resolution.

Short term inference


We conducted experiments on test samples that were sampled from the same distribution as the training dataset and had the same short-term span lengths as the training samples (t=0.1 to 0.6 for the one-dimensional case, t=0.25 to 0.75 for the two-dimensional case). The results for the one-dimensional case are presented in Figure <ref> and Table <ref>, while the results for the two-dimensional case are shown in Figure <ref>, <ref>, <ref>, <ref> and Table <ref>. Tables <ref> and <ref> provide statistics on the relative l^2 norm and l^∞ norm across 10 test samples. Due to the smaller domain size in the one-dimensional case compared to the two-dimensional case, the changes in the solutions are observed to be more dynamic. In addition, in the one-dimensional case, the solution generated by Flux NO tends to lose stability more easily because the grid size of 256 (compared to 64 in each axis for the two-dimensional case) makes it easier to violate the CFL condition. For the one-dimensional problem, the instability of Flux NO is such that only results up to t=0.2 are shown in Table <ref>. From Table <ref>, it can be seen that the norms related to velocity are the largest and increase the fastest, which can be attributed to the poorer approximation of the u_y and u_z components, while u_x is well approximated compared to these, as indicated by Figure <ref>. Consequently, although the training dataset for the one-dimensional problem is richer, consisting of 500 functions compared to 100 for the two-dimensional problem, it is expected to have lower stability and generalization power due to the reasons mentioned above. For the two-dimensional problem, the performance on test samples within the short-term training distribution appears to be satisfactory, as can be seen from Figure <ref>, <ref>, <ref>, <ref> and Table <ref>, indicating that the solution for each component has been well resolved.
























Long term inference
We conducted experiments over time spans significantly longer than the time domains of the functions used for training, with the initial conditions configured from the same distribution as training dataset. Due to instability issues, long-term experiments were not conducted for the one-dimensional case; instead, they were carried out only for the two-dimensional case, with the results presented in Figure <ref>, <ref>, <ref>, <ref> and Table <ref>. The training functions' time domain ranged from t=0.25 to t=0.75, with a fixed time interval of delta Δ t=0.005, allowing for the terminal point to be reached after 100 iterations. However, in our results, using the same time intervals, the iterations extended to much longer times: t=1.25 (200 iterations), t=1.5 (250 iterations), t=1.75 (300 iterations), and t=2.0 (350 iterations). As shown in Figures <ref> and <ref>, the error increases almost linearly with the number of iterations, affecting both the global and local characteristics of the solution and leading to progressively increasing discrepancies with the reference, as qualitatively evident in Figures <ref>, <ref> and <ref>. A notable aspect of the long-term inference for the two-dimensional problem is that, unlike the one-dimensional case, even with many more iterations, the solutions output by Flux NO did not blow up.

















High resolution inference
We conducted tests on samples at a higher resolution (96 by 96) than the resolution of the training data (64 by 64). Initial conditions were sampled from a Gaussian random field with the same power spectrum as the training data, and references were computed using the WENO-Z scheme with each snapshot spaced by a Δ t=0.005. The test results are presented in Figures <ref>, <ref>, and <ref>; trends in errors, which showed no significant difference from the original resolution, have been omitted. As seen in Figures <ref> and <ref>, while Flux NO captures the global characteristics of the solution, local instabilities gradually appear in the snapshots at t=1.5 and t=2.0. This instability, similar to that observed in the one-dimensional case, appears to arise from the higher resolution and is somewhat related to the CFL condition. Attempts to more easily satisfy the CFL condition by reducing the interval to Δ t=0.0025 for denser inference were made, yet instabilities still formed. From these experimental results, we hypothesize that as the resolution of vectorized functions increases, the domain dimension of their data distribution expands, making it challenging to approximate the data distribution itself, leading to these outcomes.
















 §.§ Tests on Out of Distribution Samples


In this section, we qualitatively evaluate the performance of our model given initial conditions outside the training distribution. We intended to apply the Brio-Wu problem, a type of Riemann problem, but since we assumed periodic boundary conditions for the one-dimensional case, we slightly modified the original Brio-Wu problem and set the initial conditions as follows:

(ρ,u_x,u_y,u_z,B_y,B_z,p)_t=0=

(1,0,0,0,1,0,1),     -0.5<x<0.5 

(0.125,0,0,0,-1,0,1),    otherwise.
 


Where B_x=1. Under these conditions, we performed calculations using the numerical flux based on Flux NO, and the results are shown in Figure <ref>. As can be seen from the figure, our model is able to approximate compound waves, rarefaction waves, and shock waves to some extent. Although Flux NO was able to capture shock waves, numerical instability issues prevented long-term inference. For the two-dimensional problem, we addressed the Orszag-Tang problem. The initial conditions for the Orszag-Tang problem are described as follows:

ρ(x,y,0)=γ^2,  v_x(x,y,0)=-siny,  v_y(x,y,0)=sinx, 

p(x,y,0)=γ,  B_x(x,y,0)=-siny,   B_y(x,y,0)=sin2x, u_z(x,y,0)=B_z(x,y,0)=0.

When given the initial conditions of the Orszag-Tang problem, the solution at t=0.25 was used as the initial condition for the Flux NO scheme, and under these initial conditions, iterations were run 50 (t=0.5), 100 (t=0.75), 350 (t=2.0), and 550 (t=3.0) times, with results shown in Figures <ref> and <ref>. Similar to the results in section 4.2, the model approximates well for short-term inference below the time length of the training sample, but accuracy decreases for long-term inference, though the trend is generally followed. It was also observed that the error magnitude increases near shock waves. The values of ρ at the cross-section of y=0.625π and how they evolve over time are shown in Figure <ref>. Given that it performs well in short-term prediction, we conducted a further experiment where we used the values at each time point of the Orszag-Tang problem as initial conditions to predict the values after a Δ t=0.5 time. The results of this are shown in Figure <ref>, and the errors compared to the reference are shown in Figure <ref>. As can be seen from the figures, when predictions are made in the short-term, the errors are significantly reduced. Considering the good generalization performance and short-term approximation ability of the Flux NO model compared to the typical Neural Operator, we can apply the Flux NO model as a surrogate for numerical schemes in short-term inference.













 §.§ Comparison with Other Methods


In this section, we demonstrate both quantitatively and qualitatively that the Flux NO model possesses superior generalization abilities compared to traditional FNO models. Like Flux NO, both FNO 2D and FNO 3D use frequencies up to 4 on each axis and have a width of 72. FNO 2D operates by predicting an 8-dimensional vector value after a delta t = 0.5 given an input function, while FNO 3D works by producing the next 25 snapshots simultaneously from a 4D image of dimensions 25 by 64 by 64 by 8, considering the time axis in the dataset. For FNO 2D, we also trained a larger model, FNO 2D (heavy), with 8 frequencies and a width of 104. The dataset was adapted from the one used for training Flux NO, structured with 10,000 function pairs for 2D FNO and 500 pairs for 3D FNO. All models used the same learning rate and weight decay settings of 1e-3 and 1e-4 respectively, with Adam as the optimizer and StepLR as the scheduler with a step_size of 200 and gamma of 0.5. Training was conducted over 1000 epochs, with a batch size of 500 for 2D FNO and 25 for 3D FNO. Experimental results on the test dataset are presented in Table <ref>, showing Flux NO's significantly superior performance, with 3D FNO being the better performer among the remaining models. To quantitatively and qualitatively compare 3D FNO and Flux NO, we tasked 3D FNO with inferring solutions for the Orszag-Tang problem. The results for Flux NO and 3D FNO are detailed in Table <ref>, with qualitative illustrations in Figures <ref>, <ref>, and <ref>. As the table and figures illustrate, unlike Flux NO, 3D FNO fails to perform effectively on out-of-distribution samples and in long-term inference scenarios.













 §.§ Memory and Time costs




In this section, we analyze the memory requirements, specifically the number of parameters, and the inference time required for each model. Experimental results for each model are presented in Table <ref>, where WENO-Z calculations were performed on a CPU, while the rest of the neural operator models were computed on a GPU basis. As indicated in the table, assuming efficient GPU computation parallelization, our Flux NO is approximately 25 times faster than the WENO-Z method. Even with higher resolution images, the computational complexity of Flux NO is inferred to depend simply on the lifting, projection layer's FCN and Fourier layer's CNN layers, and the computation of the Fast Fourier Transform due to the limited frequency usage in the Fourier layer. Because Flux NO inherently performs calculations locally, it naturally takes more time to reach the same point compared to the traditional FNO. However, unlike FNO, which can only infer fixed distributions at fixed time intervals, Flux NO offers the advantage of flexible time interval selection, allowing for continuous inference. As shown in Table <ref> and the results in Section 4.4, despite having a similar number of parameters, 2D FNO (heavy) and 3D FNO show inferior generalization capabilities compared to Flux NO. Considering these factors, Flux NO can be seen as embodying the strengths of both numerical schemes—generalization capability and robustness—and the fast computational abilities of neural operators.









 §.§ Albation Study


In this section, we conduct an ablation analysis to explore the effects of the additional loss functions we designed. Models were trained with both TVD (Total Variation Diminishing) loss and divergence-free loss selectively removed, and the impact on ∇·B and Total Variation was observed over 75 iterations across 10 test samples. The results are presented in Figure <ref>, which illustrates the influence of each loss on the corresponding values of ∇·B and Total Variation. Relaxing ∇·B is a significant factor in obtaining physical solutions and also affects the stability when computing time steps adaptively. Mitigating the increase in Total Variation is likewise related to the stability of the numerical scheme.









§ CONCLUSION


In this study, we propose a method to solve the one-dimensional and two-dimensional cases of the ideal MHD, one of the equations describing plasma, using our newly designed Flux NO technique. We have demonstrated through experiments that Flux NO can be applied not only to the one-dimensional scalar conservation law (<cit.>) but also to the more complex ideal MHD equations, exhibiting long-term inference capabilities and generalization ability on out-of-distribution samples. We implemented specific loss functions to satisfy the TVD properties and divergence-freeness of the solutions, and adapted the model architecture to allocate different models to handle various variables appropriately. Our designed Flux NO shows flexibility in these modifications, particularly in the potential for adapting or improving Neural Operator component in the scheme to suit the problem. Although Flux NO requires more iterations to achieve the same solution, resulting in longer inference times compared to other Neural Operators, it allows for continuous inference and notably shorter inference times than classical numerical schemes. The dataset we adopted in this research is quite limited, and computing resources were also constrained; nonetheless, it is noteworthy that we achieved such generalization performance. Flux NO has the potential for further advancement and broadening of its application scope with more diverse datasets, larger-scale computing resources, and modifications to the loss functions and architecture.













§ ERROR PLOTS


















§ SUPPLEMENTARY FIGURES
































§ EXPERIMENTAL DETAILS


All experiments were conducted using Pytorch 1.10.0, with Python 3.6.9.
The specifications of the hardware environment are as follows.






§ COURANT–FRIEDRICHS–LEWY CONDITION (CFL CONDITION)


The Courant-Friedrichs-Lewy (CFL) condition is a crucial criterion for the stability of numerical solutions in computational simulations. It ensures that the physical propagation speed of a wave within a given timestep is less than or equal to the numerical propagation speed, which is essential for maintaining stability. For general N-dimensional case, the CFL condition is commonly expressed as follows:


    Δ t ( ∑_i=1^Nu_i/Δ x_i) ≤ C

Here, u_i represents the wave speed, and C is the Courant number, typically set to a value less than one. This constraint ensures that the numerical method can adequately capture the physical phenomena within the constraints of the time-step and spatial resolution.




0.2in

Valiant:84,LeVeque:92,Pathak:22,shalev:88,Jin:95,Gopalani:22,Jaku:19,Kovachki:212,Wen:22,Abrahamsen:97,Alfven:42,Bar:21,Benitez:23,Bittencourt:04,Borges:08,Brackbill:80,Brio:88,Chen:22,Christlieb:14,Costa:07,Courant:52,Dai:98,Evans:88,Fu:19,Galtier:16,Godunov:59,Gottlieb:98,Gupta:21,Harten:87,Henrick:05,Hill:04,Holl:20,Holl:22,Jiang:96,Jiang:99,Kennel:85,Kim:24_2,Kovachki:21,Lax:54,Kim:24,Lee:23,Leer:74,Li:20,Li:21,Liu:94,Lu:21,Magiera:20,Mukhopadhyay:21,O'Shea:15,Pakmor:13,Pirozzoli:02,Priest:82,Ray:18,Rincon:19,Rosofsky:23,Rossmanith:06,Ruggeri:22,Shibata:11,Sweby:84,Toth:00,Vapnik:21,Vlasov:38,Wang:20,Wesson:78

