|
| 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 | +""" |
1 | 14 | import numpy as np |
2 | 15 |
|
3 | 16 |
|
4 | 17 | class LensingCosmology: |
5 | 18 | """ |
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. |
7 | 32 |
|
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. |
9 | 36 | """ |
10 | 37 |
|
11 | 38 | def arcsec_per_kpc_proper(self, z: float, xp=np) -> float: |
@@ -694,6 +721,21 @@ def Om(self, z: float, xp=np): |
694 | 721 |
|
695 | 722 |
|
696 | 723 | 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 | + """ |
697 | 739 |
|
698 | 740 | def __init__(self): |
699 | 741 |
|
|
0 commit comments