Merge pull request #242 from LSSTDESC/gaussian_sims_run4

Gaussian sims run4

examples/cosmodc2/pipeline_redmagic.yml

@@ -84,7 +84,6 @@ inputs:
     # See README for paths to download these files
     shear_catalog:  /global/projecta/projectdirs/lsst/groups/WL/users/zuntz/data/cosmoDC2-1.1.4_oneyear_unit_response/shear_catalog.hdf5
     photometry_catalog:  /global/projecta/projectdirs/lsst/groups/WL/users/zuntz/data/cosmoDC2-1.1.4_oneyear_unit_response/photometry_catalog.hdf5
-    #shear_photometry_catalog:  /global/projecta/projectdirs/lsst/groups/WL/users/zuntz/data/cosmoDC2-1.1.4_oneyear_unit_response/photometry_catalog.hdf5  # not needed anymore after metadetect merge
     lens_photometry_catalog:  /global/projecta/projectdirs/lsst/groups/WL/users/zuntz/data/cosmoDC2-1.1.4_oneyear_unit_response/photometry_catalog.hdf5
     photoz_trained_model:  /global/projecta/projectdirs/lsst/groups/WL/users/zuntz/data/cosmoDC2-1.1.4_oneyear_unit_response/cosmoDC2_trees_i25.3.npy
     fiducial_cosmology: data/fiducial_cosmology.yml

examples/gaussian_sims/config.yml

@@ -9,7 +9,7 @@ TXExposureInfo:
     dc2_name: '1.2p'
 
 TXGaussianSimsMock:
-    cat_name: '/global/projecta/projectdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/022422/12300area/shearcat_shearalm_allbins_ra_dec_g1_g2_and_fakecols_areacut.npy'
+    cat_name: '/global/cfs/cdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/051422/12300area/shearcat_allbins_ra_dec_e1_e2_and_fakecols_areacut.npy'
     visits_per_band: 16
     extra_cols: redshift_true size_true shear_1 shear_2 
     flip_g2: False # opposite of cosmodc2
@@ -89,7 +89,13 @@ TXSourceDiagnosticPlots:
 
 TXFourierGaussianCovariance:
     galaxy_bias: [1., 1., 1., 1., 1.]
-    cache_dir: ./cache_nmt
+    cache_dir: ./cache_nmt/nside2048/
+    gaussian_sims_factor: [2., 1.5, 1.25, 1.25, 1.25]
+
+TXFourierTJPCovariance:
+    galaxy_bias: [1., 1., 1., 1., 1.]
+    cache_dir: ./cache_nmt/nside2048/
+    IA: 0.
 
 TXRealGaussianCovariance:
     min_sep: 2.5
@@ -99,22 +105,24 @@ TXRealGaussianCovariance:
     nprocess: 4
     threads_per_process: 2
     nodes: 4
+    gaussian_sims_factor: [2., 1.5, 1.25, 1.25, 1.25]
     galaxy_bias: [1., 1., 1., 1., 1.]
 
 TXTwoPointFourier:
     chunk_rows: 100000
     flip_g1: False
     flip_g2: False
     apodization_size: 0.0
-    cache_dir: ./cache_nmt
+    cache_dir: ./cache_nmt/nside2048/
     true_shear: False
     n_ell: 30
-    ell_max: 6143 # nside * 3 - 1, since Namaster computes that anyway. 
+    ell_max: 6144 # nside * 3 , since Namaster computes that anyway. 
     nside: 2048
-    analytic_noise: True
+    analytic_noise: True 
+    gaussian_sims_factor: [2., 1.5, 1.25, 1.25, 1.25]
 
 TXTwoPoint:
-    bin_slop: 0.05
+    bin_slop: 0.01
     delta_gamma: 0.02
     do_pos_pos: True
     do_shear_shear: True

examples/gaussian_sims/pipeline_gaussian_sims.yml

@@ -32,7 +32,7 @@ stages:
       nodes: 2
     - name: TXPhotozPlots
     - name: TXJackknifeCenters
-    - name: TXMapPlots
+    #- name: TXMapPlots
     - name: TXTracerMetadata
     - name: TXSourceMaps
       nprocess: 12
@@ -56,45 +56,47 @@ stages:
       nprocess: 4
       nodes: 4
       threads_per_process: 32
-    #- name: TXFourierGaussianCovariance
-    #  threads_per_process: 64
-    - name: TXFourierTJPCovariance
+    - name: TXFourierGaussianCovariance
       threads_per_process: 64
+    #- name: TXFourierTJPCovariance
+    #  nodes: 6
+    #  nprocess: 6
+    #  threads_per_process: 32
     - name: TXConvergenceMaps    # Make convergence kappa maps from g1, g2 maps
       threads_per_process: 32
-    - name: TXConvergenceMapPlots # Plot the convergence map
+    #- name: TXConvergenceMapPlots # Plot the convergence map
     #- name: TXTwoPointPlotsFourier 
-    # - name: TXNullBlinding
-    #- name: TXTwoPoint
-    #  threads_per_process: 32
-    #  nprocess: 6
-    #  nodes: 6
-    #- name: TXTwoPointTheoryReal
-    #- name: TXTwoPointTheoryFourier
+    - name: TXNullBlinding
+    - name: TXTwoPoint
+      threads_per_process: 32
+      nprocess: 6
+      nodes: 6
+    - name: TXTwoPointTheoryReal
+    - name: TXTwoPointTheoryFourier
     #- name: TXTwoPointPlots
-    #- name: TXRealGaussianCovariance
+    - name: TXRealGaussianCovariance
       threads_per_process: 64
 
 
-output_dir: data/gaussian_sims/outputs_gaussian_sims/gaussian_sims_srdnzs_fullsky/022422/12300area/2022/March2/
+output_dir: data/gaussian_sims/outputs_gaussian_sims/gaussian_sims_srdnzs_fullsky/051422/12300area/2022/May23/
 config: examples/gaussian_sims/config.yml
 
 # On NERSC, set this before running:
 # export DATA=${LSST}/groups/WL/users/zuntz/data/metacal-testbed
 
 inputs:
     # See README for paths to download these files
-    lens_catalog:  /global/projecta/projectdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/022422/12300area/TXPipe_inputs/lens_catalog.hdf5
-    lens_tomography_catalog: /global/projecta/projectdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/022422/12300area/TXPipe_inputs/lens_tomography_catalog.hdf5
-    lens_photoz_stack: /global/projecta/projectdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/022422/12300area/TXPipe_inputs/lens_photoz_stack.hdf5 
-    shear_photoz_stack: /global/projecta/projectdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/022422/12300area/TXPipe_inputs/shear_photoz_stack.hdf5 
+    lens_catalog:  /global/cfs/cdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/051422/12300area/TXPipe_inputs/lens_catalog.hdf5
+    lens_tomography_catalog: /global/cfs/cdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/051422/12300area/TXPipe_inputs/lens_tomography_catalog.hdf5
+    lens_photoz_stack: /global/cfs/cdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/051422/12300area/TXPipe_inputs/lens_photoz_stack.hdf5 
+    shear_photoz_stack: /global/cfs/cdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/051422/12300area/TXPipe_inputs/shear_photoz_stack.hdf5 
     fiducial_cosmology: data/fiducial_cosmology.yml
     calibration_table:  /global/projecta/projectdirs/lsst/groups/WL/users/zuntz/data/cosmoDC2-1.1.4_oneyear_unit_response/sample_cosmodc2_w10year_errors.dat
     response_model: /global/projecta/projectdirs/lsst/groups/WL/users/zuntz/data/DESY1-R-model.hdf5 
-    mask:  /global/projecta/projectdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/022422/12300area/TXPipe_inputs/mask.hdf5
+    mask:  /global/cfs/cdirs/lsst/groups/WL/users/jprat/gaussian_sims_srdnzs_fullsky/051422/12300area/TXPipe_inputs/mask.hdf5
 
 
 resume: True
-log_dir: data/gaussian_sims/outputs_gaussian_sims/gaussian_sims_srdnzs_fullsky/
-pipeline_log: data/gaussian_sims/outputs_gaussian_sims/gaussian_sims_srdnzs_fullsky/log.txt
+log_dir: data/gaussian_sims/outputs_gaussian_sims/gaussian_sims_srdnzs_fullsky/051422/May23/
+pipeline_log: data/gaussian_sims/outputs_gaussian_sims/gaussian_sims_srdnzs_fullsky/051422/May23/log.txt
 

txpipe/base_stage.py

@@ -54,6 +54,7 @@ def memory_report(self, tag=None):
         print(
             f"{t}: Process {self.rank}:{tag} Remaining memory on {host} {avail:.1f} GB / {total:.1f} GB"
         )
+        sys.stdout.flush()
 
     def combined_iterators(self, rows, *inputs, parallel=True):
         if not len(inputs) % 3 == 0:

txpipe/covariance.py

@@ -46,6 +46,7 @@ class TXFourierGaussianCovariance(PipelineStage):
         "pickled_wigner_transform": "",
         "use_true_shear": False,
         "galaxy_bias": [0.0],
+        "gaussian_sims_factor": [1.],
     }
 
     def run(self):
@@ -114,8 +115,20 @@ def read_number_statistics(self):
         input_data = self.open_input("tracer_metadata")
 
         # per-bin quantities
-        N_eff = input_data["tracers/N_eff"][:]
+        N_eff = input_data["tracers/N_eff"][:] # for sources
         N_lens = input_data["tracers/lens_counts"][:]
+        # For the gaussian sims, lambda = nbar(1+b*delta),
+        # instead of lambda = nbar(1+delta), where delta is the density contrast field.
+        # Then, we need to scale up the shot noise term for the lenses
+        # in the covariance for the same b factor.
+        # Here we decrease the number density for this factor, since shot noise term is 1/nbar. 
+        print('N_lens:', N_lens)
+        N_lens = N_lens/np.array(self.config["gaussian_sims_factor"])**2
+
+        if self.config["gaussian_sims_factor"] != [1.]:
+            print("ATTENTION: We are dividing N_lens by the gaussian sims factor squared:", np.array(self.config["gaussian_sims_factor"])**2)
+            print("Scaled N_lens is:", N_lens)
+
         if self.config["use_true_shear"]:
             nbins = len(input_data["tracers/sigma_e"][:])
             sigma_e = np.array([0.0 for i in range(nbins)])
@@ -387,18 +400,17 @@ def compute_covariance_block(
                 lb, cov["final_b"] = bin_cov(r=ell, r_bins=ell_bins, cov=cov["final"])
         return cov
 
-    def get_angular_bins(self, two_point_data):
-        # Assume that the ell binning is the same for each of the bins.
-        # This is true in the current pipeline.
-        X = two_point_data.get_data_points("galaxy_shear_cl_ee", i=0, j=0)
-        # Further assume that the ell ranges are contiguous, so that
-        # the max value of one window is the min value of the next.
-        # So we just need the lower edges of each bin and then the
-        # final maximum value of the last bin
-        ell_edges = [x["window"].min for x in X]
-        ell_edges.append(X[-1]["window"].max)
-
-        return np.array(ell_edges)
+    def get_angular_bins(self, cl_sacc):
+        # This function replicates `choose_ell_bins` in twopoint_fourier.py
+        # TODO: Move this to txpipe/utils/nmt_utils.py
+        from .utils.nmt_utils import MyNmtBin
+
+        ell_min = cl_sacc.metadata["binning/ell_min"]
+        ell_max = cl_sacc.metadata["binning/ell_max"]
+        ell_spacing = cl_sacc.metadata["binning/ell_spacing"]
+        n_ell = cl_sacc.metadata["binning/n_ell"]
+        edges = np.unique(np.geomspace(ell_min, ell_max, n_ell).astype(int))
+        return edges
 
     def make_wigner_transform(self, meta):
         import threadpoolctl
@@ -588,6 +600,7 @@ class TXRealGaussianCovariance(TXFourierGaussianCovariance):
         "pickled_wigner_transform": "",
         "use_true_shear": False,
         "galaxy_bias": [0.0],
+        "gaussian_sims_factor": [1.],        
     }
 
     def run(self):
@@ -644,18 +657,19 @@ class TXFourierTJPCovariance(PipelineStage):
         ("twopoint_data_fourier", SACCFile),  # For the binning information
         ("tracer_metadata_yml", YamlFile),  # For metadata
         ("mask", MapsFile),  # For the lens mask
+        ("density_maps", MapsFile),  # For the clustering mask
         ("source_maps", MapsFile),  # For the sources masks
     ]
 
     outputs = [
         ("summary_statistics_fourier", SACCFile),
     ]
 
-    config_options = {"galaxy_bias": [0.0], "IA": 0.5, "cache_dir": ""}
+    config_options = {"galaxy_bias": [0.0], "IA": 0.5, "cache_dir": "",}
 
     def run(self):
         import tjpcov.main
-
+        import healpy
         # Read the metadata from earlier in the pipeline
         with self.open_input("tracer_metadata_yml", wrapper=True) as f:
             meta = f.content
@@ -692,22 +706,34 @@ def run(self):
             tjp_config[f"bias_lens_{i}"] = bias[i]
             tjp_config[f"Ngal_lens_{i}"] = meta["lens_density"][i]
 
-        for i in range(nbin_source):
-            tjp_config[f"Ngal_source_{i}"] = meta["n_eff"][i]
-            tjp_config[f"sigma_e_source_{i}"] = meta["sigma_e"][i]
-
         # Load masks
-        # Would it be better to pass a path? Masks are not that heavy, so we
-        # might save some I/O overhead reading them at once here)
+        # For clustering, we follow twopoint_fourier.py:214-219
         with self.open_input("mask", wrapper=True) as f:
             mask = f.read_map("mask")
+            mask[mask == healpy.UNSEEN] = 0.0
             if self.rank == 0:
                 print("Loaded mask")
 
+        # Set any unseen pixels to zero weight.
+        # TODO: unify this code with the code in twopoint_fourier.py
+
+        with self.open_input("density_maps", wrapper=True) as f:
+            nbin_lens = f.file["maps"].attrs["nbin_lens"]
+            d_maps = [f.read_map(f"delta_{b}") for b in range(nbin_lens)]
+            print(f"Loaded {nbin_lens} overdensity maps")
+
+        # twopoint_fourier.py:219
+        for d in d_maps:
+            mask[d == healpy.UNSEEN] = 0
+
+        # twopoint_fourier.py:225
         with self.open_input("source_maps", wrapper=True) as f:
-            lensing_weights = [
-                f.read_map(f"lensing_weight_{b}") for b in range(nbin_source)
-            ]
+            lensing_weights = []
+            for b in range(nbin_source):
+                lw = f.read_map(f"lensing_weight_{b}")
+                lw[lw == healpy.UNSEEN] = 0.0
+                lensing_weights.append(lw)
+
             if self.rank == 0:
                 print(f"Loaded {nbin_source} lensing weight maps")
 
@@ -767,7 +793,7 @@ def get_tracer_noise_from_sacc(self, cl_sacc):
             else:
                 raise KeyError(
                     "Missing n_ell_coupled metadata for tracer "
-                    + f"{trn}. Something is wrong with the input "
+                    + trn + f"Something is wrong with the input "
                     + "sacc file"
                 )
 
@@ -788,12 +814,23 @@ def get_workspaces_dict(self, cl_sacc, masks_names):
                 hashes[m] = cl_sacc.metadata[f"hash/{m}"]
         ell_hash = cl_sacc.metadata["hash/ell_hash"]
 
-        w = {}
+        w = {'00': {}, '02': {}, '22': {}}
+        # Get the number of data points per Cell
+        dtype = cl_sacc.get_data_types()[0]
+        trs = cl_sacc.get_tracer_combinations()[0]
+        ell_eff, _ = cl_sacc.get_ell_cl(dtype, *trs)
+        n_ell = ell_eff.size
         for tr1, tr2 in cl_sacc.get_tracer_combinations():
+            # This assumes that the name of the tracers will include 'lens'or
+            # 'source'
+            s1 = 0 if 'lens' in tr1 else 2
+            s2 = 0 if 'lens' in tr2 else 2
+            sk = ''.join(sorted(f'{s1}{s2}'))
             m1 = masks_names[tr1]
             m2 = masks_names[tr2]
             key = (m1, m2)
-            if (key in w) or (key[::-1] in w):
+            # checking if the combination m1,m2 or m2,m1 is already done. 
+            if (key in w[sk]) or (key[::-1] in w[sk]):
                 continue
             # Build workspace hash (twopoint_fourier.py:387-395)
             h1 = hashes[m1]
@@ -802,7 +839,7 @@ def get_workspaces_dict(self, cl_sacc, masks_names):
             if h2 != h1:
                 cache_key ^= h2
 
-            w[key] = str(cache.get_path(cache_key))
+            w[sk][key] = str(cache.get_path(cache_key))
 
         return w
 
@@ -841,7 +878,6 @@ def recover_NmtBin(self, cl_sacc):
                 "The reconstructed NmtBin object is not "
                 + "compatible with the ells in the sacc file"
             )
-
         return ell_bins
 
     def read_sacc(self):

txpipe/extensions/twopoint_scia.py

@@ -433,6 +433,10 @@ def call_treecorr(self, data, meta, i, j, k):
         return result
 
     def write_output(self, data, meta, results):
+        # This subclass only needs the root process for this task
+        if self.rank != 0:
+            return
+
         import sacc
         import treecorr
 

txpipe/input_cats.py

@@ -44,6 +44,7 @@ class TXCosmoDC2Mock(PipelineStage):
         "apply_mag_cut": False,  # used when comparing to descqa measurements
         "Mag_r_limit": -19,  # used to decide what objects to cut out
         "metadetect": True,  # Alternatively we will mock a  metacal catalog
+        "add_shape_noise": True, 
     }
 
     def data_iterator(self, gc):
@@ -522,7 +523,11 @@ def make_mock_metadetect(self, data, photo):
 
         # Use a fixed shape noise per component to generate
         # an overall
-        shape_noise = 0.26
+        if self.config["add_shape_noise"]:
+            shape_noise = 0.26
+        else:
+            shape_noise = 0.
+
         eps = np.random.normal(0, shape_noise, nobj) + 1.0j * np.random.normal(
             0, shape_noise, nobj
         )
@@ -755,7 +760,11 @@ def make_mock_metacal(self, data, photo):
 
         # Use a fixed shape noise per component to generate
         # an overall
-        shape_noise = 0.26
+        if self.config["add_shape_noise"]:
+            shape_noise = 0.26
+        else:
+            shape_noise = 0.
+
         eps = np.random.normal(0, shape_noise, nobj) + 1.0j * np.random.normal(
             0, shape_noise, nobj
         )
@@ -1003,6 +1012,7 @@ class TXGaussianSimsMock(TXCosmoDC2Mock):
         "flip_g2": False,  # this matches the metacal definition, and the treecorr/namaster one
         "apply_mag_cut": False,  # used when comparing to descqa measurements
         "metadetect": True,  # Alternatively we will mock a  metacal catalog
+        "add_shape_noise": False, # the input cats already have shape noise included
     }
 
     def data_iterator(self, cat):
@@ -1069,8 +1079,8 @@ def load_catalog(self):
 
         print(f"Loading from catalog {cat_name}")
 
-        cat = np.load(cat_name)
-
+        cat = np.load(cat_name, allow_pickle = True)
+    
         N = len(cat[0])
 
         return cat, N