Feature/cosmology jax

[Jammy2211]

This pull request refactors the cosmology module in the autogalaxy package, with full support for JAX and removing the use of astropy.

This means all cosmology calculations, such as angular diameter distances, are computing use code written in the source code and not computed via astroypy. This makes JAX integration easier and makes the code more lightweight by not needing astropy.

It removes the legacy lensing.py and wrap.py files, consolidates cosmology-related classes into model.py, and updates all references throughout the codebase. The configuration YAML has also been streamlined to use new cosmology class names. Additionally, a calculation in convert.py is improved for clarity and efficiency.

Cosmology import and docstring updates:

• Changed cosmology import in autolens/analysis/analysis/lens.py from autogalaxy.cosmology.wrap to autogalaxy.cosmology.model for Planck15.
• Updated docstrings across multiple files to remove references to "AstroPy" cosmology objects, clarifying that any cosmology object can be used. [1] [2] [3] [4] [5] [6]

Time delay calculation refactor:

• Major rewrite of the time_delays_from function in autolens/lens/tracer_util.py to:
  • Use distances in kiloparsecs (kpc) and SI units throughout.
  • Support array backends via xp for GPU/CPU compatibility.
  • Use new cosmology methods for angular diameter distances.
  • Improve clarity and maintainability of the time delay formula.

API consistency improvements:

• Passed the xp array backend parameter to the angular diameter distance function in the ray-tracing utility.

[Copilot]

Pull request overview

This pull request refactors the cosmology module in the autogalaxy package to support JAX and remove the dependency on astropy. The changes improve the time delay calculation by using explicit unit conversions (kpc) and supporting array backends (CPU/GPU) via the xp parameter.

Changes:

• Updated cosmology imports from autogalaxy.cosmology.wrap to autogalaxy.cosmology.model
• Removed references to "AstroPy" from docstrings across multiple files
• Refactored time delay calculation to use kpc units and explicit SI constants for better clarity and JAX compatibility

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

File	Description
autolens/analysis/analysis/lens.py	Updated import statement to use new cosmology module path
autolens/analysis/analysis/dataset.py	Updated docstring to remove AstroPy reference
autolens/lens/tracer.py	Updated docstrings to remove AstroPy references (2 locations)
autolens/lens/tracer_util.py	Major refactor of time delay calculation with improved clarity, JAX support via xp parameter, and proper unit handling

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.