Go back to from syntax

glass-dev · Nov 8, 2024 · 663f80a · 663f80a
1 parent 0a83b2e
commit 663f80a
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 32 deletions.
diff --git 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
@@ -63,29 +63,23 @@
 
 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]
 
 
 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]
 
 
 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
@@ -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))