Merge branch 'develop'

Author: davidesciotti

config.yaml

@@ -69,7 +69,7 @@ probe_selection:
 nz:
   # Paths to source (for WL) and lens (for GCph) redshift distributions.
   # These must have shape (z_points, zbins + 1) where z_points is the number of
-  # redshift values over which the n(z) is measured, and format .txt or .dat
+  # redshift values over which the n(z) is measured, and format .txt, .dat or .ascii
   # OR
   # as .fits files, following the LE3/Heracles/euclidlib format
   nz_sources_filename: ../common_data/Spaceborne_jobs/develop/input/nzTab-EP03-zedMin02-zedMax25-mag245.dat # Type: str
@@ -101,6 +101,12 @@ binning:
   ell_max: 3000 # Type: int | float. Maximum ell
   ell_bins: 32 # Type: int. Number of ell bins. Not used in the 'unbinned' case
 
+  # Path to user-provided ell binning scheme. Used only if 'binning_type' is set to 'from_input'.
+  # The file should have the following columns:
+  # ell, delta_ell, ell_lower_edges, ell_upper_edges  
+  # with `ell` the ell bin center, `delta_ell` its width,
+  # `ell_lower_edges` its lower edge, and `ell_upper_edges` its upper edge.
+  # The file can be in .txt, .dat or .ascii format.
   ell_bins_filename: ../common_data/Spaceborne_jobs/develop/input/ell_values_3x2pt.txt
 
   # theta binning for the real-space covariance matrix
@@ -337,20 +343,27 @@ precision:
   # see https://namaster.readthedocs.io/en/latest/api/pymaster.field.html#pymaster.field.NmtField
   n_iter_nmt: null # Type: None | int.
 
+   
+  # ======================================= JAX ========================================
+  # Whether to enable 64-bit precision for JAX computations. Setting this to False can result
+  # in noticeable differences with respect to the plain numpy implementation, since the 
+  # latter defaults to 64-bit precision.
+  jax_enable_x64: True # Type: bool. 
+
+
 misc:
   # note if you work from an interactive environment (e.g. jupyter notebook), you might
   # have to restart the kernel after changing the num_threads, and jax_platform
   num_threads: 50 # Type: int. How many threads to use for the parallel computations
 
-  # jax-specific settings
-  jax_platform: auto # Type: str. 'auto', 'cpu' or 'gpu'
-  jax_enable_x64: True # Type: bool. Whether to enable 64-bit precision for JAX computations
+  # Device used for JAX-accelerated computation. Can be one of 'auto', 'cpu' or 'gpu'
+  jax_platform: auto # Type: str. 
 
   # some sanity checks on the covariance matrix. These could take a while to run.
-  test_numpy_inversion: True # Type: bool. Test for errors in np.linalg.inv
-  test_condition_number: True # Type: bool. Check if condition number is above 1e10
-  test_cholesky_decomposition: True # Type: bool. Test if cholesky decomposition fails
-  test_symmetry: True # Type: bool. Test if covariance matrix is symmetric (cov = cov.T)
+  test_numpy_inversion: False # Type: bool. Test for errors in np.linalg.inv
+  test_condition_number: False # Type: bool. Check if condition number is above 1e10
+  test_cholesky_decomposition: False # Type: bool. Test if cholesky decomposition fails
+  test_symmetry: False # Type: bool. Test if covariance matrix is symmetric (cov = cov.T)
 
   # Whether to produce a triangle plot of the Cls, especially useful when using input
   # Cls for a quick visual comparison. Can be slow, set to False for a faster runtime.

main.py

@@ -1,9 +1,9 @@
 # ruff: noqa: E402 (ignore module import not on top of the file warnings)
 import argparse
 import contextlib
-import itertools
 import os
 import sys
+import warnings
 
 import yaml
 
@@ -39,9 +39,26 @@ def load_config(_config_path):
 
 
 cfg = load_config('config.yaml')
-# Set jax platform
+
+# JAX settings
 if cfg['misc']['jax_platform'] != 'auto':
     os.environ['JAX_PLATFORMS'] = cfg['misc']['jax_platform']
+if cfg['precision']['jax_enable_x64']:
+    os.environ['JAX_ENABLE_X64'] = '1'
+else:
+    os.environ['JAX_ENABLE_X64'] = '0'
+    warnings.warn(
+        'JAX 64-bit precision is disabled. This may lead to '
+        'noticeable differences with respect to the numpy '
+        'implementation, which uses 64-bit precision by default.',
+        stacklevel=2,
+    )
+
+# Import JAX after environment variables are set, then print device info
+import jax
+
+print(f'JAX devices: {jax.devices()}')
+print(f'JAX backend: {jax.default_backend()}')
 
 # if using the CPU, set the number of threads
 num_threads = cfg['misc']['num_threads']
@@ -65,9 +82,9 @@ def load_config(_config_path):
 # override in cfg as well
 cfg['misc']['num_threads'] = num_threads
 
+import itertools
 import pprint
 import time
-import warnings
 from functools import partial
 
 import matplotlib
@@ -2370,7 +2387,12 @@ def plot_cls():
         else 'Gauss'
     )
     cov = covs_3x2pt_2d_tosave_dict[key]
-    print(f'Testing cov {key}...')
+
+    print(
+        f'Performing sanity checks on cov {key}.\n'
+        'This can take some time for large matrices. '
+        'Please note that your files have already been saved.\n'
+    )
 
     if cfg['misc']['test_condition_number']:
         cond_number = np.linalg.cond(cov)
@@ -2418,16 +2440,16 @@ def plot_cls():
 
 # note that this is *not* compatible with %matplotlib inline in the interactive window!
 if cfg['misc']['save_figs']:
-    output_dir = f'{output_path}/figs'
-    os.makedirs(output_dir, exist_ok=True)
+    output_path_figs = f'{output_path}/figs'
+    os.makedirs(output_path_figs, exist_ok=True)
     for i, fig_num in enumerate(plt.get_fignums()):
         fig = plt.figure(fig_num)
         fig.savefig(
-            os.path.join(output_dir, f'fig_{i:03d}.pdf'),
+            os.path.join(output_path_figs, f'fig_{i:03d}.pdf'),
             bbox_inches='tight',
             pad_inches=0.1,
         )
-    print(f'Figures saved in {output_dir}\n')
+    print(f'Figures saved in {output_path_figs}\n')
 
 
 print(f'Finished in {(time.perf_counter() - script_start_time) / 60:.2f} minutes')

spaceborne/ccl_interface.py

@@ -766,7 +766,7 @@ def compute_ng_cov_3x2pt(
             symmetrize_zpairs = (probe_a, probe_b) == (probe_c, probe_d)
 
             tqdm.write(
-                f'CCL {which_ng_cov}: computing probe combination {probe_ab, probe_cd}'
+                f'CCL {which_ng_cov} cov: computing probe combination {probe_ab, probe_cd}'
             )
 
             self.cov_dict[ng_term][probe_2tpl]['4d'] = self.compute_ng_cov_probe_block(
@@ -794,7 +794,7 @@ def compute_ng_cov_3x2pt(
             nbx=len(ells),
             zbins=None,
             ind_dict=ind_dict,
-            msg=f'CCL {which_ng_cov}: ',
+            msg=f'CCL {which_ng_cov} cov: ',
         )
 
     def check_cov_blocks_symmetry(self):

spaceborne/config_checker.py

@@ -548,6 +548,9 @@ def check_types(self) -> None:
         assert isinstance(precision_cfg.get('cov_rs_int_method'), str), (
             'precision: cov_rs_int_method must be a string'
         )
+        assert isinstance(precision_cfg.get('jax_enable_x64'), bool), (
+            'precision: jax_enable_x64 must be a boolean'
+        )
 
         # misc
         assert isinstance(self.cfg.get('misc'), dict), (
@@ -578,9 +581,6 @@ def check_types(self) -> None:
         assert isinstance(misc_cfg.get('jax_platform'), str), (
             'misc: jax_platform must be a string'
         )
-        assert isinstance(misc_cfg.get('jax_enable_x64'), bool), (
-            'misc: jax_enable_x64 must be a boolean'
-        )
         assert isinstance(misc_cfg.get('cl_triangle_plot'), bool), (
             'misc: cl_triangle_plot must be a boolean'
         )

spaceborne/cov_harmonic_space.py

@@ -119,7 +119,7 @@ def consistency_checks(self):
 
     def set_gauss_cov(self, ccl_obj: CCLInterface):
         start = time.perf_counter()
-        
+
         print('\nComputing Gaussian harmonic-space covariance matrix...')
 
         # signal
@@ -154,15 +154,15 @@ def set_gauss_cov(self, ccl_obj: CCLInterface):
             )
 
         if self.cfg['precision']['cov_hs_g_ell_bin_average']:
-            # unbinned cls and noise; need the edges to compute the number of modes 
-            # (after casting them to int. n_modes is equivalent to delta_ell modulo the 
+            # unbinned cls and noise; need the edges to compute the number of modes
+            # (after casting them to int. n_modes is equivalent to delta_ell modulo the
             # fact that for delta_ell we consider non-integer ell values)
             _cl_5d = self.cl_3x2pt_unb_5d
             _noise_5d = noise_3x2pt_unb_5d
             _ell_values = self.ell_obj.ells_3x2pt_unb
             _ell_edges = self.ell_obj.ell_edges_3x2pt
         else:
-            # evaluate the covariance at the center of the ell bin and normalise by 
+            # evaluate the covariance at the center of the ell bin and normalise by
             # delta_ell
             _cl_5d = cl_3x2pt_5d
             _noise_5d = noise_3x2pt_5d

spaceborne/cov_partial_sky.py

@@ -178,7 +178,7 @@ def cl_22_list(zi, zj, spin0):
         probe_ab, probe_cd = sl.split_probe_name(probe_abcd, space='harmonic')
         probe_a, probe_b, probe_c, probe_d = list(probe_abcd)
 
-        tqdm.write(f'computing probe combination {probe_ab, probe_cd}')
+        tqdm.write(f'NaMaster G cov: computing probe combination {probe_ab, probe_cd}')
 
         s1 = spin_dict[probe_a]
         s2 = spin_dict[probe_b]
@@ -244,7 +244,7 @@ def cl_22_list(zi, zj, spin0):
         nbx=nbl,
         zbins=zbins,
         ind_dict=ind_dict,
-        msg='',
+        msg='NaMaster G cov: ',
     )
 
     return cov_dict

spaceborne/cov_ssc.py

@@ -1,6 +1,5 @@
 import time
 
-import jax
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
 import numpy as np
@@ -167,10 +166,6 @@ def __init__(self, cfg, pvt_cfg, ccl_obj, z_grid):
         self.z_grid = z_grid
         self.ccl_obj = ccl_obj
 
-        # Enable 64-bit precision if required
-        jax.config.update('jax_enable_x64', cfg['misc']['jax_enable_x64'])
-        print('JAX devices:', jax.devices())
-
         # set some useful attributes
         if self.use_ke_approx:
             self.ssc_func = ssc_integral_4D_simps_jax_ke_approx
@@ -339,7 +334,7 @@ def compute_ssc(
             nbx=nbl,
             zbins=None,
             ind_dict=self.ind_dict,
-            msg='SSC: ',
+            msg='SSC cov: ',
         )
 
         print(f'...done in {(time.perf_counter() - start):.2f} s')