CITE THIS NOTEBOOK: A 1D study of adaptive mesh refinement in neural networks for PDE solving by Gianni Tallarita. Wolfram Community FEB 28 2023.
In this post I construct an adaptive mesh refined solver for Burgers’ equation over a one-dimensional domain. This is part of a research direction I am exploring in order to implement AMR Numerical PDE solvers on GPUs, within Wolfram Language. As an introduction, I present the standard method of solving this system, along with its known failures, using a static uniform grid and a Runge-Kutta time evolver. To overcome the limitations this solver presents, we introduce adaptive mesh refinement (AMR) in the system. I then demonstrate explicitly the construction of the AMR solver in Neural Network language in order to pass the calculation to GPUs. Finally, I conclude with the limitations of this system and the natural directions of further work this topic screams for.
The theory behind the problem presented below follows chapter 1 of “Adaptive Moving Mesh Methods” by Huang and Russell https://doi.org/10.1007/978-1-4419-7916-2, to which the reader is directed for all additional details.
Burgers’ equation in 1d, a CPU AMR treatment
Burgers’ equation in 1d, a CPU AMR treatment
Burgers’ equation is a convection-diffusion equation important in several regions of physics ([1]). We are interested in this notebook (and in general) to find solutions of Burgers’ equation, posed as an initial-value problem. The equation is
u
t
u
xx
2
u
2
x
,x∈(0,1),t>0
where denotes derivatives w.r.t time, and with respect to the spatial variable . We wish to solve this equation subject to the boundary conditions and initial condition . In the above, is a free physical parameter, one of two we will be using in our solver. If is small enough, the solution quickly develops a steep front which will propagate towards the right. This present a hard problem for a numerical analysis, as the steepness of the front makes derivatives hard to calculate and eventually singular.
t
x
x
u(0,t)u(1,t)0
u(x,0)sin(2πx)+sin(πx)
1
2
ϵ
ϵ
This is clearly visible in a standard NDSolve treatment of this equation. As the solver evolves the system in time, it quickly finds a singular behaviour corresponding precisely to this steep front developing at finite time.
In[]:=
ϵ=10^-2;
sol=NDSolveValueD[u[x,t],t]==ϵD[u[x,t],{x,2}]-D,x,u[0,t]==0,u[1,t]==0,u[x,0]==Sin[2πx]+Sin[πx],u,{x,0,1},{t,0,1}
2
u[x,t]
2
1
2
13
10
Out[]=
InterpolatingFunction
In[]:=
Plot3D[sol[x,t],{x,0,1},{t,0,0.33},PlotRangeAll,AxesLabel->{"x","t"}]
Out[]=
The figure above indicates the singular behaviour of the steep front in the time evolution of the solution to Burgers’s equation.
The reader might be panicking at the sight of such horrible functional behaviour, but we can still make a lot of progress. First, let us try to understand the creation of this steep front by building our own time evolution solver, leaving the confines of NDSolve.
A Runge-Kutta solution to Burgers’ equation
A Runge-Kutta solution to Burgers’ equation
Here we write our own code to find the time evolution of the solution to Burgers’ equation. First, we will need an appropriate discretization of the spatial domain .
x
In[]:=
ClearAll@Grid1D;Grid1D[{lX_,nX_Integer}]:=Block[{dx},dx=lX/(nX-1);{dx(Range@nX-1),dx}]
In[]:=
{xs2,hx}=Grid1D[{1.,41}]
Out[]=
{{0.,0.025,0.05,0.075,0.1,0.125,0.15,0.175,0.2,0.225,0.25,0.275,0.3,0.325,0.35,0.375,0.4,0.425,0.45,0.475,0.5,0.525,0.55,0.575,0.6,0.625,0.65,0.675,0.7,0.725,0.75,0.775,0.8,0.825,0.85,0.875,0.9,0.925,0.95,0.975,1.},0.025}
This defines our grid points on the 1d line, and the effective delta that we will use in our derivatives. Now we make use of finite difference derivates.
In[]:=
D2x=NDSolve`FiniteDifferenceDerivative[Derivative[2],xs2,"DifferenceOrder"2]["DifferentiationMatrix"];Dx=NDSolve`FiniteDifferenceDerivative[Derivative[1],xs2,"DifferenceOrder"2]["DifferentiationMatrix"];
and we create our initial condition
In[]:=
uinitial=Sin[2πx]+Sin[πx]/.x->xs2;
1
2
In typical Runge-Kutta fashion, we now must define a numerical operator which acts on the r.h.s of our equation and which the time derivative will evolve in time
In[]:=
boundary=BitOr[UnitVector[41,1],UnitVector[41,41]];
In[]:=
Eq[seed_]:=Block{op},op=ϵ*D2x.seed-Dx.*;op-op*boundary
1
2
2
seed
This now calculates the r.h.s of Burgers’ equation, given an initial
ϵ
In[]:=
ϵ=10^-4;
In[]:=
Eq[uinitial]
Out[]=
{0.,-1.50016,-2.87559,-4.0124,-4.81743,-5.22619,-5.20825,-4.76954,-3.95136,-2.8263,-1.49129,-0.0587635,1.35363,2.63342,3.68317,4.42861,4.82461,4.85823,4.54874,3.94475,3.11855,2.15851,1.16017,0.216904,-0.588857,-1.19308,-1.55552,-1.66264,-1.5277,-1.18827,-0.701241,-0.136141,0.432732,0.934403,1.30877,1.5128,1.5248,1.34628,1.00127,0.533263,0.}
Note that we have included boundary conditions which enforce the preservation of the correct Dirichlet conditions at the boundaries. The rest is just the standard time evolution using a 4th order Runge Kutta procedure
ClearAll@RK4Step;RK4Step[dt_][u_]:=Block{neuralNetwork=Eq,k1,u1,k2,u2,k3,u3,k4,u4,λ1,λ2,λ3,λ4,λnew},k1=neuralNetwork[u];u1=u+k1;k2=neuralNetwork[u1];u2=u1+k2;k3=neuralNetwork[u2];u3=u2+dtk3;k4=neuralNetwork[u3];u4=u+(k1+2k2+2k3+k4);u4
dt
2
dt
2
dt
6
RK4loop[dt_,numSteps_Integer,seed_][checkpointInterval_Integer]:=Block[{step=RK4Step[dt],u=seed,start,lastCheckpoint,check},start=lastCheckpoint=AbsoluteTime[];Do[u=step[u];If[Mod[i,checkpointInterval]0,PrintTemporary["Checkpoint at step ",i,": done ",checkpointInterval," iterations in ",Round[AbsoluteTime[]-lastCheckpoint,.01]," seconds!"];lastCheckpoint=AbsoluteTime[];check={Eq[u]//Flatten//Abs//Max};Print[check];soli[i]=u;];,{i,numSteps}];Print["Finished in ",Round[AbsoluteTime[]-start,.01]," seconds!"];solnew=u;]
In[]:=
RK4loop[0.01,21,uinitial][3]
{5.65432}
{6.64951}
{8.47474}
{11.1512}
{14.6976}
{20.5989}
{18.9941}
Finished in 0.06 seconds!
Adaptive Mesh Refinement
Adaptive Mesh Refinement
and Abs denotes the absolute value.
We must now modify our r.h.s operator to include the Adaptive Mesh procedure
Let’s plot a couple of these solutions just to check that everything is ok.
What happens as we approach the singular point we encountered before? Now we see that the adaptive mesh has resolved itself to be much more dense close to the front, meaning we can resolve it without issues and push further along with time.
Conclusions
Conclusions
I have implemented a simple 1d Adaptive mesh refinement algorithm to Burgers’ equation. It is clear that this algorithm allows one to improve on singular numerical procedures. A natural extension is now to include this procedure in full 3+1 dimensional simulations. This is left for further work. However, as per most 3+1 dimensional simulations, this procedure can be extremely slow running it on CPUs. Therefore I will dedicate the next section of this work to translate the above algorithm to neural network language, which allows one to use GPUs instead. It is obvious that for this particular system, just using 41 points on a line, this procedure is cumbersome and completely inefficient. In fact, it will probably lead to slower processing times. However, it is indicative of what one would have to do to use GPUs in 3+1 d, where I have seen time gains of factors of around 130x myself.
Neural Network Implementation and GPUs
Neural Network Implementation and GPUs
In this section I will show how to pass the above algorithm to neural network language. This involves several steps.
First, let me define some global properties of the system. This DEVICE selection is what allows us to eventually use GPUs.
Next, we are going to need the equivalent of our derivative operators Dx and D2x. I obtain these by acting with convolutions on grid points of our mesh.
Which is clearly wrong at the boundary points. This is the result of using a central difference convolution on a Fixed padding scheme. The problem can be fixed by defining the correct differences at the boundary points manually in Neural network language (I have done this for the full 3+1 d system for other projects, it is a hard computational problem). This is beyond the scope of this work and it turns out we can get away with this small error in what follows. However, the reader should be very careful with this issue since it can introduce large mistakes in a numerical procedure.
We also need an equivalent of our boundary selector.
Now we are ready to pass the full r.h.s. operator to Neural Network language. Note that I pass directly to the AMR case with both x(ξ,t) and u(x(ξ,t),t). What the below network does is just what the Eqtot operator did for us above.
Let’s see what this operator looks like in Neural Network language
Here is the definition of a single 4th order Runge Kutta step.
And now let’s loop this.
Now we are ready to run. This will now work on your GPU provided you have a NVidia GPU and the latest drivers installed. Mathematica should install all relevant components for you, provided you have the right hardware.
As we can see, the Neural Network system works! It provides us with much the same answer as our previous system, however now with the ability to run on GPUs rather than CPUs. However, there are several important points which we will discuss below.
Conclusions
Conclusions
References
References
(1) - Adaptive Moving Mesh Methods - Huang and Russell - Applied Mathematical Sciences, Vol 174 - Springer