From 2455fe0622f8d99bab2a937c3e4a676d64878fb7 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Thu, 7 Nov 2024 16:51:28 +0000
Subject: [PATCH 01/15] gh-124: use `cosmology.api` over the `cosmology`
 package

Closes #124. Use the improved `cosmology.api` package throughout the
code.
---
 glass/galaxies.py     |  6 +++---
 glass/lensing.py      | 14 +++++++-------
 glass/shells.py       | 18 +++++++++---------
 pyproject.toml        |  4 +++-
 tests/test_lensing.py |  8 ++++----
 5 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/glass/galaxies.py b/glass/galaxies.py
index 16c26c9a..1384af16 100644
--- a/glass/galaxies.py
+++ b/glass/galaxies.py
@@ -33,7 +33,7 @@
 from glass.core.array import broadcast_leading_axes, cumtrapz
 
 if typing.TYPE_CHECKING:
-    from cosmology import Cosmology
+    from cosmology.api import StandardCosmology
 
     from glass.shells import RadialWindow
 
@@ -300,7 +300,7 @@ def kappa_ia_nla(  # noqa: PLR0913
     delta: npt.NDArray[np.float64],
     zeff: float,
     a_ia: float,
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
     *,
     z0: float = 0.0,
     eta: float = 0.0,
@@ -322,7 +322,7 @@ def kappa_ia_nla(  # noqa: PLR0913
     a_ia:
         Intrinsic alignments amplitude.
     cosmo:
-        Cosmology instance.
+        StandardCosmology instance.
     z0:
         Reference redshift for the redshift dependence.
     eta:
diff --git a/glass/lensing.py b/glass/lensing.py
index 4d8b19e6..627c49b0 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -40,7 +40,7 @@
 if typing.TYPE_CHECKING:
     import collections.abc
 
-    from cosmology import Cosmology
+    from cosmology.api import StandardCosmology
 
     from glass.shells import RadialWindow
 
@@ -270,7 +270,7 @@ def shear_from_convergence(
 class MultiPlaneConvergence:
     """Compute convergence fields iteratively from multiple matter planes."""
 
-    def __init__(self, cosmo: Cosmology) -> None:
+    def __init__(self, cosmo: StandardCosmology) -> None:
         """Create a new instance to iteratively compute the convergence."""
         self.cosmo = cosmo
 
@@ -334,9 +334,9 @@ def add_plane(
         t = r13 / r12
 
         # lensing weight of mass plane to be added
-        f = 3 * self.cosmo.omega_m / 2
+        f = 3 * self.cosmo.Omega_m0 / 2
         f *= x2 * self.r23
-        f *= (1 + self.z2) / self.cosmo.ef(self.z2)
+        f *= (1 + self.z2) / self.cosmo.H_over_H0(self.z2)
         f *= w2
 
         # create kappa planes on first iteration
@@ -378,7 +378,7 @@ def wlens(self) -> float:
 
 def multi_plane_matrix(
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
 ) -> npt.NDArray[np.float64]:
     """Compute the matrix of lensing contributions from each shell."""
     mpc = MultiPlaneConvergence(cosmo)
@@ -392,7 +392,7 @@ def multi_plane_matrix(
 def multi_plane_weights(
     weights: npt.NDArray[np.float64],
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
 ) -> npt.NDArray[np.float64]:
     """
     Compute effective weights for multi-plane convergence.
@@ -414,7 +414,7 @@ def multi_plane_weights(
     shells:
         Window functions of the shells.
     cosmo:
-        Cosmology instance.
+        StandardCosmology instance.
 
     """
     # ensure shape of weights ends with the number of shells
diff --git a/glass/shells.py b/glass/shells.py
index 479b589f..7a29ff38 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -54,7 +54,7 @@
 from glass.core.array import ndinterp
 
 if typing.TYPE_CHECKING:
-    from cosmology import Cosmology
+    from cosmology.api import StandardCosmology
 
 ArrayLike1D = typing.Union[collections.abc.Sequence[float], npt.NDArray[np.float64]]
 WeightFunc = typing.Callable[[ArrayLike1D], npt.NDArray[np.float64]]
@@ -62,26 +62,26 @@
 
 def distance_weight(
     z: npt.NDArray[np.float64],
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving distance."""
-    return 1 / cosmo.ef(z)  # type: ignore[no-any-return]
+    return 1 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
 
 
 def volume_weight(
     z: npt.NDArray[np.float64],
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving volume."""
-    return cosmo.xm(z) ** 2 / cosmo.ef(z)  # type: ignore[no-any-return]
+    return cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
 
 
 def density_weight(
     z: npt.NDArray[np.float64],
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in matter density."""
-    return cosmo.rho_m_z(z) * cosmo.xm(z) ** 2 / cosmo.ef(z)  # type: ignore[no-any-return]
+    return cosmo.rho_m_z(z) * cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
 
 
 class RadialWindow(typing.NamedTuple):
@@ -649,7 +649,7 @@ def redshift_grid(
 
 
 def distance_grid(
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
     zmin: float,
     zmax: float,
     *,
@@ -657,7 +657,7 @@ def distance_grid(
     num: int | None = None,
 ) -> npt.NDArray[np.float64]:
     """Redshift grid with uniform spacing in comoving distance."""
-    xmin, xmax = cosmo.dc(zmin), cosmo.dc(zmax)
+    xmin, xmax = cosmo.comoving_distance(zmin), cosmo.comoving_distance(zmax)
     if dx is not None and num is None:
         x = np.arange(xmin, np.nextafter(xmax + dx, xmax), dx)
     elif dx is None and num is not None:
diff --git a/pyproject.toml b/pyproject.toml
index 87788927..93e94ab6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ classifiers = [
     "Typing :: Typed",
 ]
 dependencies = [
-    "cosmology>=2022.10.9",
+    "cosmology.api>=0.1.0",
     "gaussiancl>=2022.10.21",
     "healpix>=2022.11.1",
     "healpy>=1.15.0",
@@ -49,6 +49,7 @@ docs = [
 ]
 examples = [
     "camb",
+    "cosmology",
     "glass.ext.camb",
     "jupyter",
     "matplotlib",
@@ -151,6 +152,7 @@ lint.isort = {known-first-party = [
 ], sections = {"cosmo" = [
     "camb",
     "cosmology",
+    "cosmology.api",
 ]}}
 lint.per-file-ignores = {"__init__.py" = [
     "F401", # unused-import
diff --git a/tests/test_lensing.py b/tests/test_lensing.py
index 356b2b78..89b87129 100644
--- a/tests/test_lensing.py
+++ b/tests/test_lensing.py
@@ -16,7 +16,7 @@
 from glass.shells import RadialWindow
 
 if typing.TYPE_CHECKING:
-    from cosmology import Cosmology
+    from cosmology.api import StandardCosmology
 
 
 @pytest.fixture
@@ -31,7 +31,7 @@ def shells() -> list[RadialWindow]:
 
 
 @pytest.fixture
-def cosmo() -> Cosmology:
+def cosmo() -> StandardCosmology:
     class MockCosmology:
         @property
         def omega_m(self) -> float:
@@ -97,7 +97,7 @@ def test_deflect_many(rng: np.random.Generator) -> None:
 
 def test_multi_plane_matrix(
     shells: list[RadialWindow],
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
     rng: np.random.Generator,
 ) -> None:
     mat = multi_plane_matrix(shells, cosmo)
@@ -119,7 +119,7 @@ def test_multi_plane_matrix(
 
 def test_multi_plane_weights(
     shells: list[RadialWindow],
-    cosmo: Cosmology,
+    cosmo: StandardCosmology,
     rng: np.random.Generator,
 ) -> None:
     w_in = np.eye(len(shells))

From 0621f3d185ca4f00dd7150edd0a9e91aaca7dc7d Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 11:22:21 +0000
Subject: [PATCH 02/15] Change docstring

---
 glass/galaxies.py | 2 +-
 glass/lensing.py  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/glass/galaxies.py b/glass/galaxies.py
index 1384af16..91ffd8e3 100644
--- a/glass/galaxies.py
+++ b/glass/galaxies.py
@@ -322,7 +322,7 @@ def kappa_ia_nla(  # noqa: PLR0913
     a_ia:
         Intrinsic alignments amplitude.
     cosmo:
-        StandardCosmology instance.
+        Cosmology instance.
     z0:
         Reference redshift for the redshift dependence.
     eta:
diff --git a/glass/lensing.py b/glass/lensing.py
index 2b40c69a..3165b631 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -413,7 +413,7 @@ def multi_plane_weights(
     shells:
         Window functions of the shells.
     cosmo:
-        StandardCosmology instance.
+        Cosmology instance.
 
     """
     # ensure shape of weights ends with the number of shells

From 88428cbb4127ab5bab59c9774587a9feb310edcc Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 12:15:56 +0000
Subject: [PATCH 03/15] Rename mock

---
 tests/test_lensing.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_lensing.py b/tests/test_lensing.py
index 89b87129..90089439 100644
--- a/tests/test_lensing.py
+++ b/tests/test_lensing.py
@@ -34,11 +34,11 @@ def shells() -> list[RadialWindow]:
 def cosmo() -> StandardCosmology:
     class MockCosmology:
         @property
-        def omega_m(self) -> float:
+        def Omega_m0(self) -> float:  # noqa: N802
             return 0.3
 
-        def ef(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
-            return (self.omega_m * (1 + z) ** 3 + 1 - self.omega_m) ** 0.5
+        def H_over_H0(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:  # noqa: N802
+            return (self.Omega_m0 * (1 + z) ** 3 + 1 - self.Omega_m0) ** 0.5
 
         def xm(
             self,

From 2b8912fd726dd477e30393babd76163daf42d480 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 12:19:29 +0000
Subject: [PATCH 04/15] Format

---
 glass/lensing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/glass/lensing.py b/glass/lensing.py
index 3165b631..4c140115 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -301,7 +301,7 @@ def add_window(self, delta: npt.NDArray[np.float64], w: RadialWindow) -> None:
                 zsrc,
                 w.za,
                 w.wa,
-            )
+            ),
         )
 
         self.add_plane(delta, zsrc, lens_weight)

From 96b7030dbed93811b7ef03603e1b8da78a7de678 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 14:23:22 +0000
Subject: [PATCH 05/15] Add typing to `StandardCosmology`

---
 glass/galaxies.py     | 2 +-
 glass/lensing.py      | 8 +++++---
 glass/shells.py       | 8 ++++----
 tests/test_lensing.py | 4 ++--
 4 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/glass/galaxies.py b/glass/galaxies.py
index 91ffd8e3..31ccaea7 100644
--- a/glass/galaxies.py
+++ b/glass/galaxies.py
@@ -300,7 +300,7 @@ def kappa_ia_nla(  # noqa: PLR0913
     delta: npt.NDArray[np.float64],
     zeff: float,
     a_ia: float,
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     *,
     z0: float = 0.0,
     eta: float = 0.0,
diff --git a/glass/lensing.py b/glass/lensing.py
index 4c140115..b2c8b2a1 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -269,7 +269,9 @@ def shear_from_convergence(
 class MultiPlaneConvergence:
     """Compute convergence fields iteratively from multiple matter planes."""
 
-    def __init__(self, cosmo: StandardCosmology) -> None:
+    def __init__(
+        self, cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]]
+    ) -> None:
         """Create a new instance to iteratively compute the convergence."""
         self.cosmo = cosmo
 
@@ -377,7 +379,7 @@ def wlens(self) -> float:
 
 def multi_plane_matrix(
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Compute the matrix of lensing contributions from each shell."""
     mpc = MultiPlaneConvergence(cosmo)
@@ -391,7 +393,7 @@ def multi_plane_matrix(
 def multi_plane_weights(
     weights: npt.NDArray[np.float64],
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """
     Compute effective weights for multi-plane convergence.
diff --git a/glass/shells.py b/glass/shells.py
index eba42521..92d9e891 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -63,7 +63,7 @@
 
 def distance_weight(
     z: npt.NDArray[np.float64],
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving distance."""
     return 1 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -71,7 +71,7 @@ def distance_weight(
 
 def volume_weight(
     z: npt.NDArray[np.float64],
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving volume."""
     return cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -79,7 +79,7 @@ def volume_weight(
 
 def density_weight(
     z: npt.NDArray[np.float64],
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in matter density."""
     return cosmo.rho_m_z(z) * cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -650,7 +650,7 @@ def redshift_grid(
 
 
 def distance_grid(
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     zmin: float,
     zmax: float,
     *,
diff --git a/tests/test_lensing.py b/tests/test_lensing.py
index 90089439..be9f7fac 100644
--- a/tests/test_lensing.py
+++ b/tests/test_lensing.py
@@ -97,7 +97,7 @@ def test_deflect_many(rng: np.random.Generator) -> None:
 
 def test_multi_plane_matrix(
     shells: list[RadialWindow],
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     rng: np.random.Generator,
 ) -> None:
     mat = multi_plane_matrix(shells, cosmo)
@@ -119,7 +119,7 @@ def test_multi_plane_matrix(
 
 def test_multi_plane_weights(
     shells: list[RadialWindow],
-    cosmo: StandardCosmology,
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     rng: np.random.Generator,
 ) -> None:
     w_in = np.eye(len(shells))

From c4221b3f22caeee88ccd62156996ebb8fac2c584 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 14:23:38 +0000
Subject: [PATCH 06/15] Change critical density

---
 glass/galaxies.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/glass/galaxies.py b/glass/galaxies.py
index 31ccaea7..66e15297 100644
--- a/glass/galaxies.py
+++ b/glass/galaxies.py
@@ -382,7 +382,7 @@ def kappa_ia_nla(  # noqa: PLR0913
 
     """
     c1 = 5e-14 / cosmo.h**2  # Solar masses per cubic Mpc
-    rho_c1 = c1 * cosmo.rho_c0
+    rho_c1 = c1 * cosmo.critical_density0
 
     prefactor = -a_ia * rho_c1 * cosmo.Om
     inverse_linear_growth = 1.0 / cosmo.gf(zeff)

From 187063a66e36f4a3eaa9e2bb84e2e0a0884d8466 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 15:46:59 +0000
Subject: [PATCH 07/15] Switch to `import cosmology.api`

---
 glass/galaxies.py     |  6 ++++--
 glass/lensing.py      | 15 +++++++++++----
 glass/shells.py       | 18 +++++++++++++-----
 tests/test_lensing.py | 14 ++++++++++----
 4 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/glass/galaxies.py b/glass/galaxies.py
index 66e15297..b622a60c 100644
--- a/glass/galaxies.py
+++ b/glass/galaxies.py
@@ -33,7 +33,7 @@
 from glass.core.array import broadcast_leading_axes, cumtrapz
 
 if typing.TYPE_CHECKING:
-    from cosmology.api import StandardCosmology
+    import cosmology.api
 
     from glass.shells import RadialWindow
 
@@ -300,7 +300,9 @@ def kappa_ia_nla(  # noqa: PLR0913
     delta: npt.NDArray[np.float64],
     zeff: float,
     a_ia: float,
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
     *,
     z0: float = 0.0,
     eta: float = 0.0,
diff --git a/glass/lensing.py b/glass/lensing.py
index b2c8b2a1..89a14fc1 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -40,7 +40,7 @@
 if typing.TYPE_CHECKING:
     import collections.abc
 
-    from cosmology.api import StandardCosmology
+    import cosmology.api
 
     from glass.shells import RadialWindow
 
@@ -270,7 +270,10 @@ class MultiPlaneConvergence:
     """Compute convergence fields iteratively from multiple matter planes."""
 
     def __init__(
-        self, cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]]
+        self,
+        cosmo: cosmology.api.StandardCosmology[
+            npt.NDArray[np.float64], npt.NDArray[np.float64]
+        ],
     ) -> None:
         """Create a new instance to iteratively compute the convergence."""
         self.cosmo = cosmo
@@ -379,7 +382,9 @@ def wlens(self) -> float:
 
 def multi_plane_matrix(
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
 ) -> npt.NDArray[np.float64]:
     """Compute the matrix of lensing contributions from each shell."""
     mpc = MultiPlaneConvergence(cosmo)
@@ -393,7 +398,9 @@ def multi_plane_matrix(
 def multi_plane_weights(
     weights: npt.NDArray[np.float64],
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
 ) -> npt.NDArray[np.float64]:
     """
     Compute effective weights for multi-plane convergence.
diff --git a/glass/shells.py b/glass/shells.py
index 92d9e891..c1dd843f 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -55,7 +55,7 @@
 from glass.core.array import ndinterp
 
 if typing.TYPE_CHECKING:
-    from cosmology.api import StandardCosmology
+    import cosmology.api
 
 ArrayLike1D = typing.Union[collections.abc.Sequence[float], npt.NDArray[np.float64]]
 WeightFunc = typing.Callable[[ArrayLike1D], npt.NDArray[np.float64]]
@@ -63,7 +63,9 @@
 
 def distance_weight(
     z: npt.NDArray[np.float64],
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving distance."""
     return 1 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -71,7 +73,9 @@ def distance_weight(
 
 def volume_weight(
     z: npt.NDArray[np.float64],
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving volume."""
     return cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -79,7 +83,9 @@ def volume_weight(
 
 def density_weight(
     z: npt.NDArray[np.float64],
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in matter density."""
     return cosmo.rho_m_z(z) * cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -650,7 +656,9 @@ def redshift_grid(
 
 
 def distance_grid(
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
     zmin: float,
     zmax: float,
     *,
diff --git a/tests/test_lensing.py b/tests/test_lensing.py
index be9f7fac..0d530c52 100644
--- a/tests/test_lensing.py
+++ b/tests/test_lensing.py
@@ -16,7 +16,7 @@
 from glass.shells import RadialWindow
 
 if typing.TYPE_CHECKING:
-    from cosmology.api import StandardCosmology
+    import cosmology.api
 
 
 @pytest.fixture
@@ -31,7 +31,9 @@ def shells() -> list[RadialWindow]:
 
 
 @pytest.fixture
-def cosmo() -> StandardCosmology:
+def cosmo() -> (
+    cosmology.api.StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]]
+):
     class MockCosmology:
         @property
         def Omega_m0(self) -> float:  # noqa: N802
@@ -97,7 +99,9 @@ def test_deflect_many(rng: np.random.Generator) -> None:
 
 def test_multi_plane_matrix(
     shells: list[RadialWindow],
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
     rng: np.random.Generator,
 ) -> None:
     mat = multi_plane_matrix(shells, cosmo)
@@ -119,7 +123,9 @@ def test_multi_plane_matrix(
 
 def test_multi_plane_weights(
     shells: list[RadialWindow],
-    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
+    cosmo: cosmology.api.StandardCosmology[
+        npt.NDArray[np.float64], npt.NDArray[np.float64]
+    ],
     rng: np.random.Generator,
 ) -> None:
     w_in = np.eye(len(shells))

From 16087c8a098439db52ad08d20f16c8bbb5fd91e5 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 15:47:28 +0000
Subject: [PATCH 08/15] Use critical density

---
 glass/shells.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/glass/shells.py b/glass/shells.py
index c1dd843f..80a28952 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -88,7 +88,12 @@ def density_weight(
     ],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in matter density."""
-    return cosmo.rho_m_z(z) * cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
+    return (  # type: ignore[no-any-return]
+        cosmo.critical_density0
+        * cosmo.Omega_m(z)
+        * cosmo.xm(z) ** 2
+        / cosmo.H_over_H0(z)
+    )
 
 
 class RadialWindow(typing.NamedTuple):

From 0a83b2e5507cbe2a9dde5a2b70e08e0e922f5676 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 15:48:38 +0000
Subject: [PATCH 09/15] Revert "NLA model for intrinsic alignments (#21)"

This reverts commit 60dc66815559193a0cafde7a4584dda91a7cd4c5.
---
 glass/__init__.py |   1 -
 glass/galaxies.py | 108 ----------------------------------------------
 2 files changed, 109 deletions(-)

diff --git a/glass/__init__.py b/glass/__init__.py
index 71b7890e..f419d9ba 100644
--- a/glass/__init__.py
+++ b/glass/__init__.py
@@ -21,7 +21,6 @@
 from glass.galaxies import (
     galaxy_shear,
     gaussian_phz,
-    kappa_ia_nla,
     redshifts,
     redshifts_from_nz,
 )
diff --git a/glass/galaxies.py b/glass/galaxies.py
index b622a60c..bb36d503 100644
--- a/glass/galaxies.py
+++ b/glass/galaxies.py
@@ -15,10 +15,6 @@
 .. autofunction:: galaxy_shear
 .. autofunction:: gaussian_phz
 
-Intrinsic alignments
---------------------
-.. autofunction:: kappa_ia_nla
-
 """  # noqa: D205, D400
 
 from __future__ import annotations
@@ -33,8 +29,6 @@
 from glass.core.array import broadcast_leading_axes, cumtrapz
 
 if typing.TYPE_CHECKING:
-    import cosmology.api
-
     from glass.shells import RadialWindow
 
 
@@ -294,105 +288,3 @@ def gaussian_phz(
             trunc = trunc[(znew < lower) | (znew > upper)]
 
     return zphot
-
-
-def kappa_ia_nla(  # noqa: PLR0913
-    delta: npt.NDArray[np.float64],
-    zeff: float,
-    a_ia: float,
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
-    *,
-    z0: float = 0.0,
-    eta: float = 0.0,
-    lbar: float = 0.0,
-    l0: float = 1e-9,
-    beta: float = 0.0,
-) -> npt.NDArray[np.float64]:
-    r"""
-    Effective convergence from intrinsic alignments using the NLA model.
-
-    Returns the effective convergence due to intrinsic alignments.
-
-    Parameters
-    ----------
-    delta:
-        Matter density contrast.
-    zeff:
-        Effective redshift of the matter field.
-    a_ia:
-        Intrinsic alignments amplitude.
-    cosmo:
-        Cosmology instance.
-    z0:
-        Reference redshift for the redshift dependence.
-    eta:
-        Power of the redshift dependence.
-    lbar:
-        Mean luminosity of the galaxy sample.
-    l0:
-        Reference luminosity for the luminosity dependence.
-    beta:
-        Power of the luminosity dependence.
-
-    Notes
-    -----
-    The Non-linear Alignments Model (NLA) describes an effective
-    convergence :math:`\kappa_{\rm IA}` that models the effect of
-    intrinsic alignments.  It is computed from the matter density
-    contrast :math:`\delta` as [1] [3]
-
-    .. math::
-
-        \kappa_{\rm IA} = f_{\rm NLA} \, \delta \;,
-
-    where the NLA factor :math:`f_{\rm NLA}` is defined as [4] [5]
-
-    .. math::
-
-        f_{\rm{NLA}}
-        = -A_{\rm IA} \, \frac{C_1 \, \bar{\rho}(z)}{D(z)} \,
-            \biggl(\frac{1+z}{1+z_0}\biggr)^\eta \,
-            \biggl(\frac{\bar{L}}{L_0}\biggr)^\beta \;,
-
-    with
-
-    * :math:`A_{\rm IA}` the intrinsic alignments amplitude,
-    * :math:`C_1` a normalisation constant [2],
-    * :math:`z` the effective redshift of the model,
-    * :math:`\bar{\rho}` the mean matter density,
-    * :math:`D` the growth factor,
-    * :math:`\eta` the power that describes the redshift-dependence with
-      respect to :math:`z_0`,
-    * :math:`\bar{L}` the mean luminosity of the galaxy sample, and
-    * :math:`\beta` the power that describes the luminosity-dependence
-      :math:`\bar{L}` with respect to :math:`L_0`.
-
-    References
-    ----------
-    * [1] Catelan P., Kamionkowski M., Blandford R. D., 2001, MNRAS,
-       320, L7. doi:10.1046/j.1365-8711.2001.04105.x
-    * [2] Hirata C. M., Seljak U., 2004, PhRvD, 70, 063526.
-       doi:10.1103/PhysRevD.70.063526
-    * [3] Bridle S., King L., 2007, NJPh, 9, 444.
-       doi:10.1088/1367-2630/9/12/444
-    * [4] Johnston, H., Georgiou, C., Joachimi, B., et al., 2019,
-        A&A, 624, A30. doi:10.1051/0004-6361/201834714
-    * [5] Tessore, N., Loureiro, A., Joachimi, B., et al., 2023,
-       OJAp, 6, 11. doi:10.21105/astro.2302.01942
-
-    """
-    c1 = 5e-14 / cosmo.h**2  # Solar masses per cubic Mpc
-    rho_c1 = c1 * cosmo.critical_density0
-
-    prefactor = -a_ia * rho_c1 * cosmo.Om
-    inverse_linear_growth = 1.0 / cosmo.gf(zeff)
-    redshift_dependence = ((1 + zeff) / (1 + z0)) ** eta
-    luminosity_dependence = (lbar / l0) ** beta
-
-    f_nla = (
-        prefactor * inverse_linear_growth * redshift_dependence * luminosity_dependence
-    )
-
-    return delta * f_nla  # type: ignore[no-any-return]

From 663f80a22f0acf86235af50abe0cdd0fbd792eae Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 16:05:32 +0000
Subject: [PATCH 10/15] Go back to `from` syntax

---
 glass/lensing.py      | 14 ++++----------
 glass/shells.py       | 16 ++++------------
 tests/test_lensing.py | 14 ++++----------
 3 files changed, 12 insertions(+), 32 deletions(-)

diff --git a/glass/lensing.py b/glass/lensing.py
index 89a14fc1..f3048d56 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -40,7 +40,7 @@
 if typing.TYPE_CHECKING:
     import collections.abc
 
-    import cosmology.api
+    from cosmology.api import StandardCosmology
 
     from glass.shells import RadialWindow
 
@@ -271,9 +271,7 @@ class MultiPlaneConvergence:
 
     def __init__(
         self,
-        cosmo: cosmology.api.StandardCosmology[
-            npt.NDArray[np.float64], npt.NDArray[np.float64]
-        ],
+        cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     ) -> None:
         """Create a new instance to iteratively compute the convergence."""
         self.cosmo = cosmo
@@ -382,9 +380,7 @@ def wlens(self) -> float:
 
 def multi_plane_matrix(
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Compute the matrix of lensing contributions from each shell."""
     mpc = MultiPlaneConvergence(cosmo)
@@ -398,9 +394,7 @@ def multi_plane_matrix(
 def multi_plane_weights(
     weights: npt.NDArray[np.float64],
     shells: collections.abc.Sequence[RadialWindow],
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """
     Compute effective weights for multi-plane convergence.
diff --git a/glass/shells.py b/glass/shells.py
index 80a28952..5b324c69 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -63,9 +63,7 @@
 
 def distance_weight(
     z: npt.NDArray[np.float64],
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving distance."""
     return 1 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -73,9 +71,7 @@ def distance_weight(
 
 def volume_weight(
     z: npt.NDArray[np.float64],
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving volume."""
     return cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
@@ -83,9 +79,7 @@ def volume_weight(
 
 def density_weight(
     z: npt.NDArray[np.float64],
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in matter density."""
     return (  # type: ignore[no-any-return]
@@ -661,9 +655,7 @@ def redshift_grid(
 
 
 def distance_grid(
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     zmin: float,
     zmax: float,
     *,
diff --git a/tests/test_lensing.py b/tests/test_lensing.py
index 0d530c52..4d5283ff 100644
--- a/tests/test_lensing.py
+++ b/tests/test_lensing.py
@@ -16,7 +16,7 @@
 from glass.shells import RadialWindow
 
 if typing.TYPE_CHECKING:
-    import cosmology.api
+    from cosmology.api import StandardCosmology
 
 
 @pytest.fixture
@@ -31,9 +31,7 @@ def shells() -> list[RadialWindow]:
 
 
 @pytest.fixture
-def cosmo() -> (
-    cosmology.api.StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]]
-):
+def cosmo() -> StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
     class MockCosmology:
         @property
         def Omega_m0(self) -> float:  # noqa: N802
@@ -99,9 +97,7 @@ def test_deflect_many(rng: np.random.Generator) -> None:
 
 def test_multi_plane_matrix(
     shells: list[RadialWindow],
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     rng: np.random.Generator,
 ) -> None:
     mat = multi_plane_matrix(shells, cosmo)
@@ -123,9 +119,7 @@ def test_multi_plane_matrix(
 
 def test_multi_plane_weights(
     shells: list[RadialWindow],
-    cosmo: cosmology.api.StandardCosmology[
-        npt.NDArray[np.float64], npt.NDArray[np.float64]
-    ],
+    cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
     rng: np.random.Generator,
 ) -> None:
     w_in = np.eye(len(shells))

From c8e5dddd6f0cbf3eddd9907f985c31dd2666ed8f Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 16:36:48 +0000
Subject: [PATCH 11/15] Use Hubble length

---
 glass/lensing.py | 10 ++++++----
 glass/shells.py  | 13 ++++++++-----
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/glass/lensing.py b/glass/lensing.py
index f3048d56..f0d7a02e 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -37,11 +37,11 @@
 import numpy as np
 import numpy.typing as npt
 
+from cosmology.api import StandardCosmology, CosmologyConstantsNamespace
+
 if typing.TYPE_CHECKING:
     import collections.abc
 
-    from cosmology.api import StandardCosmology
-
     from glass.shells import RadialWindow
 
 
@@ -330,9 +330,11 @@ def add_plane(
         w2, self.w3 = self.w3, wlens
 
         # extrapolation law
-        x2, self.x3 = self.x3, self.cosmo.xm(self.z3)
+        hubble_length = CosmologyConstantsNamespace.c / StandardCosmology.H0
+        x2 = self.x3
+        self.x3 = self.cosmo.transverse_comoving_distance(self.z3) / hubble_length
         r12 = self.r23
-        r13, self.r23 = self.cosmo.xm([z1, self.z2], self.z3) / self.x3
+        r13, self.r23 = self.cosmo.transverse_comoving_distance([z1, self.z2], self.z3) / (hubble_length * self.x3)
         t = r13 / r12
 
         # lensing weight of mass plane to be added
diff --git a/glass/shells.py b/glass/shells.py
index 5b324c69..c70aa4b1 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -52,10 +52,9 @@
 import numpy as np
 import numpy.typing as npt
 
-from glass.core.array import ndinterp
+from cosmology.api import CosmologyConstantsNamespace, StandardCosmology
 
-if typing.TYPE_CHECKING:
-    import cosmology.api
+from glass.core.array import ndinterp
 
 ArrayLike1D = typing.Union[collections.abc.Sequence[float], npt.NDArray[np.float64]]
 WeightFunc = typing.Callable[[ArrayLike1D], npt.NDArray[np.float64]]
@@ -74,7 +73,10 @@ def volume_weight(
     cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in comoving volume."""
-    return cosmo.xm(z) ** 2 / cosmo.H_over_H0(z)  # type: ignore[no-any-return]
+    hubble_length = CosmologyConstantsNamespace.c / StandardCosmology.H0
+    return (  # type: ignore[no-any-return]
+        cosmo.transverse_comoving_distance(z) / hubble_length
+    ) ** 2 / cosmo.H_over_H0(z)
 
 
 def density_weight(
@@ -82,10 +84,11 @@ def density_weight(
     cosmo: StandardCosmology[npt.NDArray[np.float64], npt.NDArray[np.float64]],
 ) -> npt.NDArray[np.float64]:
     """Uniform weight in matter density."""
+    hubble_length = CosmologyConstantsNamespace.c / StandardCosmology.H0
     return (  # type: ignore[no-any-return]
         cosmo.critical_density0
         * cosmo.Omega_m(z)
-        * cosmo.xm(z) ** 2
+        * (cosmo.transverse_comoving_distance(z) / hubble_length) ** 2
         / cosmo.H_over_H0(z)
     )
 

From 55854dc7090c6d926b4cc1a56c65ebe567b115ec Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 16:37:40 +0000
Subject: [PATCH 12/15] Add docstring

---
 tests/test_lensing.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_lensing.py b/tests/test_lensing.py
index 4d5283ff..ac1a9ced 100644
--- a/tests/test_lensing.py
+++ b/tests/test_lensing.py
@@ -45,6 +45,7 @@ def xm(
             z: npt.NDArray[np.float64],
             z2: npt.NDArray[np.float64] | None = None,
         ) -> npt.NDArray[np.float64]:
+            """Dimensionless transverse comoving distance."""
             if z2 is None:
                 return np.array(z) * 1000
             return (np.array(z2) - np.array(z)) * 1000

From d88aa4467cf0d2582c89673ff9772e8ea6671b70 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Fri, 8 Nov 2024 16:41:46 +0000
Subject: [PATCH 13/15] Fix linting

---
 glass/lensing.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/glass/lensing.py b/glass/lensing.py
index f0d7a02e..09d6336d 100644
--- a/glass/lensing.py
+++ b/glass/lensing.py
@@ -37,7 +37,7 @@
 import numpy as np
 import numpy.typing as npt
 
-from cosmology.api import StandardCosmology, CosmologyConstantsNamespace
+from cosmology.api import CosmologyConstantsNamespace, StandardCosmology
 
 if typing.TYPE_CHECKING:
     import collections.abc
@@ -334,7 +334,9 @@ def add_plane(
         x2 = self.x3
         self.x3 = self.cosmo.transverse_comoving_distance(self.z3) / hubble_length
         r12 = self.r23
-        r13, self.r23 = self.cosmo.transverse_comoving_distance([z1, self.z2], self.z3) / (hubble_length * self.x3)
+        r13, self.r23 = self.cosmo.transverse_comoving_distance(
+            [z1, self.z2], self.z3
+        ) / (hubble_length * self.x3)
         t = r13 / r12
 
         # lensing weight of mass plane to be added

From ff590a9ed0d9cc3156371dbb3aa899f87b595523 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Thu, 14 Nov 2024 09:29:35 +0000
Subject: [PATCH 14/15] Change the badge

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 42fbcbd6..2e111f6c 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@
 <!-- Essentials -->
 
 [![PyPI](https://img.shields.io/pypi/v/glass)](https://pypi.org/project/glass)
-[![Documentation](https://readthedocs.org/projects/glass/badge/?version=latest)](https://glass.readthedocs.io/latest)
+[![Documentation](https://readthedocs.org/projects/glass/badge/?version=stable)](https://glass.readthedocs.io/stable)
 [![LICENSE](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
 
 <!-- Code -->

From c7745ef7fc4c81146086b885c49785bbd4b78062 Mon Sep 17 00:00:00 2001
From: "Patrick J. Roddy" <patrickjamesroddy@gmail.com>
Date: Thu, 14 Nov 2024 09:34:55 +0000
Subject: [PATCH 15/15] Revert "Change the badge"

This reverts commit ff590a9ed0d9cc3156371dbb3aa899f87b595523.
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 2e111f6c..42fbcbd6 100644
--- a/README.md
+++ b/README.md
@@ -3,7 +3,7 @@
 <!-- Essentials -->
 
 [![PyPI](https://img.shields.io/pypi/v/glass)](https://pypi.org/project/glass)
-[![Documentation](https://readthedocs.org/projects/glass/badge/?version=stable)](https://glass.readthedocs.io/stable)
+[![Documentation](https://readthedocs.org/projects/glass/badge/?version=latest)](https://glass.readthedocs.io/latest)
 [![LICENSE](https://img.shields.io/badge/License-MIT-blue.svg)](https://opensource.org/licenses/MIT)
 
 <!-- Code -->