Merge pull request #11 from tumaer/notebooks

Notebooks

README.md

@@ -44,7 +44,7 @@ Later, you just need to `source venv/bin/activate` to activate the environment.
 If you want to use a CUDA GPU, you first need a running Nvidia driver. And then just follow the instructions [here](https://jax.readthedocs.io/en/latest/installation.html). The whole process could look like this:
 ```bash
 source .venv/bin/activate
-pip install --upgrade "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install -U "jax[cuda12]"
 ```
 
 ## Getting Started
@@ -66,17 +66,15 @@ python main.py config=cases/tgv.yaml solver.name=RIE solver.tvf=0.0
  ```
 -  Thermal diffusion
 ```bash
-python main.py config=ht.yaml
+python main.py config=cases/ht.yaml
 ```
 
-### Solver-in-the-Loop
-To train and test our Solver-in-the-Loop model, run the script in [./experiments/sitl.py](./experiments/sitl.py). This file relies on [LagrangeBench](https://github.com/tumaer/lagrangebench), which can be installed by `pip install lagrangebench`. For more information on the training and inference setup, visit the LagrangeBench website.
-
-### Inverse Problem
-The presented inverse problem of finding the initial state of a 100-step-long SPH simulation can be fully reproduced using the notebook [./experiments/inverse.ipynb](./experiments/inverse.ipynb).
-
-### Gradient Validation
-The presented validation of the gradients through the solver can be fully reproduced using the notebook [./experiments/grads.ipynb](./experiments/grads.ipynb)
+### Notebooks
+We provide four notebooks demonstrating how to use JAX-SPH:
+- [`tutorial.ipynb`](notebooks/tutorial.ipynb) [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tumaer/jax-sph/blob/main/notebooks/tutorial.ipynb), with a general overview of JAX-SPH and an example how to run the channel flow with hot bottom wall.
+- [`iclr24_grads.ipynb`](notebooks/tutorial.ipynb) [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tumaer/jax-sph/blob/main/notebooks/iclr24_grads.ipynb), with a validation of the gradients through the solver.
+- [`iclr24_inverse.ipynb`](notebooks/tutorial.ipynb) [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tumaer/jax-sph/blob/main/notebooks/iclr24_inverse.ipynb), solving the inverse problem of finding the initial state of a 100-step-long SPH simulation.
+- [`iclr24_sitl.ipynb`](notebooks/tutorial.ipynb) [![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/tumaer/jax-sph/blob/main/notebooks/iclr24_sitl.ipynb), including training and testing a Solver-in-the-Loop model using the [LagrangeBench](https://github.com/tumaer/lagrangebench) library.
 
 ## Setting up a case
 To set up a case, just add a `my_case.py` and a `my_case.yaml` file to the `cases/` directory. Every *.py case should inherit from `SimulationSetup` in `jax_sph/case_setup.py` or another case, and every *.yaml config file should either contain a complete set of parameters (see `jax_sph/defaults.py`) or extend `JAX_SPH_DEFAULTS`. Running a case in relaxation mode `case.mode=rlx` overwrites certain parts of the selected case. Passed CLI arguments overwrite any argument.

cases/ht.py

@@ -73,7 +73,7 @@ def _boundary_conditions_fn(self, state):
 
         # set incoming fluid temperature to reference_temperature
         mask_inflow = mask_fluid * (state["r"][:, 0] < 3 * self.case.dx)
-        state["T"] = jnp.where(mask_inflow, self.cfg.T_ref, state["T"])
+        state["T"] = jnp.where(mask_inflow, self.case.T_ref, state["T"])
         state["dTdt"] = jnp.where(mask_inflow, 0.0, state["dTdt"])
 
         # set the hot wall to hot_wall_temperature.
@@ -83,7 +83,7 @@ def _boundary_conditions_fn(self, state):
 
         # set the fixed wall to reference_temperature.
         mask_solid = state["tag"] == Tag.SOLID_WALL  # fixed wall
-        state["T"] = jnp.where(mask_solid, self.cfg.T_ref, state["T"])
+        state["T"] = jnp.where(mask_solid, self.case.T_ref, state["T"])
         state["dTdt"] = jnp.where(mask_solid, 0, state["dTdt"])
 
         # ensure static walls have no velocity or acceleration

cases/ht.yaml

@@ -24,4 +24,5 @@ eos:
 io:
   write_type: ["h5"]
   write_every: 100
-  data_path: "data/debug"
+  data_path: "data/debug"
+  print_props: ["Ekin", "u_max", "T_max"]

docs/requirements.txt

@@ -1,9 +1,9 @@
 h5py
-jax-md==0.2.8
-jax[cpu]==0.4.23
+-e git+https://github.com/jax-md/jax-md.git@c451353f6ddcab031f660befda256d8a4f657855#egg=jax-md
+jax[cpu]==0.4.28
+omegaconf
 pandas
 pyvista
-scipy==1.12.0
 sphinx==7.2.6
 sphinx-exec-code==0.12
 sphinx-rtd-theme==1.3.0