First milestone

[ASKabalan]

@EiffL

Milestone 1: Initial Release & Core Distributed Pipeline

Overview
This milestone marks the first major release of the package, achieving ~90% of the initial vision! The core focus of this release is delivering an end-to-end, fully distributed pipeline for cosmological simulations and lensing.

Note: Some advanced sampling utilities are excluded from this release and will be added in a future milestone.

Key Features Delivered

• LPT N-Body Simulations: Fully operational and distributed.
• Lensing Suite:
  • Born Approximation support.
  • Raytracing powered by Dorian (https://github.com/ASKabalan/Dorian).
• End-to-End Distribution: The entire pipeline from N-Body to lensing is designed to work efficiently in a distributed environment.

Upstream Dependencies (Blockers)
This release currently depends on several upstream pull requests across the ecosystem. These PRs must be merged before we can finalize this milestone:

jax_cosmo

• [PR #129 (https://github.com/Use jnp.interp DifferentiableUniverseInitiative/jax_cosmo#129): Use jnp.interp
• [PR #143 (https://github.com/Add second growth factor to jax-cosmo cache DifferentiableUniverseInitiative/jax_cosmo#143): Add second growth factor to jax-cosmo cache
• [PR #144 (https://github.com/Add massive nu DifferentiableUniverseInitiative/jax_cosmo#144): Add massive nu
• [PR #145 (https://github.com/Better cache DifferentiableUniverseInitiative/jax_cosmo#145): Better cache
• [PR #146 (https://github.com/fix error with DistrubutionNotFound and switch to PackageNotFoundError DifferentiableUniverseInitiative/jax_cosmo#146): Fix error with DistributionNotFound and switch to PackageNotFoundError

JaxPM

• [PR #53 (https://github.com/update spherical and 2d painting to work better for sharding DifferentiableUniverseInitiative/JaxPM#53): Update spherical and 2d painting to work better for sharding
• [PR #54 (https://github.com/Allow 1d and 2d normal fields in distributed mode DifferentiableUniverseInitiative/JaxPM#54): Allow 1d and 2d normal fields in distributed mode

jax-healpy

• [PR Add prelinimary example #2 (https://github.com/maintain shapes for ang2vec and vec2ang functions for batched inputs CMBSciPol/jax-healpy#2): Maintain shapes for ang2vec and vec2ang functions for batched inputs