Make Cosmology.k jitable

DifferentiableUniverseInitiative · Jan 20, 2022 · cf8500c · cf8500c
1 parent 3e45de0
commit cf8500c
Showing 1 changed file with 1 addition and 7 deletions.
diff --git a/jax_cosmo/core.py b/jax_cosmo/core.py
@@ -168,13 +168,7 @@ def Omega_k(self):
 
     @property
     def k(self):
-        if self.Omega > 1.0:  # Closed universe
-            k = 1
-        elif self.Omega == 1.0:  # Flat universe
-            k = 0
-        elif self.Omega < 1.0:  # Open Universe
-            k = -1
-        return k
+        return -np.sign(self._Omega_k).astype(np.int8)
 
     @property
     def sqrtk(self):