From 8f700abdceb14c82ae75bf3ba86bd40ccb31ecac Mon Sep 17 00:00:00 2001 From: Chaoming Wang Date: Fri, 13 Dec 2024 21:40:01 +0800 Subject: [PATCH] Update docs (#10) * update readme * update docs --- README.md | 56 +++++++++++++++++++++++++ docs/index.rst | 77 +++++++++++++++++++++++++++++----- docs/unit-examples-forward.rst | 2 +- 3 files changed, 123 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index 98dc16c..62b20c8 100644 --- a/README.md +++ b/README.md @@ -24,7 +24,63 @@ For example, it leverages ## 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) ``` diff --git a/docs/index.rst b/docs/index.rst index 933256d..d51eb22 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,14 +1,14 @@ ``pinnx`` documentation ======================== -`PINNx `_ is a library for scientific machine learning and physics-informed learning. -It is rewritten according to `DeepXDE `_ but is enhanced by our -`Brain Dynamics Programming (BDP) ecosystem `_. -For example, it leverages +`PINNx `_ is a library for scientific machine learning and physics-informed learning +in JAX. It enables to define PINN problem with explicit variables (e.g. ``x``, ``y``, ``z``) and physical units +(e.g. ``meter``, ``second``, ``kelvin``) and to solve the problem with neural networks. -- `brainstate `_ for just-in-time compilation, -- `brainunit `_ for dimensional analysis, -- `braintools `_ for checkpointing, loss functions, and other utilities. +`PINNx `_ is built on top of our `Brain Dynamics Programming (BDP) ecosystem `_. +For example, it leverages `brainstate `_ for just-in-time compilation, +`brainunit `_ for dimensional analysis, +`braintools `_ for checkpointing, loss functions, and other utilities. ---- @@ -39,7 +39,64 @@ Installation Quick Start ^^^^^^^^^^^ -To be added. + + +.. code-block:: 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) @@ -55,9 +112,7 @@ User guide .. toctree:: :maxdepth: 2 - examples-unit.rst - - + unit-examples-forward.rst See also the BDP ecosystem diff --git a/docs/unit-examples-forward.rst b/docs/unit-examples-forward.rst index 4a436ea..3a4e158 100644 --- a/docs/unit-examples-forward.rst +++ b/docs/unit-examples-forward.rst @@ -4,7 +4,7 @@ PINN Forward Examples .. toctree:: :maxdepth: 1 - unit-examples-forward/Burgers.ipynb + unit-examples-forward/burgers.ipynb unit-examples-forward/Burgers_RAR.ipynb unit-examples-forward/heat.ipynb unit-examples-forward/heat_resample.ipynb