https://github.com/ickc/python-autojax
DiRAC: revealing the nature of dark matter with the James Webb space telescope and JAX
https://github.com/ickc/python-autojax
Last synced: 2 months ago
JSON representation
DiRAC: revealing the nature of dark matter with the James Webb space telescope and JAX
- Host: GitHub
- URL: https://github.com/ickc/python-autojax
- Owner: ickc
- License: bsd-3-clause
- Created: 2025-01-23T16:54:17.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-09-17T14:08:24.000Z (10 months ago)
- Last Synced: 2026-04-07T06:38:44.237Z (3 months ago)
- Language: Python
- Homepage: https://ickc.github.io/python-autojax/
- Size: 446 KB
- Stars: 0
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# DiRAC: revealing the nature of dark matter with the James Webb space telescope and JAX
[](https://ickc.github.io/python-autojax)
[](https://github.com/ickc/python-autojax/releases)
# Introduction
This is a small self-contained repo for the 3 months project "DiRAC: revealing the nature of dark matter with the James Webb space telescope and JAX". It is organized in 3 modules. `original` is copied from [Jammy2211/dirac_rse_interferometer](https://github.com/Jammy2211/dirac_rse_interferometer) (which itself is copied from various other repos such as AutoGalaxy and AutoLens), and then they are ported to JAX in the `jax` module. While the `original` is already implemented in Numba, a `numba` module is also provided here, mainly as a starting point to port from `original` to something more vectorized first. Often time the `jax` implementation is the same as the `numba` implementation here, or a close variant of it due to differences between Numba and JAX.
As part of the goal to port to JAX implementation is to speed up, benchmark is provided to compare the 3 implementations. See instructions below to see how to run it.
For a logbook style repo that contains every notes about the project and how to run the AutoGalaxy family of softwares, see [ickc/log-PyAutoLens](https://github.com/ickc/log-PyAutoLens). There's some scripts under `external/` here that can only be run using the environment documented there, as it uses some external (to us) dependencies recorded as git submodules over there.
# Installing the project
## pip
```sh
pip install -e .[tests]
```
## conda/mamba
```sh
conda env create -f environment.yml
conda activate autojax
# update using
conda env update --name autojax --file environment.yml --prune
```
## pixi
```sh
pixi install
# prepend everything you run by pixi run, such as
pixi run pytest
```
:::{Note}
From this point on, pixi is assumed and every command is prepended with `pixi run`. For other ways of loading environment, simply load your environment as usual and disregard this prefix.
:::
# Running unit tests and benchmarks
```sh
pytest
# or if you use pixi
pixi run pytest
```
This should runs the tests and also give you benchmark information comparing different implementations.
By default, it will run the benchmark as well. If you want to run unit tests only,
```sh
pytest --benchmark-disable
# or if you use pixi
pixi run pytest --benchmark-disable
```
There are 2 benchmarks designed to test various different implementations to calculate `w_tilde` and `curvature_matrix`, `TestWTilde` and `TestCurvatureMatrix` respectively.
You can filter them like this,
```sh
pytest -k 'TestWTilde or TestCurvatureMatrix'
# or if you use pixi
pixi run pytest -k 'TestWTilde or TestCurvatureMatrix'
```