diff --git a/requirements.txt b/requirements.txt index e4d86509..6d29e11e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,6 +2,7 @@ astropy>=5.2.1,<6.0.0 graphviz==0.20.1 h5py>=3.8.0 levmarq_torch==0.0.1 +numba>=0.58.1,<0.59.0 numpy>=1.23.5 safetensors>=0.4.1 scipy>=1.8.0 diff --git a/src/caustics/cosmology/FlatLambdaCDM.py b/src/caustics/cosmology/FlatLambdaCDM.py new file mode 100644 index 00000000..ffc21a68 --- /dev/null +++ b/src/caustics/cosmology/FlatLambdaCDM.py @@ -0,0 +1,201 @@ +# mypy: disable-error-code="operator" +from typing import Optional + +import torch +from torch import Tensor + +from astropy.cosmology import default_cosmology +from scipy.special import hyp2f1 + +from ..utils import interp1d +from ..parametrized import unpack +from ..packed import Packed +from ..constants import c_Mpc_s, km_to_Mpc +from .base import ( + Cosmology, +) + +_h0_default = float(default_cosmology.get().h) +_critical_density_0_default = float( + default_cosmology.get().critical_density(0).to("solMass/Mpc^3").value +) +_Om0_default = float(default_cosmology.get().Om0) + +# Set up interpolator to speed up comoving distance calculations in Lambda-CDM +# cosmologies. Construct with float64 precision. +_comoving_distance_helper_x_grid = 10 ** torch.linspace(-3, 1, 500, dtype=torch.float64) +_comoving_distance_helper_y_grid = torch.as_tensor( + _comoving_distance_helper_x_grid + * hyp2f1(1 / 3, 1 / 2, 4 / 3, -(_comoving_distance_helper_x_grid**3)), + dtype=torch.float64, +) + +h0_default = torch.tensor(_h0_default) +critical_density_0_default = torch.tensor(_critical_density_0_default) +Om0_default = torch.tensor(_Om0_default) + + +class FlatLambdaCDM(Cosmology): + """ + Subclass of Cosmology representing a Flat Lambda Cold Dark Matter (LCDM) + cosmology with no radiation. + """ + + def __init__( + self, + h0: Optional[Tensor] = h0_default, + critical_density_0: Optional[Tensor] = critical_density_0_default, + Om0: Optional[Tensor] = Om0_default, + name: Optional[str] = None, + ): + """ + Initialize a new instance of the FlatLambdaCDM class. + + Parameters + ---------- + name: str + Name of the cosmology. + h0: Optional[Tensor] + Hubble constant over 100. Default is h0_default. + critical_density_0: (Optional[Tensor]) + Critical density at z=0. Default is critical_density_0_default. + Om0: Optional[Tensor] + Matter density parameter at z=0. Default is Om0_default. + """ + super().__init__(name) + + self.add_param("h0", h0) + self.add_param("critical_density_0", critical_density_0) + self.add_param("Om0", Om0) + + self._comoving_distance_helper_x_grid = _comoving_distance_helper_x_grid.to( + dtype=torch.float32 + ) + self._comoving_distance_helper_y_grid = _comoving_distance_helper_y_grid.to( + dtype=torch.float32 + ) + + def to( + self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None + ): + super().to(device, dtype) + self._comoving_distance_helper_y_grid = ( + self._comoving_distance_helper_y_grid.to(device, dtype) + ) + self._comoving_distance_helper_x_grid = ( + self._comoving_distance_helper_x_grid.to(device, dtype) + ) + + def hubble_distance(self, h0): + """ + Calculate the Hubble distance. + + Parameters + ---------- + h0: Tensor + Hubble constant. + + Returns + ------- + Tensor + Hubble distance. + """ + return c_Mpc_s / (100 * km_to_Mpc) / h0 + + @unpack + def critical_density( + self, + z: Tensor, + *args, + params: Optional["Packed"] = None, + h0: Optional[Tensor] = None, + critical_density_0: Optional[Tensor] = None, + Om0: Optional[Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """ + Calculate the critical density at redshift z. + + Parameters + ---------- + z: Tensor + Redshift. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + torch.Tensor + Critical density at redshift z. + """ + Ode0 = 1 - Om0 + return critical_density_0 * (Om0 * (1 + z) ** 3 + Ode0) # fmt: skip + + @unpack + def _comoving_distance_helper( + self, x: Tensor, *args, params: Optional["Packed"] = None, **kwargs + ) -> Tensor: + """ + Helper method for computing comoving distances. + + Parameters + ---------- + x: Tensor + Input tensor. + + Returns + ------- + Tensor + Computed comoving distances. + """ + return interp1d( + self._comoving_distance_helper_x_grid, + self._comoving_distance_helper_y_grid, + torch.atleast_1d(x), + ).reshape(x.shape) + + @unpack + def comoving_distance( + self, + z: Tensor, + *args, + params: Optional["Packed"] = None, + h0: Optional[Tensor] = None, + critical_density_0: Optional[Tensor] = None, + Om0: Optional[Tensor] = None, + **kwargs, + ) -> Tensor: + """ + Calculate the comoving distance to redshift z. + + Parameters + ---------- + z: Tensor + Redshift. + params: (Packed, optional) + Dynamic parameter container for the computation. + + Returns + ------- + Tensor + Comoving distance to redshift z. + """ + Ode0 = 1 - Om0 + ratio = (Om0 / Ode0) ** (1 / 3) + DH = self.hubble_distance(h0) + DC1z = self._comoving_distance_helper((1 + z) * ratio, params) + DC = self._comoving_distance_helper(ratio, params) + return DH * (DC1z - DC) / (Om0 ** (1 / 3) * Ode0 ** (1 / 6)) # fmt: skip + + @unpack + def transverse_comoving_distance( + self, + z: Tensor, + *args, + params: Optional["Packed"] = None, + h0: Optional[Tensor] = None, + critical_density_0: Optional[Tensor] = None, + Om0: Optional[Tensor] = None, + **kwargs, + ) -> Tensor: + return self.comoving_distance(z, params, **kwargs) diff --git a/src/caustics/cosmology/__init__.py b/src/caustics/cosmology/__init__.py new file mode 100644 index 00000000..c5146d0d --- /dev/null +++ b/src/caustics/cosmology/__init__.py @@ -0,0 +1,15 @@ +from .base import Cosmology +from .FlatLambdaCDM import ( + FlatLambdaCDM, + h0_default, + critical_density_0_default, + Om0_default, +) + +__all__ = [ + "Cosmology", + "FlatLambdaCDM", + "h0_default", + "critical_density_0_default", + "Om0_default", +] diff --git a/src/caustics/cosmology.py b/src/caustics/cosmology/base.py similarity index 54% rename from src/caustics/cosmology.py rename to src/caustics/cosmology/base.py index 29867136..31528f1d 100644 --- a/src/caustics/cosmology.py +++ b/src/caustics/cosmology/base.py @@ -3,38 +3,11 @@ from math import pi from typing import Optional -import torch -from astropy.cosmology import default_cosmology -from scipy.special import hyp2f1 from torch import Tensor -from .utils import interp1d -from .constants import G_over_c2, c_Mpc_s, km_to_Mpc -from .parametrized import Parametrized, unpack -from .packed import Packed - -__all__ = ( - "h0_default", - "critical_density_0_default", - "Om0_default", - "Cosmology", - "FlatLambdaCDM", -) - -h0_default = float(default_cosmology.get().h) -critical_density_0_default = float( - default_cosmology.get().critical_density(0).to("solMass/Mpc^3").value -) -Om0_default = float(default_cosmology.get().Om0) - -# Set up interpolator to speed up comoving distance calculations in Lambda-CDM -# cosmologies. Construct with float64 precision. -_comoving_distance_helper_x_grid = 10 ** torch.linspace(-3, 1, 500, dtype=torch.float64) -_comoving_distance_helper_y_grid = torch.as_tensor( - _comoving_distance_helper_x_grid - * hyp2f1(1 / 3, 1 / 2, 4 / 3, -(_comoving_distance_helper_x_grid**3)), - dtype=torch.float64, -) +from ..constants import G_over_c2 +from ..parametrized import Parametrized, unpack +from ..packed import Packed class Cosmology(Parametrized): @@ -284,169 +257,3 @@ def critical_surface_density( d_s = self.angular_diameter_distance(z_s, params) d_ls = self.angular_diameter_distance_z1z2(z_l, z_s, params) return d_s / (4 * pi * G_over_c2 * d_l * d_ls) # fmt: skip - - -class FlatLambdaCDM(Cosmology): - """ - Subclass of Cosmology representing a Flat Lambda Cold Dark Matter (LCDM) - cosmology with no radiation. - """ - - def __init__( - self, - h0: Optional[Tensor] = torch.tensor(h0_default), - critical_density_0: Optional[Tensor] = torch.tensor(critical_density_0_default), - Om0: Optional[Tensor] = torch.tensor(Om0_default), - name: Optional[str] = None, - ): - """ - Initialize a new instance of the FlatLambdaCDM class. - - Parameters - ---------- - name: str - Name of the cosmology. - h0: Optional[Tensor] - Hubble constant over 100. Default is h0_default. - critical_density_0: (Optional[Tensor]) - Critical density at z=0. Default is critical_density_0_default. - Om0: Optional[Tensor] - Matter density parameter at z=0. Default is Om0_default. - """ - super().__init__(name) - - self.add_param("h0", h0) - self.add_param("critical_density_0", critical_density_0) - self.add_param("Om0", Om0) - - self._comoving_distance_helper_x_grid = _comoving_distance_helper_x_grid.to( - dtype=torch.float32 - ) - self._comoving_distance_helper_y_grid = _comoving_distance_helper_y_grid.to( - dtype=torch.float32 - ) - - def to( - self, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None - ): - super().to(device, dtype) - self._comoving_distance_helper_y_grid = ( - self._comoving_distance_helper_y_grid.to(device, dtype) - ) - self._comoving_distance_helper_x_grid = ( - self._comoving_distance_helper_x_grid.to(device, dtype) - ) - - def hubble_distance(self, h0): - """ - Calculate the Hubble distance. - - Parameters - ---------- - h0: Tensor - Hubble constant. - - Returns - ------- - Tensor - Hubble distance. - """ - return c_Mpc_s / (100 * km_to_Mpc) / h0 - - @unpack - def critical_density( - self, - z: Tensor, - *args, - params: Optional["Packed"] = None, - h0: Optional[Tensor] = None, - critical_density_0: Optional[Tensor] = None, - Om0: Optional[Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - Calculate the critical density at redshift z. - - Parameters - ---------- - z: Tensor - Redshift. - params: (Packed, optional) - Dynamic parameter container for the computation. - - Returns - ------- - torch.Tensor - Critical density at redshift z. - """ - Ode0 = 1 - Om0 - return critical_density_0 * (Om0 * (1 + z) ** 3 + Ode0) # fmt: skip - - @unpack - def _comoving_distance_helper( - self, x: Tensor, *args, params: Optional["Packed"] = None, **kwargs - ) -> Tensor: - """ - Helper method for computing comoving distances. - - Parameters - ---------- - x: Tensor - Input tensor. - - Returns - ------- - Tensor - Computed comoving distances. - """ - return interp1d( - self._comoving_distance_helper_x_grid, - self._comoving_distance_helper_y_grid, - torch.atleast_1d(x), - ).reshape(x.shape) - - @unpack - def comoving_distance( - self, - z: Tensor, - *args, - params: Optional["Packed"] = None, - h0: Optional[Tensor] = None, - critical_density_0: Optional[Tensor] = None, - Om0: Optional[Tensor] = None, - **kwargs, - ) -> Tensor: - """ - Calculate the comoving distance to redshift z. - - Parameters - ---------- - z: Tensor - Redshift. - params: (Packed, optional) - Dynamic parameter container for the computation. - - Returns - ------- - Tensor - Comoving distance to redshift z. - """ - Ode0 = 1 - Om0 - ratio = (Om0 / Ode0) ** (1 / 3) - DH = self.hubble_distance(h0) - DC1z = self._comoving_distance_helper((1 + z) * ratio, params) - DC = self._comoving_distance_helper(ratio, params) - return DH * (DC1z - DC) / (Om0 ** (1 / 3) * Ode0 ** (1 / 6)) # fmt: skip - - @unpack - def transverse_comoving_distance( - self, - z: Tensor, - *args, - params: Optional["Packed"] = None, - h0: Optional[Tensor] = None, - critical_density_0: Optional[Tensor] = None, - Om0: Optional[Tensor] = None, - **kwargs, - ) -> Tensor: - return self.comoving_distance(z, params, **kwargs)