Merge branch 'thermal_conductivity_fix_and_paper_revision' of github.com:tluettm/PySDM into thermal_conductivity_fix_and_paper_revision

Author: tluettm

.github/workflows/tutorial_code.yml .github/workflows/docs_code_snippets.yml.github/workflows/tutorial_code.yml renamed to .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@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,23 +43,35 @@ 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@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
@@ -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
+      - 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'

.github/workflows/precommit.yml

@@ -24,7 +24,7 @@ jobs:
       - uses: actions/checkout@v4.1.6
       - 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

.github/workflows/pypi.yml

@@ -58,7 +58,7 @@ 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:

.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

.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/attributes/impl/dummy_attribute.py

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

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
 
 
@@ -72,7 +75,7 @@ def __init__(self, size, seed):
             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 @@ 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/backend_methods.py

@@ -15,3 +15,5 @@ def __init__(self):
             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/backends/impl_common/index.py

@@ -40,12 +40,10 @@ def sort_by_key(self, keys):
 
         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/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 from_ndarray(idx, array):
             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)