Jax implementation for speeding things up

flip/__init__.py

@@ -2,7 +2,17 @@
 
 import os
 
+from flip.utils import create_log
+
+log = create_log()
 from . import covariance, fisher, fitter, gridding, likelihood, power_spectra, utils
 
+try:
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+except:
+    log.add("Jax is not available, loading numpy and scipy instead")
+
 __version__ = "1.0.0"
 __flip_dir_path__ = os.path.dirname(__file__)

flip/covariance/__init__.py

@@ -1,4 +1,5 @@
 """Init file of the flip.covariance package."""
 
 from . import adamsblake17plane, adamsblake20, carreres23, lai22, ravouxcarreres, rcrk24
+from . import cov_utils
 from .covariance import CovMatrix

flip/covariance/cov_utils.py

@@ -231,6 +231,17 @@ def return_full_cov(cov):
     return full_cov
 
 
+def return_flat_cov(cov):
+    variance_val = cov[0, 0]
+    flat_cov = cov[np.triu_indices_from(cov, k=1)]
+    flat_cov = np.insert(flat_cov, 0, variance_val)
+    return flat_cov
+
+
+def return_flat_cross_cov(cov):
+    return cov.flatten()
+
+
 def return_full_cov_cross(cov, number_objects_g, number_objects_v):
     """
     The return_full_cov_cross function takes in a covariance matrix and the number of objects in each band.
@@ -245,7 +256,7 @@ def return_full_cov_cross(cov, number_objects_g, number_objects_v):
         The full covariance matrix
 
     """
-    full_cov = cov[1:].reshape((number_objects_g, number_objects_v))
+    full_cov = cov.reshape((number_objects_g, number_objects_v))
     return full_cov