split get_ana_cov_blocks.py into recipe script and compute script

Author: zatkins2

project/SO/pISO/python/get_ana_cov_blocks.py

@@ -7,7 +7,6 @@
 
 import argparse
 from os.path import join as opj
-from itertools import product
 import time
 
 import numpy as npy
@@ -30,8 +29,6 @@
 d.read_from_file(args.paramfile)
 log = log.get_logger(**d)
 
-coupling_cache_size = args.coupling_cache_size
-
 surveys = d["surveys"]
 lmax = d["lmax"]
 cov_correlation_by_noise_model = d['cov_correlation_by_noise_model']
@@ -51,6 +48,8 @@
 mcm_dir = d['mcm_dir']
 cov_dir = d['cov_dir']
 
+t0 = time.time()
+
 canonized_sn_field_info2canonized_connected_combo_2pt = npy.load(opj(cov_dir, 'canonized_sn_field_info2canonized_connected_combo_2pt.npy'), allow_pickle=True).item()
 canonized_w2s = npy.load(opj(cov_dir, 'canonized_w2s.npy'), allow_pickle=True).item()
 
@@ -61,6 +60,17 @@
 reference_sn_field_info2reference_canonized_disconnected_combo_4pt = npy.load(opj(cov_dir, 'reference_sn_field_info2reference_canonized_disconnected_combo_4pt.npy'), allow_pickle=True).item()
 canonized_wls = npy.load(opj(cov_dir, 'canonized_wls.npy'), allow_pickle=True).item()
 
+cov_block_sets2can_discon_com_4pts_and_optypes = npy.load(opj(cov_dir, 'cov_block_sets2can_discon_com_4pts_and_optypes.npy'), allow_pickle=True).item()
+cov_block2TEB_block2can_sn_alm_info2nterms = npy.load(opj(cov_dir, 'cov_block2TEB_block2can_sn_alm_info2nterms.npy'), allow_pickle=True).item()
+
+log.info(f'[Rank {so_mpi.rank}] Load metadata in {(time.time() - t0):.3f} seconds')
+
+optype2str = {
+    0: '00',
+    1: '02',
+    2: '++'
+}
+
 def update_pseudospectra_dict(f1, f2, pseudospectra_dict=None):
     if pseudospectra_dict is None:
         pseudospectra_dict = {}
@@ -154,40 +164,6 @@ def pols_disconnected_combo_4pt2ducc_optype(pol1, pol2, pol3, pol4):
         # if 2, then the spintype is ++, which is ducc optype 2
         return spin2_1 + spin2_2
 
-def update_ducc_inputs_and_nterms(sna1, sna2, sna3, sna4,
-                                  this_block_can_discon_com_4pts_and_optypes,
-                                  can_sn_alm_info2nterms):
-    # update ducc inputs with minimal unique couplings, and track their order
-    sv1, m1, TEB1, split1 = sna1
-    sv2, m2, TEB2, split2 = sna2
-    sv3, m3, TEB3, split3 = sna3
-    sv4, m4, TEB4, split4 = sna4
-
-    pol1 = TEB2pol(TEB1)
-    pol2 = TEB2pol(TEB2)
-    pol3 = TEB2pol(TEB3)
-    pol4 = TEB2pol(TEB4)
-
-    snf1 = (sv1, m1, pol1, split1)
-    snf2 = (sv2, m2, pol2, split2)
-    snf3 = (sv3, m3, pol3, split3)
-    snf4 = (sv4, m4, pol4, split4)
-    can_discon_com_4pt = get_can_discon_com_4pt(snf1, snf2, snf3, snf4)
-
-    # NOTE: although using uncanonized pol1, pol2, pol3, pol4, the optype is
-    # insensitive to disconnected 4pt canonization
-    optype = pols_disconnected_combo_4pt2ducc_optype(pol1, pol2, pol3, pol4)
-    
-    # adding to set does nothing if already in set
-    this_block_can_discon_com_4pts_and_optypes.add((can_discon_com_4pt, optype))
-
-    # track the number of times this term has appeared in this cov TEB sub-block
-    can_sn_alm_info = pspipe_list.canonize_disconnected_4pt(sna1, sna2, sna3, sna4)
-    if can_sn_alm_info not in can_sn_alm_info2nterms:
-        can_sn_alm_info2nterms[can_sn_alm_info] = 1
-    else:
-        can_sn_alm_info2nterms[can_sn_alm_info] += 1
-
 @numba.njit(parallel=True)
 def add_term_to_pseudo_cov_block(pseudo_cov_block, num_terms, w4_1234, w4_coupling, w2_12, w2_34, C12, C34, coupling):
     # important to cast the scalar to the right type before multiplication, 
@@ -217,199 +193,7 @@ def add_term_to_pseudo_cov_block(pseudo_cov_block, num_terms, w4_1234, w4_coupli
 for sv in surveys:
     nsplits[sv] = d[f'n_splits_{sv}']
 
-# first, figure out all the "shared couplings" sets of cov blocks, such that the
-# total number of couplings in each block is <= the cache size. NOTE: a nominal
-# cov_block_set might get "chopped" by blindly cutting all the cov blocks into
-# equal-length subtasks
-so_mpi.init(True)
-
-t0 = time.time()
-
-subtasks = so_mpi.taskrange(imin=0, imax=n_covs - 1)
-
-cov_block_sets2can_discon_com_4pts_and_optypes = {}
-cov_block2TEB_block2can_sn_alm_info2nterms = {}
-
-# need to initialize objects before while loop, that otherwise are re-initialized
-# in the loop
-cov_block_set = []
-can_discon_com_4pts_and_optypes = set()
-i = 0
-while True:
-    task = subtasks[i]
-    svi, mi = ni_list[task].split('&') 
-    svj, mj = nj_list[task].split('&')
-    svp, mp = np_list[task].split('&')
-    svq, mq = nq_list[task].split('&')
-    cov_block = ((svi, mi), (svj, mj),
-                 (svp, mp), (svq, mq))
-
-    # "n" holds the "noise correlation group" information: f1 and f2 have 
-    # correlated noise only if ni == nj
-    if cov_correlation_by_noise_model:
-        ni = (svi, mapnames2noise_model_tags[f'{svi}_{mi}'])
-        nj = (svj, mapnames2noise_model_tags[f'{svj}_{mj}'])
-        np = (svp, mapnames2noise_model_tags[f'{svp}_{mp}'])
-        nq = (svq, mapnames2noise_model_tags[f'{svq}_{mq}'])
-    else:
-        ni = svi
-        nj = svj
-        np = svp
-        nq = svq
-
-    # we need figure out which couplings we actually need first
-    #
-    # for each block, see which unique couplings are needed, and then try to add
-    # to existing block set of unique couplings. if resulting merged set fits in
-    # the cache, go to the next block, otherwise, end set and redo this block
-    this_block_can_discon_com_4pts_and_optypes = set()
-
-    # for each cov TEB sub-block, tracks how many times a canonical 4pt combo
-    # of (sv, m, TEB, split)s recurs, so it can be added once (times this count)
-    # rather than each time
-    TEB_block2can_sn_alm_info2nterms = {} # "alm_info" since keys are TEB instead of T and pol
-    
-    splits_cross_iterator_ij = pspipe_list.get_splits_cross_iterator(svi, nsplits[svi], svj, nsplits[svj])
-    splits_cross_iterator_pq = pspipe_list.get_splits_cross_iterator(svp, nsplits[svp], svq, nsplits[svq])        
-    for (TEBi, TEBj), (TEBp, TEBq) in product(spectra, repeat=2):    
-        if (TEBi, TEBj, TEBp, TEBq) not in TEB_block2can_sn_alm_info2nterms:
-            TEB_block2can_sn_alm_info2nterms[TEBi, TEBj, TEBp, TEBq] = {}
-        
-        can_sn_alm_info2nterms = TEB_block2can_sn_alm_info2nterms[TEBi, TEBj, TEBp, TEBq]
-
-        for (si, sj), (sp, sq) in product(splits_cross_iterator_ij, splits_cross_iterator_pq):
-
-            # ssss ipjq
-            update_ducc_inputs_and_nterms((svi, mi, TEBi, 's'), (svp, mp, TEBp, 's'),
-                                          (svj, mj, TEBj, 's'), (svq, mq, TEBq, 's'), 
-                                          this_block_can_discon_com_4pts_and_optypes,
-                                          can_sn_alm_info2nterms)
-            
-            # ssnn ipjq
-            if nj == nq and sj == sq:
-                update_ducc_inputs_and_nterms((svi, mi, TEBi, 's'), (svp, mp, TEBp, 's'),
-                                              (svj, mj, TEBj, f'n{sj}'), (svq, mq, TEBq, f'n{sj}'), 
-                                              this_block_can_discon_com_4pts_and_optypes,
-                                              can_sn_alm_info2nterms)
-
-            # nnss ipjq
-            if ni == np and si == sp:
-                update_ducc_inputs_and_nterms((svi, mi, TEBi, f'n{si}'), (svp, mp, TEBp, f'n{si}'),
-                                              (svj, mj, TEBj, 's'), (svq, mq, TEBq, 's'), 
-                                              this_block_can_discon_com_4pts_and_optypes,
-                                              can_sn_alm_info2nterms)
-                    
-            # nnnn ipjq
-            if ni == np and si == sp and nj == nq and sj == sq:
-                update_ducc_inputs_and_nterms((svi, mi, TEBi, f'n{si}'), (svp, mp, TEBp, f'n{si}'),
-                                              (svj, mj, TEBj, f'n{sj}'), (svq, mq, TEBq, f'n{sj}'), 
-                                              this_block_can_discon_com_4pts_and_optypes,
-                                              can_sn_alm_info2nterms)
-                    
-            # ssss iqjp
-            update_ducc_inputs_and_nterms((svi, mi, TEBi, 's'), (svq, mq, TEBq, 's'),
-                                          (svj, mj, TEBj, 's'), (svp, mp, TEBp, 's'), 
-                                          this_block_can_discon_com_4pts_and_optypes,
-                                          can_sn_alm_info2nterms)
-            
-            # ssnn iqjp
-            if nj == np and sj == sp:
-                update_ducc_inputs_and_nterms((svi, mi, TEBi, 's'), (svq, mq, TEBq, 's'),
-                                              (svj, mj, TEBj, f'n{sj}'), (svp, mp, TEBp, f'n{sj}'), 
-                                              this_block_can_discon_com_4pts_and_optypes,
-                                              can_sn_alm_info2nterms)
-
-            # nnss iqjp
-            if ni == nq and si == sq:
-                update_ducc_inputs_and_nterms((svi, mi, TEBi, f'n{si}'), (svq, mq, TEBq, f'n{si}'),
-                                              (svj, mj, TEBj, 's'), (svp, mp, TEBp, 's'), 
-                                              this_block_can_discon_com_4pts_and_optypes,
-                                              can_sn_alm_info2nterms)
-                    
-            # nnnn iqjp
-            if ni == nq and si == sq and nj == np and sj == sp:
-                update_ducc_inputs_and_nterms((svi, mi, TEBi, f'n{si}'), (svq, mq, TEBq, f'n{si}'),
-                                              (svj, mj, TEBj, f'n{sj}'), (svp, mp, TEBp, f'n{sj}'), 
-                                              this_block_can_discon_com_4pts_and_optypes,
-                                              can_sn_alm_info2nterms)
-
-    cov_block2TEB_block2can_sn_alm_info2nterms[cov_block] = TEB_block2can_sn_alm_info2nterms
-
-    # there are now four possibilities for what to do with this block:
-    # (a) if the current block requires more couplings than the cache size 
-    # limit, and the current cache is empty, then we have no recourse: 
-    # ending the set, resetting the cache, and redoing the block will of course
-    # never work. therefore, we first force the one block into the cache, and
-    # then end the set and reset the cache. we then go on to the next block.
-    # this *does* "violate" the cache limit, so we issue a warning
-    # (b) like (a), if adding the current block's couplings to the cache would
-    # result in a cache size more than the cache size limit, but unlike (a) if 
-    # the cache is not empty, we do have a recourse: end the set, reset the
-    # cache, and then redo this block with an empty cache. 
-    # (c) if we are on the last block of all the subtasks, but we know we are 
-    # not going to redo this block with an empty cache (i.e., not (b)), then
-    # we are also on the last task of the loop. like (a) we must force the block
-    # into the cache and end the set. unlike (a), we break the loop instead of
-    # going on to the next block. it's possible that (a) and (c) occur at the 
-    # same time, in which case (c) takes priority.
-    # (d) otherwise proceed: add this block to the current cache and go on to 
-    # the next block. hopefully this happens most of the time
-
-    single_block_set = False
-    end_set_and_redo_block = False
-    if len(can_discon_com_4pts_and_optypes & this_block_can_discon_com_4pts_and_optypes) > coupling_cache_size:
-        single_block_set = len(can_discon_com_4pts_and_optypes) == 0
-        end_set_and_redo_block = len(can_discon_com_4pts_and_optypes) > 0
-
-    end_loop = (i+1 == len(subtasks)) and not end_set_and_redo_block
-
-    if single_block_set:
-        log.warning(f"[Rank {so_mpi.rank}, Task {task}] Number of couplings for cov block {cov_block} is "
-                    f"{len(this_block_can_discon_com_4pts_and_optypes)}, which excees the coupling cache "
-                    f"size of {coupling_cache_size}. Adding to single-block-set, may result in OOM later.")
-
-    if single_block_set or end_loop:
-        cov_block_set.append(cov_block)
-        can_discon_com_4pts_and_optypes &= this_block_can_discon_com_4pts_and_optypes
-
-    if single_block_set or end_set_and_redo_block or end_loop:
-        cov_block_sets2can_discon_com_4pts_and_optypes[tuple(cov_block_set)] = can_discon_com_4pts_and_optypes
-
-        cov_block_set = []
-        can_discon_com_4pts_and_optypes = set()
-        
-        if single_block_set and not end_loop:
-            i += 1
-            continue
-        if end_set_and_redo_block:
-            continue
-        if end_loop:
-            break
-    else:
-        cov_block_set.append(cov_block)
-        can_discon_com_4pts_and_optypes &= this_block_can_discon_com_4pts_and_optypes
-        i += 1
-
-log.info(f'[Rank {so_mpi.rank}] Loop over cov block sets in {(time.time() - t0):.3f} seconds')
-
-t0 = time.time()
-
-# these may be useful to check later
-cov_block_sets2can_discon_com_4pts_and_optypes = so_mpi.gather_set_or_dict(cov_block_sets2can_discon_com_4pts_and_optypes,
-                                                                           allgather=True,
-                                                                           overlap_allowed=False)
-
-cov_block2TEB_block2can_sn_alm_info2nterms = so_mpi.gather_set_or_dict(cov_block2TEB_block2can_sn_alm_info2nterms,
-                                                                       allgather=True,
-                                                                       overlap_allowed=False)
-
-if so_mpi.rank == 0:
-    npy.save(opj(cov_dir, 'cov_block_sets2can_discon_com_4pts_and_optypes.npy'), cov_block_sets2can_discon_com_4pts_and_optypes)
-    npy.save(opj(cov_dir, 'cov_block2TEB_block2can_sn_alm_info2nterms.npy'), cov_block2TEB_block2can_sn_alm_info2nterms)
-
-log.info(f'[Rank {so_mpi.rank}] Save cov block sets in {(time.time() - t0):.3f} seconds')
-
-# now mpi over cov_block_sets
+# mpi over cov_block_sets
 cov_block_sets = list(cov_block_sets2can_discon_com_4pts_and_optypes.keys())
 n_cov_block_sets = len(cov_block_sets)
 subtasks = so_mpi.taskrange(imin=0, imax=n_cov_block_sets - 1)
@@ -441,15 +225,17 @@ def add_term_to_pseudo_cov_block(pseudo_cov_block, num_terms, w4_1234, w4_coupli
 
     optype_counts = {}
     for optype in optypes_for_ducc:
-        if optype in optype_counts:
-            optype_counts[optype] += 1
+        optypestr = optype2str[optype]
+        if optypestr in optype_counts:
+            optype_counts[optypestr] += 1
         else:
-            optype_counts[optype] = 1
+            optype_counts[optypestr] = 1
 
     specs_for_ducc = None
     optypes_for_ducc = None
 
-    log.info(f'[Rank {so_mpi.rank}, Task {task}] Calculated {optype_counts} couplings in {(time.time() - t0):.3f} seconds')            
+    optypesstr = ', '.join([f'{ct} {s}-type couplings' for s, ct in optype_counts.items()])
+    log.info(f'[Rank {so_mpi.rank}, Task {task}] Calculated {optypesstr} in {(time.time() - t0):.3f} seconds')            
 
     # now add all terms together for each cov block
     for i, cov_block in enumerate(cov_block_set):
@@ -568,6 +354,8 @@ def add_term_to_pseudo_cov_block(pseudo_cov_block, num_terms, w4_1234, w4_coupli
                                                                 dense=True, dtype=npy.float32)
 
         # finalize: need to divide the split factor from each side and cast to double
+        splits_cross_iterator_ij = pspipe_list.get_splits_cross_iterator(svi, nsplits[svi], svj, nsplits[svj])
+        splits_cross_iterator_pq = pspipe_list.get_splits_cross_iterator(svp, nsplits[svp], svq, nsplits[svq]) 
         ana_cov /= (len(splits_cross_iterator_ij) * len(splits_cross_iterator_pq))
         ana_cov = ana_cov.astype(npy.float64, copy=False)