Inverse problem for the Lorenz system
Problem setup
We will solve the Lorenz system:
with initial conditions
The reference solution is here, where the parameters \(\sigma\), \(\rho\), and \(\beta\) are to be identified from observations of the system at certain times and whose true values are 10, 15, and 8/3, respectivly.
Implementation
This description goes through the implementation of a solver for the above Lorenz system step-by-step.
First, the DeepXDE and NumPy (np) modules are imported:
import deepxde as dde
import numpy as np
We also want to define our three unknown variables, \(\sigma\), \(\rho\), and \(\beta\) which will now be called C1, C2, and C3, respectivly. These variables are given an initial guess of 1.0.
C1 = dde.Variable(1.0)
C2 = dde.Variable(1.0)
C3 = dde.Variable(1.0)
Now we can begin by creating a TimeDomain class.
geom = dde.geometry.TimeDomain(0, 3)
Next, we create the Lorenz system to solve using the dde.grad.jacobian function.
def Lorenz_system(x, y):
y1, y2, y3 = y[:, 0:1], y[:, 1:2], y[:, 2:]
dy1_x = dde.grad.jacobian(y, x, i=0)
dy2_x = dde.grad.jacobian(y, x, i=1)
dy3_x = dde.grad.jacobian(y, x, i=2)
return [
dy1_x - C1 * (y2 - y1),
dy2_x - y1 * (C2 - y3) + y2,
dy3_x - y1 * y2 + C3 * y3,
]
The first argument to Lorenz_system is the network input, i.e., the \(t\)-coordinate. The second argument is the network output, i.e., the solution \(y(x,y,z)\), but here we use y1, y2, y3 as the name of the coordinates x, y, and z, which correspond to the columns of datapoints in the 2D array, \(y\).
Next, we consider the initial conditions. We need to implement a function, which should return True for points inside the subdomain and False for the points outside.
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 (\(x\)), the second component (\(y\)), or the third component (\(z\)). Note that in our case, the point \(t\) of the initial condition is \(t = 0\).
ic1 = dde.icbc.IC(geom, lambda X: -8, boundary, component=0)
ic2 = dde.icbc.IC(geom, lambda X: 7, boundary, component=1)
ic3 = dde.icbc.IC(geom, lambda X: 27, boundary, component=2)
Now me must assign the data from Lorenz.npz to the corresponding \(t\), \(x\), \(y\), and \(z\) values for training. First we retrieve the data to train the model. The data is split into "t" and "y", which correspond to time datapoints and the cartesian coodrinate datapoints (x, y, and z), respectivly.
def gen_traindata():
data = np.load("dataset/Lorenz.npz")
return data["t"], data["y"]
Then we organize and assign the train data.
observe_t, ob_y = gen_traindata()
observe_y0 = dde.icbc.PointSetBC(observe_t, ob_y[:, 0:1], component=0)
observe_y1 = dde.icbc.PointSetBC(observe_t, ob_y[:, 1:2], component=1)
observe_y2 = dde.icbc.PointSetBC(observe_t, ob_y[:, 2:3], component=2)
Now that the problem is fully setup, we define the PDE as:
data = dde.data.PDE(
geom,
Lorenz_system,
[ic1, ic2, ic3, observe_y0, observe_y1, observe_y2],
num_domain=400,
num_boundary=2,
anchors=observe_t,
)
Where num_domain is the number of points inside the domain, and num_boundary is the number of points on the boundary. anchors are extra points beyond num_domain and num_boundary used for training.
Next, we choose the network. Here, we use a fully connected neural network of depth 4 (i.e., 3 hidden layers) and width 40:
net = dde.nn.FNN([1] + [40] * 3 + [3], "tanh", "Glorot uniform")
Now that the PDE problem and network have been created, we build a Model and choose the optimizer, learning rate, and provide the trainable variables C1, C2, and C3:
model = dde.Model(data, net)
model.compile("adam", lr=0.001, external_trainable_variables=[C1, C2, C3])
variable = dde.callbacks.VariableValue(
[C1, C2, C3], period=600, filename="variables.dat"
)
We then train the model for 60000 iterations:
losshistory, train_state = model.train(iterations=60000, callbacks=[variable])
Complete code
"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, jax"""
import deepxde as dde
import numpy as np
def gen_traindata():
data = np.load("../dataset/Lorenz.npz")
return data["t"], data["y"]
C1 = dde.Variable(1.0)
C2 = dde.Variable(1.0)
C3 = dde.Variable(1.0)
# Most backends
def Lorenz_system(x, y):
"""Lorenz system.
dy1/dx = 10 * (y2 - y1)
dy2/dx = y1 * (15 - y3) - y2
dy3/dx = y1 * y2 - 8/3 * y3
"""
y1, y2, y3 = y[:, 0:1], y[:, 1:2], y[:, 2:]
dy1_x = dde.grad.jacobian(y, x, i=0)
dy2_x = dde.grad.jacobian(y, x, i=1)
dy3_x = dde.grad.jacobian(y, x, i=2)
return [
dy1_x - C1 * (y2 - y1),
dy2_x - y1 * (C2 - y3) + y2,
dy3_x - y1 * y2 + C3 * y3,
]
# Backend JAX
# def Lorenz_system(x, y, unknowns=[C1, C2, C3]):
# C1, C2, C3 = unknowns
# y_val, y_fn = y
# y1, y2, y3 = y_val[:, 0:1], y_val[:, 1:2], y_val[:, 2:3]
# dy1_x, _ = dde.grad.jacobian(y, x, i=0)
# dy2_x, _ = dde.grad.jacobian(y, x, i=1)
# dy3_x, _ = dde.grad.jacobian(y, x, i=2)
# return [
# dy1_x - C1 * (y2 - y1),
# dy2_x - y1 * (C2 - y3) + y2,
# dy3_x - y1 * y2 + C3 * y3,
# ]
def boundary(_, on_initial):
return on_initial
geom = dde.geometry.TimeDomain(0, 3)
# Initial conditions
ic1 = dde.icbc.IC(geom, lambda X: -8, boundary, component=0)
ic2 = dde.icbc.IC(geom, lambda X: 7, boundary, component=1)
ic3 = dde.icbc.IC(geom, lambda X: 27, boundary, component=2)
# Get the train data
observe_t, ob_y = gen_traindata()
observe_y0 = dde.icbc.PointSetBC(observe_t, ob_y[:, 0:1], component=0)
observe_y1 = dde.icbc.PointSetBC(observe_t, ob_y[:, 1:2], component=1)
observe_y2 = dde.icbc.PointSetBC(observe_t, ob_y[:, 2:3], component=2)
data = dde.data.PDE(
geom,
Lorenz_system,
[ic1, ic2, ic3, observe_y0, observe_y1, observe_y2],
num_domain=400,
num_boundary=2,
anchors=observe_t,
)
net = dde.nn.FNN([1] + [40] * 3 + [3], "tanh", "Glorot uniform")
model = dde.Model(data, net)
external_trainable_variables = [C1, C2, C3]
variable = dde.callbacks.VariableValue(
external_trainable_variables, period=600, filename="variables.dat"
)
# train adam
model.compile(
"adam", lr=0.001, external_trainable_variables=external_trainable_variables
)
losshistory, train_state = model.train(iterations=20000, callbacks=[variable])
# train lbfgs (not implemented in JAX)
model.compile("L-BFGS", external_trainable_variables=external_trainable_variables)
losshistory, train_state = model.train(callbacks=[variable])
dde.saveplot(losshistory, train_state, issave=True, isplot=True)