# PINNx: Physics-Informed Neural Networks for Scientific Machine Learning in JAX
<p align="center">
<img alt="Header image of pinnx." src="https://github.com/chaobrain/pinnx/blob/main/docs/_static/pinnx.png" width=40%>
</p>
[](https://github.com/chaobrain/pinnx/actions/workflows/build.yml)
[](https://pinnx.readthedocs.io/en/latest/?badge=latest)
[](https://badge.fury.io/py/pinnx)
[](https://github.com/chaobrain/pinnx/blob/master/LICENSE)
``PINNx`` is a library for scientific machine learning and physics-informed learning in JAX.
It is inspired from [DeepXDE](https://github.com/lululxvi/deepxde) but is enhanced by our
[Brain Dynamics Programming (BDP) ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/).
For example, it leverages
- [brainstate](https://brainstate.readthedocs.io/) for just-in-time compilation,
- [brainunit](https://brainunit.readthedocs.io/) for dimensional analysis,
- [braintools](https://braintools.readthedocs.io/) for checkpointing, loss functions, and other utilities.
## Quickstart
Define a PINN with explicit variables and physical units.
```python
import brainstate as bst
import brainunit as u
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)
```
## Installation
- Install the stable version with `pip`:
``` sh
pip install pinnx --upgrade
```
## Documentation
The official documentation is hosted on Read the Docs: [https://pinnx.readthedocs.io/](https://pinnx.readthedocs.io/)
## See also the BDP ecosystem
We are building the Brain Dynamics Programming ecosystem: https://ecosystem-for-brain-dynamics.readthedocs.io/
Raw data
{
"_id": null,
"home_page": "https://github.com/chaobrain/pinnx",
"name": "pinnx",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.9",
"maintainer_email": null,
"keywords": "computational neuroscience, brain-inspired computation, brain dynamics programming",
"author": "PINNx Developers",
"author_email": "PINNx Developers <chao.brain@qq.com>",
"download_url": null,
"platform": null,
"description": "# PINNx: Physics-Informed Neural Networks for Scientific Machine Learning in JAX\r\n\r\n\r\n<p align=\"center\">\r\n \t<img alt=\"Header image of pinnx.\" src=\"https://github.com/chaobrain/pinnx/blob/main/docs/_static/pinnx.png\" width=40%>\r\n</p> \r\n\r\n\r\n[](https://github.com/chaobrain/pinnx/actions/workflows/build.yml)\r\n[](https://pinnx.readthedocs.io/en/latest/?badge=latest)\r\n[](https://badge.fury.io/py/pinnx)\r\n[](https://github.com/chaobrain/pinnx/blob/master/LICENSE)\r\n\r\n``PINNx`` is a library for scientific machine learning and physics-informed learning in JAX. \r\nIt is inspired from [DeepXDE](https://github.com/lululxvi/deepxde) but is enhanced by our \r\n[Brain Dynamics Programming (BDP) ecosystem](https://ecosystem-for-brain-dynamics.readthedocs.io/). \r\nFor example, it leverages \r\n\r\n- [brainstate](https://brainstate.readthedocs.io/) for just-in-time compilation,\r\n- [brainunit](https://brainunit.readthedocs.io/) for dimensional analysis, \r\n- [braintools](https://braintools.readthedocs.io/) for checkpointing, loss functions, and other utilities.\r\n\r\n\r\n## Quickstart\r\n\r\n\r\nDefine a PINN with explicit variables and physical units.\r\n\r\n```python\r\nimport brainstate as bst\r\nimport brainunit as u\r\nimport pinnx\r\n\r\n# geometry\r\ngeometry = pinnx.geometry.GeometryXTime(\r\n geometry=pinnx.geometry.Interval(-1, 1.),\r\n timedomain=pinnx.geometry.TimeDomain(0, 0.99)\r\n).to_dict_point(x=u.meter, t=u.second)\r\n\r\nuy = u.meter / u.second\r\nv = 0.01 / u.math.pi * u.meter ** 2 / u.second\r\n\r\n# boundary conditions\r\nbc = pinnx.icbc.DirichletBC(lambda x: {'y': 0. * uy})\r\nic = pinnx.icbc.IC(lambda x: {'y': -u.math.sin(u.math.pi * x['x'] / u.meter) * uy})\r\n\r\n# PDE equation\r\ndef pde(x, y):\r\n jacobian = approximator.jacobian(x)\r\n hessian = approximator.hessian(x)\r\n dy_x = jacobian['y']['x']\r\n dy_t = jacobian['y']['t']\r\n dy_xx = hessian['y']['x']['x']\r\n residual = dy_t + y['y'] * dy_x - v * dy_xx\r\n return residual\r\n\r\n# neural network\r\napproximator = pinnx.nn.Model(\r\n pinnx.nn.DictToArray(x=u.meter, t=u.second),\r\n pinnx.nn.FNN(\r\n [geometry.dim] + [20] * 3 + [1],\r\n \"tanh\",\r\n bst.init.KaimingUniform()\r\n ),\r\n pinnx.nn.ArrayToDict(y=uy)\r\n)\r\n\r\n# problem\r\nproblem = pinnx.problem.TimePDE(\r\n geometry,\r\n pde,\r\n [bc, ic],\r\n approximator,\r\n num_domain=2540,\r\n num_boundary=80,\r\n num_initial=160,\r\n)\r\n\r\n# training\r\ntrainer = pinnx.Trainer(problem)\r\ntrainer.compile(bst.optim.Adam(1e-3)).train(iterations=15000)\r\ntrainer.compile(bst.optim.LBFGS(1e-3)).train(2000, display_every=500)\r\ntrainer.saveplot(issave=True, isplot=True)\r\n\r\n```\r\n\r\n\r\n\r\n## Installation\r\n\r\n- Install the stable version with `pip`:\r\n\r\n``` sh\r\npip install pinnx --upgrade\r\n```\r\n\r\n\r\n## Documentation\r\n\r\nThe official documentation is hosted on Read the Docs: [https://pinnx.readthedocs.io/](https://pinnx.readthedocs.io/)\r\n\r\n\r\n## See also the BDP ecosystem\r\n\r\nWe are building the Brain Dynamics Programming ecosystem: https://ecosystem-for-brain-dynamics.readthedocs.io/\r\n\r\n",
"bugtrack_url": null,
"license": "Apache-2.0 license",
"summary": "Physics-informed.",
"version": "0.0.2.post20250105",
"project_urls": {
"Homepage": "https://github.com/chaobrain/pinnx",
"homepage": "http://github.com/chaobrain/pinnx",
"repository": "http://github.com/chaobrain/pinnx"
},
"split_keywords": [
"computational neuroscience",
" brain-inspired computation",
" brain dynamics programming"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "fd8475d2ced66a71998cc217d95319346730d281dee5e30947f21e9879c40803",
"md5": "f5b3d2c115cc157b702ae5020114e7e7",
"sha256": "872343b72c14763687e8ea90c5e7a58c03490b48b3d0cb955f34a8634e7bd9f0"
},
"downloads": -1,
"filename": "pinnx-0.0.2.post20250105-py2.py3-none-any.whl",
"has_sig": false,
"md5_digest": "f5b3d2c115cc157b702ae5020114e7e7",
"packagetype": "bdist_wheel",
"python_version": "py2.py3",
"requires_python": ">=3.9",
"size": 124050,
"upload_time": "2025-01-04T11:25:46",
"upload_time_iso_8601": "2025-01-04T11:25:46.414407Z",
"url": "https://files.pythonhosted.org/packages/fd/84/75d2ced66a71998cc217d95319346730d281dee5e30947f21e9879c40803/pinnx-0.0.2.post20250105-py2.py3-none-any.whl",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2025-01-04 11:25:46",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "chaobrain",
"github_project": "pinnx",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [
{
"name": "matplotlib",
"specs": []
},
{
"name": "numpy",
"specs": []
},
{
"name": "scikit-learn",
"specs": []
},
{
"name": "scipy",
"specs": []
},
{
"name": "brainstate",
"specs": []
},
{
"name": "brainunit",
"specs": [
[
">=",
"0.0.3"
]
]
},
{
"name": "braintools",
"specs": []
},
{
"name": "optax",
"specs": []
},
{
"name": "scikit-optimize",
"specs": []
},
{
"name": "msgpack",
"specs": []
}
],
"lcname": "pinnx"
}
JSON
