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.

Motivation

pinnx module was engineered to address a pivotal challenge prevalent in current PINN libraries: the absence of explicit physical semantics.

• Enhanced variable management with semantic clarity

Existing PINN libraries, such as DeepXDE, necessitate that users manually track the order and significance of variables. For instance, within these frameworks, variables[0] might denote amplitude, variables[1] for frequency, and so forth. This method lacks an intrinsic link between the sequence of variables and their physical meanings, thereby increasing complexity and the potential for errors. In stark contrast, pinnx empowers users to assign clear and meaningful names to variables (e.g., variables["amplitude"] and variables["frequency"]). This approach obviates the need for manual management of variable order, enhancing both code readability and maintainability.

• Simplified gradient relationship management

Another significant limitation of existing PINN libraries is their reliance on users to manage intricate Jacobian and Hessian relationships between variables. This process is not only cumbersome but also prone to mistakes. PINNx revolutionizes this aspect by employing a straightforward dictionary indexing mechanism to track intuitive gradient relationships. For example, users can effortlessly access the Hessian matrix element ${\partial }^{2}y/\partial x^2$ via hessian["y"]["x"] and the Jacobian matrix element $\partial y/\partial t$ via jacobian["y"]["t"]. This simplification streamlines the workflow, allowing users to focus on modeling rather than matrix management.

• Integration of unit-aware automatic differentiation

Another prevalent deficiency in current PINN frameworks is the lack of support for physical units, which are crucial for maintaining dimensional consistency in physical equations. Take the Burgers' equation for example, the left-hand side $\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}$ and the right-hand side $\nu \frac{{\partial }^{2}u}{\partial x^2}$ must share the same physical units (e.g., meters per second squared). To ensure such consistency, PINNx integrates brainunit.autograd module, enabling unit-aware automatic differentiation. This feature preserves unit information during the computation of first, second, or higher-order derivatives, ensuring that physical unit dimensions remain consistent throughout the differentiation process. Consequently, PINNx effectively mitigates errors arising from unit mismatches, thereby enhancing the reliability of simulations.

A quick 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.

Installation

To use the pinnx module, install it with the following command:

pip install deepxde[pinnx]

Dependency

pinnx module heavily depends on the following packages:

• brainstate: brainstate provides a state-based compilation system in JAX. It is very similar to pytorch's syntax. However, it can benefit from the JIT compilation of JAX and XLA. We recommend users to learn brainstate via its documentation: https://brainstate.readthedocs.io/quickstart/concepts-en.html
• brainunit: brainunit provides physical units and unit-aware mathematical systems tailored for Scientific AI within JAX. For the usage, please see its documentation: https://brainunit.readthedocs.io

Low-precision training support

Changing the training precision is very easy in pinnx. Simply 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)

[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]

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.

[chaoming0625]

I have updated the PR introduction to clearly explain the motivation behind the pinnx module. This should help users to understand the core purpose of the module.

[chaoming0625]

I have also added an option to install pinnx as an optional dependency, allowing users to quickly install the necessary components. Specifically, users who wish to use the pinnx module can install it with the command:

pip install deepxde[pinnx]

[chaoming0625]

I updated the PR introduction once again.

[chaoming0625]

I have maximized function reuse from deepxde to minimize API redundancy.

[chaoming0625]

I think now it is ready to merge this PR. I have tried most models. I also propose to leverage a more agile development process, in which we could consider submitting the PR in smaller increments.

[chaoming0625]

Hi, @lululxvi , could you check whether there are issues that need to be fixed? Now I have time to update the code.

[lululxvi]

Sounds great! As this PR has a lot of changes, I will gradually review each file.