backends: add basic JAX backend class (storage + moments methods implemented) + improve moments-methods unit tests

.github/workflows/readme_snippets.yml → .github/workflows/docs_code_snippets.yml

@@ -40,6 +40,17 @@ jobs:
           token: ${{ secrets.GITHUB_TOKEN }}
           files: |
             readme.png
+      - run: |
+          rm readme.png
+          sed -i -e 's/GPU/JAX/g' readme.py
+          sed -i '/1200, 2400, 3600//g' readme.py
+          python -We readme.py
+      - name: jax-artefacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: jax
+          if-no-files-found: error
+          path: readme.png
 
   julia:
     runs-on: macos-latest

PySDM/backends/__init__.py

@@ -11,11 +11,14 @@
 from numba import cuda
 
 from . import numba as _numba
+from . import jax as _jax
 
 # for pdoc
 CPU = None
 GPU = None
+JAX = None
 Numba = _numba.Numba
+Jax = _jax.Jax
 ThrustRTC = None
 
 
@@ -93,3 +96,5 @@ def _cached_backend(formulae=None, backend_class=None, **kwargs):
 
 GPU = partial(_cached_backend, backend_class=ThrustRTC)
 """ returns a cached instance of the ThrustRTC backend (cache key including formulae parameters) """
+
+JAX = partial(_cached_backend, backend_class=Jax)

PySDM/backends/impl_common/storage_utils.py

@@ -59,6 +59,10 @@ def upload(self, data):
     def fill(self, other):
         raise NotImplementedError()
 
+    @abstractmethod
+    def row_view(self, i):
+        raise NotImplementedError()
+
 
 def get_data_from_ndarray(array, storage_class: Type[StorageBase], copy_fun):
     if str(array.dtype).startswith("int"):

PySDM/backends/impl_jax/methods/__init__.py

@@ -0,0 +1,7 @@
+"""method classes of the JAX backend"""
+
+from .collisions_methods import CollisionsMethods
+from .index_methods import IndexMethods
+from .moments_methods import MomentsMethods
+from .pair_methods import PairMethods
+from .physics_methods import PhysicsMethods

PySDM/backends/impl_jax/methods/collisions_methods.py

@@ -0,0 +1,216 @@
+"""
+CPU implementation of backend methods for particle collisions
+"""
+
+from functools import cached_property
+import numpy as np
+
+import jax
+import jax.numpy as jnp
+
+from PySDM.backends.impl_common.backend_methods import BackendMethods
+from PySDM.backends.impl_numba import conf
+from PySDM.backends.impl_jax.storage import Storage
+
+
+def pair_indices(i, idx, is_first_in_pair, prob_like):
+    offset = 1 - is_first_in_pair[2 * i]
+    j = idx[2 * i + offset]
+    k = idx[2 * i + 1 + offset]
+
+    return j, k, False
+
+
+def flag_zero_multiplicity(j, k, multiplicity, healthy):
+    raise NotImplementedError()
+
+
+class CollisionsMethods(BackendMethods):
+    @cached_property
+    def _collision_coalescence_body(self):
+        @jax.jit
+        def body(
+            multiplicity,
+            idx,
+            attributes,
+            gamma,
+            is_first_in_pair,
+        ):
+            i = jnp.arange(len(multiplicity) // 2)
+
+            offset = 1 - is_first_in_pair[2 * i]
+            j = idx[2 * i + offset]
+            k = idx[2 * i + 1 + offset]
+            mj = multiplicity[j]
+            mk = multiplicity[k]
+
+            new_n = mj - gamma[i] * mk
+
+            pos_mask = new_n > 0
+            zero_mask = ~pos_mask
+
+            # CASE new_n > 0
+            mult_pos_updates = jnp.where(pos_mask, new_n, multiplicity[j])
+            # attr_pos_delta = gamma[i] * attributes[:, j]
+            # attr_pos_delta = attr_pos_delta * pos_mask
+
+            # CASE new_n <= 0
+            mj_new = mk // 2
+            mk_new = mk - mj_new
+
+            mult_j_zero = jnp.where(zero_mask, mj_new, multiplicity[j])
+            mult_k_zero = jnp.where(zero_mask, mk_new, multiplicity[k])
+
+            # new_attr = gamma[i] * attributes[:, j] + attributes[:, k]
+            # new_attr = new_attr * zero_mask
+
+            mult_j = jnp.where(pos_mask, mult_pos_updates, mult_j_zero)
+            mult_k = jnp.where(zero_mask, mult_k_zero, multiplicity[k])
+
+            multiplicity = multiplicity.at[j].set(mult_j)
+
+            multiplicity = multiplicity.at[k].set(mult_k)
+
+            # return mult_j, mult_k
+            # attributes = attributes.at[:, k].add(attr_pos_delta)
+
+            # attributes = attributes.at[:, j].set(
+            #     jnp.where(zero_mask, new_attr, attributes[:, j])
+            # )
+            # attributes = attributes.at[:, k].set(
+            #     jnp.where(zero_mask, new_attr, attributes[:, k])
+            # )
+            return multiplicity
+
+            # if new_n > 0:
+            #     multiplicity[j] = new_n
+            #     for a in range(len(attributes)):
+            #         attributes[a, k] += gamma[i] * attributes[a, j]
+            # else:  # new_n == 0
+            #     multiplicity[j] = multiplicity[k] // 2
+            #     multiplicity[k] = multiplicity[k] - multiplicity[j]
+            #     for a in range(len(attributes)):
+            #         attributes[a, j] = gamma[i] * attributes[a, j] + attributes[a, k]
+            #         attributes[a, k] = attributes[a, j]
+            # return multiplicity, attributes
+
+            # return multiplicity, attributes
+
+        return body
+
+    def collision_coalescence(
+        self,
+        *,
+        multiplicity,
+        idx,
+        attributes,
+        gamma,
+        healthy,
+        cell_id,
+        coalescence_rate,
+        is_first_in_pair,
+    ):
+        # pass
+        # indices = jnp.arange(len(multiplicity) // 2)
+        # mapped_collision_coalescence = jax.vmap(
+        #     self._collision_coalescence_body, (None, None, None, None, None, 0)
+        # )
+        # multiplicity.data, attributes.data = mapped_collision_coalescence(
+        #     multiplicity.data,
+        #     idx.data,
+        #     attributes.data,
+        #     gamma.data,
+        #     is_first_in_pair.indicator.data,
+        #     indices,
+        # )
+
+        multiplicity.data = self._collision_coalescence_body(
+            multiplicity.data,
+            idx.data,
+            attributes.data,
+            gamma.data,
+            is_first_in_pair.indicator.data,
+        )
+
+        # print(f"{mult_pos_updates=} \n {mult_j_zero=} \n {mult_k_zero=} \n {mult_k=} \n {pos_mask=}")
+
+    @cached_property
+    def _compute_gamma_body(self):
+        # pylint: disable=too-many-arguments,too-many-locals
+        @jax.jit
+        def body(prob, rand, idx, multiplicity, is_first_in_pair, i):
+            out = jnp.ceil(prob[i] - rand[i])
+            offset = 1 - is_first_in_pair[2 * i]
+            j = idx[2 * i + offset]
+            k = idx[2 * i + 1 + offset]
+
+            prop = multiplicity[j] // multiplicity[k]
+            return jnp.minimum(out.astype(int), prop)
+
+        return body
+
+    def compute_gamma(
+        self,
+        *,
+        prob,
+        rand,
+        multiplicity,
+        cell_id,
+        collision_rate_deficit,
+        collision_rate,
+        is_first_in_pair,
+        out,
+    ):
+        indices = jnp.arange(len(multiplicity) // 2)
+        mapped_compute_gamma_body = jax.vmap(
+            self._compute_gamma_body, (None, None, None, None, None, 0)
+        )
+        out.data = mapped_compute_gamma_body(
+            prob.data,
+            rand.data,
+            multiplicity.idx.data,
+            multiplicity.data,
+            is_first_in_pair.indicator.data,
+            indices,
+        )
+
+    @staticmethod
+    def make_cell_caretaker(idx_shape, idx_dtype, cell_start_len, scheme="default"):
+        class DummyCaretaker:
+            def __call__(self, *args, **kwds):
+                return
+
+        return DummyCaretaker()
+
+    @cached_property
+    def _normalize_body(self):
+        # pylint: disable=too-many-arguments
+        def body(prob, cell_start, timestep, dv, i):
+            sd_num = cell_start[1] - cell_start[0]
+            # Fixed this with (sd_num >= 2), need to add the n_cell loop, and then the normalization pass
+            norm_factor = (sd_num >= 2) * (
+                timestep / dv * sd_num * (sd_num - 1) / 2 / (sd_num // 2)
+            )
+            prob = prob.at[i].set(norm_factor)
+            return prob
+
+        return body
+
+    # pylint: disable=too-many-arguments
+    def normalize(self, prob, cell_id, cell_idx, cell_start, norm_factor, timestep, dv):
+        # FIX THIS FUNCTION!!!
+        indices = jax.numpy.arange(prob.shape[0])
+        temp_prob = jax.numpy.empty(prob.shape)
+
+        normalize_func = jax.vmap(self._normalize_body, (None, None, None, None, 0))
+
+        temp_prob = normalize_func(temp_prob, cell_start.data, timestep, dv, indices)
+
+        prob.data = jax.numpy.sum(temp_prob, axis=0)
+
+    @staticmethod
+    # pylint: disable=too-many-arguments
+    def _counting_sort_by_cell_id_and_update_cell_start(
+        new_idx, idx, cell_id, cell_idx, length, cell_start
+    ):
+        raise NotImplementedError()

PySDM/backends/impl_jax/methods/index_methods.py

@@ -0,0 +1,21 @@
+"""
+CPU implementation of shuffling and sorting backend methods
+"""
+
+from functools import cached_property
+
+import jax
+
+from PySDM.backends.impl_common.backend_methods import BackendMethods
+
+
+class IndexMethods(BackendMethods):
+
+    @cached_property
+    def shuffle_local(self):
+        @jax.jit
+        def body(idx, u01, cell_start):
+            # IMPLEMENT
+            return
+
+        return body