Merge pull request #2 from ddrous/main

Putting Dev Up To Date

.github/workflows/updec-linux.yml

@@ -0,0 +1,23 @@
+name: Updec CI/CD
+
+on: [push]
+
+jobs:
+  linux-tests:
+    runs-on: ubuntu-latest
+    strategy:
+      max-parallel: 5
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v3
+      with:
+        python-version: '3.10'
+    - name: Install Updes & Pytest
+      run: |
+        pip install --upgrade pip && pip install --upgrade "jax[cpu]" && python -m pip install -e .
+    - name: Run package-wide test
+      run: |
+        pytest --pyargs updes
+

.gitignore

@@ -8,7 +8,39 @@
 .cache
 .DS_Store
 Updec.egg-info/*
-examples/temp/*
-!examples/temp/README.md
+updes.egg-info/*
+examples/data/*
+!examples/data/README.md
 scripts/reports/*
 !scripts/reports/README.md
+runs/*
+!runs/README.md
+demos/temp/*
+*msh
+*vtk
+quadqs_statistics.json
+
+## Python wheels
+*.whl
+*.tar.gz
+
+## Quarto website
+docs/.quarto/*
+docs/_site/*
+
+!updes/tests/data/mesh.msh
+
+demos/Laplace/data/*
+!demos/Laplace/data/README.md
+!demos/Laplace/data/Comparison
+!demos/Laplace/data/TempFolder/*
+
+demos/NavierStokes/data/*
+!demos/NavierStokes/data/README.md
+!demos/NavierStokes/data/Comparison
+!demos/NavierStokes/data/TempFolder/*
+
+demos/Advection/data/*
+!demos/Advection/data/README.md
+!demos/Advection/data/Comparison
+!demos/Advection/data/TempFolder/*

README.md

@@ -1,30 +1,101 @@
-# 𝕌pdec
+# 𝕌pdes
 
-𝕌pdec is a scalable framework for meshless and data-driven PDE modelling and control. __𝕌pdec__ stands for Universal(__𝕌__) Partial Differential Equations (__PDE__) Controller(__C__).
+𝕌pdes is a general-purpose library for mesh-free PDE simulation and control. There is no faster way to explore the realm of PDEs !
 
+<table>
+  <tr>
+    <th><img src="docs/assets/laplace.png" width="180"></th>
+    <th><img src="docs/assets/adv_diff_periodic.gif" width="350"></th>
+    <!-- <th><img src="docs/assets/burgers_u.gif" width="188"></th> -->
+    <th><img src="docs/assets/ns_norm.gif" width="150"></th>
+  </tr>
+</table>
 
-## Features
-𝕌pdec support the following methods:
-- Direct Numerical simulation
-- Differentiable Physics
-- Physics-Informed Neural network
-- More to come ...
 
-𝕌pdec compares the methods in terms of:
-- accuracy and generalization error
-- performance and speed
-- robustness to perturbation
-- More to come ...
+## Features
+𝕌pdes leverages Radial Basis Functions (RBFs) and JAX to provide the following features:
+- User-centric design: no need to re-implement a solver for each new PDE
+- Lightning fast mesh-free simulation via Radial Basis Functions
+- Robust differentiable simulation via JAX, and portable across CPU, GPU, and TPU
+- Support for Dirichlet, Neumann, Robin, and Periodic boundary conditions
+- Automatic generation of normals from 2D GMSH meshes
 
-𝕌pdec offers parallel-in-time parallelization accros CPUs and GPUs, allowing users to scale and combine the above methods.
+𝕌pdes is easily extendable, with additional features added frequently.
 
 
 ## Getting started
-Check out the example notebooks and scripts in  the folder [`demos`](./demos)!
+The package is available on PyPI. You can install it with
+```
+pip install updes
+```
+
+The example below illustrates how to solve the Laplace equation with Dirichlet and Neumann boundary conditions:
+```python
+import updes
+import jax.numpy as jnp
+
+# Create a mesh-free cloud of points on a unit square
+facet_types={"South":"n", "West":"d", "North":"d", "East":"d"}
+cloud = updes.SquareCloud(Nx=30, Ny=20, facet_types=facet_types)
+
+# Define the differential operator (left-hand side of the PDE)
+def my_diff_operator(x, center, rbf, monomial, fields):
+    return updes.nodal_laplacian(x, center, rbf, monomial)
+
+# Define the right-hand side of the PDE
+def my_rhs_operator(x, centers, rbf, fields):
+    return 0.0
+
+# Set a sin function as the Dirichlet BC on the North, and zero everywhere else
+sine = lambda coord: jnp.sin(jnp.pi * coord[0])
+zero = lambda coord: 0.0
+boundary_conditions = {"South":zero, "West":zero, "North":sine, "East":zero}
+
+# Solve the Laplace equation with a JIT-compiled solver
+sol = updes.pde_solver_jit(diff_operator=my_diff_operator, 
+                    rhs_operator=my_rhs_operator, 
+                    cloud=cloud, 
+                    boundary_conditions=boundary_conditions, 
+                    rbf=updes.polyharmonic,
+                    max_degree=1)
+
+# Visualize the solution
+cloud.visualize_field(sol.vals, cmap="jet", projection="3d", title="RBF solution")
+```
+
+𝕌pdes can handle much complicated cases with little modifications to the code above. Check out the [documentation](https://ddrous.github.io/Updes/) and Python interactive notebooks in the [`demos`](./demos) folder !
+
+
+
+
+## To-Dos
+1. Logo, contributors guide, and developer documentation
+2. More introductory examples in the documentation :
+    - Integration with neural networks and [Equinox](https://github.com/patrick-kidger/equinox)
+    - Non-linear and multi-dimensional PDEs
+    - Adjoint schemes for fluid flows
+3. Better point generation with accurate geometry and normals: 
+    - USD format
+    - GMSH tutorial
+4. Support for 3D radial basis functions
+
+We welcome contributions from the community. Please feel free to open an issue or a pull request.
 
 
 ## Dependencies
-- PhiFlow: for differentiable physics
-- Diffrax: for neural ODEs
-- PyOMP: for shared-memory parallelization
-- MPI4Jax (optional): for distributed-memory parallelization
+- **Core**: [JAX](https://github.com/google/jax) - [GMSH](https://pypi.org/project/gmsh/) - [Matplotlib](https://github.com/matplotlib/matplotlib) - [Seaborn](https://github.com/mwaskom/seaborn) - [Scikit-Learn](https://github.com/scikit-learn/scikit-learn)
+- **Optional**: [PyVista](https://github.com/pyvista/pyvista) - [FFMPEG](https://github.com/kkroening/ffmpeg-python) - [QuartoDoc](https://github.com/machow/quartodoc/)
+
+See the `pyproject.toml` file the specific versions of the dependencies.
+
+
+## Cite us !
+If you use this software, please cite us with the following BibTeX entry:
+```
+@inproceedings{nzoyem2023comparison,
+  title={A comparison of mesh-free differentiable programming and data-driven strategies for optimal control under PDE constraints},
+  author={Nzoyem Ngueguin, Roussel Desmond and Barton, David AW and Deakin, Tom},
+  booktitle={Proceedings of the SC'23 Workshops of The International Conference on High Performance Computing, Network, Storage, and Analysis},
+  pages={21--28},
+  year={2023}}
+```

demos/Advection/00_advection_with_rbf.py

@@ -0,0 +1,133 @@
+# %%
+
+"""
+Test of the Updec package on the Advection-Diffusion equation with RBFs: 
+PDE here: https://en.wikipedia.org/wiki/Convection%E2%80%93diffusion_equation
+"""
+
+import time
+
+import jax
+import jax.numpy as jnp
+
+# jax.config.update('jax_platform_name', 'cpu')
+jax.config.update("jax_enable_x64", True)
+
+from updes import *
+# key = jax.random.PRNGKey(13)
+key = None
+
+# from torch.utils.tensorboard import SummaryWriter
+
+
+RUN_NAME = "TempFolder"
+DATAFOLDER = "./data/" + RUN_NAME +"/"
+make_dir(DATAFOLDER)
+
+RBF = partial(polyharmonic, a=1)
+MAX_DEGREE = 1
+
+DT = 1e-4
+NB_TIMESTEPS = 100
+PLOT_EVERY = 10
+
+VEL = jnp.array([100.0, 0.0])
+## Diffusive constant
+K = 0.08
+
+Nx = 40
+Ny = 20
+SUPPORT_SIZE = "max"
+
+facet_types={"South":"d", "West":"d", "North":"d", "East":"n"}
+cloud = SquareCloud(Nx=Nx, Ny=Ny, facet_types=facet_types, noise_key=key, support_size=SUPPORT_SIZE)
+
+cloud.visualize_cloud(s=0.1, figsize=(7,3));
+
+# print("Local supports:", cloud.local_supports[0])
+
+
+# %%
+
+def my_diff_operator(x, center=None, rbf=None, monomial=None, fields=None):
+    val = nodal_value(x, center, rbf, monomial)
+    grad = nodal_gradient(x, center, rbf, monomial)
+    lap = nodal_laplacian(x, center, rbf, monomial)
+    return (val/DT) + jnp.dot(VEL, grad) - K*lap
+
+def my_rhs_operator(x, centers=None, rbf=None, fields=None):
+    return value(x, fields[:,0], centers, RBF) / DT     ## TODO value ?
+
+d_zero = lambda x: 0.
+boundary_conditions = {"South":d_zero, "West":d_zero, "North":d_zero, "East":d_zero}
+
+
+## u0 is zero everywhere except at a point in the middle
+u0 = jnp.zeros(cloud.N)
+source_id = int(cloud.N*0.01)
+source_neighbors = jnp.array(cloud.local_supports[source_id][:cloud.N//40])
+u0 = u0.at[source_neighbors].set(0.95)
+
+## Begin timestepping for 100 steps
+
+# fig = plt.figure(figsize=(6,3))
+# ax1= fig.add_subplot(1, 1, 1, projection='3d')
+# ax = fig.add_subplot(1, 1, 1)
+
+
+u = u0.copy()
+ulist = [u]
+
+start = time.time()
+
+for i in range(1, NB_TIMESTEPS+1):
+    ufield = pde_solver_jit(diff_operator=my_diff_operator,
+                        rhs_operator = my_rhs_operator,
+                        rhs_args=[u],
+                        cloud = cloud,
+                        boundary_conditions = boundary_conditions, 
+                        rbf=RBF,
+                        max_degree=MAX_DEGREE,)
+
+    u = ufield.vals
+    ulist.append(u)
+
+    if i<=3 or i%PLOT_EVERY==0:
+        print(f"Step {i}")
+        # plt.cla()
+        # cloud.visualize_field(u, cmap="jet", projection="3d", title=f"Step {i}")
+        ax, _ = cloud.visualize_field(u, cmap="jet", title=f"Step {i}", vmin=0, vmax=1, figsize=(6,3),colorbar=False)
+        # plt.draw()
+        plt.show()
+
+
+walltime = time.time() - start
+
+minutes = walltime // 60 % 60
+seconds = walltime % 60
+print(f"Walltime: {minutes} minutes {seconds:.2f} seconds")
+
+
+
+# %%
+
+# ax = plt.gca()
+filename = DATAFOLDER + "advection_diffusion_rbf.gif"
+cloud.animate_fields([ulist], cmaps="jet", filename=filename, figsize=(7,3), titles=["Advection-Diffusion with RBFs"])
+
+
+
+# %%
+
+
+## Write stuff to tensorboard
+# run_name = str(datetime.datetime.now())[:19]        ##For tensorboard
+# writer = SummaryWriter("runs/"+run_name, comment='-Laplace')
+# hparams_dict = {"rbf":RBF.__name__, "max_degree":MAX_DEGREE, "nb_nodes":Nx*Ny, "support_size":SUPPORT_SIZE}
+# metrics_dict = {"metrics/mse_error":float(error), "metrics/wall_time":walltime}
+# writer.add_hparams(hparams_dict, metrics_dict, run_name="hp_params")
+# writer.add_figure("plots", fig)
+# writer.flush()
+# writer.close()
+
+# %%