cosmax 0.0.7

Fast and differentiable implementations of operations needed for inference and analysis in cosmology. Powered by JAX

Cosmax

Fast and differentiable tools for analysis and inference on structured and unstructured data in cosmology. Powered by JAX.

Install with pip install cosmax.

What can I do with this package?

Classical Use Cases

Generate the Matter Power Spectrum from a 3D densitiy field:

Generate the ICs for a given matter power spectrum:

Cloud in a cell mass assignment scheme to convert unstructured particles to a 3D density field:

Optimization

With gradient optimization over the cic mass assignment scheme, we can find the particle positions that best fit an observed density field:

We can also use the power spectrum as a loss function to find a conditional IC density field :

What are the limitations?

• The package is currently in development and the API is not stable.
• The package generally ONLY works with 3D square boxes and periodic boundary conditions.
• I am not a physicist, so the package may contain errors and inaccuracies.

Development

To develop, clone the repository and install the package in editable mode:

pip install -e .

To release as pip package, tests, docs and builds are handled automatically by github actions as defined in .github/workflows. To make a new release:

git tag v*.*.*
git push origin v*.*.*

and change the version number in pyproject.toml.

Test

pytest

Build

python -m build

Local Docs

With the pip package sphinx installed, run

sphinx-apidoc -o docs/source cosmax/
sphinx-build -b html docs/source docs/_build

to view locally

cd docs/_build
python -m http.server

Acknowledgements

• PowerBox was used as a reference implementation the matter power spectrum.

• Felix Köhlers implementation of Fourier Neural Operator was used as a guide to adapt the fno to jax and equinox.

• The original NeuralOperator served as a reference for the fno implementation.