Update docs (#10)

* update readme

* update docs

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)
 
 ```
 

docs/index.rst

@@ -1,14 +1,14 @@
 ``pinnx`` documentation
 ========================
 
-`PINNx <https://github.com/chaobrain/pinnx>`_ is a library for scientific machine learning and physics-informed learning.
-It is rewritten according to `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
+`PINNx <https://github.com/chaobrain/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 <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.
+`PINNx <https://github.com/chaobrain/pinnx>`_ is built on top of 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.
 
 
 ----
@@ -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

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