Add pinnx submodule

[chaoming0625]

In response to chaobrain/pinnx#19, #1904, and #1899, I am trying to merge pinnx into DeepXDE.

This pull request includes a submodule deepxde.pinnx, which enables explicit variables and physical units in physics-informed neural networks.

A code example:

import brainstate as bst
import brainunit as u
from deepxde import pinnx

# geometry
geometry = pinnx.geometry.GeometryXTime(
    geometry=pinnx.geometry.Interval(-1, 1.),
    timedomain=pinnx.geometry.TimeDomain(0, 0.99)
).to_dict_point(x=u.meter, t=u.second)

uy = u.meter / u.second
v = 0.01 / u.math.pi * u.meter ** 2 / u.second

# boundary conditions
bc = pinnx.icbc.DirichletBC(lambda x: {'y': 0. * uy})
ic = pinnx.icbc.IC(lambda x: {'y': -u.math.sin(u.math.pi * x['x'] / u.meter) * uy})

# PDE equation
def pde(x, y):
    jacobian = approximator.jacobian(x)
    hessian = approximator.hessian(x)
    dy_x = jacobian['y']['x']
    dy_t = jacobian['y']['t']
    dy_xx = hessian['y']['x']['x']
    residual = dy_t + y['y'] * dy_x - v * dy_xx
    return residual

# neural network
approximator = pinnx.nn.Model(
    pinnx.nn.DictToArray(x=u.meter, t=u.second),
    pinnx.nn.FNN(
        [geometry.dim] + [20] * 3 + [1],
        "tanh",
        bst.init.KaimingUniform()
    ),
    pinnx.nn.ArrayToDict(y=uy)
)

# problem
problem = pinnx.problem.TimePDE(
    geometry,
    pde,
    [bc, ic],
    approximator,
    num_domain=2540,
    num_boundary=80,
    num_initial=160,
)

# training
trainer = pinnx.Trainer(problem)
trainer.compile(bst.optim.Adam(1e-3)).train(iterations=15000)
trainer.compile(bst.optim.LBFGS(1e-3)).train(2000, display_every=500)
trainer.saveplot(issave=True, isplot=True)

This new submodule support most of PINN examples in DeepXDE.

The documentation is hosted in docs/pinnx_docs directory. The ample examples are included in examples/pinnx_examples.

[lululxvi]

Hi @chaoming0625 , thank you for adding the code. It is nice.

My first suggestion is that as there is some overlapped code between DeepXDE and PINNx, such as geometry, we should reduce the duplication as much as possible. You can import and reuse the code in DeepXDE.

[chaoming0625]

Hi, @lululxvi , I have made some changes. However, deepxde.pinnx.geometry is necessary, since most of the points are generated using jax.numpy rather than numpy. These changes can enable the model compatible with bfloat16.

[chaoming0625]

In deepxde.pinnx, changing the traning precision is very easy. Setting

import brainstate as bst


bst.environ.set(precision='b16')

Models will be trained using bfloat16.

Here is an example:

import brainstate as bst
import brainunit as u
import optax
from deepxde import pinnx

bst.environ.set(precision='b16')  # or, 32, 64, 16

geometry = pinnx.geometry.GeometryXTime(
    geometry=pinnx.geometry.Interval(-1, 1.),
    timedomain=pinnx.geometry.TimeDomain(0, 0.99)
).to_dict_point(x=u.meter, t=u.second)

uy = u.meter / u.second
bc = pinnx.icbc.DirichletBC(lambda x: {'y': 0. * uy})
ic = pinnx.icbc.IC(lambda x: {'y': -u.math.sin(u.math.pi * x['x'] / u.meter) * uy})

v = 0.01 / u.math.pi * u.meter ** 2 / u.second


def pde(x, y):
    jacobian = approximator.jacobian(x)
    hessian = approximator.hessian(x)
    dy_x = jacobian['y']['x']
    dy_t = jacobian['y']['t']
    dy_xx = hessian['y']['x']['x']
    residual = dy_t + y['y'] * dy_x - v * dy_xx
    return residual


approximator = pinnx.nn.Model(
    pinnx.nn.DictToArray(x=u.meter, t=u.second),
    pinnx.nn.FNN(
        [geometry.dim] + [20] * 3 + [1],
        "tanh",
    ),
    pinnx.nn.ArrayToDict(y=uy)
)

problem = pinnx.problem.TimePDE(
    geometry,
    pde,
    [bc, ic],
    approximator,
    num_domain=2540,
    num_boundary=80,
    num_initial=160,
)

trainer = pinnx.Trainer(problem, )
trainer.compile(bst.optim.OptaxOptimizer(optax.adamw(1e-3)))
trainer.train(iterations=10000)

[chaoming0625]

Therefore, hosting geometry submodule is essential for deepxde.pinnx module.

[chaoming0625]

Hi, @lululxvi , in the next, my colleague will continue to promote the integration of 'pinnx' into 'deepxde', and he will focus more on solving every problem in the integration.

[lululxvi]

Hi, @lululxvi , in the next, my colleague will continue to promote the integration of 'pinnx' into 'deepxde', and he will focus more on solving every problem in the integration.

No problem.

[lululxvi]

DeepXDE uses LGPL-2.1 license. Can we also use LGPL-2.1 license in these files?

[chaoming0625]

Sorry for this line. It should be changed to DeepXDE's license.

[lululxvi]

Then we just remove these comments

[Routhleck]

Hi @lululxvi, I have just changed to LGPL-2.1 license in these files.

[lululxvi]

How about we simply remove the license info in source code? As DeepXDE has the license file at https://github.com/lululxvi/deepxde/blob/master/LICENSE

[Routhleck]

Sure, I have removed the liccense info in source code.