diff --git a/glass/cosmology.py b/glass/cosmology.py
index 40e0cb1d..c9d8f441 100644
--- a/glass/cosmology.py
+++ b/glass/cosmology.py
@@ -16,8 +16,14 @@ class Cosmology(
     cosmology.api.HasInverseComovingDistance[AnyArray, AnyArray],  # ty: ignore[invalid-type-arguments]
     cosmology.api.HasLittleH[AnyArray],  # ty: ignore[invalid-type-arguments]
     cosmology.api.HasOmegaM0[AnyArray],  # ty: ignore[invalid-type-arguments]
-    cosmology.api.HasOmegaM[AnyArray, AnyArray],  # ty: ignore[invalid-type-arguments]
     cosmology.api.HasTransverseComovingDistance[AnyArray, AnyArray],  # ty: ignore[invalid-type-arguments]
     typing.Protocol,
 ):
     """Cosmology protocol for GLASS."""
+
+
+class CosmologyWithOmegaM(
+    Cosmology,
+    cosmology.api.HasOmegaM[AnyArray, AnyArray],  # ty: ignore[invalid-type-arguments]
+):
+    """Cosmology protocol for GLASS with OmegaM."""
diff --git a/glass/shells.py b/glass/shells.py
index 95471890..212936d9 100644
--- a/glass/shells.py
+++ b/glass/shells.py
@@ -67,7 +67,7 @@
     from types import ModuleType
 
     from glass._types import FloatArray, IntArray, UnifiedGenerator
-    from glass.cosmology import Cosmology
+    from glass.cosmology import Cosmology, CosmologyWithOmegaM
 
 
 @dataclasses.dataclass
@@ -142,11 +142,11 @@ class DensityWeight:
     Attributes
     ----------
     cosmo
-        Cosmology instance.
+        Cosmology instance with OmegaM.
 
     """
 
-    cosmo: Cosmology
+    cosmo: CosmologyWithOmegaM
 
     def __call__(self, z: FloatArray) -> FloatArray:
         """