From 982256141fd3b70cdc74be2133cd5a3aad4740f6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 28 Jan 2026 19:57:47 +0000 Subject: [PATCH 01/26] combining into single module and added method angular_diameter_distance_kpc_z1z2 --- autogalaxy/cosmology/lensing.py | 345 ---------------- autogalaxy/cosmology/model.py | 505 ++++++++++++++++++++---- autogalaxy/cosmology/wrap.py | 1 + test_autogalaxy/cosmology/test_model.py | 19 +- 4 files changed, 439 insertions(+), 431 deletions(-) diff --git a/autogalaxy/cosmology/lensing.py b/autogalaxy/cosmology/lensing.py index 5823c6d9d..d3f5a12fa 100644 --- a/autogalaxy/cosmology/lensing.py +++ b/autogalaxy/cosmology/lensing.py @@ -1,346 +1 @@ -import math -import numpy as np - -class LensingCosmology: - """ - Class containing specific functions for performing gravitational lensing cosmology calculations. - - By inheriting from the astropy `cosmo.FLRW` class this provides many additional methods for performing cosmological - calculations. - """ - - def __init__(self): - - from astropy.cosmology import FLRW - - self._cosmo = FLRW() - - def arcsec_per_kpc_from(self, redshift: float) -> float: - """ - Angular separation in arcsec corresponding to a proper kpc at redshift `z`. - - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - - Parameters - ---------- - redshift - Input redshift from which the angular separation is calculated at. - """ - return self.arcsec_per_kpc_proper(z=redshift).value - - def kpc_per_arcsec_from(self, redshift: float) -> float: - """ - Separation in transverse proper kpc corresponding to an arcminute at redshift `z`. - - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - - Parameters - ---------- - redshift - Input redshift from which the transverse proper kpc value is calculated at. - """ - return 1.0 / self.arcsec_per_kpc_proper(z=redshift).value - - def angular_diameter_distance_to_earth_in_kpc_from(self, redshift: float) -> float: - """ - Angular diameter distance from the input `redshift` to redshift zero (e.g. us, the observer on earth) in - kiloparsecs. - - This gives the proper (sometimes called 'physical') transverse distance corresponding to an angle of 1 radian - for an object at redshift `z`. - - Weinberg, 1972, pp 421-424; Weedman, 1986, pp 65-67; Peebles, 1993, pp 325-327. - - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - - Parameters - ---------- - redshift - Input redshift from which the angular diameter distance to Earth is calculated. - """ - angular_diameter_distance_kpc = self.angular_diameter_distance(z=redshift).to( - "kpc" - ) - - return angular_diameter_distance_kpc.value - - def angular_diameter_distance_between_redshifts_in_kpc_from( - self, redshift_0: float, redshift_1: float - ) -> float: - """ - Angular diameter distance from an input `redshift_0` to another input `redshift_1`. - - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - - Parameters - ---------- - redshift_0 - Redshift from which the angular diameter distance to the other redshift is calculated. - redshift_1 - Redshift from which the angular diameter distance to the other redshift is calculated. - """ - angular_diameter_distance_between_redshifts_kpc = ( - self.angular_diameter_distance_z1z2(redshift_0, redshift_1).to("kpc") - ) - - return angular_diameter_distance_between_redshifts_kpc.value - - def cosmic_average_density_from(self, redshift: float) -> float: - """ - Critical density of the Universe at an input `redshift` in units of solar masses. - - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - - Parameters - ---------- - redshift - Redshift at which the critiical density in solMass of the Universe is calculated. - """ - - cosmic_average_density_kpc = ( - self.critical_density(z=redshift).to("solMass / kpc^3").value - ) - - kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift) - - return cosmic_average_density_kpc * kpc_per_arcsec**3.0 - - def cosmic_average_density_solar_mass_per_kpc3_from(self, redshift: float) -> float: - """ - Critical density of the Universe at an input `redshift` in units of solar masses per kiloparsecs**3. - - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - - Parameters - ---------- - redshift - Redshift at which the critiical density in solMass/kpc^3 of the Universe is calculated. - """ - cosmic_average_density_kpc = ( - self.critical_density(z=redshift).to("solMass / kpc^3").value - ) - - return cosmic_average_density_kpc - - def critical_surface_density_between_redshifts_from( - self, redshift_0: float, redshift_1: float - ) -> float: - """ - The critical surface density for lensing, often written as $\sigma_{cr}$, is given by: - - critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) - - c = speed of light - G = Newton's gravity constant - D_s = angular_diameter_distance_of_source_redshift_to_earth - D_ls = angular_diameter_distance_of_lens_redshift_to_source_redshift - D_l = angular_diameter_distance_of_lens_redshift_to_earth - - This function returns the critical surface density in units of solar masses, which are convenient units for - converting the inferred masses of a model from angular units (e.g. dimensionless units inferred from - data in arcseconds) to solar masses. - - Parameters - ---------- - redshift_0 - The redshift of the first strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. - redshift_1 - The redshift of the second strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. - """ - critical_surface_density_kpc = ( - self.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_0, redshift_1=redshift_1 - ) - ) - - kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0) - - return critical_surface_density_kpc * kpc_per_arcsec**2.0 - - def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - self, redshift_0: float, redshift_1: float - ) -> float: - """ - The critical surface density for lensing, often written as $\sigma_{cr}$, is given by: - - critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) - - c = speed of light - G = Newton's gravity constant - D_s = Angular diameter distance of source redshift to earth - D_ls = Angular diameter distance of lens redshift to source redshift - D_l = Angular diameter distance of lens redshift to earth - - This function returns the critical surface density in units of solar masses / kpc^2, which are convenient - units for converting the inferred masses of a model from angular units (e.g. dimensionless units inferred - from data in arcseconds) to solar masses. - - Parameters - ---------- - redshift_0 - The redshift of the first strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. - redshift_1 - The redshift of the second strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. - """ - from astropy import constants - - const = constants.c.to("kpc / s") ** 2.0 / ( - 4 * math.pi * constants.G.to("kpc3 / (solMass s2)") - ) - - angular_diameter_distance_of_redshift_0_to_earth_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_0) - ) - - angular_diameter_distance_of_redshift_1_to_earth_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) - ) - - angular_diameter_distance_between_redshifts_kpc = ( - self.angular_diameter_distance_between_redshifts_in_kpc_from( - redshift_0=redshift_0, redshift_1=redshift_1 - ) - ) - - return ( - const - * angular_diameter_distance_of_redshift_1_to_earth_kpc - / ( - angular_diameter_distance_between_redshifts_kpc - * angular_diameter_distance_of_redshift_0_to_earth_kpc - ) - ).value - - def scaling_factor_between_redshifts_from( - self, redshift_0: float, redshift_1: float, redshift_final: float - ) -> float: - """ - For strong lens systems with more than 2 planes, the deflection angles between different planes must be scaled - by the angular diameter distances between the planes in order to properly perform multi-plane ray-tracing. This - function computes the factor to scale deflections between `redshift_0` and `reshift_final`, to deflections between - `redshift_0` and `redshift_1`. - - The second redshift should be strictly larger than the first. The scaling factor is unity when `redshift_1` - is `redshift_final`, and 0 when `redshift_0` is equal to `redshift_1`. - - For a system with a first lens galaxy l0 at `redshift_0`, second lens galaxy l1 at `redshift_1` and final - source galaxy at `redshift_final` this scaling factor is given by: - - (D_l0l1 * D_s) / (D_l1* D_l0s) - - The critical surface density for lensing, often written as $\\sigma_{cr}$, is given by: - - critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) - - D_l0l1 = Angular diameter distance of first lens redshift to second lens redshift. - D_s = Angular diameter distance of source redshift to earth - D_l1 = Angular diameter distance of second lens redshift to Earth. - D_l0s = Angular diameter distance of first lens redshift to source redshift - - For systems with more planes this scaling factor is computed multiple times for the different redshift - combinations and applied recursively when scaling the deflection angles. - - Parameters - ---------- - redshift_0 - The redshift of the first strong lens galaxy. - redshift_1 - The redshift of the second strong lens galaxy. - redshift_final - The redshift of the source galaxy. - """ - angular_diameter_distance_between_redshifts_0_and_1 = ( - self.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_1) - .to("kpc") - .value - ) - - angular_diameter_distance_to_redshift_final = ( - self.angular_diameter_distance(z=redshift_final).to("kpc").value - ) - - angular_diameter_distance_of_redshift_1_to_earth = ( - self.angular_diameter_distance(z=redshift_1).to("kpc").value - ) - - angular_diameter_distance_between_redshift_0_and_final = ( - self.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_final) - .to("kpc") - .value - ) - - return ( - angular_diameter_distance_between_redshifts_0_and_1 - * angular_diameter_distance_to_redshift_final - ) / ( - angular_diameter_distance_of_redshift_1_to_earth - * angular_diameter_distance_between_redshift_0_and_final - ) - - def velocity_dispersion_from( - self, redshift_0: float, redshift_1: float, einstein_radius: float - ) -> float: - """ - For a strong lens galaxy with an Einstien radius in arcseconds, the corresponding velocity dispersion of the - lens galaxy can be computed (assuming an isothermal mass distribution). - - The velocity dispersion is given by: - - velocity dispersion = (c * R_Ein * D_s) / (4 * pi * D_l * D_ls) - - c = speed of light - D_s = Angular diameter distance of source redshift to earth - D_ls = Angular diameter distance of lens redshift to source redshift - D_l = Angular diameter distance of lens redshift to earth - - Parameters - ---------- - redshift_0 - The redshift of the first strong lens galaxy (the lens). - redshift_1 - The redshift of the second strong lens galaxy (the source). - """ - from astropy import constants - - const = constants.c.to("kpc / s") - - angular_diameter_distance_to_redshift_0_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) - ) - - angular_diameter_distance_to_redshift_1_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) - ) - - angular_diameter_distance_between_redshifts_kpc = ( - self.angular_diameter_distance_between_redshifts_in_kpc_from( - redshift_0=redshift_0, redshift_1=redshift_1 - ) - ) - - kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0) - - einstein_radius_kpc = einstein_radius * kpc_per_arcsec - - velocity_dispersion_kpc = const * np.sqrt( - (einstein_radius_kpc * angular_diameter_distance_to_redshift_1_kpc) - / ( - 4 - * np.pi - * angular_diameter_distance_to_redshift_0_kpc - * angular_diameter_distance_between_redshifts_kpc - ) - ) - - return velocity_dispersion_kpc.to("km/s").value diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 2e30a529b..74d201e98 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -1,58 +1,346 @@ from astropy import cosmology as cosmo +import numpy as np -from autogalaxy.cosmology.lensing import LensingCosmology +import math +import numpy as np -class LambdaCDMWrap(cosmo.LambdaCDM, LensingCosmology): - def __init__( - self, - H0: float = 67.66, - Om0: float = 0.30966, - Ode0: float = 0.69034, - Tcmb0: float = 2.7255, - Neff: float = 3.046, - m_nu: float = 0.06, - Ob0: float = 0.04897, - ): +class LensingCosmology: + """ + Class containing specific functions for performing gravitational lensing cosmology calculations. + + By inheriting from the astropy `cosmo.FLRW` class this provides many additional methods for performing cosmological + calculations. + """ + def arcsec_per_kpc_from(self, redshift: float) -> float: """ - A wrapper for the astropy `LambdaCDM` cosmology class, which allows it to be used for modeling such - that the cosmological parameters are free parameters which can be fitted for. + Angular separation in arcsec corresponding to a proper kpc at redshift `z`. - The interface of this class is the same as the astropy `LambdaCDM` class, it simply overwrites the - __init__ method and inherits from it in a way that ensures **PyAutoFit** can compose a model from it - without issue. + For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. + This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - The class also inherits from `LensingCosmology`, which is a class that provides additional functionality - for calculating lensing specific quantities in the cosmology. + Parameters + ---------- + redshift + Input redshift from which the angular separation is calculated at. + """ + return self.arcsec_per_kpc_proper(z=redshift).value + + def kpc_per_arcsec_from(self, redshift: float) -> float: + """ + Separation in transverse proper kpc corresponding to an arcminute at redshift `z`. + + For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. + This function therefore returns only the value of the astropy function it wraps, omitting the units instance. Parameters ---------- - H0 - The Hubble constant at z=0. - Om0 - The total matter density at z=0. - Ode0 - The dark energy density at z=0. - Tcmb0 - The CMB temperature at z=0. - Neff - The effective number of neutrinos. - m_nu - The sum of the neutrino masses. - Ob0 - The baryon density at z=0. + redshift + Input redshift from which the transverse proper kpc value is calculated at. """ - super().__init__( - H0=H0, - Om0=Om0, - Ode0=Ode0, - Tcmb0=Tcmb0, - Neff=Neff, - m_nu=m_nu, - Ob0=Ob0, - name="FlatLambdaCDM", + return 1.0 / self.arcsec_per_kpc_proper(z=redshift).value + + def angular_diameter_distance_to_earth_in_kpc_from(self, redshift: float) -> float: + """ + Angular diameter distance from the input `redshift` to redshift zero (e.g. us, the observer on earth) in + kiloparsecs. + + This gives the proper (sometimes called 'physical') transverse distance corresponding to an angle of 1 radian + for an object at redshift `z`. + + Weinberg, 1972, pp 421-424; Weedman, 1986, pp 65-67; Peebles, 1993, pp 325-327. + + For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. + This function therefore returns only the value of the astropy function it wraps, omitting the units instance. + + Parameters + ---------- + redshift + Input redshift from which the angular diameter distance to Earth is calculated. + """ + angular_diameter_distance_kpc = self.angular_diameter_distance(z=redshift).to( + "kpc" + ) + + return angular_diameter_distance_kpc.value + + def angular_diameter_distance_between_redshifts_in_kpc_from( + self, redshift_0: float, redshift_1: float + ) -> float: + """ + Angular diameter distance from an input `redshift_0` to another input `redshift_1`. + + For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. + This function therefore returns only the value of the astropy function it wraps, omitting the units instance. + + Parameters + ---------- + redshift_0 + Redshift from which the angular diameter distance to the other redshift is calculated. + redshift_1 + Redshift from which the angular diameter distance to the other redshift is calculated. + """ + angular_diameter_distance_between_redshifts_kpc = ( + self.angular_diameter_distance_z1z2(redshift_0, redshift_1).to("kpc") ) + return angular_diameter_distance_between_redshifts_kpc.value + + def cosmic_average_density_from(self, redshift: float) -> float: + """ + Critical density of the Universe at an input `redshift` in units of solar masses. + + For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. + This function therefore returns only the value of the astropy function it wraps, omitting the units instance. + + Parameters + ---------- + redshift + Redshift at which the critiical density in solMass of the Universe is calculated. + """ + + cosmic_average_density_kpc = ( + self.critical_density(z=redshift).to("solMass / kpc^3").value + ) + + kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift) + + return cosmic_average_density_kpc * kpc_per_arcsec**3.0 + + def cosmic_average_density_solar_mass_per_kpc3_from(self, redshift: float) -> float: + """ + Critical density of the Universe at an input `redshift` in units of solar masses per kiloparsecs**3. + + For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. + This function therefore returns only the value of the astropy function it wraps, omitting the units instance. + + Parameters + ---------- + redshift + Redshift at which the critiical density in solMass/kpc^3 of the Universe is calculated. + """ + cosmic_average_density_kpc = ( + self.critical_density(z=redshift).to("solMass / kpc^3").value + ) + + return cosmic_average_density_kpc + + def critical_surface_density_between_redshifts_from( + self, redshift_0: float, redshift_1: float + ) -> float: + """ + The critical surface density for lensing, often written as $\sigma_{cr}$, is given by: + + critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) + + c = speed of light + G = Newton's gravity constant + D_s = angular_diameter_distance_of_source_redshift_to_earth + D_ls = angular_diameter_distance_of_lens_redshift_to_source_redshift + D_l = angular_diameter_distance_of_lens_redshift_to_earth + + This function returns the critical surface density in units of solar masses, which are convenient units for + converting the inferred masses of a model from angular units (e.g. dimensionless units inferred from + data in arcseconds) to solar masses. + + Parameters + ---------- + redshift_0 + The redshift of the first strong lens galaxy (E.g. the lens galaxy) for which the critical surface + density is calculated. + redshift_1 + The redshift of the second strong lens galaxy (E.g. the lens galaxy) for which the critical surface + density is calculated. + """ + critical_surface_density_kpc = ( + self.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( + redshift_0=redshift_0, redshift_1=redshift_1 + ) + ) + + kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0) + + return critical_surface_density_kpc * kpc_per_arcsec**2.0 + + def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( + self, redshift_0: float, redshift_1: float + ) -> float: + """ + The critical surface density for lensing, often written as $\sigma_{cr}$, is given by: + + critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) + + c = speed of light + G = Newton's gravity constant + D_s = Angular diameter distance of source redshift to earth + D_ls = Angular diameter distance of lens redshift to source redshift + D_l = Angular diameter distance of lens redshift to earth + + This function returns the critical surface density in units of solar masses / kpc^2, which are convenient + units for converting the inferred masses of a model from angular units (e.g. dimensionless units inferred + from data in arcseconds) to solar masses. + + Parameters + ---------- + redshift_0 + The redshift of the first strong lens galaxy (E.g. the lens galaxy) for which the critical surface + density is calculated. + redshift_1 + The redshift of the second strong lens galaxy (E.g. the lens galaxy) for which the critical surface + density is calculated. + """ + from astropy import constants + + const = constants.c.to("kpc / s") ** 2.0 / ( + 4 * math.pi * constants.G.to("kpc3 / (solMass s2)") + ) + + angular_diameter_distance_of_redshift_0_to_earth_kpc = ( + self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_0) + ) + + angular_diameter_distance_of_redshift_1_to_earth_kpc = ( + self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) + ) + + angular_diameter_distance_between_redshifts_kpc = ( + self.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=redshift_0, redshift_1=redshift_1 + ) + ) + + return ( + const + * angular_diameter_distance_of_redshift_1_to_earth_kpc + / ( + angular_diameter_distance_between_redshifts_kpc + * angular_diameter_distance_of_redshift_0_to_earth_kpc + ) + ).value + + def scaling_factor_between_redshifts_from( + self, redshift_0: float, redshift_1: float, redshift_final: float + ) -> float: + """ + For strong lens systems with more than 2 planes, the deflection angles between different planes must be scaled + by the angular diameter distances between the planes in order to properly perform multi-plane ray-tracing. This + function computes the factor to scale deflections between `redshift_0` and `reshift_final`, to deflections between + `redshift_0` and `redshift_1`. + + The second redshift should be strictly larger than the first. The scaling factor is unity when `redshift_1` + is `redshift_final`, and 0 when `redshift_0` is equal to `redshift_1`. + + For a system with a first lens galaxy l0 at `redshift_0`, second lens galaxy l1 at `redshift_1` and final + source galaxy at `redshift_final` this scaling factor is given by: + + (D_l0l1 * D_s) / (D_l1* D_l0s) + + The critical surface density for lensing, often written as $\\sigma_{cr}$, is given by: + + critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) + + D_l0l1 = Angular diameter distance of first lens redshift to second lens redshift. + D_s = Angular diameter distance of source redshift to earth + D_l1 = Angular diameter distance of second lens redshift to Earth. + D_l0s = Angular diameter distance of first lens redshift to source redshift + + For systems with more planes this scaling factor is computed multiple times for the different redshift + combinations and applied recursively when scaling the deflection angles. + + Parameters + ---------- + redshift_0 + The redshift of the first strong lens galaxy. + redshift_1 + The redshift of the second strong lens galaxy. + redshift_final + The redshift of the source galaxy. + """ + angular_diameter_distance_between_redshifts_0_and_1 = ( + self.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_1) + .to("kpc") + .value + ) + + angular_diameter_distance_to_redshift_final = ( + self.angular_diameter_distance(z=redshift_final).to("kpc").value + ) + + angular_diameter_distance_of_redshift_1_to_earth = ( + self.angular_diameter_distance(z=redshift_1).to("kpc").value + ) + + angular_diameter_distance_between_redshift_0_and_final = ( + self.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_final) + .to("kpc") + .value + ) + + return ( + angular_diameter_distance_between_redshifts_0_and_1 + * angular_diameter_distance_to_redshift_final + ) / ( + angular_diameter_distance_of_redshift_1_to_earth + * angular_diameter_distance_between_redshift_0_and_final + ) + + def velocity_dispersion_from( + self, redshift_0: float, redshift_1: float, einstein_radius: float + ) -> float: + """ + For a strong lens galaxy with an Einstien radius in arcseconds, the corresponding velocity dispersion of the + lens galaxy can be computed (assuming an isothermal mass distribution). + + The velocity dispersion is given by: + + velocity dispersion = (c * R_Ein * D_s) / (4 * pi * D_l * D_ls) + + c = speed of light + D_s = Angular diameter distance of source redshift to earth + D_ls = Angular diameter distance of lens redshift to source redshift + D_l = Angular diameter distance of lens redshift to earth + + Parameters + ---------- + redshift_0 + The redshift of the first strong lens galaxy (the lens). + redshift_1 + The redshift of the second strong lens galaxy (the source). + """ + from astropy import constants + + const = constants.c.to("kpc / s") + + angular_diameter_distance_to_redshift_0_kpc = ( + self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) + ) + + angular_diameter_distance_to_redshift_1_kpc = ( + self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) + ) + + angular_diameter_distance_between_redshifts_kpc = ( + self.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=redshift_0, redshift_1=redshift_1 + ) + ) + + kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0) + + einstein_radius_kpc = einstein_radius * kpc_per_arcsec + + velocity_dispersion_kpc = const * np.sqrt( + (einstein_radius_kpc * angular_diameter_distance_to_redshift_1_kpc) + / ( + 4 + * np.pi + * angular_diameter_distance_to_redshift_0_kpc + * angular_diameter_distance_between_redshifts_kpc + ) + ) + + return velocity_dispersion_kpc.to("km/s").value + class FlatLambdaCDMWrap(cosmo.FlatLambdaCDM, LensingCosmology): def __init__( @@ -100,53 +388,100 @@ def __init__( name="FlatLambdaCDM", ) + @staticmethod + def _simpson_1d(y, x, xp=np): + """ + Composite Simpson's rule on a 1D grid. -class FlatwCDMWrap(cosmo.FlatwCDM, LensingCosmology): - def __init__( + Requirements: + - x is 1D, evenly spaced + - len(x) is odd (i.e. number of intervals is even) + + Works with numpy or jax.numpy passed as xp. + """ + n = x.shape[0] + if (n % 2) == 0: + raise ValueError("Simpson's rule requires an odd number of samples (even number of intervals).") + + h = (x[-1] - x[0]) / (n - 1) + + # y0 + yN + s = y[0] + y[-1] + + # 4 * sum of odd indices + s = s + 4.0 * xp.sum(y[1:-1:2]) + + # 2 * sum of even indices (excluding endpoints) + s = s + 2.0 * xp.sum(y[2:-1:2]) + + return (h / 3.0) * s + + def angular_diameter_distance_kpc_z1z2( self, - H0: float = 67.66, - Om0: float = 0.30966, - w0: float = -1.0, - Tcmb0: float = 2.7255, - Neff: float = 3.046, - m_nu: float = 0.06, - Ob0: float = 0.04897, + z1: float, + z2: float, + n_steps: int = 8193, # odd by default + xp=np, ): """ - A wrapper for the astropy `FlatwCDM` cosmology class, which allows it to be used for modeling such - that the cosmological parameters are free parameters which can be fitted for. + D_A(z1,z2) in kpc for flat wCDM using Simpson's rule. + Includes radiation (photons + massless neutrinos) to better match astropy when Tcmb0 is set. - The interface of this class is the same as the astropy `FlatwCDM` class, it simply overwrites the - __init__ method and inherits from it in a way that ensures **PyAutoFit** can compose a model from it - without issue. + Assumes: + - Flat universe: Omega_k = 0 + - Dark energy equation of state constant w0 + - Neutrinos treated as radiation via Neff (massless approximation) + - H0 in km/s/Mpc + """ + # Ensure odd number of samples for Simpson (safe: n_steps is a Python int) + if (n_steps % 2) == 0: + n_steps += 1 + + z1a = xp.asarray(z1) + z2a = xp.asarray(z2) + same = z1a == z2a + + c_km_s = xp.asarray(299792.458) + + H0 = xp.asarray(self.H0) + h = H0 / xp.asarray(100.0) + + Om0 = xp.asarray(self.Om0) + w0 = xp.asarray(self.w0) + + # ---- Radiation density today (photons + massless neutrinos) ---- + # Omega_gamma * h^2 ≈ 2.469e-5 * (Tcmb/2.7255)^4 + Tcmb = xp.asarray(self.Tcmb0) + Ogamma_h2 = xp.asarray(2.469e-5) * (Tcmb / xp.asarray(2.7255)) ** 4 + Ogamma0 = Ogamma_h2 / (h * h) + + Neff = xp.asarray(self.Neff) + # Omega_nu (massless) = Omega_gamma * 0.2271 * Neff + Onu0 = Ogamma0 * xp.asarray(0.2271) * Neff + + Or0 = Ogamma0 + Onu0 + + # Flatness: Omega_de = 1 - Omega_m - Omega_r + Ode0 = xp.asarray(1.0) - Om0 - Or0 + + def E(z): + zp1 = xp.asarray(1.0) + z + # E^2 = Om(1+z)^3 + Or(1+z)^4 + Ode(1+z)^{3(1+w)} + return xp.sqrt( + Om0 * zp1**3 + + Or0 * zp1**4 + + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) + ) + + z_grid = xp.linspace(z1a, z2a, n_steps) + integrand = 1.0 / E(z_grid) + + integral = self._simpson_1d(integrand, z_grid, xp=xp) + + Dc_Mpc = (c_km_s / H0) * integral + Da_Mpc = Dc_Mpc / (xp.asarray(1.0) + z2a) + Da_kpc = Da_Mpc * xp.asarray(1.0e3) + + return xp.where(same, xp.asarray(0.0), Da_kpc) - The class also inherits from `LensingCosmology`, which is a class that provides additional functionality - for calculating lensing specific quantities in the cosmology. - Parameters - ---------- - H0 - The Hubble constant at z=0. - Om0 - The total matter density at z=0. - w0 - The dark energy equation of state at z=0. - Tcmb0 - The CMB temperature at z=0. - Neff - The effective number of neutrinos. - m_nu - The sum of the neutrino masses. - Ob0 - The baryon density at z=0. - """ - super().__init__( - H0=H0, - Om0=Om0, - w0=w0, - Tcmb0=Tcmb0, - Neff=Neff, - m_nu=m_nu, - Ob0=Ob0, - name="FlatwCDM", - ) diff --git a/autogalaxy/cosmology/wrap.py b/autogalaxy/cosmology/wrap.py index 9adbf6bd9..ac8fb9f22 100644 --- a/autogalaxy/cosmology/wrap.py +++ b/autogalaxy/cosmology/wrap.py @@ -3,6 +3,7 @@ class Planck15(cosmo.FlatLambdaCDM, LensingCosmology): + def __init__(self): Planck15_astropy = cosmo.Planck15 diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index 8154a82da..4ecfe10b3 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -1,7 +1,24 @@ import pytest +from astropy import cosmology as cosmo -def test__cosmology(Planck15): +import autogalaxy as ag + +def test__angular_diameter_distance(): + + cosmology_ap = cosmo.FlatwCDM(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + + angular_diameter_distance_ap = cosmology_ap.angular_diameter_distance_z1z2(0.1, 1.0).to("kpc").value + + cosmology = ag.cosmo.FlatwCDMWrap(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + + angular_diameter_distance = ( + cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) + ) + + assert angular_diameter_distance == pytest.approx(angular_diameter_distance_ap, 1.0e-4) + +def test__critical_surface_density(): from autogalaxy.cosmology.model import FlatwCDMWrap From 638c0b5646394fcdb58731b8a12113c0ffba480a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 28 Jan 2026 20:22:25 +0000 Subject: [PATCH 02/26] arcsec_tokpc functions added with jax support --- autogalaxy/analysis/analysis/analysis.py | 2 +- autogalaxy/analysis/analysis/dataset.py | 2 +- autogalaxy/config/priors/cosmology.yaml | 46 +------- autogalaxy/cosmology/__init__.py | 4 +- autogalaxy/cosmology/lensing.py | 1 - autogalaxy/cosmology/model.py | 95 ++++++++++------ autogalaxy/cosmology/wrap.py | 18 ---- autogalaxy/fixtures.py | 7 -- autogalaxy/imaging/model/analysis.py | 2 +- autogalaxy/interferometer/model/analysis.py | 2 +- autogalaxy/profiles/mass/dark/abstract.py | 2 +- .../profiles/mass/dark/nfw_truncated.py | 2 +- autogalaxy/quantity/model/analysis.py | 2 +- test_autogalaxy/conftest.py | 5 - test_autogalaxy/cosmology/test_lensing.py | 80 -------------- test_autogalaxy/cosmology/test_model.py | 101 ++++++++++++++++-- .../profiles/mass/dark/test_nfw_mcr.py | 20 ++-- .../mass/dark/test_nfw_truncated_mcr.py | 8 +- 18 files changed, 178 insertions(+), 221 deletions(-) delete mode 100644 autogalaxy/cosmology/lensing.py delete mode 100644 autogalaxy/cosmology/wrap.py delete mode 100644 test_autogalaxy/cosmology/test_lensing.py diff --git a/autogalaxy/analysis/analysis/analysis.py b/autogalaxy/analysis/analysis/analysis.py index 438ac013c..a79a2ba8f 100644 --- a/autogalaxy/analysis/analysis/analysis.py +++ b/autogalaxy/analysis/analysis/analysis.py @@ -7,7 +7,7 @@ from autogalaxy.galaxy.galaxy import Galaxy from autogalaxy.galaxy.galaxies import Galaxies -from autogalaxy.cosmology.lensing import LensingCosmology +from autogalaxy.cosmology.model import LensingCosmology logger = logging.getLogger(__name__) diff --git a/autogalaxy/analysis/analysis/dataset.py b/autogalaxy/analysis/analysis/dataset.py index 54e7b48bd..f2ad68159 100644 --- a/autogalaxy/analysis/analysis/dataset.py +++ b/autogalaxy/analysis/analysis/dataset.py @@ -8,7 +8,7 @@ import autoarray as aa from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages -from autogalaxy.cosmology.lensing import LensingCosmology +from autogalaxy.cosmology.model import LensingCosmology from autogalaxy.analysis.analysis.analysis import Analysis from autogalaxy.analysis.result import ResultDataset diff --git a/autogalaxy/config/priors/cosmology.yaml b/autogalaxy/config/priors/cosmology.yaml index 656941bba..e9feba055 100644 --- a/autogalaxy/config/priors/cosmology.yaml +++ b/autogalaxy/config/priors/cosmology.yaml @@ -1,54 +1,10 @@ -model.LambdaCDMWrap: +model.FlatLambdaCDM: H0: type: Constant value: 67.66 Om0: type: Constant value: 0.30966 - Ode0: - type: Constant - value: 0.69034 - Tcmb0: - type: Constant - value: 2.7255 - Neff: - type: Constant - value: 3.046 - m_nu: - type: Constant - value: 0.06 - Ob0: - type: Constant - value: 0.04897 -model.FlatLambdaCDMWrap: - H0: - type: Constant - value: 67.66 - Om0: - type: Constant - value: 0.30966 - Tcmb0: - type: Constant - value: 2.7255 - Neff: - type: Constant - value: 3.046 - m_nu: - type: Constant - value: 0.06 - Ob0: - type: Constant - value: 0.04897 -model.FlatwCDMWrap: - H0: - type: Constant - value: 67.66 - Om0: - type: Constant - value: 0.30966 - w0: - type: Constant - value: -1.0 Tcmb0: type: Constant value: 2.7255 diff --git a/autogalaxy/cosmology/__init__.py b/autogalaxy/cosmology/__init__.py index 26f887402..4c4090146 100644 --- a/autogalaxy/cosmology/__init__.py +++ b/autogalaxy/cosmology/__init__.py @@ -1,3 +1 @@ -from .lensing import LensingCosmology -from .wrap import Planck15 -from .model import FlatwCDMWrap, FlatLambdaCDMWrap +from .model import FlatLambdaCDM, Planck15, LensingCosmology diff --git a/autogalaxy/cosmology/lensing.py b/autogalaxy/cosmology/lensing.py deleted file mode 100644 index d3f5a12fa..000000000 --- a/autogalaxy/cosmology/lensing.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 74d201e98..f55fff9ba 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -1,44 +1,67 @@ -from astropy import cosmology as cosmo import numpy as np - import math -import numpy as np class LensingCosmology: """ Class containing specific functions for performing gravitational lensing cosmology calculations. - By inheriting from the astropy `cosmo.FLRW` class this provides many additional methods for performing cosmological - calculations. + This version is JAX-compatible by using an explicit `xp` backend (NumPy or jax.numpy). """ - def arcsec_per_kpc_from(self, redshift: float) -> float: + + def arcsec_per_kpc_proper(self, z: float, xp=np) -> float: + """ + Angular separation in arcsec corresponding to 1 proper kpc at redshift z. + + This matches astropy.cosmology.arcsec_per_kpc_proper. + + Proper transverse distance uses the angular diameter distance D_A(z): + + arcsec_per_kpc_proper = 206265 / D_A(z) + + where D_A(z) is in kpc. + """ + + angular_diameter_distance_kpc = self.angular_diameter_distance_kpc_z1z2( + 0.0, z, xp=xp + ) + + return xp.asarray(206265.0) / angular_diameter_distance_kpc + + def kpc_proper_per_arcsec(self, z: float, xp=np) -> float: + """ + Proper transverse separation in kpc corresponding to 1 arcsec at redshift z. + + This matches the inverse of astropy.cosmology.arcsec_per_kpc_proper: + + kpc_proper_per_arcsec = D_A(z) / 206265 + """ + + angular_diameter_distance_kpc = self.angular_diameter_distance_kpc_z1z2( + 0.0, z, xp=xp + ) + + return angular_diameter_distance_kpc / xp.asarray(206265.0) + + def arcsec_per_kpc_from(self, redshift: float, xp=np) -> float: """ Angular separation in arcsec corresponding to a proper kpc at redshift `z`. For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - Parameters - ---------- - redshift - Input redshift from which the angular separation is calculated at. + This is a thin convenience wrapper around `arcsec_per_kpc_proper`. """ - return self.arcsec_per_kpc_proper(z=redshift).value + return self.arcsec_per_kpc_proper(z=redshift, xp=xp) - def kpc_per_arcsec_from(self, redshift: float) -> float: + def kpc_per_arcsec_from(self, redshift: float, xp=np) -> float: """ - Separation in transverse proper kpc corresponding to an arcminute at redshift `z`. + Separation in transverse proper kpc corresponding to an arcsec at redshift `z`. For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. - Parameters - ---------- - redshift - Input redshift from which the transverse proper kpc value is calculated at. + This is a thin convenience wrapper around `kpc_proper_per_arcsec`. """ - return 1.0 / self.arcsec_per_kpc_proper(z=redshift).value + return self.kpc_proper_per_arcsec(z=redshift, xp=xp) def angular_diameter_distance_to_earth_in_kpc_from(self, redshift: float) -> float: """ @@ -342,7 +365,7 @@ def velocity_dispersion_from( return velocity_dispersion_kpc.to("km/s").value -class FlatLambdaCDMWrap(cosmo.FlatLambdaCDM, LensingCosmology): +class FlatLambdaCDM(LensingCosmology): def __init__( self, H0: float = 67.66, @@ -378,15 +401,15 @@ def __init__( Ob0 The baryon density at z=0. """ - super().__init__( - H0=H0, - Om0=Om0, - Tcmb0=Tcmb0, - Neff=Neff, - m_nu=m_nu, - Ob0=Ob0, - name="FlatLambdaCDM", - ) + self.H0 = H0 + self.Om0 = Om0 + self.Tcmb0 = Tcmb0 + self.Neff = Neff + self.m_nu = m_nu + self.Ob0 = Ob0 + + # Make ΛCDM a special case of wCDM + self.w0 = -1.0 @staticmethod def _simpson_1d(y, x, xp=np): @@ -485,3 +508,15 @@ def E(z): return xp.where(same, xp.asarray(0.0), Da_kpc) +class Planck15(FlatLambdaCDM): + + def __init__(self): + + super().__init__( + H0=67.74, + Om0=0.3075, + Tcmb0=2.7255, + Neff=3.046, + m_nu=0.06, + Ob0=0.0486, + ) diff --git a/autogalaxy/cosmology/wrap.py b/autogalaxy/cosmology/wrap.py deleted file mode 100644 index ac8fb9f22..000000000 --- a/autogalaxy/cosmology/wrap.py +++ /dev/null @@ -1,18 +0,0 @@ -from astropy import cosmology as cosmo -from autogalaxy.cosmology.lensing import LensingCosmology - - -class Planck15(cosmo.FlatLambdaCDM, LensingCosmology): - - def __init__(self): - Planck15_astropy = cosmo.Planck15 - - super().__init__( - H0=Planck15_astropy.H0, - Om0=Planck15_astropy.Om0, - Tcmb0=Planck15_astropy.Tcmb0, - Neff=Planck15_astropy.Neff, - m_nu=Planck15_astropy.m_nu, - Ob0=Planck15_astropy.Ob0, - name=Planck15_astropy.name, - ) diff --git a/autogalaxy/fixtures.py b/autogalaxy/fixtures.py index 569c8a805..510758a74 100644 --- a/autogalaxy/fixtures.py +++ b/autogalaxy/fixtures.py @@ -140,13 +140,6 @@ def make_galaxies_x2_inversion_7x7(): return [make_gal_x1_lp(), source_gal_inversion] -# COSMOLOGY # - - -def make_Planck15(): - return ag.cosmo.Planck15() - - # ELLIPSE FITING diff --git a/autogalaxy/imaging/model/analysis.py b/autogalaxy/imaging/model/analysis.py index 83f8b641f..2d1e6fabd 100644 --- a/autogalaxy/imaging/model/analysis.py +++ b/autogalaxy/imaging/model/analysis.py @@ -6,7 +6,7 @@ from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages from autogalaxy.analysis.analysis.dataset import AnalysisDataset -from autogalaxy.cosmology.lensing import LensingCosmology +from autogalaxy.cosmology.model import LensingCosmology from autogalaxy.imaging.model.result import ResultImaging from autogalaxy.imaging.model.visualizer import VisualizerImaging from autogalaxy.imaging.fit_imaging import FitImaging diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py index ac4a56c41..010ffae83 100644 --- a/autogalaxy/interferometer/model/analysis.py +++ b/autogalaxy/interferometer/model/analysis.py @@ -9,7 +9,7 @@ from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages from autogalaxy.analysis.analysis.dataset import AnalysisDataset -from autogalaxy.cosmology.lensing import LensingCosmology +from autogalaxy.cosmology.model import LensingCosmology from autogalaxy.interferometer.model.result import ResultInterferometer from autogalaxy.interferometer.fit_interferometer import FitInterferometer from autogalaxy.interferometer.model.visualizer import VisualizerInterferometer diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py index cc85de723..bc616eb86 100644 --- a/autogalaxy/profiles/mass/dark/abstract.py +++ b/autogalaxy/profiles/mass/dark/abstract.py @@ -4,7 +4,7 @@ import autoarray as aa from autogalaxy.profiles.mass.abstract.abstract import MassProfile -from autogalaxy.cosmology.lensing import LensingCosmology +from autogalaxy.cosmology.model import LensingCosmology from autogalaxy.profiles.mass.abstract.mge_numpy import ( MassProfileMGE, ) diff --git a/autogalaxy/profiles/mass/dark/nfw_truncated.py b/autogalaxy/profiles/mass/dark/nfw_truncated.py index 7a1ff50fb..154bda1ec 100644 --- a/autogalaxy/profiles/mass/dark/nfw_truncated.py +++ b/autogalaxy/profiles/mass/dark/nfw_truncated.py @@ -3,7 +3,7 @@ import autoarray as aa -from autogalaxy.cosmology.lensing import LensingCosmology +from autogalaxy.cosmology.model import LensingCosmology from autogalaxy.profiles.mass.dark.abstract import AbstractgNFW diff --git a/autogalaxy/quantity/model/analysis.py b/autogalaxy/quantity/model/analysis.py index 3537710f7..1889199b2 100644 --- a/autogalaxy/quantity/model/analysis.py +++ b/autogalaxy/quantity/model/analysis.py @@ -5,7 +5,7 @@ import autofit as af from autogalaxy.analysis.analysis.analysis import Analysis -from autogalaxy.cosmology.lensing import LensingCosmology +from autogalaxy.cosmology.model import LensingCosmology from autogalaxy.quantity.dataset_quantity import DatasetQuantity from autogalaxy.quantity.model.result import ResultQuantity from autogalaxy.quantity.model.visualizer import VisualizerQuantity diff --git a/test_autogalaxy/conftest.py b/test_autogalaxy/conftest.py index a017b64e7..9c8027d9c 100644 --- a/test_autogalaxy/conftest.py +++ b/test_autogalaxy/conftest.py @@ -302,11 +302,6 @@ def make_galaxies_x2_7x7(): return fixtures.make_galaxies_x2_7x7() -@pytest.fixture(name="Planck15") -def make_Planck15(): - return fixtures.make_Planck15() - - @pytest.fixture(name="adapt_galaxy_name_image_dict_7x7") def make_adapt_galaxy_name_image_dict_7x7(): return fixtures.make_adapt_galaxy_name_image_dict_7x7() diff --git a/test_autogalaxy/cosmology/test_lensing.py b/test_autogalaxy/cosmology/test_lensing.py deleted file mode 100644 index 057204932..000000000 --- a/test_autogalaxy/cosmology/test_lensing.py +++ /dev/null @@ -1,80 +0,0 @@ -import pytest -import numpy as np - - -def test__arcsec_to_kpc_conversion(Planck15): - arcsec_per_kpc = Planck15.arcsec_per_kpc_from(redshift=0.1) - - assert arcsec_per_kpc == pytest.approx(0.525060, 1e-5) - - kpc_per_arcsec = Planck15.kpc_per_arcsec_from(redshift=0.1) - - assert kpc_per_arcsec == pytest.approx(1.904544, 1e-5) - - arcsec_per_kpc = Planck15.arcsec_per_kpc_from(redshift=1.0) - - assert arcsec_per_kpc == pytest.approx(0.1214785, 1e-5) - - kpc_per_arcsec = Planck15.kpc_per_arcsec_from(redshift=1.0) - - assert kpc_per_arcsec == pytest.approx(8.231907, 1e-5) - - -def test__angular_diameter_distances(Planck15): - angular_diameter_distance_to_earth_kpc = ( - Planck15.angular_diameter_distance_to_earth_in_kpc_from(redshift=0.1) - ) - - assert angular_diameter_distance_to_earth_kpc == pytest.approx(392840, 1e-5) - - angular_diameter_distance_between_redshifts_kpc = ( - Planck15.angular_diameter_distance_between_redshifts_in_kpc_from( - redshift_0=0.1, redshift_1=1.0 - ) - ) - - assert angular_diameter_distance_between_redshifts_kpc == pytest.approx( - 1481890.4, 1e-5 - ) - - -def test__cosmic_average_densities_solar_mass_per_kpc3(Planck15): - cosmic_average_density = Planck15.cosmic_average_density_from(redshift=0.6) - - assert cosmic_average_density == pytest.approx(81280.09116133313, 1.0e-4) - - cosmic_average_density = Planck15.cosmic_average_density_solar_mass_per_kpc3_from( - redshift=0.6 - ) - - assert cosmic_average_density == pytest.approx(249.20874, 1.0e-4) - - -def test__critical_surface_mass_densities(Planck15): - critical_surface_density = Planck15.critical_surface_density_between_redshifts_from( - redshift_0=0.1, redshift_1=1.0 - ) - - assert critical_surface_density == pytest.approx(17593241668, 1e-2) - - critical_surface_density = ( - Planck15.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=0.1, redshift_1=1.0 - ) - ) - - assert critical_surface_density == pytest.approx(4.85e9, 1e-2) - - -def test__velocity_dispersion_from(Planck15): - velocity_dispersion = Planck15.velocity_dispersion_from( - redshift_0=0.5, redshift_1=1.0, einstein_radius=1.0 - ) - - assert velocity_dispersion == pytest.approx(249.03449, 1.0e-4) - - velocity_dispersion = Planck15.velocity_dispersion_from( - redshift_0=0.5, redshift_1=1.0, einstein_radius=2.0 - ) - - assert velocity_dispersion == pytest.approx(np.sqrt(2) * 249.03449, 1.0e-4) diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index 4ecfe10b3..97ddbe315 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -1,16 +1,58 @@ import pytest - -from astropy import cosmology as cosmo +import numpy as np import autogalaxy as ag +def test__arcsec_to_kpc_conversion(): + + cosmology = ag.cosmology.Planck15() + + arcsec_per_kpc = cosmology.arcsec_per_kpc_from(redshift=0.1) + + assert arcsec_per_kpc == pytest.approx(0.525060, 1e-5) + + kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=0.1) + + assert kpc_per_arcsec == pytest.approx(1.904544, 1e-5) + + arcsec_per_kpc = cosmology.arcsec_per_kpc_from(redshift=1.0) + + assert arcsec_per_kpc == pytest.approx(0.1214785, 1e-5) + + kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=1.0) + + assert kpc_per_arcsec == pytest.approx(8.231907, 1e-5) + + +def test__angular_diameter_distances(): + + cosmology = ag.cosmology.Planck15() + + angular_diameter_distance_to_earth_kpc = ( + cosmology.angular_diameter_distance_to_earth_in_kpc_from(redshift=0.1) + ) + + assert angular_diameter_distance_to_earth_kpc == pytest.approx(392840, 1e-5) + + angular_diameter_distance_between_redshifts_kpc = ( + cosmology.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=0.1, redshift_1=1.0 + ) + ) + + assert angular_diameter_distance_between_redshifts_kpc == pytest.approx( + 1481890.4, 1e-5 + ) + def test__angular_diameter_distance(): - cosmology_ap = cosmo.FlatwCDM(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + from astropy import cosmology as cosmo_ap + + cosmology_ap = cosmo_ap.FlatwCDM(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) angular_diameter_distance_ap = cosmology_ap.angular_diameter_distance_z1z2(0.1, 1.0).to("kpc").value - cosmology = ag.cosmo.FlatwCDMWrap(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + cosmology = ag.cosmology.FlatwCDMWrap(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) angular_diameter_distance = ( cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) @@ -18,23 +60,60 @@ def test__angular_diameter_distance(): assert angular_diameter_distance == pytest.approx(angular_diameter_distance_ap, 1.0e-4) -def test__critical_surface_density(): +def test__cosmic_average_densities_solar_mass_per_kpc3(): + + cosmology = ag.cosmology.Planck15() + + cosmic_average_density = cosmology.cosmic_average_density_from(redshift=0.6) + + assert cosmic_average_density == pytest.approx(81280.09116133313, 1.0e-4) + + cosmic_average_density = cosmology.cosmic_average_density_solar_mass_per_kpc3_from( + redshift=0.6 + ) - from autogalaxy.cosmology.model import FlatwCDMWrap + assert cosmic_average_density == pytest.approx(249.20874, 1.0e-4) - cosmology = FlatwCDMWrap() + +def test__critical_surface_mass_densities(): + + cosmology = ag.cosmology.Planck15() + + critical_surface_density = cosmology.critical_surface_density_between_redshifts_from( + redshift_0=0.1, redshift_1=1.0 + ) + + assert critical_surface_density == pytest.approx(17593241668, 1e-2) critical_surface_density = ( - cosmology.critical_surface_density_between_redshifts_from( + cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( redshift_0=0.1, redshift_1=1.0 ) ) - assert critical_surface_density == pytest.approx(17613991217.945473, 1.0e-4) + assert critical_surface_density == pytest.approx(4.85e9, 1e-2) - from autogalaxy.cosmology.model import FlatLambdaCDMWrap - cosmology = FlatLambdaCDMWrap() +def test__velocity_dispersion_from(): + + cosmology = ag.cosmology.Planck15() + + velocity_dispersion = cosmology.velocity_dispersion_from( + redshift_0=0.5, redshift_1=1.0, einstein_radius=1.0 + ) + + assert velocity_dispersion == pytest.approx(249.03449, 1.0e-4) + + velocity_dispersion = cosmology.velocity_dispersion_from( + redshift_0=0.5, redshift_1=1.0, einstein_radius=2.0 + ) + + assert velocity_dispersion == pytest.approx(np.sqrt(2) * 249.03449, 1.0e-4) + + +def test__critical_surface_density(): + + cosmology = ag.cosmology.FlatLambdaCDM() critical_surface_density = ( cosmology.critical_surface_density_between_redshifts_from( diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py index 47964eb69..c26204ced 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py +++ b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py @@ -7,9 +7,9 @@ def test__mass_and_concentration_consistent_with_normal_nfw(): - from autogalaxy.cosmology.model import FlatLambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = FlatLambdaCDMWrap(H0=70.0, Om0=0.3) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.NFWMCRDuffySph( centre=(1.0, 2.0), @@ -56,9 +56,9 @@ def test__mass_and_concentration_consistent_with_normal_nfw(): def test__mass_and_concentration_consistent_with_normal_nfw__scatter_0(): - from autogalaxy.cosmology.model import FlatLambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = FlatLambdaCDMWrap(H0=70.0, Om0=0.3) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.NFWMCRLudlowSph( centre=(1.0, 2.0), @@ -106,9 +106,9 @@ def test__mass_and_concentration_consistent_with_normal_nfw__scatter_0(): def test__same_as_above_but_elliptical(): - from autogalaxy.cosmology.model import FlatLambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = FlatLambdaCDMWrap(H0=70.0, Om0=0.3) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.NFWMCRLudlow( centre=(1.0, 2.0), @@ -162,9 +162,9 @@ def test__same_as_above_but_elliptical(): def test__same_as_above_but_generalized_elliptical(): - from autogalaxy.cosmology.model import FlatLambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = FlatLambdaCDMWrap(H0=70.0, Om0=0.3) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.gNFWMCRLudlow( centre=(1.0, 2.0), @@ -219,9 +219,9 @@ def test__same_as_above_but_generalized_elliptical(): def test__same_as_above_but_cored_nfw(): - from autogalaxy.cosmology.model import FlatLambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = FlatLambdaCDMWrap(H0=70.0, Om0=0.3) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.cNFWMCRLudlowSph( centre=(1.0, 2.0), diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py b/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py index 99188189b..7e22856af 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py +++ b/test_autogalaxy/profiles/mass/dark/test_nfw_truncated_mcr.py @@ -8,9 +8,9 @@ def test__duffy__mass_and_concentration_consistent_with_normal_truncated_nfw(): - from autogalaxy.cosmology.model import FlatLambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = FlatLambdaCDMWrap(H0=70.0, Om0=0.3) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.NFWTruncatedMCRDuffySph( centre=(1.0, 2.0), @@ -55,9 +55,9 @@ def test__duffy__mass_and_concentration_consistent_with_normal_truncated_nfw(): def test__ludlow__mass_and_concentration_consistent_with_normal_truncated_nfw__scatter_0(): - from autogalaxy.cosmology.model import FlatLambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = FlatLambdaCDMWrap(H0=70.0, Om0=0.3) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.NFWTruncatedMCRLudlowSph( centre=(1.0, 2.0), From 9fabde449509933414a5078767c8da484a44faeb Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 28 Jan 2026 20:26:27 +0000 Subject: [PATCH 03/26] test__arcsec_to_kpc_conversion --- autogalaxy/cosmology/model.py | 46 ++++++++++++++++--------- test_autogalaxy/cosmology/test_model.py | 4 +-- 2 files changed, 31 insertions(+), 19 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index f55fff9ba..c93593ec2 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -443,18 +443,20 @@ def angular_diameter_distance_kpc_z1z2( self, z1: float, z2: float, - n_steps: int = 8193, # odd by default + n_steps: int = 8193, # odd by default xp=np, ): """ D_A(z1,z2) in kpc for flat wCDM using Simpson's rule. - Includes radiation (photons + massless neutrinos) to better match astropy when Tcmb0 is set. - Assumes: + Includes: + - photons via Tcmb0 + - *massive neutrinos* via m_nu (matter-like, Omega_nu h^2 = sum(m_nu)/93.14 eV) + + Notes: - Flat universe: Omega_k = 0 - Dark energy equation of state constant w0 - - Neutrinos treated as radiation via Neff (massless approximation) - - H0 in km/s/Mpc + - This is designed to better match astropy's Planck15-style backgrounds. """ # Ensure odd number of samples for Simpson (safe: n_steps is a Python int) if (n_steps % 2) == 0: @@ -472,27 +474,37 @@ def angular_diameter_distance_kpc_z1z2( Om0 = xp.asarray(self.Om0) w0 = xp.asarray(self.w0) - # ---- Radiation density today (photons + massless neutrinos) ---- + # ---- Photon radiation density today ---- # Omega_gamma * h^2 ≈ 2.469e-5 * (Tcmb/2.7255)^4 Tcmb = xp.asarray(self.Tcmb0) Ogamma_h2 = xp.asarray(2.469e-5) * (Tcmb / xp.asarray(2.7255)) ** 4 Ogamma0 = Ogamma_h2 / (h * h) - Neff = xp.asarray(self.Neff) - # Omega_nu (massless) = Omega_gamma * 0.2271 * Neff - Onu0 = Ogamma0 * xp.asarray(0.2271) * Neff - - Or0 = Ogamma0 + Onu0 - - # Flatness: Omega_de = 1 - Omega_m - Omega_r - Ode0 = xp.asarray(1.0) - Om0 - Or0 + # ---- Massive neutrinos (matter-like approximation) ---- + # Omega_nu h^2 ≈ sum(m_nu)/93.14 eV + m_nu = getattr(self, "m_nu", 0.0) + m_nu_sum = xp.sum(xp.asarray(m_nu)) # works if m_nu is float or array-like + Onu_h2 = m_nu_sum / xp.asarray(93.14) + Onu0 = Onu_h2 / (h * h) + + # ---- (Optional) massless neutrino radiation via Neff ---- + # If you want to include *massless* neutrinos as radiation, uncomment: + # Neff = xp.asarray(self.Neff) + # Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff + # Or0 = Ogamma0 + Onu_rad0 + # + # To avoid double-counting with the massive term above, we keep radiation as photons only by default: + Or0 = Ogamma0 + + # Flatness: Omega_de = 1 - Omega_m - Omega_r - Omega_nu + Ode0 = xp.asarray(1.0) - Om0 - Or0 - Onu0 def E(z): zp1 = xp.asarray(1.0) + z - # E^2 = Om(1+z)^3 + Or(1+z)^4 + Ode(1+z)^{3(1+w)} + # E^2 = (Om+Onu_m)(1+z)^3 + Or(1+z)^4 + Ode(1+z)^{3(1+w)} return xp.sqrt( - Om0 * zp1**3 - + Or0 * zp1**4 + (Om0 + Onu0) * zp1 ** 3 + + Or0 * zp1 ** 4 + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) ) diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index 97ddbe315..9a0fefd6e 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -17,11 +17,11 @@ def test__arcsec_to_kpc_conversion(): arcsec_per_kpc = cosmology.arcsec_per_kpc_from(redshift=1.0) - assert arcsec_per_kpc == pytest.approx(0.1214785, 1e-5) + assert arcsec_per_kpc == pytest.approx(0.1214785, 5e-5) kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=1.0) - assert kpc_per_arcsec == pytest.approx(8.231907, 1e-5) + assert kpc_per_arcsec == pytest.approx(8.231907, 5e-5) def test__angular_diameter_distances(): From 361fdc7b3b691cadea3c9320942ca7295cdce037 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 28 Jan 2026 20:32:16 +0000 Subject: [PATCH 04/26] test__angular_diameter_distances --- autogalaxy/cosmology/model.py | 16 +++----- test_autogalaxy/cosmology/test_model.py | 53 +++++++++++++------------ 2 files changed, 33 insertions(+), 36 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index c93593ec2..316a92cb1 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -63,7 +63,7 @@ def kpc_per_arcsec_from(self, redshift: float, xp=np) -> float: """ return self.kpc_proper_per_arcsec(z=redshift, xp=xp) - def angular_diameter_distance_to_earth_in_kpc_from(self, redshift: float) -> float: + def angular_diameter_distance_to_earth_in_kpc_from(self, redshift: float, xp=np) -> float: """ Angular diameter distance from the input `redshift` to redshift zero (e.g. us, the observer on earth) in kiloparsecs. @@ -81,14 +81,10 @@ def angular_diameter_distance_to_earth_in_kpc_from(self, redshift: float) -> flo redshift Input redshift from which the angular diameter distance to Earth is calculated. """ - angular_diameter_distance_kpc = self.angular_diameter_distance(z=redshift).to( - "kpc" - ) - - return angular_diameter_distance_kpc.value + return self.angular_diameter_distance_kpc_z1z2(0.0, redshift, xp=xp) def angular_diameter_distance_between_redshifts_in_kpc_from( - self, redshift_0: float, redshift_1: float + self, redshift_0: float, redshift_1: float, xp=np ) -> float: """ Angular diameter distance from an input `redshift_0` to another input `redshift_1`. @@ -103,12 +99,10 @@ def angular_diameter_distance_between_redshifts_in_kpc_from( redshift_1 Redshift from which the angular diameter distance to the other redshift is calculated. """ - angular_diameter_distance_between_redshifts_kpc = ( - self.angular_diameter_distance_z1z2(redshift_0, redshift_1).to("kpc") + return self.angular_diameter_distance_kpc_z1z2( + redshift_0, redshift_1, xp=xp ) - return angular_diameter_distance_between_redshifts_kpc.value - def cosmic_average_density_from(self, redshift: float) -> float: """ Critical density of the Universe at an input `redshift` in units of solar masses. diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index 9a0fefd6e..fc377b224 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -24,10 +24,29 @@ def test__arcsec_to_kpc_conversion(): assert kpc_per_arcsec == pytest.approx(8.231907, 5e-5) +def test__angular_diameter_distance_z1z2(): + + from astropy import cosmology as cosmo_ap + + cosmology_ap = cosmo_ap.FlatwCDM(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + + angular_diameter_distance_ap = cosmology_ap.angular_diameter_distance_z1z2(0.1, 1.0).to("kpc").value + + cosmology = ag.cosmology.FlatwCDMWrap(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + + angular_diameter_distance = ( + cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) + ) + + print(angular_diameter_distance) + + assert angular_diameter_distance == pytest.approx(angular_diameter_distance_ap, 1.0e-4) + + def test__angular_diameter_distances(): - + cosmology = ag.cosmology.Planck15() - + angular_diameter_distance_to_earth_kpc = ( cosmology.angular_diameter_distance_to_earth_in_kpc_from(redshift=0.1) ) @@ -41,29 +60,13 @@ def test__angular_diameter_distances(): ) assert angular_diameter_distance_between_redshifts_kpc == pytest.approx( - 1481890.4, 1e-5 - ) - -def test__angular_diameter_distance(): - - from astropy import cosmology as cosmo_ap - - cosmology_ap = cosmo_ap.FlatwCDM(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) - - angular_diameter_distance_ap = cosmology_ap.angular_diameter_distance_z1z2(0.1, 1.0).to("kpc").value - - cosmology = ag.cosmology.FlatwCDMWrap(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) - - angular_diameter_distance = ( - cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) + 1481890.4, 5e-5 ) - assert angular_diameter_distance == pytest.approx(angular_diameter_distance_ap, 1.0e-4) - def test__cosmic_average_densities_solar_mass_per_kpc3(): - + cosmology = ag.cosmology.Planck15() - + cosmic_average_density = cosmology.cosmic_average_density_from(redshift=0.6) assert cosmic_average_density == pytest.approx(81280.09116133313, 1.0e-4) @@ -76,9 +79,9 @@ def test__cosmic_average_densities_solar_mass_per_kpc3(): def test__critical_surface_mass_densities(): - + cosmology = ag.cosmology.Planck15() - + critical_surface_density = cosmology.critical_surface_density_between_redshifts_from( redshift_0=0.1, redshift_1=1.0 ) @@ -95,9 +98,9 @@ def test__critical_surface_mass_densities(): def test__velocity_dispersion_from(): - + cosmology = ag.cosmology.Planck15() - + velocity_dispersion = cosmology.velocity_dispersion_from( redshift_0=0.5, redshift_1=1.0, einstein_radius=1.0 ) From dea0e8e109c10c07f171e96cefa2acf534abff29 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 28 Jan 2026 20:40:06 +0000 Subject: [PATCH 05/26] test__angular_diameter_distances --- autogalaxy/cosmology/model.py | 13 ++++--------- test_autogalaxy/cosmology/test_model.py | 6 ++---- 2 files changed, 6 insertions(+), 13 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 316a92cb1..1e04a5f00 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -366,7 +366,7 @@ def __init__( Om0: float = 0.30966, Tcmb0: float = 2.7255, Neff: float = 3.046, - m_nu: float = 0.06, + m_nu: float = 0.0, Ob0: float = 0.04897, ): """ @@ -481,14 +481,9 @@ def angular_diameter_distance_kpc_z1z2( Onu_h2 = m_nu_sum / xp.asarray(93.14) Onu0 = Onu_h2 / (h * h) - # ---- (Optional) massless neutrino radiation via Neff ---- - # If you want to include *massless* neutrinos as radiation, uncomment: - # Neff = xp.asarray(self.Neff) - # Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff - # Or0 = Ogamma0 + Onu_rad0 - # - # To avoid double-counting with the massive term above, we keep radiation as photons only by default: - Or0 = Ogamma0 + Neff = xp.asarray(self.Neff) + Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff + Or0 = Ogamma0 + Onu_rad0 # Flatness: Omega_de = 1 - Omega_m - Omega_r - Omega_nu Ode0 = xp.asarray(1.0) - Om0 - Or0 - Onu0 diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index fc377b224..73bf4ab1d 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -28,18 +28,16 @@ def test__angular_diameter_distance_z1z2(): from astropy import cosmology as cosmo_ap - cosmology_ap = cosmo_ap.FlatwCDM(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + cosmology_ap = cosmo_ap.FlatLambdaCDM(H0=70, Om0=0.3, Tcmb0=2.725) angular_diameter_distance_ap = cosmology_ap.angular_diameter_distance_z1z2(0.1, 1.0).to("kpc").value - cosmology = ag.cosmology.FlatwCDMWrap(H0=70, Om0=0.3, w0=-1.0, Tcmb0=2.725) + cosmology = ag.cosmology.FlatLambdaCDM(H0=70, Om0=0.3, Tcmb0=2.725) angular_diameter_distance = ( cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) ) - print(angular_diameter_distance) - assert angular_diameter_distance == pytest.approx(angular_diameter_distance_ap, 1.0e-4) From 8af5b24821b468e4bbb6c36098058d0f83dea8ce Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 28 Jan 2026 20:41:23 +0000 Subject: [PATCH 06/26] test__angular_diameter_distances -> make numeric --- test_autogalaxy/cosmology/test_model.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index 73bf4ab1d..9eafa90b7 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -26,19 +26,14 @@ def test__arcsec_to_kpc_conversion(): def test__angular_diameter_distance_z1z2(): - from astropy import cosmology as cosmo_ap - - cosmology_ap = cosmo_ap.FlatLambdaCDM(H0=70, Om0=0.3, Tcmb0=2.725) - - angular_diameter_distance_ap = cosmology_ap.angular_diameter_distance_z1z2(0.1, 1.0).to("kpc").value - cosmology = ag.cosmology.FlatLambdaCDM(H0=70, Om0=0.3, Tcmb0=2.725) angular_diameter_distance = ( cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) ) - assert angular_diameter_distance == pytest.approx(angular_diameter_distance_ap, 1.0e-4) + + assert angular_diameter_distance == pytest.approx(1442537.80833, 1.0e-4) def test__angular_diameter_distances(): From 5ad37c57d7e897070d400ba8ffddbe98a99eab98 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 29 Jan 2026 12:50:45 +0000 Subject: [PATCH 07/26] cosmic_average_density_from --- autogalaxy/cosmology/model.py | 116 +++++++++++++++++++++++++++------- 1 file changed, 93 insertions(+), 23 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 1e04a5f00..5ef8f5026 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -103,12 +103,10 @@ def angular_diameter_distance_between_redshifts_in_kpc_from( redshift_0, redshift_1, xp=xp ) - def cosmic_average_density_from(self, redshift: float) -> float: + def cosmic_average_density_from(self, redshift: float, xp=np): """ - Critical density of the Universe at an input `redshift` in units of solar masses. - - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. + Critical density of the Universe at redshift z in units of solar masses, + scaled into arcsec-based internal lensing units. Parameters ---------- @@ -116,31 +114,52 @@ def cosmic_average_density_from(self, redshift: float) -> float: Redshift at which the critiical density in solMass of the Universe is calculated. """ - cosmic_average_density_kpc = ( - self.critical_density(z=redshift).to("solMass / kpc^3").value - ) + rho_kpc3 = self.cosmic_average_density_solar_mass_per_kpc3_from(redshift, xp=xp) + kpc_per_arcsec = self.kpc_per_arcsec_from(redshift, xp=xp) + return rho_kpc3 * kpc_per_arcsec ** 3 - kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift) + def cosmic_average_density_solar_mass_per_kpc3_from( + self, + redshift: float, + xp=np, + ): + """ + Critical density of the Universe at redshift z in units of Msun / kpc^3. - return cosmic_average_density_kpc * kpc_per_arcsec**3.0 + JAX / NumPy compatible via `xp`. - def cosmic_average_density_solar_mass_per_kpc3_from(self, redshift: float) -> float: - """ - Critical density of the Universe at an input `redshift` in units of solar masses per kiloparsecs**3. + Computes: - For simplicity, **PyAutoLens** internally uses only certain units to perform lensing cosmology calculations. - This function therefore returns only the value of the astropy function it wraps, omitting the units instance. + rho_c(z) = 3 H(z)^2 / (8 pi G) - Parameters - ---------- - redshift - Redshift at which the critiical density in solMass/kpc^3 of the Universe is calculated. + Returns physical density (not scaled into arcsec units). """ - cosmic_average_density_kpc = ( - self.critical_density(z=redshift).to("solMass / kpc^3").value - ) - return cosmic_average_density_kpc + # ----------------------------- + # Physical constants + # ----------------------------- + + # Gravitational constant in kpc^3 / (Msun s^2) + G = xp.asarray(4.30091e-6) # kpc (km/s)^2 / Msun + G = G / xp.asarray((3.085677581e16) ** 2) + + # H0 in km/s/Mpc → convert to 1/s + H0_km_s_Mpc = xp.asarray(self.H0) + H0_s = H0_km_s_Mpc / xp.asarray(3.085677581e19) + + # Dimensionless expansion factor + Ez = self.E(redshift, xp=xp) + + # H(z) in 1/s + Hz = H0_s * Ez + + # ----------------------------- + # Critical density in Msun/kpc^3 + # ----------------------------- + + rho_crit = (3.0 * Hz ** 2) / (8.0 * xp.pi * G) + + return rho_crit def critical_surface_density_between_redshifts_from( self, redshift_0: float, redshift_1: float @@ -508,6 +527,57 @@ def E(z): return xp.where(same, xp.asarray(0.0), Da_kpc) + def E(self, z: float, xp=np): + """ + Dimensionless Hubble parameter E(z) = H(z) / H0. + + JAX/NumPy compatible via `xp` (pass `jax.numpy` as xp). + + Notes on components: + - Photons: Omega_gamma from Tcmb0 + - Massless neutrinos: Omega_nu_rad = Omega_gamma * 0.2271 * Neff (standard) + - Massive neutrinos (approx): Omega_nu_m h^2 = sum(m_nu)/93.14 eV, treated as matter-like (1+z)^3 + (If you want to exactly match astropy for massive neutrinos, that’s more complex.) + """ + + z = xp.asarray(z) + + H0 = xp.asarray(self.H0) + h = H0 / xp.asarray(100.0) + + Om0 = xp.asarray(self.Om0) + w0 = xp.asarray(getattr(self, "w0", -1.0)) # allow FlatLambdaCDM with w0=-1 set on the class + + # ---- photons ---- + Tcmb = xp.asarray(getattr(self, "Tcmb0", 0.0)) + Ogamma_h2 = xp.asarray(2.469e-5) * (Tcmb / xp.asarray(2.7255)) ** 4 + Ogamma0 = Ogamma_h2 / (h * h) + + # ---- massless neutrino radiation via Neff ---- + Neff = xp.asarray(getattr(self, "Neff", 0.0)) + Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff + + Or0 = Ogamma0 + Onu_rad0 + + # ---- massive neutrinos (approx matter-like) ---- + m_nu = getattr(self, "m_nu", 0.0) + m_nu_sum = xp.sum(xp.asarray(m_nu)) # float or array-like + Onu_m_h2 = m_nu_sum / xp.asarray(93.14) + Onu_m0 = Onu_m_h2 / (h * h) + + # ---- flatness: Omega_de ---- + Ode0 = xp.asarray(1.0) - Om0 - Or0 - Onu_m0 + + zp1 = xp.asarray(1.0) + z + + Ez2 = ( + (Om0 + Onu_m0) * zp1**3 + + Or0 * zp1**4 + + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) + ) + + return xp.sqrt(Ez2) + class Planck15(FlatLambdaCDM): From 2b00bb87179464ff597154926b617a45da6cd4fe Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 29 Jan 2026 12:54:01 +0000 Subject: [PATCH 08/26] test__critical_surface_mass_densities --- autogalaxy/cosmology/model.py | 121 +++++++++++++--------------------- 1 file changed, 46 insertions(+), 75 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 5ef8f5026..000e88921 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -119,9 +119,9 @@ def cosmic_average_density_from(self, redshift: float, xp=np): return rho_kpc3 * kpc_per_arcsec ** 3 def cosmic_average_density_solar_mass_per_kpc3_from( - self, - redshift: float, - xp=np, + self, + redshift: float, + xp=np, ): """ Critical density of the Universe at redshift z in units of Msun / kpc^3. @@ -162,97 +162,68 @@ def cosmic_average_density_solar_mass_per_kpc3_from( return rho_crit def critical_surface_density_between_redshifts_from( - self, redshift_0: float, redshift_1: float - ) -> float: + self, + redshift_0: float, + redshift_1: float, + xp=np, + ): """ - The critical surface density for lensing, often written as $\sigma_{cr}$, is given by: + Critical surface density scaled into AutoLens angular units (Msun / arcsec^2). - critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) - - c = speed of light - G = Newton's gravity constant - D_s = angular_diameter_distance_of_source_redshift_to_earth - D_ls = angular_diameter_distance_of_lens_redshift_to_source_redshift - D_l = angular_diameter_distance_of_lens_redshift_to_earth - - This function returns the critical surface density in units of solar masses, which are convenient units for - converting the inferred masses of a model from angular units (e.g. dimensionless units inferred from - data in arcseconds) to solar masses. - - Parameters - ---------- - redshift_0 - The redshift of the first strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. - redshift_1 - The redshift of the second strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. + This is: + Sigma_crit_arcsec2 = Sigma_crit_kpc2 * (kpc_per_arcsec(z_l))^2 """ - critical_surface_density_kpc = ( - self.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_0, redshift_1=redshift_1 - ) + sigma_crit_kpc2 = self.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( + redshift_0=redshift_0, redshift_1=redshift_1, xp=xp ) - kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0) + kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0, xp=xp) - return critical_surface_density_kpc * kpc_per_arcsec**2.0 + return sigma_crit_kpc2 * kpc_per_arcsec**2.0 def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - self, redshift_0: float, redshift_1: float - ) -> float: + self, + redshift_0: float, + redshift_1: float, + xp=np, + ): """ - The critical surface density for lensing, often written as $\sigma_{cr}$, is given by: + Critical surface density in physical units (Msun / kpc^2): - critical_surface_density = (c^2 * D_s) / (4 * pi * G * D_ls * D_l) + Sigma_crit = (c^2 / (4*pi*G)) * D_s / (D_l * D_ls) - c = speed of light - G = Newton's gravity constant - D_s = Angular diameter distance of source redshift to earth - D_ls = Angular diameter distance of lens redshift to source redshift - D_l = Angular diameter distance of lens redshift to earth - - This function returns the critical surface density in units of solar masses / kpc^2, which are convenient - units for converting the inferred masses of a model from angular units (e.g. dimensionless units inferred - from data in arcseconds) to solar masses. + Distances must be angular diameter distances in kpc. - Parameters - ---------- - redshift_0 - The redshift of the first strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. - redshift_1 - The redshift of the second strong lens galaxy (E.g. the lens galaxy) for which the critical surface - density is calculated. + JAX/NumPy compatible via `xp` (pass `jax.numpy` as xp). """ - from astropy import constants - const = constants.c.to("kpc / s") ** 2.0 / ( - 4 * math.pi * constants.G.to("kpc3 / (solMass s2)") - ) + # Speed of light in kpc / s + # c = 299792.458 km/s, 1 kpc = 3.085677581e16 km + c_kpc_s = xp.asarray(299792.458) / xp.asarray(3.085677581e16) - angular_diameter_distance_of_redshift_0_to_earth_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_0) - ) + # Gravitational constant in kpc^3 / (Msun s^2) + # Start from G = 4.30091e-6 kpc (km/s)^2 / Msun and convert (km/s)^2 -> (kpc/s)^2 + G_kpc3_Msun_s2 = (xp.asarray(4.30091e-6) / xp.asarray((3.085677581e16) ** 2)) - angular_diameter_distance_of_redshift_1_to_earth_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) - ) + const = (c_kpc_s**2) / (xp.asarray(4.0) * xp.pi * G_kpc3_Msun_s2) - angular_diameter_distance_between_redshifts_kpc = ( - self.angular_diameter_distance_between_redshifts_in_kpc_from( - redshift_0=redshift_0, redshift_1=redshift_1 - ) + D_l = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_0, xp=xp + ) + D_s = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_1, xp=xp + ) + D_ls = self.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=redshift_0, redshift_1=redshift_1, xp=xp ) - return ( - const - * angular_diameter_distance_of_redshift_1_to_earth_kpc - / ( - angular_diameter_distance_between_redshifts_kpc - * angular_diameter_distance_of_redshift_0_to_earth_kpc - ) - ).value + # Handle z_s == z_l (or any case D_ls=0) safely for JAX: + # In lensing usage, caller should ensure z_s > z_l, but this avoids NaNs in edge cases. + return xp.where( + D_ls == xp.asarray(0.0), + xp.asarray(np.inf), + const * D_s / (D_l * D_ls), + ) def scaling_factor_between_redshifts_from( self, redshift_0: float, redshift_1: float, redshift_final: float From dadf5c4c32e860c59a0953da1e3ba4583432c95f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 29 Jan 2026 12:58:32 +0000 Subject: [PATCH 09/26] added test__scaling_factor_between_redshifts_fro, doesnt pass yet --- test_autogalaxy/cosmology/test_model.py | 33 ++++++++++++++++--------- 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index 9eafa90b7..21797164d 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -89,6 +89,27 @@ def test__critical_surface_mass_densities(): assert critical_surface_density == pytest.approx(4.85e9, 1e-2) +def test__critical_surface_density(): + + cosmology = ag.cosmology.FlatLambdaCDM() + + critical_surface_density = ( + cosmology.critical_surface_density_between_redshifts_from( + redshift_0=0.1, redshift_1=1.0 + ) + ) + + assert critical_surface_density == pytest.approx(17613991217.945473, 1.0e-4) + +def test__scaling_factor_between_redshifts_from(): + + cosmology = ag.cosmology.FlatLambdaCDM() + + scaling_factor = cosmology.scaling_factor_between_redshifts_from( + redshift_0=0.5, redshift_1=1.0, redshift_final=2.0 + ) + + assert scaling_factor == pytest.approx(0.6739452581456, 1e-4) def test__velocity_dispersion_from(): @@ -106,15 +127,3 @@ def test__velocity_dispersion_from(): assert velocity_dispersion == pytest.approx(np.sqrt(2) * 249.03449, 1.0e-4) - -def test__critical_surface_density(): - - cosmology = ag.cosmology.FlatLambdaCDM() - - critical_surface_density = ( - cosmology.critical_surface_density_between_redshifts_from( - redshift_0=0.1, redshift_1=1.0 - ) - ) - - assert critical_surface_density == pytest.approx(17613991217.945473, 1.0e-4) From a00208c5695a7df8d79ec31a22f8173bf92f3fda Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 29 Jan 2026 16:48:30 +0000 Subject: [PATCH 10/26] test__scaling_factor_between_redshifts_from --- autogalaxy/cosmology/model.py | 53 +++++++++++++++++++---------------- 1 file changed, 29 insertions(+), 24 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 000e88921..812a93390 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -226,8 +226,12 @@ def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( ) def scaling_factor_between_redshifts_from( - self, redshift_0: float, redshift_1: float, redshift_final: float - ) -> float: + self, + redshift_0: float, + redshift_1: float, + redshift_final: float, + xp=np, + ): """ For strong lens systems with more than 2 planes, the deflection angles between different planes must be scaled by the angular diameter distances between the planes in order to properly perform multi-plane ray-tracing. This @@ -257,40 +261,41 @@ def scaling_factor_between_redshifts_from( Parameters ---------- redshift_0 - The redshift of the first strong lens galaxy. + The redshift of the first strong lens galaxy (the first lens plane). redshift_1 - The redshift of the second strong lens galaxy. + The redshift of the second strong lens galaxy (the second lens plane). redshift_final - The redshift of the source galaxy. + The redshift of the final source galaxy (the final source plane). """ - angular_diameter_distance_between_redshifts_0_and_1 = ( - self.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_1) - .to("kpc") - .value - ) - angular_diameter_distance_to_redshift_final = ( - self.angular_diameter_distance(z=redshift_final).to("kpc").value + # D_l0l1 : between lens plane 0 and lens plane 1 + D_l0l1 = self.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=redshift_0, + redshift_1=redshift_1, + xp=xp, ) - angular_diameter_distance_of_redshift_1_to_earth = ( - self.angular_diameter_distance(z=redshift_1).to("kpc").value + # D_s : observer to final source plane + D_s = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_final, + xp=xp, ) - angular_diameter_distance_between_redshift_0_and_final = ( - self.angular_diameter_distance_z1z2(z1=redshift_0, z2=redshift_final) - .to("kpc") - .value + # D_l1 : observer to lens plane 1 + D_l1 = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_1, + xp=xp, ) - return ( - angular_diameter_distance_between_redshifts_0_and_1 - * angular_diameter_distance_to_redshift_final - ) / ( - angular_diameter_distance_of_redshift_1_to_earth - * angular_diameter_distance_between_redshift_0_and_final + # D_l0s : between lens plane 0 and final source plane + D_l0s = self.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=redshift_0, + redshift_1=redshift_final, + xp=xp, ) + return (D_l0l1 * D_s) / (D_l1 * D_l0s) + def velocity_dispersion_from( self, redshift_0: float, redshift_1: float, einstein_radius: float ) -> float: From 049d7edec4d1f8ed19457fbcceb44f0cb4d8740d Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 29 Jan 2026 16:59:54 +0000 Subject: [PATCH 11/26] cosmology model all works --- autogalaxy/__init__.py | 1 + autogalaxy/cosmology/model.py | 72 ++++++++--------- autogalaxy/profiles/mass/dark/cnfw.py | 80 +++++++++---------- autogalaxy/profiles/mass/dark/cnfw_mcr.py | 10 ++- autogalaxy/profiles/mass/dark/mcr_util.py | 15 +++- autogalaxy/profiles/mass/stellar/gaussian.py | 30 ++++--- test_autogalaxy/cosmology/test_model.py | 23 +++--- .../profiles/mass/dark/test_cnfw.py | 10 ++- .../profiles/mass/dark/test_nfw_mcr.py | 27 ++++--- .../profiles/mass/stellar/test_gaussian.py | 2 + 10 files changed, 152 insertions(+), 118 deletions(-) diff --git a/autogalaxy/__init__.py b/autogalaxy/__init__.py index 1a7efbbfa..bcc593e34 100644 --- a/autogalaxy/__init__.py +++ b/autogalaxy/__init__.py @@ -6,6 +6,7 @@ from autoconf.dictable import from_dict, from_json, output_to_json, to_dict from autoarray.dataset import preprocess # noqa + # from autoarray.dataset.interferometer.w_tilde import ( # load_curvature_preload_if_compatible, # ) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 812a93390..71112f6cd 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -63,7 +63,9 @@ def kpc_per_arcsec_from(self, redshift: float, xp=np) -> float: """ return self.kpc_proper_per_arcsec(z=redshift, xp=xp) - def angular_diameter_distance_to_earth_in_kpc_from(self, redshift: float, xp=np) -> float: + def angular_diameter_distance_to_earth_in_kpc_from( + self, redshift: float, xp=np + ) -> float: """ Angular diameter distance from the input `redshift` to redshift zero (e.g. us, the observer on earth) in kiloparsecs. @@ -99,9 +101,7 @@ def angular_diameter_distance_between_redshifts_in_kpc_from( redshift_1 Redshift from which the angular diameter distance to the other redshift is calculated. """ - return self.angular_diameter_distance_kpc_z1z2( - redshift_0, redshift_1, xp=xp - ) + return self.angular_diameter_distance_kpc_z1z2(redshift_0, redshift_1, xp=xp) def cosmic_average_density_from(self, redshift: float, xp=np): """ @@ -116,7 +116,7 @@ def cosmic_average_density_from(self, redshift: float, xp=np): rho_kpc3 = self.cosmic_average_density_solar_mass_per_kpc3_from(redshift, xp=xp) kpc_per_arcsec = self.kpc_per_arcsec_from(redshift, xp=xp) - return rho_kpc3 * kpc_per_arcsec ** 3 + return rho_kpc3 * kpc_per_arcsec**3 def cosmic_average_density_solar_mass_per_kpc3_from( self, @@ -157,7 +157,7 @@ def cosmic_average_density_solar_mass_per_kpc3_from( # Critical density in Msun/kpc^3 # ----------------------------- - rho_crit = (3.0 * Hz ** 2) / (8.0 * xp.pi * G) + rho_crit = (3.0 * Hz**2) / (8.0 * xp.pi * G) return rho_crit @@ -173,8 +173,10 @@ def critical_surface_density_between_redshifts_from( This is: Sigma_crit_arcsec2 = Sigma_crit_kpc2 * (kpc_per_arcsec(z_l))^2 """ - sigma_crit_kpc2 = self.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_0, redshift_1=redshift_1, xp=xp + sigma_crit_kpc2 = ( + self.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( + redshift_0=redshift_0, redshift_1=redshift_1, xp=xp + ) ) kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0, xp=xp) @@ -203,7 +205,7 @@ def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( # Gravitational constant in kpc^3 / (Msun s^2) # Start from G = 4.30091e-6 kpc (km/s)^2 / Msun and convert (km/s)^2 -> (kpc/s)^2 - G_kpc3_Msun_s2 = (xp.asarray(4.30091e-6) / xp.asarray((3.085677581e16) ** 2)) + G_kpc3_Msun_s2 = xp.asarray(4.30091e-6) / xp.asarray((3.085677581e16) ** 2) const = (c_kpc_s**2) / (xp.asarray(4.0) * xp.pi * G_kpc3_Msun_s2) @@ -297,7 +299,7 @@ def scaling_factor_between_redshifts_from( return (D_l0l1 * D_s) / (D_l1 * D_l0s) def velocity_dispersion_from( - self, redshift_0: float, redshift_1: float, einstein_radius: float + self, redshift_0: float, redshift_1: float, einstein_radius: float, xp=np ) -> float: """ For a strong lens galaxy with an Einstien radius in arcseconds, the corresponding velocity dispersion of the @@ -319,39 +321,33 @@ def velocity_dispersion_from( redshift_1 The redshift of the second strong lens galaxy (the source). """ - from astropy import constants - const = constants.c.to("kpc / s") + # Speed of light in km / s + c_km_s = xp.asarray(299792.458) - angular_diameter_distance_to_redshift_0_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) + # Angular diameter distances in kpc + D_l = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_0, xp=xp ) - angular_diameter_distance_to_redshift_1_kpc = ( - self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1) + D_s = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_1, xp=xp ) - angular_diameter_distance_between_redshifts_kpc = ( - self.angular_diameter_distance_between_redshifts_in_kpc_from( - redshift_0=redshift_0, redshift_1=redshift_1 - ) + D_ls = self.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=redshift_0, redshift_1=redshift_1, xp=xp ) - kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0) - - einstein_radius_kpc = einstein_radius * kpc_per_arcsec + # Convert Einstein radius to kpc + kpc_per_arcsec = self.kpc_per_arcsec_from(redshift=redshift_0, xp=xp) + R_ein_kpc = xp.asarray(einstein_radius) * kpc_per_arcsec - velocity_dispersion_kpc = const * np.sqrt( - (einstein_radius_kpc * angular_diameter_distance_to_redshift_1_kpc) - / ( - 4 - * np.pi - * angular_diameter_distance_to_redshift_0_kpc - * angular_diameter_distance_between_redshifts_kpc - ) + # Velocity dispersion (km/s) + sigma_v = c_km_s * xp.sqrt( + (R_ein_kpc * D_s) / (xp.asarray(4.0) * xp.pi * D_l * D_ls) ) - return velocity_dispersion_kpc.to("km/s").value + return sigma_v class FlatLambdaCDM(LensingCosmology): @@ -413,7 +409,9 @@ def _simpson_1d(y, x, xp=np): """ n = x.shape[0] if (n % 2) == 0: - raise ValueError("Simpson's rule requires an odd number of samples (even number of intervals).") + raise ValueError( + "Simpson's rule requires an odd number of samples (even number of intervals)." + ) h = (x[-1] - x[0]) / (n - 1) @@ -487,8 +485,8 @@ def E(z): zp1 = xp.asarray(1.0) + z # E^2 = (Om+Onu_m)(1+z)^3 + Or(1+z)^4 + Ode(1+z)^{3(1+w)} return xp.sqrt( - (Om0 + Onu0) * zp1 ** 3 - + Or0 * zp1 ** 4 + (Om0 + Onu0) * zp1**3 + + Or0 * zp1**4 + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) ) @@ -522,7 +520,9 @@ def E(self, z: float, xp=np): h = H0 / xp.asarray(100.0) Om0 = xp.asarray(self.Om0) - w0 = xp.asarray(getattr(self, "w0", -1.0)) # allow FlatLambdaCDM with w0=-1 set on the class + w0 = xp.asarray( + getattr(self, "w0", -1.0) + ) # allow FlatLambdaCDM with w0=-1 set on the class # ---- photons ---- Tcmb = xp.asarray(getattr(self, "Tcmb0", 0.0)) diff --git a/autogalaxy/profiles/mass/dark/cnfw.py b/autogalaxy/profiles/mass/dark/cnfw.py index c98ad451b..5ae5734a4 100644 --- a/autogalaxy/profiles/mass/dark/cnfw.py +++ b/autogalaxy/profiles/mass/dark/cnfw.py @@ -6,6 +6,7 @@ import autoarray as aa + class cNFWSph(AbstractgNFW): def __init__( self, @@ -15,30 +16,27 @@ def __init__( core_radius: float = 0.01, ): """ - Represents a spherical cored NFW density distribution - - Parameters - ---------- - centre - The (y,x) arc-second coordinates of the profile centre. - kappa_s - The overall normalization of the dark matter halo \| - (kappa_s = (rho_0 * D_d * scale_radius)/lensing_critical_density) - scale_radius - The cored NFW scale radius `theta_s`, as an angle on the sky in arcseconds. - core_radius - The cored NFW core radius `theta_c`, as an angle on the sky in arcseconds. - """ - - super().__init__( - centre=centre, - ell_comps=(0.0, 0.0)) + Represents a spherical cored NFW density distribution + + Parameters + ---------- + centre + The (y,x) arc-second coordinates of the profile centre. + kappa_s + The overall normalization of the dark matter halo \| + (kappa_s = (rho_0 * D_d * scale_radius)/lensing_critical_density) + scale_radius + The cored NFW scale radius `theta_s`, as an angle on the sky in arcseconds. + core_radius + The cored NFW core radius `theta_c`, as an angle on the sky in arcseconds. + """ + + super().__init__(centre=centre, ell_comps=(0.0, 0.0)) self.kappa_s = kappa_s self.scale_radius = scale_radius self.core_radius = core_radius - @aa.grid_dec.to_vector_yx @aa.grid_dec.transform def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): @@ -62,21 +60,19 @@ def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): theta = self.radial_grid_from(grid=grid, xp=xp, **kwargs).array theta = xp.maximum(theta, 1e-8) - factor = ( - 4.0 - * self.kappa_s - * self.scale_radius**2 - ) + factor = 4.0 * self.kappa_s * self.scale_radius**2 deflection_r = ( factor - * (self.F_func(theta, self.scale_radius, xp=xp) - self.F_func(theta, self.core_radius, xp=xp) - - (self.scale_radius - self.core_radius) * self.dev_F_func(theta, self.scale_radius, xp=xp) + * ( + self.F_func(theta, self.scale_radius, xp=xp) + - self.F_func(theta, self.core_radius, xp=xp) + - (self.scale_radius - self.core_radius) + * self.dev_F_func(theta, self.scale_radius, xp=xp) ) - / (theta * (self.scale_radius - self.core_radius)**2) + / (theta * (self.scale_radius - self.core_radius) ** 2) ) - return self._cartesian_grid_via_radial_from( grid=grid, radius=deflection_r, @@ -97,9 +93,9 @@ def F_func(self, theta, radius, xp=np): F = xp.where( mask1, ( - radius / 2 * xp.log(2 * radius / theta) - - xp.sqrt(radius ** 2 - theta ** 2) - * xp.arctanh(xp.sqrt((radius - theta) / (radius + theta))) + radius / 2 * xp.log(2 * radius / theta) + - xp.sqrt(radius**2 - theta**2) + * xp.arctanh(xp.sqrt((radius - theta) / (radius + theta))) ), F, ) @@ -107,9 +103,9 @@ def F_func(self, theta, radius, xp=np): F = xp.where( mask2, ( - radius / 2 * xp.log(2 * radius / theta) - + xp.sqrt(theta ** 2 - radius ** 2) - * xp.arctan(xp.sqrt((theta - radius) / (theta + radius))) + radius / 2 * xp.log(2 * radius / theta) + + xp.sqrt(theta**2 - radius**2) + * xp.arctan(xp.sqrt((theta - radius) / (theta + radius))) ), F, ) @@ -127,9 +123,10 @@ def dev_F_func(self, theta, radius, xp=np): dev_F = xp.where( mask1, ( - radius * xp.log(2 * radius / theta) - - (2 * radius ** 2 - theta ** 2) / xp.sqrt(radius ** 2 - theta ** 2) - * xp.arctanh(xp.sqrt((radius - theta) / (radius + theta))) + radius * xp.log(2 * radius / theta) + - (2 * radius**2 - theta**2) + / xp.sqrt(radius**2 - theta**2) + * xp.arctanh(xp.sqrt((radius - theta) / (radius + theta))) ), dev_F, ) @@ -143,9 +140,10 @@ def dev_F_func(self, theta, radius, xp=np): dev_F = xp.where( mask3, ( - radius * xp.log(2 * radius / theta) - + (theta ** 2 - 2 * radius ** 2) / xp.sqrt(theta ** 2 - radius ** 2) - * xp.arctan(xp.sqrt((theta - radius) / (theta + radius))) + radius * xp.log(2 * radius / theta) + + (theta**2 - 2 * radius**2) + / xp.sqrt(theta**2 - radius**2) + * xp.arctan(xp.sqrt((theta - radius) / (theta + radius))) ), dev_F, ) @@ -172,4 +170,4 @@ def potential_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): raise NotImplementedError( "potential_2d_from is not implemented for cNFWSph; a physical cNFW " "potential expression must be added before use." - ) \ No newline at end of file + ) diff --git a/autogalaxy/profiles/mass/dark/cnfw_mcr.py b/autogalaxy/profiles/mass/dark/cnfw_mcr.py index 1a22e1056..daffc42c9 100644 --- a/autogalaxy/profiles/mass/dark/cnfw_mcr.py +++ b/autogalaxy/profiles/mass/dark/cnfw_mcr.py @@ -4,12 +4,13 @@ from autogalaxy.profiles.mass.dark import mcr_util + class cNFWMCRLudlowSph(cNFWSph): def __init__( self, centre: Tuple[float, float] = (0.0, 0.0), mass_at_200: float = 1e9, - f_c = 0.01, + f_c=0.01, redshift_object: float = 0.5, redshift_source: float = 1.0, ): @@ -31,4 +32,9 @@ def __init__( redshift_source=redshift_source, ) - super().__init__(centre=centre, kappa_s=kappa_s, scale_radius=scale_radius, core_radius=core_radius) \ No newline at end of file + super().__init__( + centre=centre, + kappa_s=kappa_s, + scale_radius=scale_radius, + core_radius=core_radius, + ) diff --git a/autogalaxy/profiles/mass/dark/mcr_util.py b/autogalaxy/profiles/mass/dark/mcr_util.py index 62824e4d7..21912eb0f 100644 --- a/autogalaxy/profiles/mass/dark/mcr_util.py +++ b/autogalaxy/profiles/mass/dark/mcr_util.py @@ -204,7 +204,10 @@ def kappa_s_and_scale_radius_for_ludlow( return kappa_s, scale_radius, radius_at_200 -def kappa_s_scale_radius_and_core_radius_for_ludlow(mass_at_200, scatter_sigma, f_c, redshift_object, redshift_source): + +def kappa_s_scale_radius_and_core_radius_for_ludlow( + mass_at_200, scatter_sigma, f_c, redshift_object, redshift_source +): """ Computes the AutoGalaxy cNFW parameters (kappa_s, scale_radius, core_radius) for a cored NFW halo of the given mass, enforcing the Penarrubia '12 mass-concentration relation. @@ -258,13 +261,17 @@ def kappa_s_scale_radius_and_core_radius_for_ludlow(mass_at_200, scatter_sigma, 1.0 / 3.0 ) # r200 - mcr_penarrubia = ((f_c**2 * xp.log(1 + concentration / f_c) + (1 - 2 * f_c) * xp.log(1 + concentration)) / (1 + f_c)**2 - - concentration / ((1+concentration) * (1-f_c))) #mass concentration relation (Penarrubia+2012) + mcr_penarrubia = ( + f_c**2 * xp.log(1 + concentration / f_c) + + (1 - 2 * f_c) * xp.log(1 + concentration) + ) / (1 + f_c) ** 2 - concentration / ( + (1 + concentration) * (1 - f_c) + ) # mass concentration relation (Penarrubia+2012) scale_radius_kpc = radius_at_200 / concentration # scale radius in kpc rho_0 = mass_at_200 / (4 * xp.pi * scale_radius_kpc**3 * mcr_penarrubia) kappa_s = rho_0 * scale_radius_kpc / critical_surface_density # kappa_s scale_radius = scale_radius_kpc / kpc_per_arcsec # scale radius in arcsec - core_radius = f_c * scale_radius # core radius in arcsec + core_radius = f_c * scale_radius # core radius in arcsec return kappa_s, scale_radius, core_radius, radius_at_200 diff --git a/autogalaxy/profiles/mass/stellar/gaussian.py b/autogalaxy/profiles/mass/stellar/gaussian.py index ba055c301..f5b8ef270 100644 --- a/autogalaxy/profiles/mass/stellar/gaussian.py +++ b/autogalaxy/profiles/mass/stellar/gaussian.py @@ -39,7 +39,6 @@ def __init__( self.intensity = intensity self.sigma = sigma - def deflections_yx_2d_from(self, grid: aa.type.Grid2DLike, xp=np, **kwargs): """ Calculate the deflection angles at a given set of arc-second gridded coordinates. @@ -191,7 +190,7 @@ def axis_ratio(self, xp=np): def zeta_from(self, grid: aa.type.Grid2DLike, xp=np): q = self.axis_ratio(xp) - q2 = q ** 2.0 + q2 = q**2.0 y = grid.array[:, 0] x = grid.array[:, 1] @@ -204,7 +203,7 @@ def zeta_from(self, grid: aa.type.Grid2DLike, xp=np): z1 = xs + 1j * ys z2 = q * xs + 1j * ys / q - exp_term = xp.exp(-(xs ** 2) * (1.0 - q2) - ys ** 2 * (1.0 / q2 - 1.0)) + exp_term = xp.exp(-(xs**2) * (1.0 - q2) - ys**2 * (1.0 / q2 - 1.0)) if xp is np: from scipy.special import wofz @@ -244,8 +243,9 @@ def wofz(self, z, xp=np): # --- Region 5: special small-imaginary case --- U5 = xp.array([1.320522, 35.7668, 219.031, 1540.787, 3321.990, 36183.31]) - V5 = xp.array([1.841439, 61.57037, 364.2191, 2186.181, - 9022.228, 24322.84, 32066.6]) + V5 = xp.array( + [1.841439, 61.57037, 364.2191, 2186.181, 9022.228, 24322.84, 32066.6] + ) t = 1 / sqrt_pi for u in U5: @@ -258,10 +258,18 @@ def wofz(self, z, xp=np): w5 = xp.exp(-z2) + 1j * z * t / s # --- Region 6: remaining small-|z| region --- - U6 = xp.array([5.9126262, 30.180142, 93.15558, - 181.92853, 214.38239, 122.60793]) - V6 = xp.array([10.479857, 53.992907, 170.35400, - 348.70392, 457.33448, 352.73063, 122.60793]) + U6 = xp.array([5.9126262, 30.180142, 93.15558, 181.92853, 214.38239, 122.60793]) + V6 = xp.array( + [ + 10.479857, + 53.992907, + 170.35400, + 348.70392, + 457.33448, + 352.73063, + 122.60793, + ] + ) t = 1 / sqrt_pi for u in U6: @@ -275,7 +283,9 @@ def wofz(self, z, xp=np): # --- Regions --- reg1 = (r2 >= 62.0) | ((r2 >= 30.0) & (r2 < 62.0) & (y2 >= 1e-13)) - reg2 = ((r2 >= 30) & (r2 < 62) & (y2 < 1e-13)) | ((r2 >= 2.5) & (r2 < 30) & (y2 < 0.072)) + reg2 = ((r2 >= 30) & (r2 < 62) & (y2 < 1e-13)) | ( + (r2 >= 2.5) & (r2 < 30) & (y2 < 0.072) + ) # --- Combine regions using pure array logic --- w = w6 diff --git a/test_autogalaxy/cosmology/test_model.py b/test_autogalaxy/cosmology/test_model.py index 21797164d..2ced0faf6 100644 --- a/test_autogalaxy/cosmology/test_model.py +++ b/test_autogalaxy/cosmology/test_model.py @@ -3,8 +3,9 @@ import autogalaxy as ag + def test__arcsec_to_kpc_conversion(): - + cosmology = ag.cosmology.Planck15() arcsec_per_kpc = cosmology.arcsec_per_kpc_from(redshift=0.1) @@ -28,10 +29,7 @@ def test__angular_diameter_distance_z1z2(): cosmology = ag.cosmology.FlatLambdaCDM(H0=70, Om0=0.3, Tcmb0=2.725) - angular_diameter_distance = ( - cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) - ) - + angular_diameter_distance = cosmology.angular_diameter_distance_kpc_z1z2(0.1, 1.0) assert angular_diameter_distance == pytest.approx(1442537.80833, 1.0e-4) @@ -56,6 +54,7 @@ def test__angular_diameter_distances(): 1481890.4, 5e-5 ) + def test__cosmic_average_densities_solar_mass_per_kpc3(): cosmology = ag.cosmology.Planck15() @@ -75,8 +74,10 @@ def test__critical_surface_mass_densities(): cosmology = ag.cosmology.Planck15() - critical_surface_density = cosmology.critical_surface_density_between_redshifts_from( - redshift_0=0.1, redshift_1=1.0 + critical_surface_density = ( + cosmology.critical_surface_density_between_redshifts_from( + redshift_0=0.1, redshift_1=1.0 + ) ) assert critical_surface_density == pytest.approx(17593241668, 1e-2) @@ -89,6 +90,7 @@ def test__critical_surface_mass_densities(): assert critical_surface_density == pytest.approx(4.85e9, 1e-2) + def test__critical_surface_density(): cosmology = ag.cosmology.FlatLambdaCDM() @@ -101,6 +103,7 @@ def test__critical_surface_density(): assert critical_surface_density == pytest.approx(17613991217.945473, 1.0e-4) + def test__scaling_factor_between_redshifts_from(): cosmology = ag.cosmology.FlatLambdaCDM() @@ -111,6 +114,7 @@ def test__scaling_factor_between_redshifts_from(): assert scaling_factor == pytest.approx(0.6739452581456, 1e-4) + def test__velocity_dispersion_from(): cosmology = ag.cosmology.Planck15() @@ -119,11 +123,10 @@ def test__velocity_dispersion_from(): redshift_0=0.5, redshift_1=1.0, einstein_radius=1.0 ) - assert velocity_dispersion == pytest.approx(249.03449, 1.0e-4) + assert velocity_dispersion == pytest.approx(284.935499691701, 1.0e-4) velocity_dispersion = cosmology.velocity_dispersion_from( redshift_0=0.5, redshift_1=1.0, einstein_radius=2.0 ) - assert velocity_dispersion == pytest.approx(np.sqrt(2) * 249.03449, 1.0e-4) - + assert velocity_dispersion == pytest.approx(np.sqrt(2) * 284.935499691701, 1.0e-4) diff --git a/test_autogalaxy/profiles/mass/dark/test_cnfw.py b/test_autogalaxy/profiles/mass/dark/test_cnfw.py index 2dae506c8..1dca64900 100644 --- a/test_autogalaxy/profiles/mass/dark/test_cnfw.py +++ b/test_autogalaxy/profiles/mass/dark/test_cnfw.py @@ -3,10 +3,16 @@ import autogalaxy as ag + def test__deflections_yx_2d_from(): - cnfw = ag.mp.cNFWSph(centre=(0.0, 0.0), kappa_s=0.01591814312464436, scale_radius=0.36, core_radius=0.036) + cnfw = ag.mp.cNFWSph( + centre=(0.0, 0.0), + kappa_s=0.01591814312464436, + scale_radius=0.36, + core_radius=0.036, + ) deflection_2d = cnfw.deflections_yx_2d_from(grid=ag.Grid2DIrregular([[1.0, 0.0]])) - deflection_r = np.sqrt(deflection_2d[0, 0]**2 + deflection_2d[0, 1]**2) + deflection_r = np.sqrt(deflection_2d[0, 0] ** 2 + deflection_2d[0, 1] ** 2) assert deflection_r == pytest.approx(0.006034319441107217, 1.0e-8) diff --git a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py index c26204ced..b2aff11ca 100644 --- a/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py +++ b/test_autogalaxy/profiles/mass/dark/test_nfw_mcr.py @@ -217,6 +217,7 @@ def test__same_as_above_but_generalized_elliptical(): assert (deflections_ludlow == deflections).all() + def test__same_as_above_but_cored_nfw(): from autogalaxy.cosmology.model import FlatLambdaCDM @@ -224,12 +225,12 @@ def test__same_as_above_but_cored_nfw(): cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3) mp = ag.mp.cNFWMCRLudlowSph( - centre=(1.0, 2.0), - mass_at_200=1.0e9, - f_c=0.01, - redshift_object=0.6, - redshift_source=2.5, - ) + centre=(1.0, 2.0), + mass_at_200=1.0e9, + f_c=0.01, + redshift_object=0.6, + redshift_source=2.5, + ) mass_at_200_via_mass = mp.mass_at_200_solar_masses( redshift_object=0.6, redshift_source=2.5, cosmology=cosmology @@ -239,17 +240,17 @@ def test__same_as_above_but_cored_nfw(): ) cnfw_kappa_s = ag.mp.cNFWSph( - centre=(1.0, 2.0), - kappa_s=mp.kappa_s, - scale_radius=mp.scale_radius, - core_radius=mp.core_radius, - ) + centre=(1.0, 2.0), + kappa_s=mp.kappa_s, + scale_radius=mp.scale_radius, + core_radius=mp.core_radius, + ) mass_at_200_via_kappa_s = cnfw_kappa_s.mass_at_200_solar_masses( - redshift_object=0.6, redshift_source=2.5, cosmology=cosmology + redshift_object=0.6, redshift_source=2.5, cosmology=cosmology ) concentration_via_kappa_s = cnfw_kappa_s.concentration( - redshift_profile=0.6, redshift_source=2.5, cosmology=cosmology + redshift_profile=0.6, redshift_source=2.5, cosmology=cosmology ) # We are using the NFWTruncatedSph to check the mass gives a consistent kappa_s, given certain radii. diff --git a/test_autogalaxy/profiles/mass/stellar/test_gaussian.py b/test_autogalaxy/profiles/mass/stellar/test_gaussian.py index 7a91b2a79..af8b26e5b 100644 --- a/test_autogalaxy/profiles/mass/stellar/test_gaussian.py +++ b/test_autogalaxy/profiles/mass/stellar/test_gaussian.py @@ -67,6 +67,7 @@ def test__deflections_2d_via_analytic_from(): assert deflections[0, 0] == pytest.approx(1.10812, 1.0e-4) assert deflections[0, 1] == pytest.approx(0.35467, 1.0e-4) + @pytest.mark.skip(reason="Not JAX compatible") def test__deflections_2d_via_integral_from(): mp = ag.mp.Gaussian( @@ -137,6 +138,7 @@ def test__deflections_2d_via_integral_from(): assert deflections == pytest.approx(deflections_via_analytic.array, 1.0e-3) + @pytest.mark.skip(reason="Not JAX compatible") def test__deflections_yx_2d_from(): mp = ag.mp.Gaussian() From dc8a0c596a5a489cda3746439586c8371d72b404 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 10:15:49 +0000 Subject: [PATCH 12/26] added critical_density method --- autogalaxy/analysis/analysis/analysis.py | 2 +- autogalaxy/cosmology/model.py | 35 +++++++++++++++++++ autogalaxy/profiles/mass/dark/abstract.py | 10 +++--- autogalaxy/profiles/mass/dark/mcr_util.py | 4 +-- .../profiles/mass/dark/nfw_truncated.py | 2 +- .../mass/dark/nfw_virial_mass_conc.py | 2 +- autogalaxy/profiles/mass/point/smbh.py | 2 +- 7 files changed, 46 insertions(+), 11 deletions(-) diff --git a/autogalaxy/analysis/analysis/analysis.py b/autogalaxy/analysis/analysis/analysis.py index a79a2ba8f..5d94bba4c 100644 --- a/autogalaxy/analysis/analysis/analysis.py +++ b/autogalaxy/analysis/analysis/analysis.py @@ -36,7 +36,7 @@ def __init__( The Cosmology assumed for this analysis. """ - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 self.cosmology = cosmology or Planck15() self.preloads = preloads diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 71112f6cd..e7055f0c0 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -161,6 +161,41 @@ def cosmic_average_density_solar_mass_per_kpc3_from( return rho_crit + def critical_density(self, z: float, xp=np): + """ + Critical density of the Universe at redshift z, returned in Msun / kpc^3. + + This is an xp (NumPy / JAX) drop-in for the Astropy method: + astropy.cosmology.FLRW.critical_density(z) + + Astropy returns g/cm^3, but in AutoLens you were immediately converting to: + solMass / kpc^3 + + So this function returns Msun / kpc^3 directly. + + Requires: + self.H0 (km/s/Mpc) + self.E(z, xp=xp) where E(z) = H(z)/H0 + """ + + z = xp.asarray(z) + + # H0 in km/s/Mpc -> 1/s + # 1 Mpc = 3.085677581e19 km + H0_s = xp.asarray(self.H0) / xp.asarray(3.085677581e19) + + # H(z) in 1/s + Hz = H0_s * self.E(z, xp=xp) + + # Gravitational constant G in kpc^3 / (Msun s^2) + # Start from: G = 4.30091e-6 kpc (km/s)^2 / Msun + # Convert (km/s)^2 -> (kpc/s)^2 by dividing by (km per kpc)^2 + km_per_kpc = xp.asarray(3.085677581e16) + G = xp.asarray(4.30091e-6) / (km_per_kpc**2) + + # rho_c = 3 H(z)^2 / (8 pi G) [Msun / kpc^3] + return (xp.asarray(3.0) * Hz**2) / (xp.asarray(8.0) * xp.pi * G) + def critical_surface_density_between_redshifts_from( self, redshift_0: float, diff --git a/autogalaxy/profiles/mass/dark/abstract.py b/autogalaxy/profiles/mass/dark/abstract.py index bc616eb86..9af7bf795 100644 --- a/autogalaxy/profiles/mass/dark/abstract.py +++ b/autogalaxy/profiles/mass/dark/abstract.py @@ -213,7 +213,7 @@ def rho_at_scale_radius_solar_mass_per_kpc3( The Cosmic average density is defined at the redshift of the profile. """ - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = cosmology or Planck15() @@ -237,7 +237,7 @@ def delta_concentration( cosmology: LensingCosmology = None, ): - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = cosmology or Planck15() @@ -274,7 +274,7 @@ def concentration( xp=np, ): - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = cosmology or Planck15() @@ -321,7 +321,7 @@ def radius_at_200( """ Returns `r_{200m}` for this halo in **arcseconds** """ - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = cosmology or Planck15() @@ -343,7 +343,7 @@ def mass_at_200_solar_masses( cosmology: LensingCosmology = None, xp=np, ): - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = cosmology or Planck15() diff --git a/autogalaxy/profiles/mass/dark/mcr_util.py b/autogalaxy/profiles/mass/dark/mcr_util.py index 21912eb0f..3cb7fccc4 100644 --- a/autogalaxy/profiles/mass/dark/mcr_util.py +++ b/autogalaxy/profiles/mass/dark/mcr_util.py @@ -12,7 +12,7 @@ def kappa_s_and_scale_radius_for_duffy(mass_at_200, redshift_object, redshift_so from astropy import units - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = Planck15() @@ -70,7 +70,7 @@ def _ludlow16_cosmology_callback( from astropy import units from colossus.cosmology import cosmology as col_cosmology from colossus.halo.concentration import concentration as col_concentration - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 # ----------------------- # Colossus cosmology diff --git a/autogalaxy/profiles/mass/dark/nfw_truncated.py b/autogalaxy/profiles/mass/dark/nfw_truncated.py index 154bda1ec..6894f3f43 100644 --- a/autogalaxy/profiles/mass/dark/nfw_truncated.py +++ b/autogalaxy/profiles/mass/dark/nfw_truncated.py @@ -114,7 +114,7 @@ def mass_at_truncation_radius_solar_mass( cosmology: LensingCosmology = None, xp=np, ): - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = cosmology or Planck15() diff --git a/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py index d11cf8acd..ce86ebf22 100644 --- a/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py +++ b/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py @@ -10,7 +10,7 @@ def kappa_s_and_scale_radius( ): from astropy import units - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = Planck15() diff --git a/autogalaxy/profiles/mass/point/smbh.py b/autogalaxy/profiles/mass/point/smbh.py index 40fc054ec..9fde58255 100644 --- a/autogalaxy/profiles/mass/point/smbh.py +++ b/autogalaxy/profiles/mass/point/smbh.py @@ -29,7 +29,7 @@ def __init__( redshift_source The redshift of the source galaxy, which is used to convert the mass of the SMBH to an Einstein radius. """ - from autogalaxy.cosmology.wrap import Planck15 + from autogalaxy.cosmology.model import Planck15 cosmology = Planck15() From 3c0b102afc844466208cb47628c201b2e3bb6e67 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 10:31:58 +0000 Subject: [PATCH 13/26] remove astropy units in dark calculations for mcr --- autogalaxy/profiles/mass/dark/mcr_util.py | 54 +++++++++++------------ 1 file changed, 27 insertions(+), 27 deletions(-) diff --git a/autogalaxy/profiles/mass/dark/mcr_util.py b/autogalaxy/profiles/mass/dark/mcr_util.py index 3cb7fccc4..efb1c192b 100644 --- a/autogalaxy/profiles/mass/dark/mcr_util.py +++ b/autogalaxy/profiles/mass/dark/mcr_util.py @@ -16,29 +16,30 @@ def kappa_s_and_scale_radius_for_duffy(mass_at_200, redshift_object, redshift_so cosmology = Planck15() - cosmic_average_density = ( - cosmology.critical_density(redshift_object).to(units.solMass / units.kpc**3) - ).value + # Msun / kpc^3 (no units conversion needed) + cosmic_average_density = cosmology.critical_density(redshift_object, xp=np) + # Msun / kpc^2 critical_surface_density = ( cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_object, redshift_1=redshift_source + redshift_0=redshift_object, + redshift_1=redshift_source, + xp=np, ) ) - kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object) + # kpc / arcsec + kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object, xp=np) + # r200 in kpc radius_at_200 = ( mass_at_200 / (200.0 * cosmic_average_density * (4.0 * np.pi / 3.0)) - ) ** ( - 1.0 / 3.0 - ) # r200 - coefficient = 5.71 * (1.0 + redshift_object) ** ( - -0.47 - ) # The coefficient of Duffy mass-concentration (Duffy+2008) - concentration = coefficient * (mass_at_200 / 2.952465309e12) ** ( - -0.084 - ) # mass-concentration relation. (Duffy+2008) + ) ** (1.0 / 3.0) + + # Duffy+2008 mass–concentration (as in your code) + coefficient = 5.71 * (1.0 + redshift_object) ** (-0.47) + concentration = coefficient * (mass_at_200 / 2.952465309e12) ** (-0.084) + de_c = ( 200.0 / 3.0 @@ -46,12 +47,12 @@ def kappa_s_and_scale_radius_for_duffy(mass_at_200, redshift_object, redshift_so concentration**3 / (np.log(1.0 + concentration) - concentration / (1.0 + concentration)) ) - ) # rho_c + ) - scale_radius_kpc = radius_at_200 / concentration # scale radius in kpc - rho_s = cosmic_average_density * de_c # rho_s - kappa_s = rho_s * scale_radius_kpc / critical_surface_density # kappa_s - scale_radius = scale_radius_kpc / kpc_per_arcsec # scale radius in arcsec + scale_radius_kpc = radius_at_200 / concentration + rho_s = cosmic_average_density * de_c # Msun / kpc^3 + kappa_s = rho_s * scale_radius_kpc / critical_surface_density # dimensionless + scale_radius = scale_radius_kpc / kpc_per_arcsec # arcsec return kappa_s, scale_radius, radius_at_200 @@ -67,7 +68,6 @@ def _ludlow16_cosmology_callback( """ import numpy as np - from astropy import units from colossus.cosmology import cosmology as col_cosmology from colossus.halo.concentration import concentration as col_concentration from autogalaxy.cosmology.model import Planck15 @@ -86,24 +86,24 @@ def _ludlow16_cosmology_callback( ) # ----------------------- - # Astropy cosmology + # AutoGalaxy cosmology (no astropy.units) # ----------------------- cosmology = Planck15() - cosmic_average_density = ( - cosmology.critical_density(redshift_object) - .to(units.solMass / units.kpc**3) - .value - ) + # Msun / kpc^3 (your xp drop-in should return this directly) + cosmic_average_density = cosmology.critical_density(redshift_object, xp=np) + # Msun / kpc^2 critical_surface_density = ( cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( redshift_0=redshift_object, redshift_1=redshift_source, + xp=np, ) ) - kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object) + # kpc / arcsec + kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object, xp=np) return ( np.float64(concentration), From 39d3c9c2f6e473a315f1dccf019da2ba325cb49b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 14:10:18 +0000 Subject: [PATCH 14/26] lots of fixes to make tests in test__values_of_quantities_for_real_cosmology pass --- autogalaxy/cosmology/model.py | 216 ++++++++++++------ .../mass/dark/gnfw_virial_mass_conc.py | 44 ++-- .../mass/dark/gnfw_virial_mass_gnfw_conc.py | 48 ++-- .../profiles/mass/dark/test_abstract.py | 40 ++-- 4 files changed, 203 insertions(+), 145 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index e7055f0c0..937866e3f 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -17,7 +17,7 @@ def arcsec_per_kpc_proper(self, z: float, xp=np) -> float: Proper transverse distance uses the angular diameter distance D_A(z): - arcsec_per_kpc_proper = 206265 / D_A(z) + arcsec_per_kpc_proper = 206264.806247 / D_A(z) where D_A(z) is in kpc. """ @@ -26,7 +26,7 @@ def arcsec_per_kpc_proper(self, z: float, xp=np) -> float: 0.0, z, xp=xp ) - return xp.asarray(206265.0) / angular_diameter_distance_kpc + return xp.asarray(206264.806247) / angular_diameter_distance_kpc def kpc_proper_per_arcsec(self, z: float, xp=np) -> float: """ @@ -34,14 +34,14 @@ def kpc_proper_per_arcsec(self, z: float, xp=np) -> float: This matches the inverse of astropy.cosmology.arcsec_per_kpc_proper: - kpc_proper_per_arcsec = D_A(z) / 206265 + kpc_proper_per_arcsec = D_A(z) / 206264.806247 """ angular_diameter_distance_kpc = self.angular_diameter_distance_kpc_z1z2( 0.0, z, xp=xp ) - return angular_diameter_distance_kpc / xp.asarray(206265.0) + return angular_diameter_distance_kpc / xp.asarray(206264.806247) def arcsec_per_kpc_from(self, redshift: float, xp=np) -> float: """ @@ -234,34 +234,32 @@ def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( JAX/NumPy compatible via `xp` (pass `jax.numpy` as xp). """ - # Speed of light in kpc / s - # c = 299792.458 km/s, 1 kpc = 3.085677581e16 km - c_kpc_s = xp.asarray(299792.458) / xp.asarray(3.085677581e16) + # Distances in kpc + D_l_kpc = self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_0, xp=xp) + D_s_kpc = self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1, xp=xp) + D_ls_kpc = self.angular_diameter_distance_between_redshifts_in_kpc_from( + redshift_0=redshift_0, redshift_1=redshift_1, xp=xp + ) - # Gravitational constant in kpc^3 / (Msun s^2) - # Start from G = 4.30091e-6 kpc (km/s)^2 / Msun and convert (km/s)^2 -> (kpc/s)^2 - G_kpc3_Msun_s2 = xp.asarray(4.30091e-6) / xp.asarray((3.085677581e16) ** 2) + # kpc -> m + kpc_to_m = xp.asarray(3.085677581491367e19) + D_l = D_l_kpc * kpc_to_m + D_s = D_s_kpc * kpc_to_m + D_ls = D_ls_kpc * kpc_to_m - const = (c_kpc_s**2) / (xp.asarray(4.0) * xp.pi * G_kpc3_Msun_s2) + # SI constants + c = xp.asarray(299792458.0) # m/s + G = xp.asarray(6.67430e-11) # m^3/(kg s^2) + Msun = xp.asarray(1.98847e30) # kg - D_l = self.angular_diameter_distance_to_earth_in_kpc_from( - redshift=redshift_0, xp=xp - ) - D_s = self.angular_diameter_distance_to_earth_in_kpc_from( - redshift=redshift_1, xp=xp - ) - D_ls = self.angular_diameter_distance_between_redshifts_in_kpc_from( - redshift_0=redshift_0, redshift_1=redshift_1, xp=xp - ) + # Sigma_crit in kg / m^2 + prefac = (c * c) / (xp.asarray(4.0) * xp.pi * G) + sigma_SI = prefac * D_s / (D_l * D_ls) - # Handle z_s == z_l (or any case D_ls=0) safely for JAX: - # In lensing usage, caller should ensure z_s > z_l, but this avoids NaNs in edge cases. - return xp.where( - D_ls == xp.asarray(0.0), - xp.asarray(np.inf), - const * D_s / (D_l * D_ls), - ) + # kg/m^2 -> Msun/kpc^2 + sigma = sigma_SI * (kpc_to_m * kpc_to_m) / Msun + return xp.where(D_ls_kpc == xp.asarray(0.0), xp.asarray(np.inf), sigma) def scaling_factor_between_redshifts_from( self, redshift_0: float, @@ -461,24 +459,69 @@ def _simpson_1d(y, x, xp=np): return (h / 3.0) * s + def _m_nu_sum_eV(self, xp=np): + m = xp.asarray(getattr(self, "m_nu", 0.0)) + if getattr(m, "ndim", 0) == 0: + # scalar means per-species mass (astropy convention) + n = int(np.floor(float(getattr(self, "Neff", 3.046)))) + return m * xp.asarray(float(n)) + return xp.sum(m) + + import numpy as np + + def _radiation_and_massive_nu_densities(self, h, xp=np): + """ + Returns (Or0, Onu_m0) where: + - Or0 is photons + *massless* neutrino radiation (dimensionless today) + - Onu_m0 is massive neutrino density today treated as matter-like (dimensionless today) + + Uses a simple split: N_eff_massless = Neff - n_massive. + """ + Tcmb = xp.asarray(self.Tcmb0) + Ogamma_h2 = xp.asarray(2.469e-5) * (Tcmb / xp.asarray(2.7255)) ** 4 + Ogamma0 = Ogamma_h2 / (h * h) + + # Massive neutrinos: Omega_nu h^2 ≈ sum(m_nu)/93.14 eV + m_nu = getattr(self, "m_nu", 0.0) + m_arr = xp.asarray(m_nu) + if getattr(m_arr, "ndim", 0) == 0: + # interpret scalar as total sum (your current convention) + m_sum = m_arr + n_massive = xp.asarray(1.0) if m_arr > 0 else xp.asarray(0.0) + else: + m_sum = xp.sum(m_arr) + n_massive = xp.sum(m_arr > xp.asarray(0.0)) + + Onu_m_h2 = m_sum / xp.asarray(93.14) + Onu_m0 = Onu_m_h2 / (h * h) + + # Only the *massless* share of Neff contributes as radiation if we're separately adding massive nu matter + Neff = xp.asarray(self.Neff) + Neff_massless = xp.maximum(Neff - n_massive, xp.asarray(0.0)) + Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff_massless + + Or0 = Ogamma0 + Onu_rad0 + return Or0, Onu_m0 + def angular_diameter_distance_kpc_z1z2( - self, - z1: float, - z2: float, - n_steps: int = 8193, # odd by default - xp=np, + self, + z1: float, + z2: float, + n_steps: int = 8193, # odd by default + xp=np, ): """ D_A(z1,z2) in kpc for flat wCDM using Simpson's rule. Includes: - photons via Tcmb0 - - *massive neutrinos* via m_nu (matter-like, Omega_nu h^2 = sum(m_nu)/93.14 eV) + - neutrinos split into: + * massive part from m_nu (approx matter-like today) + * massless part from Neff, with the number of massive species subtracted to avoid double-counting Notes: - Flat universe: Omega_k = 0 - Dark energy equation of state constant w0 - - This is designed to better match astropy's Planck15-style backgrounds. """ # Ensure odd number of samples for Simpson (safe: n_steps is a Python int) if (n_steps % 2) == 0: @@ -496,37 +539,45 @@ def angular_diameter_distance_kpc_z1z2( Om0 = xp.asarray(self.Om0) w0 = xp.asarray(self.w0) - # ---- Photon radiation density today ---- - # Omega_gamma * h^2 ≈ 2.469e-5 * (Tcmb/2.7255)^4 - Tcmb = xp.asarray(self.Tcmb0) + # ---- photons ---- + Tcmb = xp.asarray(getattr(self, "Tcmb0", 0.0)) Ogamma_h2 = xp.asarray(2.469e-5) * (Tcmb / xp.asarray(2.7255)) ** 4 Ogamma0 = Ogamma_h2 / (h * h) - # ---- Massive neutrinos (matter-like approximation) ---- - # Omega_nu h^2 ≈ sum(m_nu)/93.14 eV + # ---- massive neutrinos (approx matter-like today) ---- + m_nu_sum = self._m_nu_sum_eV(xp=xp) + Onu_m_h2 = m_nu_sum / xp.asarray(93.14) + Onu_m0 = Onu_m_h2 / (h * h) + + # ---- massless neutrino radiation via Neff (avoid double-counting) ---- + Neff = xp.asarray(getattr(self, "Neff", 0.0)) + m_nu = getattr(self, "m_nu", 0.0) - m_nu_sum = xp.sum(xp.asarray(m_nu)) # works if m_nu is float or array-like - Onu_h2 = m_nu_sum / xp.asarray(93.14) - Onu0 = Onu_h2 / (h * h) + m_arr = xp.asarray(m_nu) - Neff = xp.asarray(self.Neff) - Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff + if getattr(m_arr, "ndim", 0) == 0: + n_massive = xp.where(m_arr > xp.asarray(0.0), xp.asarray(1.0), xp.asarray(0.0)) + else: + n_massive = xp.sum(m_arr > xp.asarray(0.0)) + + Neff_massless = xp.maximum(Neff - n_massive, xp.asarray(0.0)) + + Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff_massless Or0 = Ogamma0 + Onu_rad0 - # Flatness: Omega_de = 1 - Omega_m - Omega_r - Omega_nu - Ode0 = xp.asarray(1.0) - Om0 - Or0 - Onu0 + # ---- flatness: Omega_de ---- + Ode0 = xp.asarray(1.0) - Om0 - Or0 - Onu_m0 - def E(z): + def E_local(z): zp1 = xp.asarray(1.0) + z - # E^2 = (Om+Onu_m)(1+z)^3 + Or(1+z)^4 + Ode(1+z)^{3(1+w)} return xp.sqrt( - (Om0 + Onu0) * zp1**3 - + Or0 * zp1**4 + (Om0 + Onu_m0) * zp1 ** 3 + + Or0 * zp1 ** 4 + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) ) z_grid = xp.linspace(z1a, z2a, n_steps) - integrand = 1.0 / E(z_grid) + integrand = xp.asarray(1.0) / E_local(z_grid) integral = self._simpson_1d(integrand, z_grid, xp=xp) @@ -542,11 +593,12 @@ def E(self, z: float, xp=np): JAX/NumPy compatible via `xp` (pass `jax.numpy` as xp). - Notes on components: + Components: - Photons: Omega_gamma from Tcmb0 - - Massless neutrinos: Omega_nu_rad = Omega_gamma * 0.2271 * Neff (standard) - - Massive neutrinos (approx): Omega_nu_m h^2 = sum(m_nu)/93.14 eV, treated as matter-like (1+z)^3 - (If you want to exactly match astropy for massive neutrinos, that’s more complex.) + - Neutrinos: + * massive part from m_nu (approx matter-like today) + * massless part from Neff, but we subtract the number of massive species to avoid double-counting + - Dark energy: constant w0, with Omega_de set by flatness """ z = xp.asarray(z) @@ -555,26 +607,35 @@ def E(self, z: float, xp=np): h = H0 / xp.asarray(100.0) Om0 = xp.asarray(self.Om0) - w0 = xp.asarray( - getattr(self, "w0", -1.0) - ) # allow FlatLambdaCDM with w0=-1 set on the class + w0 = xp.asarray(getattr(self, "w0", -1.0)) # ---- photons ---- Tcmb = xp.asarray(getattr(self, "Tcmb0", 0.0)) Ogamma_h2 = xp.asarray(2.469e-5) * (Tcmb / xp.asarray(2.7255)) ** 4 Ogamma0 = Ogamma_h2 / (h * h) - # ---- massless neutrino radiation via Neff ---- - Neff = xp.asarray(getattr(self, "Neff", 0.0)) - Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff + # ---- massive neutrinos (approx matter-like today) ---- + # NOTE: _m_nu_sum_eV should follow YOUR convention (scalar=total, list=sum, etc.) + m_nu_sum = self._m_nu_sum_eV(xp=xp) + Onu_m_h2 = m_nu_sum / xp.asarray(93.14) + Onu_m0 = Onu_m_h2 / (h * h) - Or0 = Ogamma0 + Onu_rad0 + # ---- massless neutrino radiation via Neff (avoid double-counting) ---- + # Approx: subtract number of massive species from Neff + Neff = xp.asarray(getattr(self, "Neff", 0.0)) - # ---- massive neutrinos (approx matter-like) ---- m_nu = getattr(self, "m_nu", 0.0) - m_nu_sum = xp.sum(xp.asarray(m_nu)) # float or array-like - Onu_m_h2 = m_nu_sum / xp.asarray(93.14) - Onu_m0 = Onu_m_h2 / (h * h) + m_arr = xp.asarray(m_nu) + + if getattr(m_arr, "ndim", 0) == 0: + n_massive = xp.where(m_arr > xp.asarray(0.0), xp.asarray(1.0), xp.asarray(0.0)) + else: + n_massive = xp.sum(m_arr > xp.asarray(0.0)) + + Neff_massless = xp.maximum(Neff - n_massive, xp.asarray(0.0)) + + Onu_rad0 = Ogamma0 * xp.asarray(0.2271) * Neff_massless + Or0 = Ogamma0 + Onu_rad0 # ---- flatness: Omega_de ---- Ode0 = xp.asarray(1.0) - Om0 - Or0 - Onu_m0 @@ -582,13 +643,28 @@ def E(self, z: float, xp=np): zp1 = xp.asarray(1.0) + z Ez2 = ( - (Om0 + Onu_m0) * zp1**3 - + Or0 * zp1**4 - + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) + (Om0 + Onu_m0) * zp1 ** 3 + + Or0 * zp1 ** 4 + + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) ) return xp.sqrt(Ez2) + def Om(self, z: float, xp=np): + """ + Matter density parameter at redshift z: Omega_m(z). + + JAX / NumPy compatible via `xp`. + + For flat models using your E(z): + Om(z) = Om0 (1+z)^3 / E(z)^2 + """ + z = xp.asarray(z) + zp1 = xp.asarray(1.0) + z + + Ez = self.E(z, xp=xp) + return xp.asarray(self.Om0) * zp1**3 / (Ez * Ez) + class Planck15(FlatLambdaCDM): @@ -599,6 +675,6 @@ def __init__(self): Om0=0.3075, Tcmb0=2.7255, Neff=3.046, - m_nu=0.06, + m_nu=[0.0, 0.0, 0.06], Ob0=0.0486, ) diff --git a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py index 7a2292fbd..e8a7ba1cf 100644 --- a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py +++ b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py @@ -9,51 +9,47 @@ def kappa_s_and_scale_radius( cosmology, virial_mass, c_2, overdens, redshift_object, redshift_source, inner_slope ): - from astropy import units from scipy.integrate import quad - concentration = (2 - inner_slope) * c_2 # gNFW concentration + concentration = (2.0 - inner_slope) * c_2 # gNFW concentration - critical_density = ( - cosmology.critical_density(redshift_object).to(units.solMass / units.kpc**3) - ).value + critical_density = cosmology.critical_density(redshift_object, xp=np) # Msun / kpc^3 critical_surface_density = ( cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_object, redshift_1=redshift_source + redshift_0=redshift_object, + redshift_1=redshift_source, + xp=np, ) - ) + ) # Msun / kpc^2 - kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object) + kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object, xp=np) # kpc / arcsec if overdens == 0: - x = cosmology.Om(redshift_object) - 1 - overdens = 18 * np.pi**2 + 82 * x - 39 * x**2 # Bryan & Norman (1998) + x = cosmology.Om(redshift_object, xp=np) - 1.0 + overdens = 18.0 * np.pi**2 + 82.0 * x - 39.0 * x**2 # Bryan & Norman (1998) + # r_vir in kpc virial_radius = ( virial_mass / (overdens * critical_density * (4.0 * np.pi / 3.0)) - ) ** ( - 1.0 / 3.0 - ) # r_vir + ) ** (1.0 / 3.0) - scale_radius_kpc = ( - virial_radius / concentration - ) # scale radius of gNFW profile in kpc + # scale radius in kpc + scale_radius_kpc = virial_radius / concentration - ############################## + # Normalization integral for gNFW def integrand(r): - return (r**2 / r**inner_slope) * (1 + r / scale_radius_kpc) ** (inner_slope - 3) + return (r**2 / r**inner_slope) * (1.0 + r / scale_radius_kpc) ** (inner_slope - 3.0) de_c = ( (overdens / 3.0) * (virial_radius**3 / scale_radius_kpc**inner_slope) - / quad(integrand, 0, virial_radius)[0] - ) # rho_c - ############################## + / quad(integrand, 0.0, virial_radius)[0] + ) - rho_s = critical_density * de_c # rho_s - kappa_s = rho_s * scale_radius_kpc / critical_surface_density # kappa_s - scale_radius = scale_radius_kpc / kpc_per_arcsec # scale radius in arcsec + rho_s = critical_density * de_c # Msun / kpc^3 + kappa_s = rho_s * scale_radius_kpc / critical_surface_density # dimensionless + scale_radius = scale_radius_kpc / kpc_per_arcsec # arcsec return kappa_s, scale_radius, virial_radius, overdens diff --git a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py index e30205982..c463dc983 100644 --- a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py +++ b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py @@ -15,51 +15,49 @@ def kappa_s_and_scale_radius( redshift_source, inner_slope, ): - from astropy import units from scipy.integrate import quad - # gNFW concentration imported - - critical_density = ( - cosmology.critical_density(redshift_object).to(units.solMass / units.kpc**3) - ).value + # Msun / kpc^3 (no astropy.units) + critical_density = cosmology.critical_density(redshift_object, xp=np) + # Msun / kpc^2 critical_surface_density = ( cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_object, redshift_1=redshift_source + redshift_0=redshift_object, + redshift_1=redshift_source, + xp=np, ) ) - kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object) + # kpc / arcsec + kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object, xp=np) if overdens == 0: - x = cosmology.Om(redshift_object) - 1 - overdens = 18 * np.pi**2 + 82 * x - 39 * x**2 # Bryan & Norman (1998) + x = cosmology.Om(redshift_object, xp=np) - 1.0 + overdens = 18.0 * np.pi ** 2 + 82.0 * x - 39.0 * x ** 2 # Bryan & Norman (1998) + # r_vir in kpc virial_radius = ( - virial_mass / (overdens * critical_density * (4.0 * np.pi / 3.0)) - ) ** ( - 1.0 / 3.0 - ) # r_vir + virial_mass / (overdens * critical_density * (4.0 * np.pi / 3.0)) + ) ** (1.0 / 3.0) - scale_radius_kpc = ( - virial_radius / concentration - ) # scale radius of gNFW profile in kpc + # scale radius of gNFW in kpc + scale_radius_kpc = virial_radius / concentration ############################## def integrand(r): - return (r**2 / r**inner_slope) * (1 + r / scale_radius_kpc) ** (inner_slope - 3) + return (r ** 2 / r ** inner_slope) * (1.0 + r / scale_radius_kpc) ** (inner_slope - 3.0) de_c = ( - (overdens / 3.0) - * (virial_radius**3 / scale_radius_kpc**inner_slope) - / quad(integrand, 0, virial_radius)[0] - ) # rho_c + (overdens / 3.0) + * (virial_radius ** 3 / scale_radius_kpc ** inner_slope) + / quad(integrand, 0.0, virial_radius)[0] + ) ############################## - rho_s = critical_density * de_c # rho_s - kappa_s = rho_s * scale_radius_kpc / critical_surface_density # kappa_s - scale_radius = scale_radius_kpc / kpc_per_arcsec # scale radius in arcsec + rho_s = critical_density * de_c # Msun / kpc^3 + kappa_s = rho_s * scale_radius_kpc / critical_surface_density # dimensionless + scale_radius = scale_radius_kpc / kpc_per_arcsec # arcsec return kappa_s, scale_radius, virial_radius, overdens diff --git a/test_autogalaxy/profiles/mass/dark/test_abstract.py b/test_autogalaxy/profiles/mass/dark/test_abstract.py index e3cfecc2f..a058381bc 100644 --- a/test_autogalaxy/profiles/mass/dark/test_abstract.py +++ b/test_autogalaxy/profiles/mass/dark/test_abstract.py @@ -253,24 +253,12 @@ def test__mass_at_200__unit_conversions_work(): ) assert mass_at_200 == pytest.approx(mass_calc, 1.0e-5) - # cosmology = ag.m.MockCosmology(arcsec_per_kpc=0.5, kpc_per_arcsec=2.0, critical_surface_density=2.0, - # cosmic_average_density=1.0) - # - # radius_at_200 = mp.radius_at_200_for_units(unit_length='arcsec', redshift_galaxy=0.5, redshift_source=1.0, - # cosmology=cosmology) - # - # mass_at_200 = mp.mass_at_200(cosmology=cosmology, redshift_galaxy=0.5, redshift_source=1.0, unit_length='arcsec', - # unit_mass='solMass') - # - # mass_calc = 200.0 * ((4.0 / 3.0) * np.pi) * cosmology.cosmic_average_density * (radius_at_200 ** 3.0) - # assert mass_at_200 == pytest.approx(mass_calc, 1.0e-5) - def test__values_of_quantities_for_real_cosmology(): - from autogalaxy.cosmology.model import LambdaCDMWrap + from autogalaxy.cosmology.model import FlatLambdaCDM - cosmology = LambdaCDMWrap(H0=70.0, Om0=0.3, Ode0=0.7) + cosmology = FlatLambdaCDM(H0=70.0, Om0=0.3, m_nu=0.06) mp = ag.mp.NFWTruncatedSph(kappa_s=0.5, scale_radius=5.0, truncation_radius=10.0) @@ -313,12 +301,12 @@ def test__values_of_quantities_for_real_cosmology(): cosmology=cosmology, ) - assert rho == pytest.approx(29086405.474155396, 1.0e-4) - assert delta_concentration == pytest.approx(213881.71990418772, 1.0e-4) - assert concentration == pytest.approx(18.67527773217461, 1.0e-4) - assert radius_at_200 == pytest.approx(93.37638866087305, 1.0e-4) - assert mass_at_200 == pytest.approx(27641476612470.0, 1.0e-4) - assert mass_at_truncation_radius == pytest.approx(14871675423873.854, 1.0e-4) + assert rho == pytest.approx(29161077.6138, rel=1e-4) + assert delta_concentration == pytest.approx(214430.445767, 1.0e-4) + assert concentration == pytest.approx(18.694005966122, 1.0e-4) + assert radius_at_200 == pytest.approx(93.4700298306, 1.0e-4) + assert mass_at_200 == pytest.approx(27664107754813.824, 1.0e-4) + assert mass_at_truncation_radius == pytest.approx(14883851437772.34, 1.0e-4) rho = mp.rho_at_scale_radius_solar_mass_per_kpc3( redshift_object=0.6, redshift_source=2.5, cosmology=cosmology @@ -359,9 +347,9 @@ def test__values_of_quantities_for_real_cosmology(): cosmology=cosmology, ) - assert rho == pytest.approx(29086405.474155396, 1.0e-4) - assert delta_concentration == pytest.approx(110527.8022275522, 1.0e-4) - assert concentration == pytest.approx(14.394455391541, 1.0e-4) - assert radius_at_200 == pytest.approx(71.97227695770, 1.0e-4) - assert mass_at_200 == pytest.approx(24493332582335.38, 1.0e-4) - assert mass_at_truncation_radius == pytest.approx(13177910041504.6, 1.0e-4) + assert rho == pytest.approx(29161077.6138651, 1.0e-4) + assert delta_concentration == pytest.approx(110450.279912950, 1.0e-4) + assert concentration == pytest.approx(14.3904385787937, 1.0e-4) + assert radius_at_200 == pytest.approx(71.9521928939688, 1.0e-4) + assert mass_at_200 == pytest.approx(24499163428482.89, 1.0e-4) + assert mass_at_truncation_radius == pytest.approx(13181047155073.828, 1.0e-4) From 3b70800d290893021457fcaf5656df9128f5769f Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 14:12:02 +0000 Subject: [PATCH 15/26] fix skip --- test_autogalaxy/profiles/mass/stellar/test_gaussian.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/test_autogalaxy/profiles/mass/stellar/test_gaussian.py b/test_autogalaxy/profiles/mass/stellar/test_gaussian.py index af8b26e5b..af9d24ce5 100644 --- a/test_autogalaxy/profiles/mass/stellar/test_gaussian.py +++ b/test_autogalaxy/profiles/mass/stellar/test_gaussian.py @@ -68,7 +68,6 @@ def test__deflections_2d_via_analytic_from(): assert deflections[0, 1] == pytest.approx(0.35467, 1.0e-4) -@pytest.mark.skip(reason="Not JAX compatible") def test__deflections_2d_via_integral_from(): mp = ag.mp.Gaussian( centre=(0.0, 0.0), @@ -139,7 +138,6 @@ def test__deflections_2d_via_integral_from(): assert deflections == pytest.approx(deflections_via_analytic.array, 1.0e-3) -@pytest.mark.skip(reason="Not JAX compatible") def test__deflections_yx_2d_from(): mp = ag.mp.Gaussian() From 7548dc95f4250440bdf6271f6e6d3302926c7d9d Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 14:16:03 +0000 Subject: [PATCH 16/26] more units removal --- autogalaxy/convert.py | 11 +-- .../ellipse/plot/fit_ellipse_plot_util.py | 3 +- autogalaxy/profiles/mass/dark/mcr_util.py | 3 - .../mass/dark/nfw_virial_mass_conc.py | 67 +++++++++++++------ 4 files changed, 57 insertions(+), 27 deletions(-) diff --git a/autogalaxy/convert.py b/autogalaxy/convert.py index 5bda71120..ff1c2a17d 100644 --- a/autogalaxy/convert.py +++ b/autogalaxy/convert.py @@ -331,9 +331,12 @@ def multipole_comps_from( ------- The multipole component parameters. """ - from astropy import units + # Degrees -> radians conversion + deg_to_rad = xp.asarray(np.pi / 180.0) - multipole_comp_0 = k_m * xp.sin(phi_m * float(m) * units.deg.to(units.rad)) - multipole_comp_1 = k_m * xp.cos(phi_m * float(m) * units.deg.to(units.rad)) + angle = phi_m * float(m) * deg_to_rad - return (multipole_comp_0, multipole_comp_1) + multipole_comp_0 = k_m * xp.sin(angle) + multipole_comp_1 = k_m * xp.cos(angle) + + return multipole_comp_0, multipole_comp_1 diff --git a/autogalaxy/ellipse/plot/fit_ellipse_plot_util.py b/autogalaxy/ellipse/plot/fit_ellipse_plot_util.py index 7455a96d2..cea3c97ab 100644 --- a/autogalaxy/ellipse/plot/fit_ellipse_plot_util.py +++ b/autogalaxy/ellipse/plot/fit_ellipse_plot_util.py @@ -1,4 +1,3 @@ -from astropy import units import itertools import matplotlib.pyplot as plt import numpy as np @@ -6,6 +5,8 @@ def plot_ellipse_residuals(array, fit_list, colors, output, for_subplot: bool = False): + from astropy import units + color = itertools.cycle(colors) if for_subplot: diff --git a/autogalaxy/profiles/mass/dark/mcr_util.py b/autogalaxy/profiles/mass/dark/mcr_util.py index efb1c192b..1fe0e2c31 100644 --- a/autogalaxy/profiles/mass/dark/mcr_util.py +++ b/autogalaxy/profiles/mass/dark/mcr_util.py @@ -9,9 +9,6 @@ def kappa_s_and_scale_radius_for_duffy(mass_at_200, redshift_object, redshift_so Interprets mass as *`M_{200c}`*, not `M_{200m}`. """ - - from astropy import units - from autogalaxy.cosmology.model import Planck15 cosmology = Planck15() diff --git a/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py index ce86ebf22..d82541df2 100644 --- a/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py +++ b/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py @@ -1,36 +1,57 @@ from typing import Tuple +import numpy as np -from autogalaxy.profiles.mass.dark.nfw import NFWSph -import numpy as np +from autogalaxy.profiles.mass.dark.nfw import NFWSph +from autogalaxy.cosmology.model import Planck15 def kappa_s_and_scale_radius( - virial_mass, concentration, virial_overdens, redshift_object, redshift_source + virial_mass, + concentration, + virial_overdens, + redshift_object, + redshift_source, ): + """ + Computes kappa_s and scale radius for an NFW halo given virial mass and concentration. - from astropy import units - from autogalaxy.cosmology.model import Planck15 + No astropy.units required. + + Assumes: + - virial_mass in Msun + - critical_density in Msun / kpc^3 + - critical_surface_density in Msun / kpc^2 + - kpc_per_arcsec in kpc / arcsec + """ cosmology = Planck15() - cosmic_average_density = ( - cosmology.critical_density(redshift_object).to(units.solMass / units.kpc**3) - ).value + # Msun / kpc^3 + cosmic_average_density = cosmology.critical_density(redshift_object, xp=np) + # Msun / kpc^2 critical_surface_density = ( cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( - redshift_0=redshift_object, redshift_1=redshift_source + redshift_0=redshift_object, + redshift_1=redshift_source, + xp=np, ) ) - kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object) + # kpc / arcsec + kpc_per_arcsec = cosmology.kpc_per_arcsec_from( + redshift=redshift_object, + xp=np, + ) + # Virial radius r_vir in kpc virial_radius = ( - virial_mass / (virial_overdens * cosmic_average_density * (4.0 * np.pi / 3.0)) - ) ** ( - 1.0 / 3.0 - ) # r_vir + virial_mass + / (virial_overdens * cosmic_average_density * (4.0 * np.pi / 3.0)) + ) ** (1.0 / 3.0) + + # Characteristic overdensity factor de_c = ( virial_overdens / 3.0 @@ -38,16 +59,24 @@ def kappa_s_and_scale_radius( concentration**3 / (np.log(1.0 + concentration) - concentration / (1.0 + concentration)) ) - ) # rho_c + ) + + # Scale radius in kpc + scale_radius_kpc = virial_radius / concentration - scale_radius_kpc = virial_radius / concentration # scale radius in kpc - rho_s = cosmic_average_density * de_c # rho_s - kappa_s = rho_s * scale_radius_kpc / critical_surface_density # kappa_s - scale_radius = scale_radius_kpc / kpc_per_arcsec # scale radius in arcsec + # Characteristic density rho_s in Msun / kpc^3 + rho_s = cosmic_average_density * de_c + + # Dimensionless lensing normalization + kappa_s = rho_s * scale_radius_kpc / critical_surface_density + + # Scale radius in arcsec + scale_radius = scale_radius_kpc / kpc_per_arcsec return kappa_s, scale_radius, virial_radius + class NFWVirialMassConcSph(NFWSph): def __init__( self, From 43bedee61d6173d3a83446e83229351dee9d7ba9 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 14:18:06 +0000 Subject: [PATCH 17/26] remove another multiupole comps --- autogalaxy/profiles/mass/total/power_law_multipole.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/autogalaxy/profiles/mass/total/power_law_multipole.py b/autogalaxy/profiles/mass/total/power_law_multipole.py index c0f652af3..63a4f4653 100644 --- a/autogalaxy/profiles/mass/total/power_law_multipole.py +++ b/autogalaxy/profiles/mass/total/power_law_multipole.py @@ -147,12 +147,10 @@ def k_m_and_angle_m_from(self, xp=np) -> Tuple[float, float]: angle_m The multipole orientation angle in radians. """ - from astropy import units - k_m, angle_m = convert.multipole_k_m_and_phi_m_from( multipole_comps=self.multipole_comps, m=self.m, xp=xp ) - angle_m *= units.deg.to(units.rad) + angle_m *= xp.asarray(np.pi / 180.0) return k_m, angle_m From 37d2e5919d6cad8cf6d578aa661a31c6de6f7e53 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 14:19:36 +0000 Subject: [PATCH 18/26] remove astropy from mock cosmology --- autogalaxy/util/mock/mock_cosmology.py | 27 ++++++++++++++++++----- test_autogalaxy/imaging/test_simulator.py | 6 ++++- 2 files changed, 26 insertions(+), 7 deletions(-) diff --git a/autogalaxy/util/mock/mock_cosmology.py b/autogalaxy/util/mock/mock_cosmology.py index d41f918fb..ac13b9638 100644 --- a/autogalaxy/util/mock/mock_cosmology.py +++ b/autogalaxy/util/mock/mock_cosmology.py @@ -34,13 +34,28 @@ def kpc_per_arcsec_proper(self, z): def angular_diameter_distance(self, z): return Value(value=1.0) - def angular_diameter_distance_z1z2(self, z1, z2): - from astropy import constants + def angular_diameter_distance_z1z2(self, z1, z2, xp=np): + """ + Replacement for the astropy-constants version. - const = constants.c.to("kpc / s") ** 2.0 / ( - 4 * math.pi * constants.G.to("kpc3 / (solMass s2)") - ) - return Value(value=self.critical_surface_density * const.value) + Computes the factor: + + const = c^2 / (4*pi*G) + + in units of Msun / kpc. + + No astropy.constants required. + """ + + # Speed of light in kpc/s + c_kpc_s = xp.asarray(299792.458) / xp.asarray(3.085677581e16) + + # Gravitational constant in kpc^3 / (Msun s^2) + G_kpc3_Msun_s2 = xp.asarray(4.30091e-6) / xp.asarray((3.085677581e16) ** 2) + + const = (c_kpc_s ** 2) / (xp.asarray(4.0) * xp.pi * G_kpc3_Msun_s2) + + return Value(value=self.critical_surface_density * const) def critical_density(self, z): return Value(value=self.cosmic_average_density) diff --git a/test_autogalaxy/imaging/test_simulator.py b/test_autogalaxy/imaging/test_simulator.py index 64a64707c..1b9aacd9e 100644 --- a/test_autogalaxy/imaging/test_simulator.py +++ b/test_autogalaxy/imaging/test_simulator.py @@ -1,4 +1,3 @@ -from astropy.io import fits import numpy as np import os from os import path @@ -9,6 +8,9 @@ def create_fits(fits_path, array): + + from astropy.io import fits + file_dir = os.path.split(fits_path)[0] if not os.path.exists(file_dir): @@ -61,6 +63,8 @@ def test__from_fits__all_imaging_data_structures_are_flipped_for_ds9(): overwrite=True, ) + from astropy.io import fits + hdu_list = fits.open(image_path) image = np.array(hdu_list[0].data).astype("float64") assert (image == np.array([[1.0, 0.0], [0.0, 0.0]])).all() From 9151b992e992592d4dfc069eea92d52a4eb21507 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 14:23:22 +0000 Subject: [PATCH 19/26] black --- autogalaxy/cosmology/model.py | 37 ++++++++++++------- .../mass/dark/gnfw_virial_mass_conc.py | 12 ++++-- .../mass/dark/gnfw_virial_mass_gnfw_conc.py | 16 ++++---- autogalaxy/profiles/mass/dark/mcr_util.py | 4 +- .../mass/dark/nfw_virial_mass_conc.py | 4 +- autogalaxy/util/mock/mock_cosmology.py | 3 +- 6 files changed, 46 insertions(+), 30 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 937866e3f..0d1322e0a 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -235,8 +235,12 @@ def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( """ # Distances in kpc - D_l_kpc = self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_0, xp=xp) - D_s_kpc = self.angular_diameter_distance_to_earth_in_kpc_from(redshift=redshift_1, xp=xp) + D_l_kpc = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_0, xp=xp + ) + D_s_kpc = self.angular_diameter_distance_to_earth_in_kpc_from( + redshift=redshift_1, xp=xp + ) D_ls_kpc = self.angular_diameter_distance_between_redshifts_in_kpc_from( redshift_0=redshift_0, redshift_1=redshift_1, xp=xp ) @@ -260,6 +264,7 @@ def critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( sigma = sigma_SI * (kpc_to_m * kpc_to_m) / Msun return xp.where(D_ls_kpc == xp.asarray(0.0), xp.asarray(np.inf), sigma) + def scaling_factor_between_redshifts_from( self, redshift_0: float, @@ -504,11 +509,11 @@ def _radiation_and_massive_nu_densities(self, h, xp=np): return Or0, Onu_m0 def angular_diameter_distance_kpc_z1z2( - self, - z1: float, - z2: float, - n_steps: int = 8193, # odd by default - xp=np, + self, + z1: float, + z2: float, + n_steps: int = 8193, # odd by default + xp=np, ): """ D_A(z1,z2) in kpc for flat wCDM using Simpson's rule. @@ -556,7 +561,9 @@ def angular_diameter_distance_kpc_z1z2( m_arr = xp.asarray(m_nu) if getattr(m_arr, "ndim", 0) == 0: - n_massive = xp.where(m_arr > xp.asarray(0.0), xp.asarray(1.0), xp.asarray(0.0)) + n_massive = xp.where( + m_arr > xp.asarray(0.0), xp.asarray(1.0), xp.asarray(0.0) + ) else: n_massive = xp.sum(m_arr > xp.asarray(0.0)) @@ -571,8 +578,8 @@ def angular_diameter_distance_kpc_z1z2( def E_local(z): zp1 = xp.asarray(1.0) + z return xp.sqrt( - (Om0 + Onu_m0) * zp1 ** 3 - + Or0 * zp1 ** 4 + (Om0 + Onu_m0) * zp1**3 + + Or0 * zp1**4 + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) ) @@ -628,7 +635,9 @@ def E(self, z: float, xp=np): m_arr = xp.asarray(m_nu) if getattr(m_arr, "ndim", 0) == 0: - n_massive = xp.where(m_arr > xp.asarray(0.0), xp.asarray(1.0), xp.asarray(0.0)) + n_massive = xp.where( + m_arr > xp.asarray(0.0), xp.asarray(1.0), xp.asarray(0.0) + ) else: n_massive = xp.sum(m_arr > xp.asarray(0.0)) @@ -643,9 +652,9 @@ def E(self, z: float, xp=np): zp1 = xp.asarray(1.0) + z Ez2 = ( - (Om0 + Onu_m0) * zp1 ** 3 - + Or0 * zp1 ** 4 - + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) + (Om0 + Onu_m0) * zp1**3 + + Or0 * zp1**4 + + Ode0 * zp1 ** (xp.asarray(3.0) * (xp.asarray(1.0) + w0)) ) return xp.sqrt(Ez2) diff --git a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py index e8a7ba1cf..cae42f239 100644 --- a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py +++ b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_conc.py @@ -13,7 +13,9 @@ def kappa_s_and_scale_radius( concentration = (2.0 - inner_slope) * c_2 # gNFW concentration - critical_density = cosmology.critical_density(redshift_object, xp=np) # Msun / kpc^3 + critical_density = cosmology.critical_density( + redshift_object, xp=np + ) # Msun / kpc^3 critical_surface_density = ( cosmology.critical_surface_density_between_redshifts_solar_mass_per_kpc2_from( @@ -23,7 +25,9 @@ def kappa_s_and_scale_radius( ) ) # Msun / kpc^2 - kpc_per_arcsec = cosmology.kpc_per_arcsec_from(redshift=redshift_object, xp=np) # kpc / arcsec + kpc_per_arcsec = cosmology.kpc_per_arcsec_from( + redshift=redshift_object, xp=np + ) # kpc / arcsec if overdens == 0: x = cosmology.Om(redshift_object, xp=np) - 1.0 @@ -39,7 +43,9 @@ def kappa_s_and_scale_radius( # Normalization integral for gNFW def integrand(r): - return (r**2 / r**inner_slope) * (1.0 + r / scale_radius_kpc) ** (inner_slope - 3.0) + return (r**2 / r**inner_slope) * (1.0 + r / scale_radius_kpc) ** ( + inner_slope - 3.0 + ) de_c = ( (overdens / 3.0) diff --git a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py index c463dc983..c07d290b3 100644 --- a/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py +++ b/autogalaxy/profiles/mass/dark/gnfw_virial_mass_gnfw_conc.py @@ -34,24 +34,26 @@ def kappa_s_and_scale_radius( if overdens == 0: x = cosmology.Om(redshift_object, xp=np) - 1.0 - overdens = 18.0 * np.pi ** 2 + 82.0 * x - 39.0 * x ** 2 # Bryan & Norman (1998) + overdens = 18.0 * np.pi**2 + 82.0 * x - 39.0 * x**2 # Bryan & Norman (1998) # r_vir in kpc virial_radius = ( - virial_mass / (overdens * critical_density * (4.0 * np.pi / 3.0)) - ) ** (1.0 / 3.0) + virial_mass / (overdens * critical_density * (4.0 * np.pi / 3.0)) + ) ** (1.0 / 3.0) # scale radius of gNFW in kpc scale_radius_kpc = virial_radius / concentration ############################## def integrand(r): - return (r ** 2 / r ** inner_slope) * (1.0 + r / scale_radius_kpc) ** (inner_slope - 3.0) + return (r**2 / r**inner_slope) * (1.0 + r / scale_radius_kpc) ** ( + inner_slope - 3.0 + ) de_c = ( - (overdens / 3.0) - * (virial_radius ** 3 / scale_radius_kpc ** inner_slope) - / quad(integrand, 0.0, virial_radius)[0] + (overdens / 3.0) + * (virial_radius**3 / scale_radius_kpc**inner_slope) + / quad(integrand, 0.0, virial_radius)[0] ) ############################## diff --git a/autogalaxy/profiles/mass/dark/mcr_util.py b/autogalaxy/profiles/mass/dark/mcr_util.py index 1fe0e2c31..2ba54bcb7 100644 --- a/autogalaxy/profiles/mass/dark/mcr_util.py +++ b/autogalaxy/profiles/mass/dark/mcr_util.py @@ -47,9 +47,9 @@ def kappa_s_and_scale_radius_for_duffy(mass_at_200, redshift_object, redshift_so ) scale_radius_kpc = radius_at_200 / concentration - rho_s = cosmic_average_density * de_c # Msun / kpc^3 + rho_s = cosmic_average_density * de_c # Msun / kpc^3 kappa_s = rho_s * scale_radius_kpc / critical_surface_density # dimensionless - scale_radius = scale_radius_kpc / kpc_per_arcsec # arcsec + scale_radius = scale_radius_kpc / kpc_per_arcsec # arcsec return kappa_s, scale_radius, radius_at_200 diff --git a/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py b/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py index d82541df2..2fc1bf511 100644 --- a/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py +++ b/autogalaxy/profiles/mass/dark/nfw_virial_mass_conc.py @@ -47,8 +47,7 @@ def kappa_s_and_scale_radius( # Virial radius r_vir in kpc virial_radius = ( - virial_mass - / (virial_overdens * cosmic_average_density * (4.0 * np.pi / 3.0)) + virial_mass / (virial_overdens * cosmic_average_density * (4.0 * np.pi / 3.0)) ) ** (1.0 / 3.0) # Characteristic overdensity factor @@ -76,7 +75,6 @@ def kappa_s_and_scale_radius( return kappa_s, scale_radius, virial_radius - class NFWVirialMassConcSph(NFWSph): def __init__( self, diff --git a/autogalaxy/util/mock/mock_cosmology.py b/autogalaxy/util/mock/mock_cosmology.py index ac13b9638..17fee4abd 100644 --- a/autogalaxy/util/mock/mock_cosmology.py +++ b/autogalaxy/util/mock/mock_cosmology.py @@ -1,4 +1,5 @@ import math +import numpy as np # Mock Cosmology # @@ -53,7 +54,7 @@ def angular_diameter_distance_z1z2(self, z1, z2, xp=np): # Gravitational constant in kpc^3 / (Msun s^2) G_kpc3_Msun_s2 = xp.asarray(4.30091e-6) / xp.asarray((3.085677581e16) ** 2) - const = (c_kpc_s ** 2) / (xp.asarray(4.0) * xp.pi * G_kpc3_Msun_s2) + const = (c_kpc_s**2) / (xp.asarray(4.0) * xp.pi * G_kpc3_Msun_s2) return Value(value=self.critical_surface_density * const) From 78e771bd90f0d86ef801c50ac1964653284a19b3 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 30 Jan 2026 14:42:28 +0000 Subject: [PATCH 20/26] minor --- autogalaxy/cosmology/model.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 0d1322e0a..f81c1eb5b 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -472,8 +472,6 @@ def _m_nu_sum_eV(self, xp=np): return m * xp.asarray(float(n)) return xp.sum(m) - import numpy as np - def _radiation_and_massive_nu_densities(self, h, xp=np): """ Returns (Or0, Onu_m0) where: From 3cf4f81f16ecb9bf9fa75060a3300d33418f40dc Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 30 Jan 2026 14:49:59 +0000 Subject: [PATCH 21/26] Update autogalaxy/cosmology/model.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- autogalaxy/cosmology/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index f81c1eb5b..343641e00 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -345,7 +345,7 @@ def velocity_dispersion_from( The velocity dispersion is given by: - velocity dispersion = (c * R_Ein * D_s) / (4 * pi * D_l * D_ls) + sigma_v = c * sqrt((R_Ein * D_s) / (4 * pi * D_l * D_ls)) c = speed of light D_s = Angular diameter distance of source redshift to earth From 50636a06ad18592641cec86cc383c1e95a9f9aab Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 30 Jan 2026 14:50:11 +0000 Subject: [PATCH 22/26] Update autogalaxy/cosmology/model.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- autogalaxy/cosmology/model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 343641e00..4cccf234e 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -340,7 +340,7 @@ def velocity_dispersion_from( self, redshift_0: float, redshift_1: float, einstein_radius: float, xp=np ) -> float: """ - For a strong lens galaxy with an Einstien radius in arcseconds, the corresponding velocity dispersion of the + For a strong lens galaxy with an Einstein radius in arcseconds, the corresponding velocity dispersion of the lens galaxy can be computed (assuming an isothermal mass distribution). The velocity dispersion is given by: From bedcf9a18ecb65884aae5388282fc27b9f551312 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 30 Jan 2026 14:53:00 +0000 Subject: [PATCH 23/26] Initial plan From 562d1b92c5979a8580d318d93fd1d553e094275c Mon Sep 17 00:00:00 2001 From: James Nightingale Date: Fri, 30 Jan 2026 14:53:02 +0000 Subject: [PATCH 24/26] Update autogalaxy/cosmology/model.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- autogalaxy/cosmology/model.py | 1 - 1 file changed, 1 deletion(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 4cccf234e..6fd78729a 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -1,5 +1,4 @@ import numpy as np -import math class LensingCosmology: From 0f13b2f03ec63ceace02454a03a7c35d800cebcc Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 30 Jan 2026 14:54:32 +0000 Subject: [PATCH 25/26] Update FlatLambdaCDM docstring to reflect JAX-compatible implementation Co-authored-by: Jammy2211 <23455639+Jammy2211@users.noreply.github.com> --- autogalaxy/cosmology/model.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 4cccf234e..0c452892e 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -399,15 +399,15 @@ def __init__( Ob0: float = 0.04897, ): """ - A wrapper for the astropy `FlatLambdaCDM` cosmology class, which allows it to be used for modeling such - that the cosmological parameters are free parameters which can be fitted for. + A JAX-compatible implementation of the FlatLambdaCDM cosmology model, which allows it to be used for + modeling such that the cosmological parameters are free parameters which can be fitted for. - The interface of this class is the same as the astropy `FlatLambdaCDM` class, it simply overwrites the - __init__ method and inherits from it in a way that ensures **PyAutoFit** can compose a model from it - without issue. + This is a custom implementation designed to work with both NumPy and JAX backends, enabling automatic + differentiation and GPU acceleration when needed. The interface and parameter naming conventions follow + the astropy `FlatLambdaCDM` class for compatibility. - The class also inherits from `LensingCosmology`, which is a class that provides additional functionality - for calculating lensing specific quantities in the cosmology. + The class inherits from `LensingCosmology`, which provides additional functionality for calculating + lensing specific quantities in the cosmology. Parameters ---------- From 1042a7193d18b6d911b5c5ddd101cf9299849a6a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 30 Jan 2026 14:56:26 +0000 Subject: [PATCH 26/26] Remove trailing whitespace from docstring Co-authored-by: Jammy2211 <23455639+Jammy2211@users.noreply.github.com> --- autogalaxy/cosmology/model.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/autogalaxy/cosmology/model.py b/autogalaxy/cosmology/model.py index 0c452892e..bfbce142b 100644 --- a/autogalaxy/cosmology/model.py +++ b/autogalaxy/cosmology/model.py @@ -399,14 +399,14 @@ def __init__( Ob0: float = 0.04897, ): """ - A JAX-compatible implementation of the FlatLambdaCDM cosmology model, which allows it to be used for + A JAX-compatible implementation of the FlatLambdaCDM cosmology model, which allows it to be used for modeling such that the cosmological parameters are free parameters which can be fitted for. - This is a custom implementation designed to work with both NumPy and JAX backends, enabling automatic - differentiation and GPU acceleration when needed. The interface and parameter naming conventions follow + This is a custom implementation designed to work with both NumPy and JAX backends, enabling automatic + differentiation and GPU acceleration when needed. The interface and parameter naming conventions follow the astropy `FlatLambdaCDM` class for compatibility. - The class inherits from `LensingCosmology`, which provides additional functionality for calculating + The class inherits from `LensingCosmology`, which provides additional functionality for calculating lensing specific quantities in the cosmology. Parameters