Diffusion equation with hard initial and boundary conditions
Problem setup
We will solve a diffusion equation with hard initial and boundary conditions:
with the initial condition
and the Dirichlet boundary condition
The reference solution is \(y = e^{-t} \sin(\pi x)\).
Implementation
This description goes through the implementation of a solver for the above described diffusion equation step-by-step.
First, the DeepXDE, NumPy (np
), and TensorFlow (tf
) modules are imported:
import deepxde as dde
import numpy as np
from deepxde.backend import tf
We begin by defining computational geometries. We can use a built-in class Interval
and TimeDomain
and we combine both the domains using GeometryXTime
as follows
geom = dde.geometry.Interval(-1, 1)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)
Next, we express the PDE residual of the diffusion equation:
def pde(x, y):
dy_t = dde.grad.jacobian(y, x, j=1)
dy_xx = dde.grad.hessian(y, x, j=0)
return (
dy_t
- dy_xx
+ tf.exp(-x[:, 1:])
* (tf.sin(np.pi * x[:, 0:1]) - np.pi ** 2 * tf.sin(np.pi * x[:, 0:1]))
)
The first argument to pde
is 2-dimensional vector where the first component(x[:,0:1]
) is \(x\)-coordinate and the second component (x[:,1:]
) is the \(t\)-coordinate. The second argument is the network output, i.e., the solution \(y(x, t)\).
The reference solution func
is defined as:
def func(x):
return np.sin(np.pi * x[:, 0:1]) * np.exp(-x[:, 1:])
Now, we have specified the geometry and the PDE residual. However, in order to apply hard boundary and initial conditions, they are not specified and excluded from the loss function. We then define the TimePDE
problem as
data = dde.data.TimePDE(geomtime, pde, [], num_domain=40, solution=func, num_test=10000)
The number 40 is the number of training residual points sampled inside the domain. 10000 points for testing the PDE residual.
Next, we choose the network. Here, we use a fully connected neural network of depth 4 (i.e., 3 hidden layers) and width 32:
layer_size = [2] + [32] * 3 + [1]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)
Then we construct a function that spontaneously satisfies both the initial and the boundary conditions to transform the network output. In this case, \(t(1-x^2)y + sin(\pi x)\) is used. When \(t\) is equal to 0, the initial condition \(sin(\pi x)\) is recovered. When \(x\) is equal to -1 or 1, the boundary condition \(y(-1, t) = y(1, t) = 0\) is recovered. Hence the initial and boundary conditions are both hard conditions.
net.apply_output_transform(
lambda x, y: x[:, 1:2] * (1 - x[:, 0:1] ** 2) * y + tf.sin(np.pi * x[:, 0:1])
)
Now, we have the PDE problem and the network. We build a Model
and choose the optimizer and learning rate. We then train the model for 10000 iterations.
model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["l2 relative error"])
losshistory, train_state = model.train(iterations=10000)
We also save and plot the best trained result and loss history.
dde.saveplot(losshistory, train_state, issave=True, isplot=True)
Complete code
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, jax, paddle"""
import deepxde as dde
import numpy as np
# Backend tensorflow.compat.v1 or tensorflow
from deepxde.backend import tf
# Backend pytorch
# import torch
# Backend jax
# import jax.numpy as jnp
# Backend paddle
# import paddle
def pde(x, y):
# Most backends
dy_t = dde.grad.jacobian(y, x, i=0, j=1)
dy_xx = dde.grad.hessian(y, x, i=0, j=0)
# Backend jax
# dy_t, _ = dde.grad.jacobian(y, x, i=0, j=1)
# dy_xx, _ = dde.grad.hessian(y, x, i=0, j=0)
# Backend tensorflow.compat.v1 or tensorflow
return (
dy_t
- dy_xx
+ tf.exp(-x[:, 1:])
* (tf.sin(np.pi * x[:, 0:1]) - np.pi ** 2 * tf.sin(np.pi * x[:, 0:1]))
)
# Backend pytorch
# return (
# dy_t
# - dy_xx
# + torch.exp(-x[:, 1:])
# * (torch.sin(np.pi * x[:, 0:1]) - np.pi ** 2 * torch.sin(np.pi * x[:, 0:1]))
# )
# Backend jax
# return (
# dy_t
# - dy_xx
# + jnp.exp(-x[:, 1:])
# * (jnp.sin(np.pi * x[..., 0:1]) - np.pi ** 2 * jnp.sin(np.pi * x[..., 0:1]))
# )
# Backend paddle
# return (
# dy_t
# - dy_xx
# + paddle.exp(-x[:, 1:])
# * (paddle.sin(np.pi * x[:, 0:1]) - np.pi ** 2 * paddle.sin(np.pi * x[:, 0:1]))
# )
def func(x):
return np.sin(np.pi * x[:, 0:1]) * np.exp(-x[:, 1:])
geom = dde.geometry.Interval(-1, 1)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)
data = dde.data.TimePDE(geomtime, pde, [], num_domain=40, solution=func, num_test=10000)
layer_size = [2] + [32] * 3 + [1]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)
net.apply_output_transform(
# Backend tensorflow.compat.v1 or tensorflow
lambda x, y: x[:, 1:2] * (1 - x[:, 0:1] ** 2) * y + tf.sin(np.pi * x[:, 0:1])
# Backend pytorch
# lambda x, y: x[:, 1:2] * (1 - x[:, 0:1] ** 2) * y + torch.sin(np.pi * x[:, 0:1])
# Backend jax
# lambda x, y: x[..., 1:2] * (1 - x[..., 0:1] ** 2) * y + jnp.sin(np.pi * x[..., 0:1])
# Backend paddle
# lambda x, y: x[:, 1:2] * (1 - x[:, 0:1] ** 2) * y + paddle.sin(np.pi * x[:, 0:1])
)
model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["l2 relative error"])
losshistory, train_state = model.train(iterations=10000)
dde.saveplot(losshistory, train_state, issave=True, isplot=True)