Merge pull request #21 from tumaer/upgrade_to_jax_062

Upgrade to jax>=0.6.2

Author: arturtoshev

.github/workflows/publish.yml

@@ -10,11 +10,11 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
     - name: Set up Python
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v6
       with:
-        python-version: '3.10'
+        python-version: '3.12'
     - name: Install Poetry
       run: |
         python -m pip install --upgrade pip

.github/workflows/ruff.yml

@@ -4,5 +4,7 @@ jobs:
   ruff:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
-      - uses: chartboost/ruff-action@v1
+      - uses: actions/checkout@v5
+      - uses: astral-sh/ruff-action@v3
+      - run: ruff check --fix
+      - run: ruff format

.github/workflows/tests.yml

@@ -15,12 +15,12 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10"]  # Add more Python versions here, e.g. "3.9", "3.11"
+        python-version: ["3.10", "3.12"]
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v6
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install Poetry
@@ -35,7 +35,7 @@ jobs:
       run: |
         .venv/bin/pytest --cov-report=xml
     - name: Upload coverage report to Codecov
-      uses: codecov/codecov-action@v4
+      uses: codecov/codecov-action@v5
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
         file: ./coverage.xml

.pre-commit-config.yaml

@@ -8,7 +8,7 @@ exclude: |
   )$
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.4.0
+    rev: v6.0.0
     hooks:
       - id: check-merge-conflict
       - id: check-added-large-files
@@ -18,8 +18,8 @@ repos:
       - id: check-yaml
       - id: requirements-txt-fixer
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.1.8'
+    rev: 'v0.14.0'
     hooks:
-      - id: ruff
+      - id: ruff-check
         args: [ --fix ]
       - id: ruff-format

README.md

@@ -31,15 +31,15 @@
 Install the `jax-sph` library from PyPI as
 
 ```bash
-python3.10 -m venv venv
+python3.12 -m venv venv
 source venv/bin/activate
 pip install jax-sph
 ```
 
 By default `jax-sph` is installed without GPU support. If you have a CUDA-capable GPU, follow the instructions in the [GPU support](#gpu-support) section.
 
 ### Clone
-We recommend using a `poetry` or `python3-venv` environment.
+We recommend using a `poetry` or `python3-venv` environment. Last tested with python3.10 + jax[cpu]==0.6.2 and python3.12 + jax[cpu]==0.7.2.
 
 **Using Poetry**
 ```bash
@@ -51,7 +51,7 @@ Later, you just need to `source .venv/bin/activate` to activate the environment.
 
 **Using `python3-venv`**
 ```bash
-python3 -m venv venv
+python3.12 -m venv venv
 source venv/bin/activate
 pip install -r requirements.txt
 pip install -e . # to install jax_sph in interactive mode
@@ -62,7 +62,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 -U "jax[cuda12]==0.4.29"
+pip install -U "jax[cuda12]" # specify version as "jax[cuda12]==0.6.2"
 ```
 
 ## Getting Started

cases/ht.py

@@ -1,6 +1,5 @@
 "Heat Transfer over a flat plate setup"
 
-
 import jax.numpy as jnp
 import numpy as np
 from omegaconf import DictConfig

cases/ldc.py

@@ -1,6 +1,5 @@
 """Lid-driven cavity case setup"""
 
-
 import jax.numpy as jnp
 import numpy as np
 from omegaconf import DictConfig

cases/tgv.py

@@ -1,6 +1,5 @@
 """Taylor-Green case setup"""
 
-
 import jax.numpy as jnp
 import numpy as np
 from omegaconf import DictConfig

docs/requirements.txt

@@ -1,6 +1,6 @@
 absl-py>=2.1.0
 h5py
-jax[cpu]==0.4.29
+jax[cpu]==0.6.2
 jraph>=0.0.6.dev0
 omegaconf
 pandas

jax_sph/case_setup.py

@@ -191,7 +191,7 @@ def initialize(self):
                 assert state[k][mask].shape == _state[k][_mask].shape, ValueError(
                     f"Shape mismatch for key {k} in state0 file."
                 )
-                state[k][mask] = _state[k][_mask]
+                state[k] = state[k].at[mask].set(_state[k][_mask])
 
         # the following arguments are needed for dataset generation
         cfg.case.c_ref, cfg.case.p_ref, cfg.case.p_bg = c_ref, p_ref, p_bg