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 module refactor and cleanup:

• Removed autogalaxy/cosmology/lensing.py and autogalaxy/cosmology/wrap.py, consolidating all cosmology classes (LensingCosmology, Planck15, etc.) into autogalaxy/cosmology/model.py. All imports throughout the codebase now reference the new module and class names. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13]

• Updated the cosmology configuration file (cosmology.yaml) to use new class names (e.g., FlatLambdaCDM instead of legacy wrappers) and removed redundant or obsolete configuration blocks.

General code improvements:

• Simplified the multipole component calculation in convert.py by removing the dependency on astropy.units and using a direct degrees-to-radians conversion, improving clarity and performance.

Minor improvements:

• Deferred the import of astropy.units in fit_ellipse_plot_util.py to within the plotting function, reducing unnecessary imports at the module level.

• Cleaned up the autogalaxy/__init__.py and removed unused cosmology fixture code. [1] [2]

[Copilot]

Pull request overview

This pull request refactors the cosmology module in autogalaxy to be JAX-compatible by removing dependency on astropy and implementing custom cosmological calculations. The changes consolidate cosmology classes from multiple files into a single model.py file, remove legacy wrapper classes, and update all references throughout the codebase.

Changes:

• Removed cosmology/lensing.py and cosmology/wrap.py, consolidating all cosmology functionality into cosmology/model.py with custom JAX-compatible implementations
• Updated imports from autogalaxy.cosmology.lensing and autogalaxy.cosmology.wrap to autogalaxy.cosmology.model throughout the codebase
• Removed pytest skip markers for JAX-incompatible tests in mass profile modules, indicating these are now JAX-compatible
• Simplified convert.py to use direct degree-to-radian conversion instead of astropy.units
• Updated configuration file to reference the new FlatLambdaCDM class name
• Deferred astropy imports to function scope where appropriate

Reviewed changes

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

Show a summary per file

File	Description
autogalaxy/cosmology/model.py	Added complete JAX-compatible cosmology implementation with LensingCosmology, FlatLambdaCDM, and Planck15 classes
autogalaxy/cosmology/lensing.py	Deleted - functionality moved to model.py
autogalaxy/cosmology/wrap.py	Deleted - functionality moved to model.py
autogalaxy/cosmology/init.py	Updated exports to reflect new module structure
autogalaxy/convert.py	Replaced astropy.units with direct numpy conversion for degrees to radians
autogalaxy/profiles/mass/total/power_law_multipole.py	Updated to use direct conversion instead of astropy.units
autogalaxy/profiles/mass/dark/*.py	Updated imports from cosmology.lensing/wrap to cosmology.model, removed astropy.units dependencies
autogalaxy/profiles/mass/point/smbh.py	Updated import from cosmology.wrap to cosmology.model
autogalaxy/profiles/mass/stellar/gaussian.py	Code formatting improvements and removed blank line
autogalaxy/analysis/analysis/*.py	Updated cosmology imports to use new module structure
autogalaxy/imaging/model/analysis.py	Updated cosmology imports
autogalaxy/interferometer/model/analysis.py	Updated cosmology imports
autogalaxy/quantity/model/analysis.py	Updated cosmology imports
autogalaxy/util/mock/mock_cosmology.py	Updated to remove astropy.constants dependency
autogalaxy/ellipse/plot/fit_ellipse_plot_util.py	Moved astropy.units import to function scope
autogalaxy/fixtures.py	Removed unused Planck15 fixture
autogalaxy/config/priors/cosmology.yaml	Simplified to only include FlatLambdaCDM configuration
test_autogalaxy/cosmology/test_model.py	Consolidated and expanded cosmology tests with updated expected values for new implementation
test_autogalaxy/cosmology/test_lensing.py	Deleted - tests moved to test_model.py
test_autogalaxy/conftest.py	Removed Planck15 fixture
test_autogalaxy/profiles/mass/stellar/test_gaussian.py	Removed JAX-incompatibility skip markers
test_autogalaxy/profiles/mass/dark/test_*.py	Updated imports and test expected values, formatting improvements
test_autogalaxy/imaging/test_simulator.py	Moved astropy.io.fits import to function scope

---

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

[Copilot]

Breaking change in default parameter value. The old FlatLambdaCDMWrap had a default value of m_nu=0.06, but the new FlatLambdaCDM has m_nu=0.0. This will change cosmological calculations for any code that instantiates FlatLambdaCDM without explicitly specifying m_nu. Consider either: 1) Restoring the default to m_nu=0.06 for backward compatibility, or 2) Documenting this as a breaking change.

Suggested change

	m_nu: float = 0.0,
	m_nu: float = 0.06,

[Copilot]

The docstring incorrectly states this is "a wrapper for the astropy FlatLambdaCDM cosmology class", but this implementation no longer wraps astropy - it's a standalone JAX-compatible implementation. The docstring should be updated to reflect that this is a custom JAX-compatible implementation of the FlatLambdaCDM cosmology model.

[Jammy2211]

@copilot open a new pull request to apply changes based on this feedback

[Copilot]

@Jammy2211 I've opened a new pull request, #270, to work on those changes. Once the pull request is ready, I'll request review from you.