Inverse PDE in jax

[outspeckle]

Hello Dr. Lu Lu and Deepxde community,

What do you think would be the best way to use dde.Variable in jax? I tried adding adding a flax.core.variables.Variable for the free parameters and include it in self.net.params in model.py (line 328), but have problems initalizing it correctly. Or would it be better to handle it separately from flax and include it in the loss later?

[lululxvi]

Good question. I haven't looked into this yet, and am not sure what the best solution is.

@ZongrenZou Any ideas?

[ZongrenZou]

Soon. I didn't take inverse problem on priority for backend JAX before. As said, logic behind it for JAX is different from other backends. We'll work this out soon (perhaps in two weeks).

[ZongrenZou]

@outspeckle We just added the support for inverse problem in JAX. You may want to check the "Lorenz_inverse.py" example in "pinn_inverse" folder. Let us know if you have any other questions.

As for the design behind this, we didn't use flax.core.variables.Variable. We simply treat unknowns in inverse problems as JAX arrays. The reason behind it is that, unlike other backends, JAX traces (trainable) variables through functions and their arguments, while, for example, in Tensorflow, all variables, once created, are traced automatically and hence people don't need to trace them explicitly.

[outspeckle]

This is really great, thank you very much! The approach via implicit tracing works very well, its easily possible to access them via model.external_trainable_variables.

[bonneted]

Good morning @ZongrenZou @lululxvi,
I have a problem running Lorenz_inverse.py with JAX. The trainable variables are not updated during the training with JAX but it works fine with Pytorch. Here is the code I run for JAX (I just modified the training for a quick check) :

"""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
)
print(f"C1: {C1}, C2: {C2}, C3: {C3}")
losshistory, train_state = model.train(iterations=1000, callbacks=[variable])
print(f"C1: {C1}, C2: {C2}, C3: {C3}")

# train lbfgs
# 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)

Here is the output.
Pytorch backend :

Using backend: pytorch
Other supported backends: tensorflow.compat.v1, tensorflow, jax, paddle.
paddle supports more examples now and is recommended.

Compiling model...
'compile' took 0.469550 s

C1: 1.0, C2: 1.0, C3: 1.0
Training model...

Step      Train loss                                                                                    Test loss                                                                                     Test metric
0         [9.98e-02, 8.76e-02, 9.65e-03, 6.40e+01, 4.90e+01, 7.29e+02, 3.52e+01, 3.84e+01, 2.56e+02]    [9.98e-02, 8.76e-02, 9.65e-03, 6.40e+01, 4.90e+01, 7.29e+02, 3.52e+01, 3.84e+01, 2.56e+02]    []  
1000      [6.80e+00, 2.55e+00, 2.63e+00, 2.97e+01, 4.73e-03, 1.70e+00, 3.29e+01, 7.45e+00, 1.03e+01]    [6.80e+00, 2.55e+00, 2.63e+00, 2.97e+01, 4.73e-03, 1.70e+00, 3.29e+01, 7.45e+00, 1.03e+01]    []  

Best model at step 1000:
  train loss: 9.41e+01
  test loss: 9.41e+01
  test metric: []

'train' took 9.868791 s

C1: 0.2607625424861908, C2: 1.7225128412246704, C3: 0.32752102613449097

And with JAX (no variable update):

Using backend: jax
Other supported backends: tensorflow.compat.v1, tensorflow, pytorch, paddle.
paddle supports more examples now and is recommended.
Enable just-in-time compilation with XLA.

Compiling model...
'compile' took 0.673463 s

C1: 1.0, C2: 1.0, C3: 1.0
Training model...

Step      Train loss                                                                                    Test loss                                                                                     Test metric
0         [3.80e+00, 3.47e+00, 1.94e-01, 6.40e+01, 4.90e+01, 7.29e+02, 2.31e+01, 3.92e+01, 2.64e+02]    [3.80e+00, 3.47e+00, 1.94e-01, 6.40e+01, 4.90e+01, 7.29e+02, 2.31e+01, 3.92e+01, 2.64e+02]    []  
1000      [5.50e+00, 2.71e+00, 2.99e+00, 3.23e+01, 1.75e-03, 1.68e+00, 3.31e+01, 7.23e+00, 1.06e+01]    [5.50e+00, 2.71e+00, 2.99e+00, 3.23e+01, 1.75e-03, 1.68e+00, 3.31e+01, 7.23e+00, 1.06e+01]    []  

Best model at step 1000:
  train loss: 9.61e+01
  test loss: 9.61e+01
  test metric: []

'train' took 5.746234 s

C1: 1.0, C2: 1.0, C3: 1.0

Do you have the same issue?

[bonneted]

I just realized that I need to access the values using model.external_trainable_variables as explained by @outspeckle and it works fine.
I still have a problem with the callback though, the file is empty (only the iterations column).

[bonneted]

The callback VariableValue was simply not implemented for jax.
I have implemented it but it's not the cleanest because it reset the model at each iteration. I have opened a pull request.