Skip to content

Commit

Permalink
Add PI-DeepONet example for 1D Poisson equation (#1311)
Browse files Browse the repository at this point in the history
  • Loading branch information
samuelburbulla authored Aug 23, 2023
1 parent de933f4 commit 683682c
Show file tree
Hide file tree
Showing 3 changed files with 205 additions and 0 deletions.
Binary file added docs/demos/operator/poisson_1d_pideeponet.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
114 changes: 114 additions & 0 deletions docs/demos/operator/poisson_1d_pideeponet.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,114 @@
Poisson equation in 1D
======================

Problem setup
-------------

We will learn the solution operator

.. math:: G: f \mapsto u

for the one-dimensional Poisson problem

.. math:: u''(x) = f(x), \qquad x \in [0, 1],

with zero Dirichlet boundary conditions :math:`u(0) = u(1) = 0`.

The source term :math:`f` is supposed to be an arbitrary continuous function.


Implementation
--------------

The solution operator can be learned by training a physics-informed DeepONet.

First, we define the PDE with boundary conditions and the domain:

.. code-block:: python
def equation(x, y, f):
dy_xx = dde.grad.hessian(y, x)
return -dy_xx - f
geom = dde.geometry.Interval(0, 1)
def u_boundary(_):
return 0
def boundary(_, on_boundary):
return on_boundary
bc = dde.icbc.DirichletBC(geom, u_boundary, boundary)
pde = dde.data.PDE(geom, equation, bc, num_domain=100, num_boundary=2)
Next, we specify the function space for :math:`f` and the corresponding evaluation points.
For this example, we use the ``dde.data.PowerSeries`` to get the function space
of polynomials of degree three.
Together with the PDE, the function space is used to define a
PDEOperator ``dde.data.PDEOperatorCartesianProd`` that incorporates the PDE into
the loss function.

.. code-block:: python
degree = 3
space = dde.data.PowerSeries(N=degree + 1)
num_eval_points = 10
evaluation_points = geom.uniform_points(num_eval_points, boundary=True)
pde_op = dde.data.PDEOperatorCartesianProd(
pde,
space,
evaluation_points,
num_function=100,
)
The DeepONet can be defined using ``dde.nn.DeepONetCartesianProd``.
The branch net is chosen as a fully connected neural network of size ``[m, 32, p]`` where ``p=32``
and the trunk net is a fully connected neural network of size ``[dim_x, 32, p]``.

.. code-block:: python
dim_x = 1
p = 32
net = dde.nn.DeepONetCartesianProd(
[num_eval_points, 32, p],
[dim_x, 32, p],
activation="tanh",
kernel_initializer="Glorot normal",
)
We define the ``Model`` and train it with L-BFGS:

.. code-block:: python
model = dde.Model(pde_op, net)
dde.optimizers.set_LBFGS_options(maxiter=1000)
model.compile("L-BFGS")
model.train()
Finally, the trained model can be used to predict the solution of the Poisson
equation. We sample the solution for three random representations of :math:`f`.

.. code-block:: python
n = 3
features = space.random(n)
fx = space.eval_batch(features, evaluation_points)
x = geom.uniform_points(100, boundary=True)
y = model.predict((fx, x))
![](pideeponet_poisson1d.png)


Complete code
-------------

.. literalinclude:: ../../../examples/operator/pideeponet_1d_poisson.py
:language: python
91 changes: 91 additions & 0 deletions examples/operator/poisson_1d_pideeponet.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch"""
import deepxde as dde
import matplotlib.pyplot as plt
import numpy as np


# Poisson equation: -u_xx = f
def equation(x, y, f):
dy_xx = dde.grad.hessian(y, x)
return -dy_xx - f


# Domain is interval [0, 1]
geom = dde.geometry.Interval(0, 1)


# Zero Dirichlet BC
def u_boundary(_):
return 0


def boundary(_, on_boundary):
return on_boundary


bc = dde.icbc.DirichletBC(geom, u_boundary, boundary)

# Define PDE
pde = dde.data.PDE(geom, equation, bc, num_domain=100, num_boundary=2)

# Function space for f(x) are polynomials
degree = 3
space = dde.data.PowerSeries(N=degree + 1)

# Choose evaluation points
num_eval_points = 10
evaluation_points = geom.uniform_points(num_eval_points, boundary=True)

# Define PDE operator
pde_op = dde.data.PDEOperatorCartesianProd(
pde,
space,
evaluation_points,
num_function=100,
)

# Setup DeepONet
dim_x = 1
p = 32
net = dde.nn.DeepONetCartesianProd(
[num_eval_points, 32, p],
[dim_x, 32, p],
activation="tanh",
kernel_initializer="Glorot normal",
)

# Define and train model
model = dde.Model(pde_op, net)
dde.optimizers.set_LBFGS_options(maxiter=1000)
model.compile("L-BFGS")
model.train()

# Plot realisations of f(x)
n = 3
features = space.random(n)
fx = space.eval_batch(features, evaluation_points)

x = geom.uniform_points(100, boundary=True)
y = model.predict((fx, x))

# Setup figure
fig = plt.figure(figsize=(7, 8))
plt.subplot(2, 1, 1)
plt.title("Poisson equation: Source term f(x) and solution u(x)")
plt.ylabel("f(x)")
z = np.zeros_like(x)
plt.plot(x, z, "k-", alpha=0.1)

# Plot source term f(x)
for i in range(n):
plt.plot(evaluation_points, fx[i], "--")

# Plot solution u(x)
plt.subplot(2, 1, 2)
plt.ylabel("u(x)")
plt.plot(x, z, "k-", alpha=0.1)
for i in range(n):
plt.plot(x, y[i], "-")
plt.xlabel("x")

plt.show()

0 comments on commit 683682c

Please sign in to comment.