Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add pinnx submodule #1932

Open
wants to merge 22 commits into
base: master
Choose a base branch
from
Open

Add pinnx submodule #1932

wants to merge 22 commits into from

Conversation

chaoming0625
Copy link

@chaoming0625 chaoming0625 commented Jan 15, 2025

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 2 y / x 2 via hessian["y"]["x"] and the Jacobian matrix element y / 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 u t + u u x and the right-hand side ν 2 u 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:

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)

Verified

This commit was created on GitHub.com and signed with GitHub’s verified signature.
@lululxvi
Copy link
Owner

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
Copy link
Author

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
Copy link
Author

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

@chaoming0625
Copy link
Author

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
Copy link
Owner

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.

Routhleck and others added 4 commits January 23, 2025 10:13
we can install pinnx module using the following command:

pip install deepxde[pinnx]
@chaoming0625
Copy link
Author

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
Copy link
Author

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
Copy link
Author

I updated the PR introduction once again.

@chaoming0625
Copy link
Author

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

@chaoming0625
Copy link
Author

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.

fix
fix
@chaoming0625
Copy link
Author

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

@lululxvi
Copy link
Owner

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants