Updated example notebook, tested jit and grad on the GPU

Author: TonyZhou729

ABCMB/main.py

@@ -157,9 +157,9 @@ def run_cosmology(self, params : dict):
         ells = self.SS.ells
 
         if self.return_PTBG:
-            return ells, Cls, PT, BG
+            return Cls, ells, PT, BG
         else:
-            return ells, Cls
+            return Cls, ells
 
     # @jit
     @eqx.filter_jit

Module_Tests/spectrum_test.ipynb

@@ -44,7 +44,7 @@ def test_accuracy_checker(h = 0.6762):
         }
 
         model = Model(specs) # ZZ: model now takes ellmin, ellmax for Cls, and want_lensing
-        ABC_ell, ABC_Cls = model.run_cosmology(params)
+        ABC_Cls, ABC_ell = model.run_cosmology(params)
         ABC_tt = ABC_Cls[0] 
         ABC_te = ABC_Cls[1] 
         ABC_ee = ABC_Cls[2]