jaxdecomp proto

.github/workflows/formatting.yml

@@ -0,0 +1,21 @@
+name: Code Formatting
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v3
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip isort
+        python -m pip install pre-commit
+    - name: Run pre-commit
+      run: python -m pre_commit run --all-files

.github/workflows/tests.yml

@@ -0,0 +1,45 @@
+name: Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  run_tests:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10" , "3.11" , "3.12"]
+
+    steps:
+    - name: Checkout Source
+      uses: actions/[email protected]
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install dependencies
+      run: |
+        sudo apt-get install -y libopenmpi-dev
+        python -m pip install --upgrade pip
+        pip install jax==0.4.35
+        pip install numpy setuptools cython wheel
+        pip install git+https://github.com/MP-Gadget/pfft-python
+        pip install git+https://github.com/MP-Gadget/pmesh
+        pip install git+https://github.com/ASKabalan/fastpm-python --no-build-isolation
+        pip install .[test]
+
+    - name: Run Single Device Tests
+      run: |
+        cd tests
+        pytest -v -m "not distributed"
+    - name: Run Distributed tests
+      run: |
+        pytest -v -m distributed

.gitignore

@@ -98,6 +98,11 @@ __pypackages__/
 celerybeat-schedule
 celerybeat.pid
 
+
+out
+traces
+*.npy
+*.out
 # SageMath parsed files
 *.sage.py
 

.pre-commit-config.yaml

@@ -14,4 +14,4 @@ repos:
     rev: 5.13.2
     hooks:
       - id: isort
-        name: isort (python)
+        name: isort (python)

README.md

@@ -4,16 +4,15 @@
 <!-- ALL-CONTRIBUTORS-BADGE:END -->
 JAX-powered Cosmological Particle-Mesh N-body Solver
 
-**This project is currently in an early design phase. All inputs are welcome on the [design document](https://github.com/DifferentiableUniverseInitiative/JaxPM/blob/main/design.md)**
-
 ## Goals
 
 Provide a modern infrastructure to support differentiable PM N-body simulations using JAX:
 - Keep implementation simple and readable, in pure NumPy API
-- Transparent distribution using builtin `xmap`
 - Any order forward and backward automatic differentiation
 - Support automated batching using `vmap`
 - Compatibility with external optimizer libraries like `optax`
+- Now fully distributable on **multi-GPU and multi-node** systems using [jaxDecomp](https://github.com/DifferentiableUniverseInitiative/jaxDecomp) working with`JAX v0.4.35`
+
 
 ## Open development and use
 
@@ -23,6 +22,10 @@ Current expectations are:
 - Everyone is welcome to contribute, and can join the JOSS publication (until it is submitted to the journal).
 - Anyone (including main contributors) can use this code as a framework to build and publish their own applications, with no expectation that they *need* to extend authorship to all jaxpm developers.
 
+## Getting Started
+
+To dive into JaxPM’s capabilities, please explore the **notebook section** for detailed tutorials and examples on various setups, from single-device simulations to multi-host configurations. You can find the notebooks' [README here](notebooks/README.md) for a structured guide through each tutorial.
+
 
 ## Contributors ✨