add tests for LAL cosmology use

gwinferno/preprocess/data_collection.py

@@ -12,7 +12,7 @@
 from jax.scipy.integrate import trapezoid
 from tqdm import trange
 
-from ..cosmology import PLANCK_2015_Cosmology as cosmo
+from ..cosmology import PLANCK_2015_LVK_Cosmology as cosmo
 from .conversions import chieff_from_q_component_spins
 from .conversions import chip_from_q_component_spins
 from .priors import chi_effective_prior_from_isotropic_spins

tests/models/bsplines/single_test.py

@@ -3,7 +3,7 @@
 
 import jax.numpy as jnp
 import xarray as xr
-from astropy.cosmology import Planck15
+from astropy.cosmology import Planck15 as Planck15_astropy
 
 from gwinferno.models.bsplines.single import Base1DBSplineModel
 from gwinferno.models.bsplines.single import BSplineChiEffective
@@ -15,6 +15,8 @@
 from gwinferno.models.bsplines.single import BSplineSpinTilt
 from gwinferno.models.bsplines.single import BSplineSymmetricChiEffective
 from gwinferno.preprocess.data_collection import load_injection_dataset
+from gwinferno.models.bsplines.single import Planck15
+from gwinferno.cosmology import PLANCK_2015_LVK_Cosmology
 
 
 class TestBase1DBSplineModel(unittest.TestCase):
@@ -54,8 +56,8 @@ def load_data(self):
     def spline_shape(self, model, pe_x, inj_x, redshift=False):
 
         if redshift:
-            pe_dvcdz = jnp.array(Planck15.differential_comoving_volume(pe_x).value * 4 * jnp.pi)
-            inj_dvcdz = jnp.array(Planck15.differential_comoving_volume(inj_x).value * 4 * jnp.pi)
+            pe_dvcdz = jnp.array(Planck15_astropy.differential_comoving_volume(pe_x).value * 4 * jnp.pi)
+            inj_dvcdz = jnp.array(Planck15_astropy.differential_comoving_volume(inj_x).value * 4 * jnp.pi)
             bspline = model(self.nsplines, pe_x, inj_x, pe_dvcdz, inj_dvcdz)
             pe_pdf = bspline(self.coefs, pe_samples=True)
             inj_pdf = bspline(self.coefs, pe_samples=False)
@@ -79,3 +81,6 @@ def test_functions(self):
         self.spline_shape(BSplineRatio, self.pedict["mass_ratio"], self.injdict["mass_ratio"])
         self.spline_shape(BSplineMass, self.pedict["mass_1"], self.injdict["mass_1"])
         self.spline_shape(BSplineRedshift, self.pedict["redshift"], self.injdict["redshift"], redshift=True)
+
+    def test_cosmology(self):
+        self.assertIs(Planck15, PLANCK_2015_LVK_Cosmology)

tests/models/parametric/parametric_test.py

@@ -7,6 +7,8 @@
 
 from gwinferno.models.parametric.parametric import PowerlawRedshiftModel
 from gwinferno.preprocess.data_collection import load_injection_dataset
+from gwinferno.models.parametric.parametric import Planck15
+from gwinferno.cosmology import PLANCK_2015_LVK_Cosmology
 
 
 class TestPowerlawRedshift(unittest.TestCase):
@@ -67,3 +69,6 @@ def test_bounds(self):
 
         self.assertEqual(z_nonzero_pe, 0, msg="PE PDF not properly truncated")
         self.assertEqual(z_nonzero_inj, 0, msg="Inj PDF not properly truncated")
+
+    def test_cosmology(self):
+        self.assertIs(Planck15, PLANCK_2015_LVK_Cosmology)

tests/preprocess/data_collection_test.py

@@ -0,0 +1,9 @@
+import unittest
+
+from gwinferno.preprocess.data_collection import cosmo
+from gwinferno.cosmology import PLANCK_2015_LVK_Cosmology
+
+
+class TestBase1DBSplineModel(unittest.TestCase):
+    def test_cosmology(self):
+        self.assertIs(cosmo, PLANCK_2015_LVK_Cosmology)