A simple ODE system

Problem setup

We will solve a simple ODE system:

\[\frac{dy_1}{dt} = y_2, \qquad \frac{dy_2}{dt} = - y_1, \qquad \text{where} \quad t \in [0,10],\]

with the initial conditions

\[y_1(0) = 0, \quad y_2(0) = 1.\]

The reference solution is \(y_1 = \sin(t), \quad y_2 = \cos(t)\).

Implementation

This description goes through the implementation of a solver for the above ODE system step-by-step.

First, the DeepXDE and NumPy (np) modules are imported:

import deepxde as dde
import numpy as np

We begin by defining a computational geometry. We can use a built-in class TimeDomain to define a time domain as follows

geom = dde.geometry.TimeDomain(0, 10)

Next, we express the ODE system:

def ode_system(x, y):
    y1, y2 = y[:, 0:1], y[:, 1:]
    dy1_x = dde.grad.jacobian(y, x, i=0)
    dy2_x = dde.grad.jacobian(y, x, i=1)
    return [dy1_x - y2, dy2_x + y1]

The first argument to ode_system is the network input, i.e., the \(t\)-coordinate, and here we represent it as x. The second argument to ode_system is the network output, which is a 2-dimensional vector where the first component(y[:, 0:1]) is \(y_1\) and the second component (y[:, 1:]) is \(y_2\).

Next, we consider the initial condition. We need to implement a function, which should return True for points inside the subdomain and False for the points outside. In our case, the point \(t\) of the initial condition is \(t = 0\). (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, on_initial):
    return dde.utils.isclose(x[0], 0)

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_initial is used as the second argument. If the point \(t = 0\), then on_initial is True, otherwise, on_initial is False. Thus, we can also define boundary in a simpler way:

def boundary(_, on_initial):
    return on_initial

Then the initial conditions are specified using the computational domain, initial function and boundary. The argument component refers to if this IC is for the first component or the second component.

ic1 = dde.icbc.IC(geom, lambda x: 0, boundary, component=0)
ic2 = dde.icbc.IC(geom, lambda x: 1, boundary, component=1)

Now, we have specified the geometry, ODEs, and initial conditions. Since PDE is also an ODE solver, we then define the ODE problem as

data = dde.data.PDE(geom, ode_system, [ic1, ic2], 35, 2, solution=func, num_test=100)

The number 35 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 left and right endpoints of the time domain in this case). We use 100 points for testing the ODE residual. The argument solution=func is the reference solution to compute the error of our solution, and we define it as follows:

def func(x):
    return np.hstack((np.sin(x), np.cos(x)))

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 + [2]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)

Now, we have the ODE 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 then train the model for 20000 iterations:

losshistory, train_state = model.train(iterations=20000)

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


def ode_system(x, y):
    """ODE system.
    dy1/dx = y2
    dy2/dx = -y1
    """
    # Most backends
    y1, y2 = y[:, 0:1], y[:, 1:]
    dy1_x = dde.grad.jacobian(y, x, i=0)
    dy2_x = dde.grad.jacobian(y, x, i=1)
    # Backend jax
    # y_val, y_fn = y
    # y1, y2 = y_val[:, 0:1], y_val[:, 1:]
    # dy1_x, _ = dde.grad.jacobian(y, x, i=0)
    # dy2_x, _ = dde.grad.jacobian(y, x, i=1)
    return [dy1_x - y2, dy2_x + y1]


def boundary(_, on_initial):
    return on_initial


def func(x):
    """
    y1 = sin(x)
    y2 = cos(x)
    """
    return np.hstack((np.sin(x), np.cos(x)))


geom = dde.geometry.TimeDomain(0, 10)
ic1 = dde.icbc.IC(geom, lambda x: 0, boundary, component=0)
ic2 = dde.icbc.IC(geom, lambda x: 1, boundary, component=1)
data = dde.data.PDE(geom, ode_system, [ic1, ic2], 35, 2, solution=func, num_test=100)

layer_size = [1] + [50] * 3 + [2]
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=20000)

dde.saveplot(losshistory, train_state, issave=True, isplot=True)