Heat equation with training points resampling

Heat equation with training points resampling#

Problem setup#

We will solve a heat equation:

\[ \frac{\partial u}{\partial t}=\alpha \frac{\partial^2u}{\partial x^2}, \qquad x \in [0, 1], \quad t \in [0, 1] \]

where \(alpha=0.4\) is the thermal diffusivity constant.

With Dirichlet boundary conditions:

\[ u(0,t) = u(1,t)=0, \]

and periodic(sinusoidal) inital condition:

\[ u(x,0) = \sin (\frac{n\pi x}{L}),\qquad 0<x<L, \quad n = 1,2,..... \]

where \(L=1\) is the length of the bar, \(n=1\) is the frequency of the sinusoidal initial conditions.

The exact solution is \(u(x,t) = e^{\frac{-n ^2\pi ^2 \alpha t}{L^2}}\sin (\frac{n\pi x}{L})\).

Dimensional Analysis#

Step 1: Assign Dimensions to Variables#

  1. Spatial Coordinate \(x\):

    • The dimension of \(x\) is length:

      \[ [x] = L. \]

  2. Time \(t\):

    • The dimension of time is:

      \[ [t] = T. \]

  3. Temperature \(u\):

    • Temperature has dimensions of temperature:

      \[ [u] = \Theta. \]

  4. Thermal Diffusivity \(\alpha\):

    • The term \(\alpha \frac{\partial^2 u}{\partial x^2}\) involves the second spatial derivative of temperature, which must have the same dimensions as the time derivative \(\frac{\partial u}{\partial t}\).

Step 2: Analyze the Dimensions of Each Term#

  1. Time Derivative Term:

    • The time derivative \(\frac{\partial u}{\partial t}\) has dimensions:

      \[ \left[\frac{\partial u}{\partial t}\right] = \frac{[u]}{[t]} = \frac{\Theta}{T}. \]

  2. Diffusion Term:

    • The diffusion term \(\alpha \frac{\partial^2 u}{\partial x^2}\) involves the second spatial derivative of temperature:

      \[ \left[\frac{\partial^2 u}{\partial x^2}\right] = \frac{[u]}{[x]^2} = \frac{\Theta}{L^2}. \]

    • Therefore, the diffusion term has dimensions:

      \[ \left[\alpha \frac{\partial^2 u}{\partial x^2}\right] = [\alpha] \cdot \frac{\Theta}{L^2} = \frac{\Theta}{T}. \]

Step 3: Determine the Dimensions of \(\alpha\)#

  • The diffusion term \(\alpha \frac{\partial^2 u}{\partial x^2}\) must have the same dimensions as the time derivative \(\frac{\partial u}{\partial t}\):

    \[ [\alpha] \cdot \frac{\Theta}{L^2} = \frac{\Theta}{T} \implies [\alpha] = \frac{L^2}{T}. \]

  • Therefore, the thermal diffusivity \(\alpha\) has dimensions of thermal diffusivity:

    \[ [\alpha] = \frac{L^2}{T}. \]

Step 4: Summary of Dimensions#

Variable/Parameter

Physical Meaning

Dimensions

\(x\)

Spatial coordinate

\(L\)

\(t\)

Time

\(T\)

\(u\)

Temperature

\(\Theta\)

\(\alpha\)

Thermal diffusivity

\(L^2 / T\)

Step 5: Initial and Boundary Conditions#

  1. Boundary Conditions:

    • The boundary conditions \(u(0,t) = u(1,t) = 0\) are given in temperature units:

      \[ [u(0,t)] = [u(1,t)] = \Theta. \]

  2. Initial Condition:

    • The initial condition \(u(x,0) = \sin \left(\frac{n\pi x}{L}\right)\) is given in temperature units:

      \[ [u(x,0)] = \Theta. \]

    • The term \(\sin \left(\frac{n\pi x}{L}\right)\) is dimensionless because \(\frac{n\pi x}{L}\) is dimensionless (\(n\) and \(\pi\) are dimensionless constants, and \(x\) and \(L\) have the same dimensions).

Implementation#

This description goes through the implementation of a solver for the above described Heat equation step-by-step.

First, import the libraries we need:

import brainstate
import braintools
import numpy as np
import optax
import brainunit as u

import pinnx

We begin by defining the parameters of the equation:

a = 0.4 * u.meter2 / u.second # Thermal diffusivity
L = 1 * u.meter  # Length of the bar
n = 1 * u.Hz # Frequency of the sinusoidal initial conditions

Next, we define a computational geometry and time domain. We can use a built-in class Interval and TimeDomain and we combine both the domains using GeometryXTime as follows

geomtime = pinnx.geometry.GeometryXTime(
    geometry=pinnx.geometry.Interval(0., 1.),
    timedomain=pinnx.geometry.TimeDomain(0., 1.)
).to_dict_point(x=u.meter, t=u.second)

Next, we express the PDE residual of the Heat equation:

@brainstate.transform.jit
def pde(x, y):
    jacobian = approximator.jacobian(x)
    hessian = approximator.hessian(x)
    dy_t = jacobian['y']['t']
    dy_xx = hessian['y']['x']['x']
    return dy_t - a * dy_xx

The first argument to pde is 2-dimensional vector where the first component(x[:,0]) is x-coordinate and the second componenet (x[:,1]) is the t-coordinate. The second argument is the network output, i.e., the solution u(x,t), but here we use y as the name of the variable.

Next, we consider the boundary/initial condition. We include the geomtime space, time geometry created above. We also define IC which is the inital condition for the burgers equation and we use the computational domain and initial function.

uy = u.kelvin / u.second
# Initial and boundary conditions:
bc = pinnx.icbc.DirichletBC(
    lambda x : {'y': 0. * uy}
)
ic = pinnx.icbc.IC(
    lambda x: {'y': u.math.sin(n * u.math.pi * x['x'][:] / L, unit_to_scale=u.becquerel) * uy},
)

Next, we choose the network. Here, we use a fully connected neural network of depth 4 (i.e., 3 hidden layers) and width 20:

approximator = pinnx.nn.Model(
    pinnx.nn.DictToArray(x=u.meter, t=u.second),
    pinnx.nn.FNN(
        [2] + [20] * 3 + [1],
        "tanh",
        braintools.init.KaimingUniform()
    ),
    pinnx.nn.ArrayToDict(y=uy)
)

Now, we have specified the geometry, PDE residual, boundary/initial condition and approximator. We then define the TimePDE problem as

problem = pinnx.problem.TimePDE(
    geomtime,
    pde,
    [bc, ic],
    approximator,
    num_domain=2540,
    num_boundary=80,
    num_initial=160,
    num_test=2540,
)

Warning: 2540 points required, but 2550 points sampled.

The number 2540 is the number of training residual points sampled inside the domain, and the number 80 is the number of training points sampled on the boundary. We also include 160 initial residual points for the initial conditions.

Now, we have the PDE problem and the network. We build a Model and choose the optimizer and learning rate:

trainer = pinnx.Trainer(problem)
trainer.compile(braintools.optim.Adam(1e-3))

Compiling trainer...
'compile' took 0.055764 s

<pinnx.Trainer at 0x148ceacd0>

The following code is to apply mini-batch gradient descent sampling method. The period is the period of resamping. Here, the training points in the domain will be resampled every 10 iterations.

pde_resampler = pinnx.callbacks.PDEPointResampler(period=10)

We then train the model for 20000 iterations:

trainer.train(iterations=20000, callbacks=[pde_resampler])

Training trainer...

Step      Train loss                                                            Test loss                                                           Test metric 
10000     [2.6207555e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,          [2.11663e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,          []          
           {'ibc0': {'y': 6.1262945e-06 * kelvin / second}},                     {'ibc0': {'y': 6.1262945e-06 * kelvin / second}},                              
           {'ibc1': {'y': 4.302944e-06 * kelvin / second}}]                      {'ibc1': {'y': 4.302944e-06 * kelvin / second}}]                               
11000     [1.9524387e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,          [1.6359398e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,        []          
           {'ibc0': {'y': 4.0165432e-06 * kelvin / second}},                     {'ibc0': {'y': 4.0165432e-06 * kelvin / second}},                              
           {'ibc1': {'y': 3.168525e-06 * kelvin / second}}]                      {'ibc1': {'y': 3.168525e-06 * kelvin / second}}]                               
12000     [1.5614694e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,          [1.34647535e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,       []          
           {'ibc0': {'y': 2.8513743e-06 * kelvin / second}},                     {'ibc0': {'y': 2.8513743e-06 * kelvin / second}},                              
           {'ibc1': {'y': 2.5143677e-06 * kelvin / second}}]                     {'ibc1': {'y': 2.5143677e-06 * kelvin / second}}]                              
13000     [1.5810221e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,          [1.08873655e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,       []          
           {'ibc0': {'y': 2.6305897e-06 * kelvin / second}},                     {'ibc0': {'y': 2.6305897e-06 * kelvin / second}},                              
           {'ibc1': {'y': 2.057085e-06 * kelvin / second}}]                      {'ibc1': {'y': 2.057085e-06 * kelvin / second}}]                               
14000     [1.0092881e-05 * 10.0^0 * ((kelvin / second) / second) ** 2,          [8.765763e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,         []          
           {'ibc0': {'y': 1.7242603e-06 * kelvin / second}},                     {'ibc0': {'y': 1.7242603e-06 * kelvin / second}},                              
           {'ibc1': {'y': 1.4643564e-06 * kelvin / second}}]                     {'ibc1': {'y': 1.4643564e-06 * kelvin / second}}]                              
15000     [9.48981e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,            [8.021057e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,         []          
           {'ibc0': {'y': 1.8745583e-06 * kelvin / second}},                     {'ibc0': {'y': 1.8745583e-06 * kelvin / second}},                              
           {'ibc1': {'y': 1.6735984e-06 * kelvin / second}}]                     {'ibc1': {'y': 1.6735984e-06 * kelvin / second}}]                              
16000     [8.69476e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,            [6.61613e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          []          
           {'ibc0': {'y': 1.2975498e-06 * kelvin / second}},                     {'ibc0': {'y': 1.2975498e-06 * kelvin / second}},                              
           {'ibc1': {'y': 8.5685264e-07 * kelvin / second}}]                     {'ibc1': {'y': 8.5685264e-07 * kelvin / second}}]                              
17000     [6.9987154e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          [5.5678443e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,        []          
           {'ibc0': {'y': 1.1903414e-06 * kelvin / second}},                     {'ibc0': {'y': 1.1903414e-06 * kelvin / second}},                              
           {'ibc1': {'y': 1.0031221e-06 * kelvin / second}}]                     {'ibc1': {'y': 1.0031221e-06 * kelvin / second}}]                              
18000     [5.6217564e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          [4.7939957e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,        []          
           {'ibc0': {'y': 7.706846e-07 * kelvin / second}},                      {'ibc0': {'y': 7.706846e-07 * kelvin / second}},                               
           {'ibc1': {'y': 7.064452e-07 * kelvin / second}}]                      {'ibc1': {'y': 7.064452e-07 * kelvin / second}}]                               
19000     [5.385209e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,           [4.3364116e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,        []          
           {'ibc0': {'y': 6.759552e-07 * kelvin / second}},                      {'ibc0': {'y': 6.759552e-07 * kelvin / second}},                               
           {'ibc1': {'y': 5.627439e-07 * kelvin / second}}]                      {'ibc1': {'y': 5.627439e-07 * kelvin / second}}]                               
20000     [5.322141e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,           [4.2963425e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,        []          
           {'ibc0': {'y': 7.620727e-07 * kelvin / second}},                      {'ibc0': {'y': 7.620727e-07 * kelvin / second}},                               
           {'ibc1': {'y': 5.87999e-07 * kelvin / second}}]                       {'ibc1': {'y': 5.87999e-07 * kelvin / second}}]                                

Best trainer at step 19000:
  train loss: 6.62e-06
  test loss: 5.58e-06
  test metric: []

'train' took 47.065386 s

<pinnx.Trainer at 0x148ceacd0>

After we train the network using Adam, we continue to train the network using L-BFGS to achieve a smaller loss:

trainer.compile(braintools.optim.OptaxOptimizer(optax.lbfgs(1e-3, linesearch=None)))
trainer.train(iterations=2000)

Compiling trainer...
'compile' took 0.140131 s

Training trainer...

Step      Train loss                                                            Test loss                                                             Test metric 
20000     [5.1988213e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          [4.2963425e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          []          
           {'ibc0': {'y': 7.620727e-07 * kelvin / second}},                      {'ibc0': {'y': 7.620727e-07 * kelvin / second}},                                 
           {'ibc1': {'y': 5.87999e-07 * kelvin / second}}]                       {'ibc1': {'y': 5.87999e-07 * kelvin / second}}]                                  
21000     [4.5359056e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          [3.7426814e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          []          
           {'ibc0': {'y': 6.414695e-07 * kelvin / second}},                      {'ibc0': {'y': 6.414695e-07 * kelvin / second}},                                 
           {'ibc1': {'y': 4.3834766e-07 * kelvin / second}}]                     {'ibc1': {'y': 4.3834766e-07 * kelvin / second}}]                                
22000     [4.472149e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,           [3.7037253e-06 * 10.0^0 * ((kelvin / second) / second) ** 2,          []          
           {'ibc0': {'y': 6.1141054e-07 * kelvin / second}},                     {'ibc0': {'y': 6.1141054e-07 * kelvin / second}},                                
           {'ibc1': {'y': 4.3972565e-07 * kelvin / second}}]                     {'ibc1': {'y': 4.3972565e-07 * kelvin / second}}]                                

Best trainer at step 22000:
  train loss: 5.52e-06
  test loss: 4.75e-06
  test metric: []

'train' took 10.094850 s

<pinnx.Trainer at 0x148ceacd0>

Let’s visualize and save the data.

trainer.saveplot(issave=True, isplot=True)

Saving loss history to /Users/sichaohe/Documents/GitHub/pinnx/docs/examples-pinn-forward/loss.dat ...
Saving checkpoint into /Users/sichaohe/Documents/GitHub/pinnx/docs/examples-pinn-forward/loss.dat
Saving training data to /Users/sichaohe/Documents/GitHub/pinnx/docs/examples-pinn-forward/train.dat ...
Saving checkpoint into /Users/sichaohe/Documents/GitHub/pinnx/docs/examples-pinn-forward/train.dat
Saving test data to /Users/sichaohe/Documents/GitHub/pinnx/docs/examples-pinn-forward/test.dat ...
Saving checkpoint into /Users/sichaohe/Documents/GitHub/pinnx/docs/examples-pinn-forward/test.dat
../_images/8555552f880f48416bec8382f52509f02ecf6980779f7487ee4782bed24ca273.png ../_images/218297b4f9554a4283301fc8ec1915add9be08dd600b7f760b6b24c0c59ae813.png