Antiderivative operator from an aligned dataset

Problem setup

We will learn the antiderivative operator

\[G: v \mapsto u\]

defined by an ODE

\[\frac{du(x)}{dx} = v(x), \qquad x \in [0, 1],\]

with IC \(u(0) = 0\).

We learn \(G\) from a dataset. Each data point in the dataset is one pair of \((v, u)\), generated as follows:

  1. A random function \(v\) is sampled from a Gaussian random field (GRF) with the resolution \(m = 100\).

  2. Solve \(u\) for \(v\) numerically. We assume that for each \(u\), we have the values of \(u(x)\) in the same \(N_u = 100\) locations. Because we have the values of \(u(x)\) in the same locations, we call this dataset as “aligned data”.

The datasets can be found at here.

  • The training dataset has size 150.

  • The testing dataset has size 1000.

Implementation

To use DeepONet, we need to organize the dataset in the following format:

  • Input of the branch net: the functions \(v\). It is a matrix of shape (dataset size, \(m\)), e.g., (150, 100) for the training dataset.

  • Input of the trunk net: the locations \(x\) of \(u(x)\) values. It is a matrix of shape (\(N_u\), dimension), i.e., (100, 1) for both training and testing datasets.

  • Output: The values of \(u(x)\) in different locations for different \(v\). It is a matrix of shape (dataset size, \(N_u\)), e.g., (150, 100) for the training dataset.

We first load the training dataset. The input X_train is a tuple; the first element is the branch net input, and the second element is the trunk net input.

d = np.load("antiderivative_aligned_train.npz", allow_pickle=True)
X_train, y_train = (d["X"][0], d["X"][1]), d["y"]

We also load the testing dataset, and then define the data dde.data.TripleCartesianProd:

data = dde.data.TripleCartesianProd(
    X_train=X_train, y_train=y_train, X_test=X_test, y_test=y_test
)

Next we define a DeepONet dde.nn.DeepONetCartesianProd. The branch net is chosen as a fully connected neural network of size [m, 40, 40], and the the trunk net is a fully connected neural network of size [dim_x, 40, 40].

m = 100
dim_x = 1
net = dde.nn.DeepONetCartesianProd(
    [m, 40, 40],
    [dim_x, 40, 40],
    "relu",
    "Glorot normal",
)

We define a Model, and then train the model with Adam and learning rate 0.001 for 10000 iterations:

model = dde.Model(data, net)
model.compile("adam", lr=0.001, metrics=["mean l2 relative error"])
losshistory, train_state = model.train(iterations=10000)

Complete code

"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
import deepxde as dde
import matplotlib.pyplot as plt
import numpy as np

# Load dataset
d = np.load("antiderivative_aligned_train.npz", allow_pickle=True)
X_train = (d["X"][0].astype(np.float32), d["X"][1].astype(np.float32))
y_train = d["y"].astype(np.float32)
d = np.load("antiderivative_aligned_test.npz", allow_pickle=True)
X_test = (d["X"][0].astype(np.float32), d["X"][1].astype(np.float32))
y_test = d["y"].astype(np.float32)

data = dde.data.TripleCartesianProd(
    X_train=X_train, y_train=y_train, X_test=X_test, y_test=y_test
)

# Choose a network
m = 100
dim_x = 1
net = dde.nn.DeepONetCartesianProd(
    [m, 40, 40],
    [dim_x, 40, 40],
    "relu",
    "Glorot normal",
)

# Define a Model
model = dde.Model(data, net)

# Compile and Train
model.compile("adam", lr=0.001, metrics=["mean l2 relative error"])
losshistory, train_state = model.train(iterations=10000)

# Plot the loss trajectory
dde.utils.plot_loss_history(losshistory)
plt.show()