PySDM 3.0.0rc37__tar.gz → 3.0.0rc38__tar.gz

pysdm-3.0.0rc37/.github/workflows/tutorial_code.yml → pysdm-3.0.0rc38/.github/workflows/docs_code_snippets.yml

@@ -1,4 +1,4 @@
-name: tutorial_code
+name: docs_code
 
 defaults:
   run:
@@ -16,22 +16,25 @@ jobs:
   python:
     strategy:
       matrix:
-        platform: [ubuntu-latest, macos-15-intel, macos-14, windows-latest]
+        platform: [ubuntu-latest, macos-14, windows-latest]
     runs-on: ${{ matrix.platform }}
     steps:
       - uses: actions/checkout@v4.1.6
-      - uses: actions/setup-python@v6.2.0
+      - uses: actions/setup-python@v7.0.0
         with:
-          python-version: "3.10"
+          python-version: "3.11"
       - run: python -m pip install $PIP_INSTALL_ARGS -e .[CI_version_pins]
       - run: python -m pip install $PIP_INSTALL_ARGS pytest-codeblocks pytest
       - run: |
-          python -c "import os,pytest_codeblocks; code=pytest_codeblocks.extract_from_file('docs/markdown/pysdm_landing.md'); f=open('tutorial_code.py', 'w', encoding='utf-8'); f.write('# coding: utf-8'+os.linesep); f.writelines(block.code for block in code if block.syntax=='Python'); f.close()"
-      - run: cat -n tutorial_code.py
-      - run: |
-          python -We tutorial_code.py
-          sed -i -e 's/CPU/GPU/g' tutorial_code.py
-          python -We tutorial_code.py
+          python -c "import os,pytest_codeblocks; code=pytest_codeblocks.extract_from_file('docs/markdown/pysdm_landing.md'); f=open('docs_code.py', 'w', encoding='utf-8'); f.write('# coding: utf-8'+os.linesep); f.writelines(block.code for block in code if block.syntax=='Python'); f.close()"
+      - run: cat -n docs_code.py
+      - name: run-on-cpu-backend
+        run: |
+          python -We docs_code.py
+      - name: run-on-gpu-backend
+        run: |
+          sed -i -e 's/CPU/GPU/g' docs_code.py
+          python -We docs_code.py
       
       - name: artefacts
         if: github.ref == 'refs/heads/main' && matrix.platform == 'ubuntu-latest'
@@ -40,29 +43,41 @@ jobs:
           token: ${{ secrets.GITHUB_TOKEN }}
           files: |
             readme.png
+      - name: run-with-jax-backend
+        run: |
+          sed -i -e 's/GPU/JAX/g' docs_code.py
+          sed -i -e 's/Coalescence(/Coalescence(adaptive=False,/g' docs_code.py
+          awk "{print}; /readme.png/ {exit}" docs_code.py > docs_code_jax.py
+          python -We docs_code_jax.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
     steps:
       - uses: actions/checkout@v4.1.6
-      - uses: actions/setup-python@v6.2.0
+      - uses: actions/setup-python@v7.0.0
         with:
-          python-version: "3.10"
+          python-version: "3.11"
       - run: pip install -e .[CI_version_pins]
       - run: pip install pytest-codeblocks pytest
-      - run: python -c "import pytest_codeblocks; code=pytest_codeblocks.extract_from_file('docs/markdown/pysdm_landing.md'); f=open('tutorial_code.jl', 'w'); f.writelines(block.code for block in code if block.syntax=='Julia'); f.close()"
+      - run: python -c "import pytest_codeblocks; code=pytest_codeblocks.extract_from_file('docs/markdown/pysdm_landing.md'); f=open('docs_code.jl', 'w'); f.writelines(block.code for block in code if block.syntax=='Julia'); f.close()"
       - uses: julia-actions/setup-julia@v3
-      - run: cat -n tutorial_code.jl
+      - run: cat -n docs_code.jl
       - run: echo PYTHON=`which python` >> $GITHUB_ENV
-      - run: julia tutorial_code.jl
-      - run: sed -i -e 's/CPU/GPU/g' tutorial_code.jl
-      - run: julia tutorial_code.jl
+      - run: julia docs_code.jl
+      - run: sed -i -e 's/CPU/GPU/g' docs_code.jl
+      - run: julia docs_code.jl
 
   matlab:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4.1.6
-      - uses: actions/setup-python@v6.2.0
+      - uses: actions/setup-python@v7.0.0
         with:
           python-version: "3.12"
 
@@ -70,19 +85,19 @@ jobs:
 
       - run: pip install -e .[CI_version_pins]
       - run: pip install pytest-codeblocks pytest
-      - run: python -c "import pytest_codeblocks; code=pytest_codeblocks.extract_from_file('docs/markdown/pysdm_landing.md'); f=open('tutorial_code.m', 'w'); f.writelines(block.code for block in code if block.syntax=='Matlab'); f.close()"
-      - run: cat -n tutorial_code.m
-      - uses: matlab-actions/setup-matlab@v3.0.1
+      - run: python -c "import pytest_codeblocks; code=pytest_codeblocks.extract_from_file('docs/markdown/pysdm_landing.md'); f=open('docs_code.m', 'w'); f.writelines(block.code for block in code if block.syntax=='Matlab'); f.close()"
+      - run: cat -n docs_code.m
+      - uses: matlab-actions/setup-matlab@v3.0.2
         with:
           release: R2024b
           cache: true
       - uses: matlab-actions/run-command@v3.2
         with:
-          command: tutorial_code
-      - run: sed -i -e 's/CPU/GPU/g' tutorial_code.m
+          command: docs_code
+      - run: sed -i -e 's/CPU/GPU/g' docs_code.m
       - uses: matlab-actions/run-command@v3.2
         with:
-          command: tutorial_code
+          command: docs_code
 
       - name: artefacts
         if: github.ref == 'refs/heads/main'

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/.github/workflows/pdoc.yml

@@ -23,7 +23,7 @@ jobs:
         with:
           submodules: recursive
           persist-credentials: false
-      - uses: actions/setup-python@v6.2.0
+      - uses: actions/setup-python@v7.0.0
         with:
           python-version: "3.12"
 

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/.github/workflows/precommit.yml

@@ -22,9 +22,9 @@ jobs:
     runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v4.1.6
-      - uses: actions/setup-python@v6.2.0
+      - uses: actions/setup-python@v7.0.0
         with:
-          python-version: "3.10"
+          python-version: "3.11"
       - run: |
           pip install pre-commit
           pre-commit run --all-files --show-diff-on-failure
@@ -38,7 +38,9 @@ jobs:
         with:
           submodules: recursive
           fetch-depth: 0  # https://github.com/pypa/setuptools_scm/issues/480
-      - run: pip install -r tests/devops_tests/requirements.txt
+      - run: |
+          pip install -r tests/devops_tests/requirements.txt
+          pip install "ghapi<2.0.0" "fastcore<2"
       - env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: pytest  ${{ env.pytest_options }} -k "not test_run_notebooks" tests/devops_tests

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/.github/workflows/pypi.yml

@@ -26,7 +26,7 @@ jobs:
         with:
           submodules: recursive
           fetch-depth: 0  # https://github.com/pypa/setuptools_scm/issues/480
-      - uses: actions/setup-python@v6.2.0
+      - uses: actions/setup-python@v7.0.0
         with:
           python-version: "3.12"
 
@@ -58,11 +58,11 @@ jobs:
     strategy:
       matrix:
         platform: [ ubuntu-latest, macos-latest, windows-latest ]
-        python-version: [ "3.10", "3.12" ]
+        python-version: [ "3.11", "3.12" ]
     runs-on: ${{ matrix.platform }}
     needs: [ package ]
     steps:
-      - uses: actions/setup-python@v6.2.0
+      - uses: actions/setup-python@v7.0.0
         with:
           python-version: ${{ matrix.python-version }}
       - uses: actions/download-artifact@v8

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/.github/workflows/stale.yml

@@ -18,7 +18,7 @@ jobs:
       pull-requests: write
 
     steps:
-    - uses: actions/stale@v10.3.0
+    - uses: actions/stale@v11.0.0
       with:
         repo-token: ${{ secrets.GITHUB_TOKEN }}
         stale-issue-message: 'Stale issue message'

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/.github/workflows/tests.yml

@@ -25,8 +25,8 @@ jobs:
   env-unit-tests-setup:
     strategy:
       matrix:
-        platform: [ ubuntu-24.04, macos-15-intel, macos-14, windows-latest ]
-        python-version: [ "3.10", "3.12" ]
+        platform: [ ubuntu-24.04, macos-14, windows-latest ]
+        python-version: [ "3.11", "3.12" ]
     runs-on: ${{ matrix.platform }}
     steps:
       - uses: actions/checkout@v4.1.6
@@ -45,7 +45,7 @@ jobs:
         python-version: ["3.12"]
         include: # for pylint
           - platform: ubuntu-24.04
-            python-version: "3.10"
+            python-version: "3.11"
     runs-on: ${{ matrix.platform }}
     steps:
       - uses: actions/checkout@v4.1.6
@@ -62,7 +62,7 @@ jobs:
     strategy:
       matrix:
         platform: [ubuntu-24.04]
-        python-version: ["3.10"]
+        python-version: ["3.11"]
     runs-on: ${{ matrix.platform }}
     steps:
     - uses: actions/checkout@v4.1.6
@@ -88,8 +88,8 @@ jobs:
     needs: env-unit-tests-setup
     strategy:
       matrix:
-        platform: [ ubuntu-24.04, macos-15-intel, macos-14, windows-latest ]
-        python-version: [ "3.10", "3.12" ]
+        platform: [ ubuntu-24.04, macos-14, windows-latest ]
+        python-version: [ "3.11", "3.12" ]
     runs-on: ${{ matrix.platform }}
     env:
       NUMBA_DISABLE_JIT: 1

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/.gitignore

@@ -146,3 +146,12 @@ dmypy.json
 
 #Jetbrains
 .idea/
+
+#JAX work
+tmp/
+__ptx_cache__.db
+.venv-ubuntu/
+prof/
+jax_test.ipynb
+jax_test.py
+readme.png

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pysdm
-Version: 3.0.0rc37
+Version: 3.0.0rc38
 Summary: Pythonic particle-based (super-droplet) warm-rain/aqueous-chemistry cloud microphysics package with box, parcel & 1D/2D prescribed-flow examples in Python, Julia and Matlab
 Author-email: "https://github.com/open-atmos/PySDM/graphs/contributors" <sylwester.arabas@agh.edu.pl>
 License-Expression: GPL-3.0-only
@@ -23,6 +23,7 @@ License-File: LICENSE
 Requires-Dist: ThrustRTC>=0.3.20
 Requires-Dist: CURandRTC>=0.1.2
 Requires-Dist: numba>=0.51.2
+Requires-Dist: jax>=0.8.0
 Requires-Dist: scipy>=1.15.0
 Requires-Dist: numpy
 Requires-Dist: Pint
@@ -40,10 +41,11 @@ Requires-Dist: PyPartMC; extra == "nonunit-tests"
 Provides-Extra: ci-version-pins
 Requires-Dist: PyPartMC==2.1.0; extra == "ci-version-pins"
 Requires-Dist: numba==0.60.0; extra == "ci-version-pins"
-Requires-Dist: numpy==1.26.4; extra == "ci-version-pins"
-Requires-Dist: scipy==1.15.0; extra == "ci-version-pins"
+Requires-Dist: numpy==2.0.2; extra == "ci-version-pins"
+Requires-Dist: scipy==1.15.3; extra == "ci-version-pins"
 Requires-Dist: ThrustRTC==0.3.20; extra == "ci-version-pins"
 Requires-Dist: CURandRTC==0.1.7; extra == "ci-version-pins"
+Requires-Dist: jax==0.8.0; extra == "ci-version-pins"
 Requires-Dist: Pint==0.24.4; extra == "ci-version-pins"
 Requires-Dist: chempy==0.10.1; extra == "ci-version-pins"
 Requires-Dist: pyevtk==1.7.0; extra == "ci-version-pins"
@@ -55,6 +57,7 @@ Dynamic: license-file
 [![Python 3](https://img.shields.io/static/v1?label=Python&logo=Python&color=3776AB&message=3)](https://www.python.org/)
 [![LLVM](https://img.shields.io/static/v1?label=LLVM&logo=LLVM&color=gold&message=Numba)](https://numba.pydata.org)
 [![CUDA](https://img.shields.io/static/v1?label=CUDA&logo=nVidia&color=87ce3e&message=ThrustRTC)](https://pypi.org/project/ThrustRTC/)
+[![JAX](https://img.shields.io/badge/Google-JAX-blue?style=flat&logo=google)](https://docs.jax.dev/en/latest/index.html)
 [![Linux OK](https://img.shields.io/static/v1?label=Linux&logo=Linux&color=yellow&message=%E2%9C%93)](https://en.wikipedia.org/wiki/Linux)
 [![macOS OK](https://img.shields.io/static/v1?label=macOS&logo=Apple&color=silver&message=%E2%9C%93)](https://en.wikipedia.org/wiki/macOS)
 [![Windows OK](https://img.shields.io/static/v1?label=Windows&logo=Windows&color=white&message=%E2%9C%93)](https://en.wikipedia.org/wiki/Windows)
@@ -102,15 +105,17 @@ Each tutorial also comes with a set of questions at the end that can be used as
 Like the examples, these tutorials can be executed or viewed "in the cloud" making it an especially 
   easy way for students to get started.
 
-PySDM has two alternative parallel number-crunching backends 
+PySDM has three alternative parallel number-crunching backends 
   available: multi-threaded CPU backend based on [Numba](http://numba.pydata.org/) 
   and GPU-resident backend built on top of [ThrustRTC](https://pypi.org/project/ThrustRTC/).
+  Additionally, a new hybrid backend built with [Jax](https://docs.jax.dev/en/latest/index.html) is being developed.
 The [`Numba`](https://open-atmos.github.io/PySDM/PySDM/backends/numba.html) backend (aliased ``CPU``) features multi-threaded parallelism for 
   multi-core CPUs, it uses the just-in-time compilation technique based on the LLVM infrastructure.
 The [`ThrustRTC`](https://open-atmos.github.io/PySDM/PySDM/backends/thrust_rtc.html) backend (aliased ``GPU``) offers GPU-resident operation of PySDM
   leveraging the [SIMT](https://en.wikipedia.org/wiki/Single_instruction,_multiple_threads) 
   parallelisation model. 
 Using the ``GPU`` backend requires nVidia hardware and [CUDA driver](https://developer.nvidia.com/cuda-downloads).
+The [`Jax`](https://docs.jax.dev/en/latest/index.html) backend (aliased ``JAX``) is flexible and works with any device configurations (CPU/GPU/TPU), utilising the same backend API. As It is in early development, only simple coagulation cases in single-cell simulations are supported. 
 
 For an overview of PySDM features (and the preferred way to cite PySDM in papers), please refer to our JOSS papers:
 - [Bartman et al. 2022](https://doi.org/10.21105/joss.03219) (PySDM v1).
@@ -121,7 +126,7 @@ For a list of talks and other materials on PySDM as well as a list of published
 
 ## Dependencies and Installation
 
-PySDM dependencies are: [Numpy](https://numpy.org/), [Numba](http://numba.pydata.org/), [SciPy](https://scipy.org/), 
+PySDM dependencies are: [Numpy](https://numpy.org/), [Numba](http://numba.pydata.org/), [Jax](https://docs.jax.dev/en/latest/index.html), [SciPy](https://scipy.org/), 
 [Pint](https://pint.readthedocs.io/), [chempy](https://pypi.org/project/chempy/), 
 [pyevtk](https://pypi.org/project/pyevtk/),
 [ThrustRTC](https://fynv.github.io/ThrustRTC/) and [CURandRTC](https://github.com/fynv/CURandRTC).
@@ -156,6 +161,7 @@ only the supporting .py files).
 [![Numba](https://img.shields.io/static/v1?label=+&logo=Numba&color=orange&message=Numba)](https://numba.pydata.org)
 [![LLVM](https://img.shields.io/static/v1?label=+&logo=LLVM&color=gold&message=LLVM)](https://llvm.org)
 [![CUDA](https://img.shields.io/static/v1?label=+&logo=nVidia&color=darkgreen&message=ThrustRTC/CUDA)](https://pypi.org/project/ThrustRTC/)
+[![JAX](https://img.shields.io/badge/Google-JAX-blue?style=flat&logo=google)](https://docs.jax.dev/en/latest/index.html)
 [![NumPy](https://img.shields.io/static/v1?label=+&logo=numpy&color=blue&message=NumPy)](https://numpy.org/)
 [![pytest](https://img.shields.io/static/v1?label=+&logo=pytest&color=purple&message=pytest)](https://pytest.org/)   
 [![Colab](https://img.shields.io/static/v1?label=+&logo=googlecolab&color=darkred&message=Colab)](https://colab.research.google.com/)

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/PySDM/attributes/impl/dummy_attribute.py

@@ -12,7 +12,7 @@ class DummyAttribute(Attribute):
 
     def allocate(self, idx):
         super().allocate(idx)
-        self.data[:] = np.nan
+        self.data.fill(np.nan)
 
     def get(self):
         return self.data

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/PySDM/backends/__init__.py

@@ -11,11 +11,14 @@ import warnings
 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
 
 
@@ -72,7 +75,7 @@ else:
             super().__init__(size, seed)
             self.generator = np.random.default_rng(seed)
 
-        def __call__(self, storage):
+        def u01(self, storage):
             # pylint: disable=unsupported-assignment-operation
             storage.data.ndarray[:] = self.generator.uniform(0, 1, storage.shape)
 
@@ -93,3 +96,5 @@ CPU = partial(_cached_backend, backend_class=Numba)
 
 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-3.0.0rc37 → pysdm-3.0.0rc38}/PySDM/backends/impl_common/backend_methods.py

@@ -15,3 +15,5 @@ class BackendMethods:
             self.Storage = None
         if not hasattr(self, "default_jit_flags"):
             self.default_jit_flags = {}
+        if not hasattr(self, "block_until_ready"):
+            self.block_until_ready = False

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/PySDM/backends/impl_common/index.py

@@ -40,12 +40,10 @@ def make_Index(backend):
 
         def shuffle(self, temporary, parts=None):
             if parts is None:
-                backend.shuffle_global(
-                    idx=self.data, length=self.length, u01=temporary.data
-                )
+                backend.shuffle_global(idx=self, u01=temporary.data)
             else:
                 backend.shuffle_local(
-                    idx=self.data, u01=temporary.data, cell_start=parts.data
+                    idx=self, u01=temporary.data, cell_start=parts.data
                 )
 
         def remove_zero_n_or_flagged(self, indexed_storage):

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/PySDM/backends/impl_common/indexed_storage.py

@@ -2,6 +2,8 @@
 attribute storage class featuring particle permutation logic
 """
 
+from PySDM.backends.jax import Jax
+
 from .storage_utils import StorageSignature
 
 
@@ -39,6 +41,13 @@ def make_IndexedStorage(backend):
             result = IndexedStorage.indexed(idx, storage)
             return result
 
+        def row_view(self, i):
+            # TODO #1913: Implement IndexedStorage row_view for Jax (check on Numba/ThrustRTC)
+            if isinstance(backend, Jax):
+                return super().row_view(i)
+
+            return IndexedStorage.indexed(self.idx, super().row_view(i))
+
         def to_ndarray(self, *, raw=False):
             result = backend.Storage.to_ndarray(self)
             dim = len(self.shape)

{pysdm-3.0.0rc37 → pysdm-3.0.0rc38}/PySDM/backends/impl_common/storage_utils.py

@@ -48,15 +48,15 @@ class StorageBase:
         raise NotImplementedError()
 
     @abstractmethod
-    def urand(self, generator):
+    def upload(self, data):
         raise NotImplementedError()
 
     @abstractmethod
-    def upload(self, data):
+    def fill(self, other):
         raise NotImplementedError()
 
     @abstractmethod
-    def fill(self, other):
+    def row_view(self, i):
         raise NotImplementedError()
 
 

pysdm-3.0.0rc38/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-3.0.0rc38/PySDM/backends/impl_jax/methods/collisions_methods.py

@@ -0,0 +1,199 @@
+"""
+JAX implementation of backend methods for particle collisions
+"""
+
+from functools import cached_property
+
+import jax
+import jax.numpy as jnp
+
+from PySDM.backends.impl_common.backend_methods import BackendMethods
+from PySDM.backends.impl_jax.storage import Storage
+
+
+def flag_zero_multiplicity(j, k, multiplicity, healthy):
+    raise NotImplementedError()
+
+
+class CollisionsMethods(BackendMethods):
+    @cached_property
+    def _collision_coalescence_body(self):
+        @jax.jit
+        def body(  # pylint: disable=too-many-locals
+            multiplicity,
+            idx,
+            attributes,
+            gamma,
+            is_first_in_pair,
+        ):
+            i = jnp.arange(len(multiplicity) // 2)
+            offset = 1 - is_first_in_pair[2 * i]
+            # offset = 0
+            j = idx[2 * i + offset]
+            k = idx[2 * i + 1 + offset]
+            mj = multiplicity[j]
+            mk = multiplicity[k]
+
+            new_n = mj - jnp.int64(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
+            new_attr = gamma[i] * attributes[:, j] + attributes[:, k]
+
+            new_attr = new_attr * zero_mask
+
+            mult_j = jnp.where(
+                pos_mask,
+                mult_pos_updates,
+                jnp.where(zero_mask, mk // 2, multiplicity[j]),
+            )
+            mult_k = jnp.where(
+                zero_mask,
+                jnp.where(zero_mask, mk - (mk // 2), multiplicity[k]),
+                multiplicity[k],
+            )
+
+            multiplicity = multiplicity.at[j].set(mult_j)
+
+            multiplicity = multiplicity.at[k].set(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, attributes
+
+        return body
+
+    def collision_coalescence(
+        self,
+        *,
+        multiplicity,
+        idx,
+        attributes,
+        gamma,
+        healthy,  # pylint: disable=unused-argument
+        cell_id,  # pylint: disable=unused-argument
+        coalescence_rate,  # pylint: disable=unused-argument
+        is_first_in_pair,
+        # TODO #1913: implement rates, multi-cell collisions and healthy
+    ):
+        multiplicity.data, attributes.data = self._collision_coalescence_body(
+            multiplicity.data,
+            idx.data,
+            attributes.data,
+            gamma.data,
+            is_first_in_pair.indicator.data,
+        )
+        multiplicity.data.block_until_ready()
+
+    @cached_property
+    def _compute_gamma_body(self):
+        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 jax.jit(jax.vmap(body, (None, None, None, None, None, 0)))
+
+    def compute_gamma(
+        self,
+        *,
+        prob,
+        rand,
+        multiplicity,
+        cell_id,  # pylint: disable=unused-argument
+        collision_rate_deficit,  # pylint: disable=unused-argument
+        collision_rate,  # pylint: disable=unused-argument
+        is_first_in_pair,
+        out,
+        # TODO #1913: implement rates, multi-cell collisions and healthy
+    ):
+        indices = jnp.arange(len(multiplicity) // 2)
+        out.data = self._compute_gamma_body(
+            prob.data,
+            rand.data,
+            multiplicity.idx.data,
+            multiplicity.data,
+            is_first_in_pair.indicator.data,
+            indices,
+        ).block_until_ready()
+
+    @staticmethod
+    def make_cell_caretaker(
+        idx_shape, idx_dtype, cell_start_len, scheme="default"
+    ):  # pylint: disable=unused-argument
+        # cell_start_len present for API compliance
+        class CellCaretaker:  # pylint: disable=too-few-public-methods
+            def __init__(self, idx_shape, idx_dtype, scheme):
+                assert scheme == "default"
+                self.tmp_idx = Storage.empty(idx_shape, idx_dtype)
+
+            def __call__(self, cell_id, cell_idx, cell_start, idx):
+                length = len(idx)
+                cell_start.data = (
+                    CollisionsMethods._counting_sort_by_cell_id_and_update_cell_start(
+                        self.tmp_idx.data,
+                        idx.data,
+                        cell_id.data,
+                        cell_idx.data,
+                        length,
+                        cell_start.data,
+                    )
+                )
+
+        return CellCaretaker(idx_shape, idx_dtype, scheme)
+
+    @cached_property
+    def _normalize_body(self):
+        @jax.jit
+        def body(prob, cell_start, timestep, dv):
+            sd_num = cell_start[1] - cell_start[0]
+
+            def loop_body(i, prob):
+                norm_factor = (sd_num >= 2) * (
+                    timestep / dv * sd_num * (sd_num - 1) / 2 / (sd_num // 2)
+                )
+                prob = prob.at[i].multiply(norm_factor)
+                return prob
+
+            return jax.lax.fori_loop(0, prob.shape[0], loop_body, prob)
+
+        return body
+
+    # pylint: disable=too-many-arguments,disable=unused-argument
+    def normalize(self, prob, cell_id, cell_idx, cell_start, norm_factor, timestep, dv):
+        prob.data = self._normalize_body(
+            prob.data, cell_start.data, timestep, dv
+        ).block_until_ready()
+
+    @staticmethod
+    # pylint: disable=too-many-arguments
+    def _counting_sort_by_cell_id_and_update_cell_start(
+        new_idx,  # pylint: disable=unused-argument
+        idx,  # pylint: disable=unused-argument
+        cell_id,  # pylint: disable=unused-argument
+        cell_idx,  # pylint: disable=unused-argument
+        length,  # pylint: disable=unused-argument
+        cell_start,  # pylint: disable=unused-argument
+    ):
+        # TODO #1913:  implement sorting
+        cell_start = cell_start.at[0].set(0)
+        return cell_start

pysdm-3.0.0rc38/PySDM/backends/impl_jax/methods/index_methods.py

@@ -0,0 +1,31 @@
+"""
+JAX 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_global_body(self):
+    #     @jax.jit
+    #     def body(idx, u01):
+    #         idx = idx.at[:].set(jax.numpy.argsort(u01, stable=False))
+    #         return idx
+
+    #     return body
+
+    def shuffle_global(self, idx, u01):
+        # TODO #1913: decide whether to use u01 argsort or random.permute
+        # TODO #1913: https://github.com/jax-ml/jax/issues/5328
+        # idx.data = self._shuffle_global_body(idx.data, u01) (requires actual u01)
+        u01.permute(idx)
+
+    def shuffle_local(self, idx, u01, cell_start):  # pylint: disable=unused-argument
+        # TODO #1913: implement shuffle_local
+        self.shuffle_global(idx, u01)