Update READMEs

README.md

@@ -1,6 +1,8 @@
-# Torch CFD
+# Computational Fluid Dynamics in PyTorch
 
-This is a native PyTorch port of [Google's Computational Fluid Dynamics package in Jax](https://github.com/google/jax-cfd). The main changes are documented in the `README.md` under the [`torch_cfd` directory](torch_cfd/README.md). The biggest change is many routines that rely on the functional programming of Jax have been rewritten to be a more PyTorch-friendly tensor-in to tensor-out style.
+A native PyTorch port of [Google's Computational Fluid Dynamics package in Jax](https://github.com/google/jax-cfd). The main changes are documented in the `README.md` under the [`torch_cfd` directory](torch_cfd). Biggest changes in all routines:
+- Routines that rely on the functional programming of Jax have been rewritten to be a more debugger-friendly PyTorch tensor-in-tensor-out style.
+- Functions and operators are in general implemented as `nn.Module`.
 
 ## Installation
 
@@ -9,11 +11,14 @@ pip install torch-cfd
 ```
 
 ## Contributions
-PR welcome. Current the port only includes:
+PR welcome. Currently, the port only includes:
 - Pseudospectral methods for vorticity which use anti-aliasing filtering techniques for non-linear terms to maintain stability.
-- Temporal discretization: Currently only RK4 temporal discretization, using explicit time-stepping for advection and either implicit or explicit time-stepping for diffusion.
+- Temporal discretization: Currently only RK4 temporal discretization using explicit time-stepping for advection and either implicit or explicit time-stepping for diffusion.
 - Boundary conditions: only periodic boundary conditions.
 
 ## Examples
 - Demos of different simulation setups:
-  - [2D simulation with a psuedo-spectral solver](example_Kolmogrov2d_rk4_cn_forced_turbulence.ipynb)
+  - [2D simulation with a pseudo-spectral solver](example_Kolmogrov2d_rk4_cn_forced_turbulence.ipynb)
+
+## Acknowledgments
+SC appreciates the support from the National Science Foundation DMS-2309778.

torch_cfd/README.md

@@ -3,7 +3,31 @@
 - [ ] add native PyTorch implementation for applying `torch.linalg` and `torch.fft` function directly on `GridArray`.
 - [ ] add discrete Helmholtz decomposition in both spatial and spectral domains.
 - [ ] adjust the function to act on `(batch, time, *spatial)` tensor, currently only `(*spatial)` is supported.
+- [x] add native vorticity computation, instead of taking FDM for pseudo-spectral.
 
 ## Changelog
+
+### 0.0.4
+- The forcing functions are now implemented as `nn.Module` and utilize a wrapper decorator for the potential function.
+- Added some common time stepping schemes, additional ones that Jax-CFD did not have includes the commonly used Crank-Nicholson IMEX.
+- Combined the implementation for step size satisfying the CFL condition.
+
+
+### 0.0.1
 - `grids.GridArray` is implemented as a subclass of `torch.Tensor`, not the original jax implentation uses the inheritance from `np.lib.mixins.NDArrayOperatorsMixin`. `__array_ufunc__()` is replaced by `__torch_function__()`.
-- The padding of `torch.nn.functional.pad()` is different from `jax.numpy.pad()`, PyTorch's pad starts from the last dimension, while Jax's pad starts from the first dimension. For example, `F.pad(x, (0, 0, 1, 0, 1, 1))` is equivalent to `jax.numpy.pad(x, ((1, 1), (1, 0), (0, 0)))` for an array of size `(*, t, h, w)`.
+- The padding of `torch.nn.functional.pad()` is different from `jax.numpy.pad()`, PyTorch's pad starts from the last dimension, while Jax's pad starts from the first dimension. For example, `F.pad(x, (0, 0, 1, 0, 1, 1))` is equivalent to `jax.numpy.pad(x, ((1, 1), (1, 0), (0, 0)))` for an array of size `(*, t, h, w)`.
+- A handy outer sum, which is usefully in getting the n-dimensional Laplacian in the frequency domain, is implemented as follows to replace `reduce(np.add.outer, eigenvalues)`
+    ```python
+    def outer_sum(x: Union[List[Array], Tuple[Array]]) -> Array:
+        """
+        Returns the outer sum of a list of one dimensional arrays
+        Example:
+        x = [a, b, c]
+        out = a[..., None, None] + b[..., None] + c
+        """
+
+        def _sum(a, b):
+            return a[..., None] + b
+
+        return reduce(_sum, x)
+    ```