├── .clang-format ├── .github └── workflows │ ├── formatting.yml │ ├── github-deploy.yml │ ├── joss-paper-pdf.yml │ └── tests.yml ├── .gitignore ├── .gitmodules ├── .pre-commit-config.yaml ├── .readthedocs.yml ├── CHANGELOG.md ├── CITATION.cff ├── CMakeLists.txt ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── benchmarks ├── jaxfft.py └── mpi4jaxfft.py ├── docs ├── 01-basic_usage.md ├── 02-decomposition.md ├── 03-distributed_ffts.md ├── 04-halo_exchange.md ├── 05-transpositions.md ├── 06-caveats.md ├── Makefile ├── api.rst ├── conf.py ├── index.rst ├── installation.md ├── make.bat └── requirements.txt ├── examples ├── README.md ├── lpt_nbody_demo.py ├── scatter.py ├── submit_rusty.sbatch └── visualizer.ipynb ├── joss-paper ├── assets │ ├── LPT_density_field_z0_2048.png │ ├── benchmark.png │ ├── decomp2d.gif │ ├── fft.svg │ ├── halo-exchange.svg │ ├── strong_scaling.png │ └── weak_scaling.png ├── paper.bib └── paper.md ├── pyproject.toml ├── slurms ├── README.md └── template.slurm ├── src ├── csrc │ ├── fft.cu │ ├── grid_descriptor_mgr.cc │ ├── halo.cu │ ├── include │ │ ├── checks.h │ │ ├── fft.h │ │ ├── grid_descriptor_mgr.h │ │ ├── halo.h │ │ ├── helpers.h │ │ ├── jaxdecomp.h │ │ ├── logger.hpp │ │ ├── perfostep.hpp │ │ └── transpose.h │ ├── jaxdecomp.cc │ └── transpose.cu └── jaxdecomp │ ├── __init__.py │ ├── _src │ ├── __init__.py │ ├── cudecomp │ │ ├── __init__.py │ │ ├── fft.py │ │ ├── halo.py │ │ └── transpose.py │ ├── error.py │ ├── fft_utils.py │ ├── jax │ │ ├── __init__.py │ │ ├── fft.py │ │ ├── fftfreq.py │ │ ├── halo.py │ │ └── transpose.py │ ├── pencil_utils.py │ └── spmd_ops.py │ ├── fft.py │ ├── halo.py │ ├── transpose.py │ └── typing.py └── tests ├── conftest.py ├── run_all_tests.sh ├── test_fft.py ├── test_halo.py └── test_transpose.py /.clang-format: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.clang-format -------------------------------------------------------------------------------- /.github/workflows/formatting.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.github/workflows/formatting.yml -------------------------------------------------------------------------------- /.github/workflows/github-deploy.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.github/workflows/github-deploy.yml -------------------------------------------------------------------------------- /.github/workflows/joss-paper-pdf.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.github/workflows/joss-paper-pdf.yml -------------------------------------------------------------------------------- /.github/workflows/tests.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.github/workflows/tests.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.gitignore -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.gitmodules -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /.readthedocs.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/.readthedocs.yml -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/CHANGELOG.md -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/CITATION.cff -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/CMakeLists.txt -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/README.md -------------------------------------------------------------------------------- /benchmarks/jaxfft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/benchmarks/jaxfft.py -------------------------------------------------------------------------------- /benchmarks/mpi4jaxfft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/benchmarks/mpi4jaxfft.py -------------------------------------------------------------------------------- /docs/01-basic_usage.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/01-basic_usage.md -------------------------------------------------------------------------------- /docs/02-decomposition.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/02-decomposition.md -------------------------------------------------------------------------------- /docs/03-distributed_ffts.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/03-distributed_ffts.md -------------------------------------------------------------------------------- /docs/04-halo_exchange.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/04-halo_exchange.md -------------------------------------------------------------------------------- /docs/05-transpositions.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/05-transpositions.md -------------------------------------------------------------------------------- /docs/06-caveats.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/06-caveats.md -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/api.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/api.rst -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/conf.py -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/index.rst -------------------------------------------------------------------------------- /docs/installation.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/installation.md -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/make.bat -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/docs/requirements.txt -------------------------------------------------------------------------------- /examples/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/examples/README.md -------------------------------------------------------------------------------- /examples/lpt_nbody_demo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/examples/lpt_nbody_demo.py -------------------------------------------------------------------------------- /examples/scatter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/examples/scatter.py -------------------------------------------------------------------------------- /examples/submit_rusty.sbatch: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/examples/submit_rusty.sbatch -------------------------------------------------------------------------------- /examples/visualizer.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/examples/visualizer.ipynb -------------------------------------------------------------------------------- /joss-paper/assets/LPT_density_field_z0_2048.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/assets/LPT_density_field_z0_2048.png -------------------------------------------------------------------------------- /joss-paper/assets/benchmark.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/assets/benchmark.png -------------------------------------------------------------------------------- /joss-paper/assets/decomp2d.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/assets/decomp2d.gif -------------------------------------------------------------------------------- /joss-paper/assets/fft.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/assets/fft.svg -------------------------------------------------------------------------------- /joss-paper/assets/halo-exchange.svg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/assets/halo-exchange.svg -------------------------------------------------------------------------------- /joss-paper/assets/strong_scaling.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/assets/strong_scaling.png -------------------------------------------------------------------------------- /joss-paper/assets/weak_scaling.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/assets/weak_scaling.png -------------------------------------------------------------------------------- /joss-paper/paper.bib: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/paper.bib -------------------------------------------------------------------------------- /joss-paper/paper.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/joss-paper/paper.md -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/pyproject.toml -------------------------------------------------------------------------------- /slurms/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/slurms/README.md -------------------------------------------------------------------------------- /slurms/template.slurm: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/slurms/template.slurm -------------------------------------------------------------------------------- /src/csrc/fft.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/fft.cu -------------------------------------------------------------------------------- /src/csrc/grid_descriptor_mgr.cc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/grid_descriptor_mgr.cc -------------------------------------------------------------------------------- /src/csrc/halo.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/halo.cu -------------------------------------------------------------------------------- /src/csrc/include/checks.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/checks.h -------------------------------------------------------------------------------- /src/csrc/include/fft.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/fft.h -------------------------------------------------------------------------------- /src/csrc/include/grid_descriptor_mgr.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/grid_descriptor_mgr.h -------------------------------------------------------------------------------- /src/csrc/include/halo.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/halo.h -------------------------------------------------------------------------------- /src/csrc/include/helpers.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/helpers.h -------------------------------------------------------------------------------- /src/csrc/include/jaxdecomp.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/jaxdecomp.h -------------------------------------------------------------------------------- /src/csrc/include/logger.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/logger.hpp -------------------------------------------------------------------------------- /src/csrc/include/perfostep.hpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/perfostep.hpp -------------------------------------------------------------------------------- /src/csrc/include/transpose.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/include/transpose.h -------------------------------------------------------------------------------- /src/csrc/jaxdecomp.cc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/jaxdecomp.cc -------------------------------------------------------------------------------- /src/csrc/transpose.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/csrc/transpose.cu -------------------------------------------------------------------------------- /src/jaxdecomp/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/__init__.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/__init__.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/cudecomp/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/jaxdecomp/_src/cudecomp/fft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/cudecomp/fft.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/cudecomp/halo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/cudecomp/halo.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/cudecomp/transpose.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/cudecomp/transpose.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/error.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/error.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/fft_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/fft_utils.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/jax/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/jaxdecomp/_src/jax/fft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/jax/fft.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/jax/fftfreq.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/jax/fftfreq.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/jax/halo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/jax/halo.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/jax/transpose.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/jax/transpose.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/pencil_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/pencil_utils.py -------------------------------------------------------------------------------- /src/jaxdecomp/_src/spmd_ops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/_src/spmd_ops.py -------------------------------------------------------------------------------- /src/jaxdecomp/fft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/fft.py -------------------------------------------------------------------------------- /src/jaxdecomp/halo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/halo.py -------------------------------------------------------------------------------- /src/jaxdecomp/transpose.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/transpose.py -------------------------------------------------------------------------------- /src/jaxdecomp/typing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/src/jaxdecomp/typing.py -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/tests/conftest.py -------------------------------------------------------------------------------- /tests/run_all_tests.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/tests/run_all_tests.sh -------------------------------------------------------------------------------- /tests/test_fft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/tests/test_fft.py -------------------------------------------------------------------------------- /tests/test_halo.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/tests/test_halo.py -------------------------------------------------------------------------------- /tests/test_transpose.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DifferentiableUniverseInitiative/jaxDecomp/HEAD/tests/test_transpose.py --------------------------------------------------------------------------------