Heat equation with training points resampling
Problem setup
We will solve a heat equation with training points resampling:
where \(\alpha=0.4\) is the thermal diffusivity constant.
With Dirichlet boundary conditions:
and periodic(sinusoidal) inital condition:
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})\).
Implementation
This description goes through the implementation of a solver for the above described Heat equation step-by-step.
First, the DeepXDE are imported:
import deepxde as dde
We begin by defining the parameters of the equation:
a = 0.4
L = 1
n = 1
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
geom = dde.geometry.Interval(0, L)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)
Next, we express the PDE residual of the Heat equation:
def pde(x, y):
dy_t = dde.grad.jacobian(y, x, i=0, j=1)
dy_xx = dde.grad.hessian(y, x, i=0, j=0)
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. on_boundary
is chosen here to use the whole boundary of the computational domain in considered as the boundary condition. We include the geomtime
space, time geometry created above and on_boundary
as the BCs in the DirichletBC
function of DeepXDE. We also define IC
which is the inital condition for the burgers equation and we use the computational domain, initial function, and on_initial
to specify the IC.
bc = dde.icbc.DirichletBC(geomtime, lambda x: 0, lambda _, on_boundary: on_boundary)
ic = dde.icbc.IC(
geomtime,
lambda x: np.sin(n * np.pi * x[:, 0:1] / L),
lambda _, on_initial: on_initial,
)
Now, we have specified the geometry, PDE residual, and boundary/initial condition. We then define the TimePDE
problem as
data = dde.data.TimePDE(
geomtime,
pde,
[bc, ic],
num_domain=2540,
num_boundary=80,
num_initial=160,
num_test=2540,
)
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.
Next, we choose the network. Here, we use a fully connected neural network of depth 4 (i.e., 3 hidden layers) and width 20:
net = dde.nn.FNN([2] + [20] * 3 + [1], "tanh", "Glorot normal")
Now, we have the PDE problem and the network. We build a Model
and choose the optimizer and learning rate:
model = dde.Model(data, net)
model.compile("adam", lr=1e-3)
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 = dde.callbacks.PDEPointResampler(period=10)
We then train the model for 20000 iterations:
losshistory, train_state = model.train(iterations=200000, callbacks=[pde_resampler])
After we train the network using Adam, we continue to train the network using L-BFGS to achieve a smaller loss:
model.compile("L-BFGS-B")
losshistory, train_state = model.train()
Complete code
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
import deepxde as dde
import numpy as np
def heat_eq_exact_solution(x, t):
"""Returns the exact solution for a given x and t (for sinusoidal initial conditions).
Parameters
----------
x : np.ndarray
t : np.ndarray
"""
return np.exp(-(n**2 * np.pi**2 * a * t) / (L**2)) * np.sin(n * np.pi * x / L)
def gen_exact_solution():
"""Generates exact solution for the heat equation for the given values of x and t."""
# Number of points in each dimension:
x_dim, t_dim = (256, 201)
# Bounds of 'x' and 't':
x_min, t_min = (0, 0.0)
x_max, t_max = (L, 1.0)
# Create tensors:
t = np.linspace(t_min, t_max, num=t_dim).reshape(t_dim, 1)
x = np.linspace(x_min, x_max, num=x_dim).reshape(x_dim, 1)
usol = np.zeros((x_dim, t_dim)).reshape(x_dim, t_dim)
# Obtain the value of the exact solution for each generated point:
for i in range(x_dim):
for j in range(t_dim):
usol[i][j] = heat_eq_exact_solution(x[i], t[j])
# Save solution:
np.savez("heat_eq_data", x=x, t=t, usol=usol)
# Load solution:
data = np.load("heat_eq_data.npz")
def gen_testdata():
"""Import and preprocess the dataset with the exact solution."""
# Load the data:
data = np.load("heat_eq_data.npz")
# Obtain the values for t, x, and the excat solution:
t, x, exact = data["t"], data["x"], data["usol"].T
# Process the data and flatten it out (like labels and features):
xx, tt = np.meshgrid(x, t)
X = np.vstack((np.ravel(xx), np.ravel(tt))).T
y = exact.flatten()[:, None]
return X, y
# Problem parameters:
a = 0.4 # Thermal diffusivity
L = 1 # Length of the bar
n = 1 # Frequency of the sinusoidal initial conditions
# Generate a dataset with the exact solution (if you dont have one):
gen_exact_solution()
def pde(x, y):
"""Expresses the PDE residual of the heat equation."""
dy_t = dde.grad.jacobian(y, x, i=0, j=1)
dy_xx = dde.grad.hessian(y, x, i=0, j=0)
return dy_t - a * dy_xx
# Computational geometry:
geom = dde.geometry.Interval(0, L)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)
# Initial and boundary conditions:
bc = dde.icbc.DirichletBC(geomtime, lambda x: 0, lambda _, on_boundary: on_boundary)
ic = dde.icbc.IC(
geomtime,
lambda x: np.sin(n * np.pi * x[:, 0:1] / L),
lambda _, on_initial: on_initial,
)
pde_resampler = dde.callbacks.PDEPointResampler(period=10)
# Define the PDE problem and configurations of the network:
data = dde.data.TimePDE(
geomtime,
pde,
[bc, ic],
num_domain=2540,
num_boundary=80,
num_initial=160,
num_test=2540,
)
net = dde.nn.FNN([2] + [20] * 3 + [1], "tanh", "Glorot normal")
model = dde.Model(data, net)
# Build and train the model:
model.compile("adam", lr=1e-3)
model.train(iterations=200000, callbacks=[pde_resampler])
model.compile("L-BFGS")
losshistory, train_state = model.train(callbacks=[pde_resampler])
# Plot/print the results
dde.saveplot(losshistory, train_state, issave=True, isplot=True)
X, y_true = gen_testdata()
y_pred = model.predict(X)
f = model.predict(X, operator=pde)
print("Mean residual:", np.mean(np.absolute(f)))
print("L2 relative error:", dde.metrics.l2_relative_error(y_true, y_pred))
np.savetxt("test.dat", np.hstack((X, y_true, y_pred)))