Skip to content

Commit 683682c

Browse files
author
Samuel Burbulla
authored
Add PI-DeepONet example for 1D Poisson equation (#1311)
1 parent de933f4 commit 683682c

File tree

3 files changed

+205
-0
lines changed

3 files changed

+205
-0
lines changed
238 KB
Loading
Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
Poisson equation in 1D
2+
======================
3+
4+
Problem setup
5+
-------------
6+
7+
We will learn the solution operator
8+
9+
.. math:: G: f \mapsto u
10+
11+
for the one-dimensional Poisson problem
12+
13+
.. math:: u''(x) = f(x), \qquad x \in [0, 1],
14+
15+
with zero Dirichlet boundary conditions :math:`u(0) = u(1) = 0`.
16+
17+
The source term :math:`f` is supposed to be an arbitrary continuous function.
18+
19+
20+
Implementation
21+
--------------
22+
23+
The solution operator can be learned by training a physics-informed DeepONet.
24+
25+
First, we define the PDE with boundary conditions and the domain:
26+
27+
.. code-block:: python
28+
29+
def equation(x, y, f):
30+
dy_xx = dde.grad.hessian(y, x)
31+
return -dy_xx - f
32+
33+
geom = dde.geometry.Interval(0, 1)
34+
35+
def u_boundary(_):
36+
return 0
37+
38+
def boundary(_, on_boundary):
39+
return on_boundary
40+
41+
bc = dde.icbc.DirichletBC(geom, u_boundary, boundary)
42+
43+
pde = dde.data.PDE(geom, equation, bc, num_domain=100, num_boundary=2)
44+
45+
46+
Next, we specify the function space for :math:`f` and the corresponding evaluation points.
47+
For this example, we use the ``dde.data.PowerSeries`` to get the function space
48+
of polynomials of degree three.
49+
Together with the PDE, the function space is used to define a
50+
PDEOperator ``dde.data.PDEOperatorCartesianProd`` that incorporates the PDE into
51+
the loss function.
52+
53+
.. code-block:: python
54+
55+
degree = 3
56+
space = dde.data.PowerSeries(N=degree + 1)
57+
58+
num_eval_points = 10
59+
evaluation_points = geom.uniform_points(num_eval_points, boundary=True)
60+
61+
pde_op = dde.data.PDEOperatorCartesianProd(
62+
pde,
63+
space,
64+
evaluation_points,
65+
num_function=100,
66+
)
67+
68+
69+
The DeepONet can be defined using ``dde.nn.DeepONetCartesianProd``.
70+
The branch net is chosen as a fully connected neural network of size ``[m, 32, p]`` where ``p=32``
71+
and the trunk net is a fully connected neural network of size ``[dim_x, 32, p]``.
72+
73+
.. code-block:: python
74+
75+
dim_x = 1
76+
p = 32
77+
net = dde.nn.DeepONetCartesianProd(
78+
[num_eval_points, 32, p],
79+
[dim_x, 32, p],
80+
activation="tanh",
81+
kernel_initializer="Glorot normal",
82+
)
83+
84+
85+
We define the ``Model`` and train it with L-BFGS:
86+
87+
.. code-block:: python
88+
89+
model = dde.Model(pde_op, net)
90+
dde.optimizers.set_LBFGS_options(maxiter=1000)
91+
model.compile("L-BFGS")
92+
model.train()
93+
94+
Finally, the trained model can be used to predict the solution of the Poisson
95+
equation. We sample the solution for three random representations of :math:`f`.
96+
97+
.. code-block:: python
98+
99+
n = 3
100+
features = space.random(n)
101+
fx = space.eval_batch(features, evaluation_points)
102+
103+
x = geom.uniform_points(100, boundary=True)
104+
y = model.predict((fx, x))
105+
106+
107+
![](pideeponet_poisson1d.png)
108+
109+
110+
Complete code
111+
-------------
112+
113+
.. literalinclude:: ../../../examples/operator/pideeponet_1d_poisson.py
114+
:language: python
Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,91 @@
1+
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch"""
2+
import deepxde as dde
3+
import matplotlib.pyplot as plt
4+
import numpy as np
5+
6+
7+
# Poisson equation: -u_xx = f
8+
def equation(x, y, f):
9+
dy_xx = dde.grad.hessian(y, x)
10+
return -dy_xx - f
11+
12+
13+
# Domain is interval [0, 1]
14+
geom = dde.geometry.Interval(0, 1)
15+
16+
17+
# Zero Dirichlet BC
18+
def u_boundary(_):
19+
return 0
20+
21+
22+
def boundary(_, on_boundary):
23+
return on_boundary
24+
25+
26+
bc = dde.icbc.DirichletBC(geom, u_boundary, boundary)
27+
28+
# Define PDE
29+
pde = dde.data.PDE(geom, equation, bc, num_domain=100, num_boundary=2)
30+
31+
# Function space for f(x) are polynomials
32+
degree = 3
33+
space = dde.data.PowerSeries(N=degree + 1)
34+
35+
# Choose evaluation points
36+
num_eval_points = 10
37+
evaluation_points = geom.uniform_points(num_eval_points, boundary=True)
38+
39+
# Define PDE operator
40+
pde_op = dde.data.PDEOperatorCartesianProd(
41+
pde,
42+
space,
43+
evaluation_points,
44+
num_function=100,
45+
)
46+
47+
# Setup DeepONet
48+
dim_x = 1
49+
p = 32
50+
net = dde.nn.DeepONetCartesianProd(
51+
[num_eval_points, 32, p],
52+
[dim_x, 32, p],
53+
activation="tanh",
54+
kernel_initializer="Glorot normal",
55+
)
56+
57+
# Define and train model
58+
model = dde.Model(pde_op, net)
59+
dde.optimizers.set_LBFGS_options(maxiter=1000)
60+
model.compile("L-BFGS")
61+
model.train()
62+
63+
# Plot realisations of f(x)
64+
n = 3
65+
features = space.random(n)
66+
fx = space.eval_batch(features, evaluation_points)
67+
68+
x = geom.uniform_points(100, boundary=True)
69+
y = model.predict((fx, x))
70+
71+
# Setup figure
72+
fig = plt.figure(figsize=(7, 8))
73+
plt.subplot(2, 1, 1)
74+
plt.title("Poisson equation: Source term f(x) and solution u(x)")
75+
plt.ylabel("f(x)")
76+
z = np.zeros_like(x)
77+
plt.plot(x, z, "k-", alpha=0.1)
78+
79+
# Plot source term f(x)
80+
for i in range(n):
81+
plt.plot(evaluation_points, fx[i], "--")
82+
83+
# Plot solution u(x)
84+
plt.subplot(2, 1, 2)
85+
plt.ylabel("u(x)")
86+
plt.plot(x, z, "k-", alpha=0.1)
87+
for i in range(n):
88+
plt.plot(x, y[i], "-")
89+
plt.xlabel("x")
90+
91+
plt.show()

0 commit comments

Comments
 (0)