Skip to content

Commit 73721bd

Browse files
Jammy2211claude
authored andcommitted
docs: add module docstring and improve class docstrings in cosmology package
- Add module-level docstring to model.py explaining the JAX-compatible cosmology classes and the default Planck15 cosmology - Expand LensingCosmology class docstring to list all categories of methods it provides and note that xp=numpy is the default - Add full docstring to Planck15 class with the Planck 2015 parameter values and the reference to Planck Collaboration XIII (2016) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent d99d74b commit 73721bd

1 file changed

Lines changed: 44 additions & 2 deletions

File tree

autogalaxy/cosmology/model.py

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,38 @@
1+
"""
2+
JAX-compatible cosmology models for gravitational lensing calculations.
3+
4+
This module provides `LensingCosmology`, an abstract mixin that adds lensing-specific convenience methods
5+
(arcsec/kpc conversions, critical surface densities, scaling factors) on top of any flat ΛCDM cosmology, and
6+
`FlatLambdaCDM`, a concrete JAX-compatible implementation of a flat Lambda-CDM cosmology.
7+
8+
The default cosmology used throughout **PyAutoGalaxy** and **PyAutoLens** is `Planck15`, a `FlatLambdaCDM`
9+
subclass initialized with Planck 2015 parameter values.
10+
11+
All methods accept an `xp` keyword argument (defaulting to `numpy`) so they can be traced by JAX when
12+
automatic differentiation or GPU acceleration is required.
13+
"""
114
import numpy as np
215

316

417
class LensingCosmology:
518
"""
6-
Class containing specific functions for performing gravitational lensing cosmology calculations.
19+
Mixin class providing gravitational lensing cosmology calculations on top of a flat ΛCDM cosmology.
20+
21+
All methods are JAX-compatible: pass `xp=jax.numpy` to any method to trace it through JAX for automatic
22+
differentiation or GPU/TPU acceleration. The default is `xp=numpy` so there is no JAX dependency at
23+
import time.
24+
25+
The methods cover:
26+
27+
- Unit conversions: arcseconds ↔ kpc at a given redshift.
28+
- Physical distances: luminosity distance, angular diameter distances.
29+
- Lensing-specific quantities: critical surface density Σ_cr, multi-plane scaling factors, velocity
30+
dispersion from Einstein radius.
31+
- Background cosmology: cosmic average density, critical density.
732
8-
This version is JAX-compatible by using an explicit `xp` backend (NumPy or jax.numpy).
33+
This class is inherited by `FlatLambdaCDM` (and therefore `Planck15`). It does not define the background
34+
expansion function `E(z)` or `angular_diameter_distance_kpc_z1z2` — those are implemented by the concrete
35+
subclass.
936
"""
1037

1138
def arcsec_per_kpc_proper(self, z: float, xp=np) -> float:
@@ -694,6 +721,21 @@ def Om(self, z: float, xp=np):
694721

695722

696723
class Planck15(FlatLambdaCDM):
724+
"""
725+
The Planck 2015 flat ΛCDM cosmology.
726+
727+
This is the default cosmology used throughout **PyAutoGalaxy** and **PyAutoLens**. It is initialized with
728+
the best-fit parameters from Planck Collaboration XIII (2016):
729+
730+
- H₀ = 67.74 km/s/Mpc
731+
- Ω_m = 0.3075
732+
- T_CMB = 2.7255 K
733+
- N_eff = 3.046
734+
- m_ν = [0, 0, 0.06] eV (sum = 0.06 eV)
735+
- Ω_b = 0.0486
736+
737+
These values match the `astropy.cosmology.Planck15` cosmology object.
738+
"""
697739

698740
def __init__(self):
699741

0 commit comments

Comments
 (0)