Lotka-Volterra equation
Problem setup
We will solve a Lotka-Volterra equation:
with the initial condition
and two user-specified parameters
the first of which approximates the upper bound of the range, and the second is the right bound of the domain. These two will be used for scaling.
The reference solution is generated using integrate.solve_ivp() from scipy.
Implementation
This description goes through the implementation of a solver for the above described Lotka-Volterra equation step-by-step.
First, the DeepXDE, numpy, matplotlib, and scipy packages are imported:
import deepxde as dde
import numpy as np
import matplotlib.pyplot as plt
from scipy import integrate
from deepxde.backend import tf
We begin by defining the approximate upper bound of the range and the right bound of the domain. Later, we scale by these factors to obtain a graph between 0 and 1 when graphing population vs. time.
ub = 200
rb = 20
We now define a time domain. We do this by using the built-in class TimeDomain:
geom = dde.geometry.TimeDomain(0.0, 1.0)
Next, we express the ODE system:
def ode_system(x, y):
r = y[:, 0:1]
p = y[:, 1:2]
dr_t = dde.grad.jacobian(y, x, i=0)
dp_t = dde.grad.jacobian(y, x, i=1)
return [
dr_t - 1 / ub * rb * (2.0 * ub * r - 0.04 * ub * r * ub * p),
dp_t - 1 / ub * rb * (0.02 * r * ub * p * ub - 1.06 * p * ub),
]
The first argument to ode_system is the \(t\)-coordinate, represented by x. The second argument is a 2-dimensional vector, represented as y, which contains \(r(t)\) and \(p(t)\).
Now, we define the ODE problem as
data = dde.data.PDE(geom, ode_system, [], 3000, 2, num_test = 3000)
Note that when solving this equation, we want to have hard constraints on the initial conditions, so we define this later when creating the network rather than as part of the PDE.
We have 3000 training residual points inside the domain and 2 points on the boundary. We use 3000 points for testing the ODE residual. We now create the network:
layer_size = [1] + [64] * 6 + [2]
activation = "tanh"
initializer = "Glorot normal"
net = dde.nn.FNN(layer_size, activation, initializer)
This is a neural network of depth 7 with 6 hidden layers of width 50. We use \(\tanh\) as the activation function. Since we expect to have periodic behavior in the Lotka-Volterra equation, we add a feature layer with \(\sin(kt)\). This forces the prediction to be periodic and therefore more accurate.
def input_transform(t):
return tf.concat(
(
t,
tf.sin(t),
tf.sin(2 * t),
tf.sin(3 * t),
tf.sin(4 * t),
tf.sin(5 * t),
tf.sin(6 * t),
),
axis=1,
)
As mentioned earlier, we want the initial conditions \(r(0)=\frac{100}{U}\) and \(p(0)=\frac{15}{U}\) to be hard constraints, so we transform the output:
def output_transform(t, y):
y1 = y[:, 0:1]
y2 = y[:, 1:2]
return tf.concat(
[y1 * tf.tanh(t) + 100 / ub, y2 * tf.tanh(t) + 15 / ub], axis=1
)
We add these layers:
net.apply_feature_transform(input_transform)
net.apply_output_transform(output_transform)
Now that we have defined the neural network, we build a Model, choose the optimizer and learning rate, and train it for 50000 iterations:
model = dde.Model(data, net)
model.compile("adam", lr=0.001)
losshistory, train_state = model.train(iterations=50000)
After training with Adam, we continue with L-BFGS to have an even smaller loss:
model.compile("L-BFGS")
losshistory, train_state = model.train()
dde.saveplot(losshistory, train_state, issave=True, isplot=True)
Complete code
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle, jax"""
import deepxde as dde
import matplotlib.pyplot as plt
import numpy as np
from scipy import integrate
# Import tf if using backend tensorflow.compat.v1 or tensorflow
from deepxde.backend import tf
# Import torch if using backend pytorch
# import torch
# Import paddle if using backend paddle
# import paddle
# Import jax.numpy if using backend jax
# import jax.numpy as jnp
ub = 200
rb = 20
def func(t, r):
x, y = r
dx_t = 1 / ub * rb * (2.0 * ub * x - 0.04 * ub * x * ub * y)
dy_t = 1 / ub * rb * (0.02 * ub * x * ub * y - 1.06 * ub * y)
return dx_t, dy_t
def gen_truedata():
t = np.linspace(0, 1, 100)
sol = integrate.solve_ivp(func, (0, 10), (100 / ub, 15 / ub), t_eval=t)
x_true, y_true = sol.y
x_true = x_true.reshape(100, 1)
y_true = y_true.reshape(100, 1)
return x_true, y_true
def ode_system(x, y):
# Most backends
r = y[:, 0:1]
p = y[:, 1:2]
dr_t = dde.grad.jacobian(y, x, i=0)
dp_t = dde.grad.jacobian(y, x, i=1)
# Backend jax
# y_val, y_fn = y
# r = y_val[:, 0:1]
# p = y_val[:, 1:2]
# dr_t, _ = dde.grad.jacobian(y, x, i=0)
# dp_t, _ = dde.grad.jacobian(y, x, i=1)
return [
dr_t - 1 / ub * rb * (2.0 * ub * r - 0.04 * ub * r * ub * p),
dp_t - 1 / ub * rb * (0.02 * r * ub * p * ub - 1.06 * p * ub),
]
geom = dde.geometry.TimeDomain(0, 1.0)
data = dde.data.PDE(geom, ode_system, [], 3000, 2, num_test=3000)
layer_size = [1] + [64] * 6 + [2]
activation = "tanh"
initializer = "Glorot normal"
net = dde.nn.FNN(layer_size, activation, initializer)
# Backend tensorflow.compat.v1 or tensorflow
def input_transform(t):
return tf.concat(
(
t,
tf.sin(t),
tf.sin(2 * t),
tf.sin(3 * t),
tf.sin(4 * t),
tf.sin(5 * t),
tf.sin(6 * t),
),
axis=1,
)
# Backend pytorch
# def input_transform(t):
# return torch.cat(
# [
# torch.sin(t),
# ],
# dim=1,
# )
# Backend paddle
# def input_transform(t):
# return paddle.concat(
# (
# paddle.sin(t),
# ),
# axis=1,
# )
# Backend jax
# def input_transform(t):
# if t.ndim == 1:
# t = t[None]
#
# return jnp.concatenate(
# [
# jnp.sin(t),
# ],
# axis=1
# )
# hard constraints: x(0) = 100, y(0) = 15
# Backend tensorflow.compat.v1 or tensorflow
def output_transform(t, y):
y1 = y[:, 0:1]
y2 = y[:, 1:2]
return tf.concat([y1 * tf.tanh(t) + 100 / ub, y2 * tf.tanh(t) + 15 / ub], axis=1)
# Backend pytorch
# def output_transform(t, y):
# y1 = y[:, 0:1]
# y2 = y[:, 1:2]
# return torch.cat([y1 * torch.tanh(t) + 100 / ub, y2 * torch.tanh(t) + 15 / ub], dim=1)
# Backend paddle
# def output_transform(t, y):
# y1 = y[:, 0:1]
# y2 = y[:, 1:2]
# return paddle.concat([y1 * paddle.tanh(t) + 100 / ub, y2 * paddle.tanh(t) + 15 / ub], axis=1)
# Backend jax
# def output_transform(t, y):
# y1 = y[:, 0:1]
# y2 = y[:, 1:2]
# return jnp.concatenate(
# [y1 * jnp.tanh(t) + 100 / ub, y2 * jnp.tanh(t) + 15 / ub],
# axis=1
# ).squeeze()
net.apply_feature_transform(input_transform)
net.apply_output_transform(output_transform)
model = dde.Model(data, net)
model.compile("adam", lr=0.001)
losshistory, train_state = model.train(iterations=50000)
# Most backends except jax can have a second fine tuning of the solution
model.compile("L-BFGS")
losshistory, train_state = model.train()
dde.saveplot(losshistory, train_state, issave=True, isplot=True)
plt.xlabel("t")
plt.ylabel("population")
t = np.linspace(0, 1, 100)
x_true, y_true = gen_truedata()
plt.plot(t, x_true, color="black", label="x_true")
plt.plot(t, y_true, color="blue", label="y_true")
t = t.reshape(100, 1)
sol_pred = model.predict(t)
x_pred = sol_pred[:, 0:1]
y_pred = sol_pred[:, 1:2]
plt.plot(t, x_pred, color="red", linestyle="dashed", label="x_pred")
plt.plot(t, y_pred, color="orange", linestyle="dashed", label="y_pred")
plt.legend()
plt.show()