From b86dcbc9bcf1a90a01c4727991ebe35626db12d5 Mon Sep 17 00:00:00 2001
From: Matthew Mould <mattdmould@gmail.com>
Date: Mon, 28 Apr 2025 15:47:07 -0700
Subject: [PATCH 1/2] update data collection cosmology and add tests for
 cosmology

---
 gwinferno/preprocess/data_collection.py    |  2 +-
 tests/models/bsplines/single_test.py       | 11 ++++++++---
 tests/models/parametric/parametric_test.py |  5 +++++
 tests/preprocess/data_collection_test.py   |  9 +++++++++
 4 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/gwinferno/preprocess/data_collection.py b/gwinferno/preprocess/data_collection.py
index 2d2db721..588d3b20 100644
--- a/gwinferno/preprocess/data_collection.py
+++ b/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
diff --git a/tests/models/bsplines/single_test.py b/tests/models/bsplines/single_test.py
index 3f6b0d1d..e613bab2 100644
--- a/tests/models/bsplines/single_test.py
+++ b/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_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)
diff --git a/tests/models/parametric/parametric_test.py b/tests/models/parametric/parametric_test.py
index e021f9df..6d5ab398 100644
--- a/tests/models/parametric/parametric_test.py
+++ b/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)
diff --git a/tests/preprocess/data_collection_test.py b/tests/preprocess/data_collection_test.py
index e69de29b..67d10f88 100644
--- a/tests/preprocess/data_collection_test.py
+++ b/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)

From 443b8b17f1c99392bb8d22cae3f92376384c0b73 Mon Sep 17 00:00:00 2001
From: Matthew Mould <mattdmould@gmail.com>
Date: Mon, 28 Apr 2025 17:17:33 -0700
Subject: [PATCH 2/2] fix and run tests

---
 tests/models/bsplines/single_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/models/bsplines/single_test.py b/tests/models/bsplines/single_test.py
index e613bab2..65873f7b 100644
--- a/tests/models/bsplines/single_test.py
+++ b/tests/models/bsplines/single_test.py
@@ -3,7 +3,7 @@
 
 import jax.numpy as jnp
 import xarray as xr
-from astropy.cosmology import Planck15_astropy
+from astropy.cosmology import Planck15 as Planck15_astropy
 
 from gwinferno.models.bsplines.single import Base1DBSplineModel
 from gwinferno.models.bsplines.single import BSplineChiEffective