diff --git a/.github/workflows/coverage.yaml b/.github/workflows/coverage.yaml index 49171d46..6dc13a07 100644 --- a/.github/workflows/coverage.yaml +++ b/.github/workflows/coverage.yaml @@ -30,6 +30,14 @@ jobs: echo github.event_name is: ${{ github.event_name }} echo github workspace: ${{ github.workspace }} pip --version + - name: Build Boost + id: boost + uses: egor-tensin/build-boost@v1 + with: + version: 1.83.0 + libraries: filesystem program_options system python + platform: x64 + configuration: Release - name: Install dependencies run: | python -m pip install --upgrade pip @@ -44,6 +52,37 @@ jobs: pip install -e . pip show caustic shell: bash + - name: Install phytorch dependencies + run: | + cd $GITHUB_WORKSPACE + pip install ninja + pip install git+https://github.com/rtqichen/torchdiffeq.git + pip show ninja + pip show torchdiffeq + shell: bash + + - name: Install phytorch + run: | + cd $GITHUB_WORKSPACE + git clone https://github.com/mjyb16/phytorch_forked.git + cd phytorch_forked + pip install -e . + pip show phytorch + shell: bash + + - name: Modify pytorch + run: | + complex_utils_path=$(LOCAL_VARIABLE="complex_utils.h" && echo $(cd $(dirname "$LOCAL_VARIABLE") && pwd -P)/$(basename "$LOCAL_VARIABLE")) + complex_utils_folder=$(LOCAL_VARIABLE="complex_utils.h" && echo $(cd $(dirname "$LOCAL_VARIABLE") && pwd -P)) + rm -rf complex_utils_path + mv phytorch_forked/complex_utils.h complex_utils_folder + shell: bash + + - name: Build phytorch extensions + run: | + cd $GITHUB_WORKSPACE/phytorch_forked/phytorch-extensions + python setup.py build_ext -b ../phytorch/extensions + shell: bash - name: Test with pytest run: | cd $GITHUB_WORKSPACE diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index 779a9949..736da8e1 100644 --- a/.github/workflows/python-app.yml +++ b/.github/workflows/python-app.yml @@ -21,7 +21,7 @@ jobs: strategy: matrix: python-version: ["3.9", "3.10"] - os: [ubuntu-latest, windows-latest, macOS-latest] + os: [ubuntu-latest, macOS-latest] steps: - name: Checkout caustic @@ -31,6 +31,22 @@ jobs: uses: actions/setup-python@v4 with: python-version: ${{ matrix.python-version }} + - name: Build Boost + id: boost + uses: egor-tensin/build-boost@v1 + with: + version: 1.83.0 + libraries: filesystem program_options system python + platform: x64 + configuration: Release + + - name: Show paths + run: | + printf 'Boost has been unpacked to: %s\n' '${{ steps.boost.outputs.root }}' + printf 'Libraries can be found here: %s\n' '${{ steps.boost.outputs.librarydir }}' + + shell: bash + - name: Record State run: | pwd @@ -54,11 +70,99 @@ jobs: cd $GITHUB_WORKSPACE pip install -e . pip show caustic + cd $GITHUB_WORKSPACE + pwd + shell: bash + + - name: Install phytorch dependencies + run: | + cd $GITHUB_WORKSPACE + pip install ninja + pip install git+https://github.com/rtqichen/torchdiffeq.git + pip show ninja + pip show torchdiffeq + shell: bash + + - name: Install phytorch + run: | + cd $GITHUB_WORKSPACE + git clone https://github.com/mjyb16/phytorch_forked.git + cd phytorch_forked + pip install -e . + pip show phytorch + shell: bash + + - name: Modify pytorch + run: | + complex_utils_path=$(LOCAL_VARIABLE="complex_utils.h" && echo $(cd $(dirname "$LOCAL_VARIABLE") && pwd -P)/$(basename "$LOCAL_VARIABLE")) + complex_utils_folder=$(LOCAL_VARIABLE="complex_utils.h" && echo $(cd $(dirname "$LOCAL_VARIABLE") && pwd -P)) + rm -rf complex_utils_path + mv phytorch_forked/complex_utils.h complex_utils_folder shell: bash + - name: Build phytorch extensions + run: | + cd $GITHUB_WORKSPACE/phytorch_forked/phytorch-extensions + python setup.py build_ext -b ../phytorch/extensions + shell: bash + + - name: Test with pytest run: | cd $GITHUB_WORKSPACE pytest test shell: bash + + build_windows: + + runs-on: ${{matrix.os}} + strategy: + matrix: + python-version: ["3.9", "3.10"] + os: [windows-latest] + + steps: + - name: Checkout caustic + uses: actions/checkout@v3 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + + - name: Record State + run: | + pwd + echo github.ref is: ${{ github.ref }} + echo GITHUB_SHA is: $GITHUB_SHA + echo github.event_name is: ${{ github.event_name }} + echo github workspace: ${{ github.workspace }} + pip --version + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install pytest pytest-cov torch + # Install deps + cd $GITHUB_WORKSPACE + pip install -r requirements.txt + shell: bash + + - name: Install Caustic + run: | + cd $GITHUB_WORKSPACE + pip install -e . + pip show caustic + cd $GITHUB_WORKSPACE + pwd + shell: bash + + + + - name: Test with pytest + run: | + cd $GITHUB_WORKSPACE + pytest test --ignore-glob='*cosmology.py' + shell: bash + diff --git a/README.md b/README.md index bc17ab01..08587ce9 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,8 @@ Simply install caustic from PyPI: pip install caustic ``` +- If you would like to take gradients through cosmological parameters, see [here](https://github.com/Ciela-Institute/caustic/blob/differentiable_cosmo/phytorch_instructions.md) for additional installation instructions. + ## Contributing Please reach out to one of us if you're interested in contributing! diff --git a/caustic/cosmology/FlatLambdaCDM.py b/caustic/cosmology/FlatLambdaCDM.py new file mode 100644 index 00000000..70c1be52 --- /dev/null +++ b/caustic/cosmology/FlatLambdaCDM.py @@ -0,0 +1,132 @@ +from .base import Cosmology +from typing import Any, 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 + +__all__ = ( + "FlatLambdaCDM", +) + +H0_default = float(default_cosmology.get().h)*100 +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, +) + +class FlatLambdaCDM(Cosmology): + """ + Subclass of Cosmology representing a Flat Lambda Cold Dark Matter (LCDM) cosmology with no radiation. + NOTE: THIS IS THE VERSION WITHOUT PRECISE COSMOLOGICAL GRADIENTS AND WHICH USES AN APPROXIMATION SCHEME. It is kept for backwards compatibility. To use autograd through cosmological calculations, follow the installation instructions for phytorch. + """ + + 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: str = None, + ): + """ + Initialize a new instance of the FlatLambdaCDM class. + + Args: + 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 hubble_distance(self, H0): + """ + Calculate the Hubble distance. + + Args: + h0 (Tensor): Hubble constant. + + Returns: + Tensor: Hubble distance. + """ + return c_Mpc_s / (km_to_Mpc) / H0 + + @unpack(1) + def critical_density(self, z: Tensor, H0, central_critical_density, Om0, *args, params: Optional["Packed"] = None) -> torch.Tensor: + """ + Calculate the critical density at redshift z. + + Args: + z (Tensor): Redshift. + params (Packed, optional): Dynamic parameter container for the computation. + + Returns: + torch.Tensor: Critical density at redshift z. + """ + Ode0 = 1 - Om0 + return central_critical_density * (Om0 * (1 + z) ** 3 + Ode0) + + @unpack(1) + def _comoving_distance_helper(self, x: Tensor, *args, params: Optional["Packed"] = None) -> Tensor: + """ + Helper method for computing comoving distances. + + Args: + 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(1) + def comoving_distance(self, z: Tensor, H0, central_critical_density, Om0, *args, params: Optional["Packed"] = None) -> Tensor: + """ + Calculate the comoving distance to redshift z. + + Args: + 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) + return ( + self.hubble_distance(H0) + * ( + self._comoving_distance_helper((1 + z) * ratio, params) + - self._comoving_distance_helper(ratio, params) + ) + / (Om0 ** (1 / 3) * Ode0 ** (1 / 6)) + ) diff --git a/caustic/cosmology/LambdaCDM.py b/caustic/cosmology/LambdaCDM.py new file mode 100644 index 00000000..6fde9afe --- /dev/null +++ b/caustic/cosmology/LambdaCDM.py @@ -0,0 +1,101 @@ +from .base import Cosmology +from typing import Any, Optional +from phytorch.cosmology.drivers.analytic import LambdaCDM as phy_LCDM +from phytorch.units import si, astro +from astropy.cosmology import default_cosmology + +import torch +from torch import Tensor + +from ..constants import G_over_c2, c_Mpc_s, km_to_Mpc +from ..parametrized import unpack + +__all__ = ( + "LambdaCDM", +) + +#Default values for parameters, from astropy +H0_default = float(default_cosmology.get().h)*100. +Om0_default = float(default_cosmology.get().Om0) +Ob0_default = float(default_cosmology.get().Ob0) +Ode0_default = float(default_cosmology.get().Ode0) + +class LambdaCDM(Cosmology): + """ + Subclass of Cosmology representing a Lambda Cold Dark Matter (LCDM) cosmology with no radiation. + """ + + def __init__( + self, + H0: Optional[Tensor] = torch.tensor(H0_default), + Om0: Optional[Tensor] = torch.tensor(Om0_default), + Ob0: Optional[Tensor] = torch.tensor(Ob0_default), + Ode0: Optional[Tensor] = torch.tensor(Ode0_default), + name: str = None, + ): + """ + Initialize a new instance of the LambdaCDM class. + + Args: + name (str): Name of the cosmology. + H0 (Optional[Tensor]): Hubble constant over 100. Default is H0_default from Astropy. + Om0 (Optional[Tensor]): Matter density parameter at z=0. Default is Om0_default from Astropy. + Ob0 (Optional[Tensor]): Baryon density parameter at z=0. Default is Ob0_default from Astropy. + Ode0 (Optional[Tensor]): Dark Energy density parameter at z=0. Default is Ode0_default from Astropy. + """ + super().__init__(name = name) + + self.add_param("H0", H0) + self.add_param("Om0", Om0) + self.add_param("Ob0", Ob0) + self.add_param("Ode0", Ode0) + self.cosmo = phy_LCDM(H0 = H0, Om0 = Om0, Ob0 = Ob0, Ode0 = Ode0) + + def hubble_distance(self, H0): + """ + Calculate the Hubble distance in Mpc. + + Args: + h0 (Tensor): Hubble constant. + + Returns: + Tensor: Hubble distance in Mpc. + """ + self.cosmo.H0 = H0 * si.km / si.s / astro.Mpc + return self.cosmo.hubble_distance.to(astro.Mpc) + + @unpack(1) + def critical_density(self, z: Tensor, H0, Om0, Ob0, Ode0, *args, params: Optional["Packed"] = None) -> torch.Tensor: + """ + Calculate the critical density at redshift z. + + Args: + z (Tensor): Redshift. + params (Packed, optional): Dynamic parameter container for the computation. + + Returns: + torch.Tensor: Critical density at redshift z in g/cm^3. + """ + self.cosmo.H0 = H0 * si.km / si.s / astro.Mpc + self.cosmo.Om0 = Om0 + self.cosmo.Ob0 = Ob0 + self.cosmo.Ode0 = Ode0 + return self.cosmo.critical_density(z).to(si.g / si.cm**3) + + @unpack(1) + def comoving_distance(self, z: Tensor, H0, Om0, Ob0, Ode0, *args, params: Optional["Packed"] = None) -> Tensor: + """ + Calculate the comoving distance to redshift z. + + Args: + z (Tensor): Redshift. + params (Packed, optional): Dynamic parameter container for the computation. + + Returns: + Tensor: Comoving distance to redshift z. + """ + self.cosmo.H0 = H0 * si.km / si.s / astro.Mpc + self.cosmo.Om0 = Om0 + self.cosmo.Ob0 = Ob0 + self.cosmo.Ode0 = Ode0 + return self.cosmo.comoving_distance_dimless(z)*self.hubble_distance(H0) diff --git a/caustic/cosmology/__init__.py b/caustic/cosmology/__init__.py new file mode 100644 index 00000000..37459fd6 --- /dev/null +++ b/caustic/cosmology/__init__.py @@ -0,0 +1,2 @@ +from .base import * +from .FlatLambdaCDM import * diff --git a/caustic/cosmology.py b/caustic/cosmology/base.py similarity index 52% rename from caustic/cosmology.py rename to caustic/cosmology/base.py index 024cdaf9..9dac18a6 100644 --- a/caustic/cosmology.py +++ b/caustic/cosmology/base.py @@ -1,37 +1,13 @@ from abc import abstractmethod from math import pi from typing import Any, 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 ..constants import G_over_c2, c_Mpc_s, km_to_Mpc +from ..parametrized import Parametrized, unpack __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, ) @@ -174,106 +150,4 @@ def critical_surface_density( d_l = self.angular_diameter_distance(z_l, params) d_s = self.angular_diameter_distance(z_s, params) d_ls = self.angular_diameter_distance_z1z2(z_l, z_s, params) - return d_s / d_l / d_ls / (4 * pi * G_over_c2) - - -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: str = None, - ): - """ - Initialize a new instance of the FlatLambdaCDM class. - - Args: - 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 hubble_distance(self, h0): - """ - Calculate the Hubble distance. - - Args: - h0 (Tensor): Hubble constant. - - Returns: - Tensor: Hubble distance. - """ - return c_Mpc_s / (100 * km_to_Mpc) / h0 - - @unpack(1) - def critical_density(self, z: Tensor, h0, central_critical_density, Om0, *args, params: Optional["Packed"] = None) -> torch.Tensor: - """ - Calculate the critical density at redshift z. - - Args: - z (Tensor): Redshift. - params (Packed, optional): Dynamic parameter container for the computation. - - Returns: - torch.Tensor: Critical density at redshift z. - """ - Ode0 = 1 - Om0 - return central_critical_density * (Om0 * (1 + z) ** 3 + Ode0) - - @unpack(1) - def _comoving_distance_helper(self, x: Tensor, *args, params: Optional["Packed"] = None) -> Tensor: - """ - Helper method for computing comoving distances. - - Args: - 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(1) - def comoving_distance(self, z: Tensor, h0, central_critical_density, Om0, *args, params: Optional["Packed"] = None) -> Tensor: - """ - Calculate the comoving distance to redshift z. - - Args: - 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) - return ( - self.hubble_distance(h0) - * ( - self._comoving_distance_helper((1 + z) * ratio, params) - - self._comoving_distance_helper(ratio, params) - ) - / (Om0 ** (1 / 3) * Ode0 ** (1 / 6)) - ) + return d_s / d_l / d_ls / (4 * pi * G_over_c2) \ No newline at end of file diff --git a/phytorch_instructions.md b/phytorch_instructions.md new file mode 100644 index 00000000..07684678 --- /dev/null +++ b/phytorch_instructions.md @@ -0,0 +1,39 @@ +In order to use phytorch (which enables gradients through cosmological parameters), do the following (assumes PyTorch > 2.0): + +- Install Boost: https://www.boost.org/ +- `pip install ninja` +- git clone [torchdiffeq](https://github.com/rtqichen/torchdiffeq), cd into torchdiffeq, and run `pip install -e .` +- git clone [phytorch](https://github.com/kosiokarchev/phytorch), cd into phytorch, and run `pip install -e .` +- In `(your local environment here)/lib/python3.9/site-packages/torch/include/c10/util/complex_utils.h`, change the following: +``` +template +bool isnan(const c10::complex& v) { + return std::isnan(v.real()) || std::isnan(v.imag()); +} + +} // namespace std +``` +To: +``` +template C10_HOST_DEVICE +bool isnan(const c10::complex& v) { + return std::isnan(v.real()) || std::isnan(v.imag()); +} + +} // namespace std +``` + +- In `phytorch/phytorch-extensions/common/complex.h`, change +``` +DEF_COMPLEX_CHECK isnan(complex a) { return isnan(a.real()) or isnan(a.imag()); } +``` +To: +``` +//DEF_COMPLEX_CHECK isnan(complex a) { return isnan(a.real()) or isnan(a.imag()); } +``` +- Run the following: +``` +cd phytorch-extensions +python setup.py build_ext -b ../phytorch/extensions +``` + diff --git a/test/test_cosmology.py b/test/test_cosmology.py index 4e8794cc..216dbd69 100644 --- a/test/test_cosmology.py +++ b/test/test_cosmology.py @@ -4,11 +4,15 @@ import torch from astropy.cosmology import Cosmology as Cosmology_AP from astropy.cosmology import FlatLambdaCDM as AstropyFlatLambdaCDM +from astropy.cosmology import LambdaCDM as AstropyLambdaCDM -from caustic.cosmology import Cosmology, FlatLambdaCDM as CausticFlatLambdaCDM, Om0_default, h0_default +from caustic.cosmology import Cosmology +from caustic.cosmology import FlatLambdaCDM as CausticFlatLambdaCDM +from caustic.cosmology.LambdaCDM import LambdaCDM as CausticLambdaCDM +from caustic.cosmology.FlatLambdaCDM import Om0_default, H0_default -def get_cosmologies() -> List[Tuple[Cosmology, Cosmology_AP]]: +def get_cosmologies_flat_lcdm() -> List[Tuple[Cosmology, Cosmology_AP]]: """ Gets caustic cosmologies and corresponding astropy ones. """ @@ -16,22 +20,59 @@ def get_cosmologies() -> List[Tuple[Cosmology, Cosmology_AP]]: cosmologies.append( ( CausticFlatLambdaCDM(name="cosmo"), - AstropyFlatLambdaCDM(h0_default, Om0_default, Tcmb0=0), + AstropyFlatLambdaCDM(H0_default, Om0_default, Tcmb0=0), ) ) return cosmologies -def test_comoving_dist(): +def test_comoving_dist_flat_lcdm(): rtol = 1e-3 atol = 0 zs = torch.linspace(0.05, 3, 10) - for cosmology, cosmology_ap in get_cosmologies(): + for cosmology, cosmology_ap in get_cosmologies_flat_lcdm(): vals = cosmology.comoving_distance(zs).numpy() - vals_ref = cosmology_ap.comoving_distance(zs).value / 1e2 # type: ignore + vals_ref = cosmology_ap.comoving_distance(zs).value # type: ignore + assert np.allclose(vals, vals_ref, rtol, atol) + + +def get_cosmologies_lcdm() -> List[Tuple[Cosmology, Cosmology_AP]]: + """ + Gets caustic cosmologies and corresponding astropy ones. + """ + cosmologies = [] + cosmologies.append( + ( + CausticLambdaCDM(name="cosmo_phytorch"), + AstropyFlatLambdaCDM(H0_default, Om0_default, Tcmb0=0), + ) + ) + return cosmologies + + +def test_comoving_dist_lcdm(): + rtol = 1e-3 + atol = 0 + + zs = torch.linspace(0.05, 3, 10) + for cosmology, cosmology_ap in get_cosmologies_lcdm(): + vals = cosmology.comoving_distance(zs).numpy() + vals_ref = cosmology_ap.comoving_distance(zs).value # type: ignore + assert np.allclose(vals, vals_ref, rtol, atol) + +def test_critical_density_lcdm(): + rtol = 1e-2 #Not sure why, but there is a bit of a difference at the 1e-3 level + atol = 0 + + zs = torch.linspace(0.05, 3, 10) + for cosmology, cosmology_ap in get_cosmologies_lcdm(): + vals = cosmology.critical_density(zs).numpy() + vals_ref = cosmology_ap.critical_density(zs).value # type: ignore assert np.allclose(vals, vals_ref, rtol, atol) if __name__ == "__main__": - test_comoving_dist() + test_comoving_dist_flat_lcdm() + test_comoving_dist_lcdm() + test_critical_density_lcdm() diff --git a/test/test_multiplane.py b/test/test_multiplane.py index 0e36f6a4..99889da6 100644 --- a/test/test_multiplane.py +++ b/test/test_multiplane.py @@ -47,7 +47,7 @@ def test(): ) # Use same cosmology - cosmo_ap = FlatLambdaCDM_ap(cosmology.h0.value, cosmology.Om0.value, Tcmb0=0) + cosmo_ap = FlatLambdaCDM_ap(cosmology.H0.value, cosmology.Om0.value, Tcmb0=0) lens_ls = LensModel( lens_model_list=["SIE" for _ in range(len(xs))], z_source=z_s.item(), diff --git a/test/test_parametrized.py b/test/test_parametrized.py index e8f4882f..972a35d1 100644 --- a/test/test_parametrized.py +++ b/test/test_parametrized.py @@ -11,7 +11,7 @@ def test_params(): class Sim(Simulator): def __init__(self): super().__init__() - self.cosmo = FlatLambdaCDM(h0=None) + self.cosmo = FlatLambdaCDM(H0=None) self.epl = EPL(self.cosmo) self.sersic = Sersic() self.add_param("z_s", 1.0) @@ -111,7 +111,7 @@ def test_parametrized_name_collision(): class Sim(Simulator): def __init__(self): super().__init__() - self.cosmo = FlatLambdaCDM(h0=None) + self.cosmo = FlatLambdaCDM(H0=None) # These two module are identical and will create a name collision self.lens1 = EPL(self.cosmo) self.lens2 = EPL(self.cosmo) @@ -123,7 +123,7 @@ def __init__(self): assert sim.lens2.name == "EPL_1" # Case 2: name collision in parents of a module - cosmo = FlatLambdaCDM(h0=None) + cosmo = FlatLambdaCDM(H0=None) lens = EPL(cosmo) class Sim(Simulator): def __init__(self): diff --git a/test/utils.py b/test/utils.py index 51b9bc14..e26768b3 100644 --- a/test/utils.py +++ b/test/utils.py @@ -21,7 +21,7 @@ def __init__(self, name="simulator"): else: self.add_param("z_s", None) z_l = 0.5 - self.cosmo = FlatLambdaCDM(h0=0.7 if cosmo_static else None, name="cosmo") + self.cosmo = FlatLambdaCDM(H0=70 if cosmo_static else None, name="cosmo") if use_nfw: self.lens = NFW(self.cosmo, z_l=z_l, name="lens") # NFW wactually depend on cosmology, so a better test for Parametrized else: @@ -41,8 +41,8 @@ def forward(self, params): z_s = torch.tensor([1.0, 1.5]) sim_params = [z_s] # default cosmo params - h0 = torch.tensor([0.68, 0.75]) - cosmo_params = [h0] + H0 = torch.tensor([68, 75]) + cosmo_params = [H0] # default lens params if use_nfw: x0 = torch.tensor([0., 0.1]) @@ -84,7 +84,7 @@ def __init__(self, name="test"): pixel_scale = 0.04 z_l = 0.5 self.z_s = torch.tensor(1.0) - self.cosmo = FlatLambdaCDM(h0=0.7 if cosmo_static else None, name="cosmo") + self.cosmo = FlatLambdaCDM(H0=70 if cosmo_static else None, name="cosmo") self.epl = EPL(self.cosmo, z_l=z_l, name="lens") self.kappa = PixelatedConvergence(pixel_scale, n_pix, self.cosmo, z_l=z_l, shape=(n_pix, n_pix), name="kappa") self.source = Pixelated(x0=0., y0=0., pixelscale=pixel_scale/2, shape=(n_pix, n_pix), name="source") @@ -99,7 +99,7 @@ def forward(self, params): return self.source.brightness(bx, by, params) # default cosmo params - h0 = torch.tensor([0.68, 0.75]) + H0 = torch.tensor([68, 75]) # default lens params x0 = torch.tensor([0, 0.1]) y0 = torch.tensor([0, 0.1]) @@ -111,7 +111,7 @@ def forward(self, params): kappa = torch.randn([2, n_pix, n_pix]) source = torch.randn([2, n_pix, n_pix]) - cosmo_params = [h0] + cosmo_params = [H0] lens_params = [x0, y0, q, phi, b, t] if not batched_params: cosmo_params = [_x[0] for _x in cosmo_params]