Poisson equation in 1D with Dirichlet boundary conditions
Problem setup
We will solve a Poisson equation:
with the Dirichlet boundary conditions
The exact solution is \(u(x) = \sin(\pi x)\).
Implementation
This description goes through the implementation of a solver for the above described Poisson equation step-by-step.
First, the DeepXDE and TensorFlow (tf) modules are imported:
import deepxde as dde
from deepxde.backend import tf
We begin by defining a computational geometry. We can use a built-in class Interval as follows
geom = dde.geometry.Interval(-1, 1)
Next, we express the PDE residual of the Poisson equation:
def pde(x, y):
dy_xx = dde.grad.hessian(y, x)
return -dy_xx - np.pi ** 2 * tf.sin(np.pi * x)
The first argument to pde is the network input, i.e., the \(x\)-coordinate. The second argument is the network output, i.e., the solution \(u(x)\), but here we use y as the name of the variable.
Next, we consider the Dirichlet boundary condition. A simple Python function, returning a boolean, is used to define the subdomain for the Dirichlet boundary condition (\(\{-1, 1\}\)). The function should return True for those points inside the subdomain and False for the points outside. In our case, the points \(x\) of the Dirichlet boundary condition are \(x=-1\) and \(x=1\). (Note that because of rounding-off errors, it is often wise to use dde.utils.isclose to test whether two floating point values are equivalent.)
def boundary(x, _):
return dde.utils.isclose(x[0], -1) or dde.utils.isclose(x[0], 1)
The argument x to boundary is the network input and is a \(d\)-dim vector, where \(d\) is the dimension and \(d=1\) in this case. To facilitate the implementation of boundary, a boolean on_boundary is used as the second argument. If the point x (the first argument) is on the entire boundary of the geometry (the left and right endpoints of the interval in this case), then on_boundary is True, otherwise, on_boundary is False. Thus, we can also define boundary in a simpler way:
def boundary(x, on_boundary):
return on_boundary
Next, we define a function to return the value of \(u(x)\) for the points \(x\) on the boundary. In this case, it is \(u(x)=0\).
def func(x):
return 0
If the function value is not a constant, we can also use NumPy to compute. For example, \(\sin(\pi x)\) is 0 on the boundary, and thus we can also use
def func(x):
return np.sin(np.pi * x)
Then, the Dirichlet boundary condition is
bc = dde.icbc.DirichletBC(geom, func, boundary)
Now, we have specified the geometry, PDE residual, and Dirichlet boundary condition. We then define the PDE problem as
data = dde.data.PDE(geom, pde, bc, 16, 2, solution=func, num_test=100)
The number 16 is the number of training residual points sampled inside the domain, and the number 2 is the number of training points sampled on the boundary. The argument solution=func is the reference solution to compute the error of our solution, and can be ignored if we don’t have a reference solution. We use 100 residual 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 50:
layer_size = [1] + [50] * 3 + [1]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)
Now, we have the PDE problem and the network. We bulid a Model and choose the optimizer and learning rate:
model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["l2 relative error"])
We also compute the \(L^2\) relative error as a metric during training. We can also use callbacks to save the model and the movie during training, which is optional.
checkpointer = dde.callbacks.ModelCheckpoint(
"model/model.ckpt", verbose=1, save_better_only=True
)
# ImageMagick (https://imagemagick.org/) is required to generate the movie.
movie = dde.callbacks.MovieDumper(
"model/movie", [-1], [1], period=100, save_spectrum=True, y_reference=func
)
We then train the model for 10000 iterations:
losshistory, train_state = model.train(
iterations=10000, callbacks=[checkpointer, movie]
)
Complete code
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, jax, paddle"""
import deepxde as dde
import matplotlib.pyplot as plt
import numpy as np
# Import tf if using backend tensorflow.compat.v1 or tensorflow
from deepxde.backend import tf
# Import torch if using backend pytorch
# import torch
# Import jax.numpy if using backend jax
# import jax.numpy as jnp
# Import paddle if using backend paddle
# import paddle
def pde(x, y):
# Most backends
dy_xx = dde.grad.hessian(y, x)
# Backend jax
# dy_xx, _ = dde.grad.hessian(y, x)
# Use tf.sin for backend tensorflow.compat.v1 or tensorflow
return -dy_xx - np.pi ** 2 * tf.sin(np.pi * x)
# Use torch.sin for backend pytorch
# return -dy_xx - np.pi ** 2 * torch.sin(np.pi * x)
# Use jax.numpy.sin for backend jax
# return -dy_xx - np.pi ** 2 * jnp.sin(np.pi * x)
# Use paddle.sin for backend paddle
# return -dy_xx - np.pi ** 2 * paddle.sin(np.pi * x)
def boundary(x, on_boundary):
return on_boundary
def func(x):
return np.sin(np.pi * x)
geom = dde.geometry.Interval(-1, 1)
bc = dde.icbc.DirichletBC(geom, func, boundary)
data = dde.data.PDE(geom, pde, bc, 16, 2, solution=func, num_test=100)
layer_size = [1] + [50] * 3 + [1]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)
model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["l2 relative error"])
losshistory, train_state = model.train(iterations=10000)
# Optional: Save the model during training.
# checkpointer = dde.callbacks.ModelCheckpoint(
# "model/model", verbose=1, save_better_only=True
# )
# Optional: Save the movie of the network solution during training.
# ImageMagick (https://imagemagick.org/) is required to generate the movie.
# movie = dde.callbacks.MovieDumper(
# "model/movie", [-1], [1], period=100, save_spectrum=True, y_reference=func
# )
# losshistory, train_state = model.train(iterations=10000, callbacks=[checkpointer, movie])
dde.saveplot(losshistory, train_state, issave=True, isplot=True)
# Optional: Restore the saved model with the smallest training loss
# model.restore(f"model/model-{train_state.best_step}.ckpt", verbose=1)
# Plot PDE residual
x = geom.uniform_points(1000, True)
y = model.predict(x, operator=pde)
plt.figure()
plt.plot(x, y)
plt.xlabel("x")
plt.ylabel("PDE residual")
plt.show()