different subset names fot cat and zcat

redmapper/galaxy.py

@@ -129,7 +129,8 @@ def __init__(self, *arrays, **kwargs):
         self.depth = 10 if 'depth' not in kwargs else kwargs['depth']
 
     @classmethod
-    def load_merged_or_in_parallel(cls, indices, tab, filename, zredfile, kw_border, columns=None, use_zred=False):
+    def load_merged_or_in_parallel(cls, indices, tab, filename, zredfile, kw_border, cat_subset=None, use_zred=False,
+                                   zcat_subset=None):
 
         """
         A helper function created by @jacobic to speed things up.
@@ -142,7 +143,7 @@ def load_merged_or_in_parallel(cls, indices, tab, filename, zredfile, kw_border,
                     fname_merged = fn_merged.decode()
                 except AttributeError:
                     fname_merged = fn_merged
-                cat = fitsio.read(fname_merged, ext=1, lower=True, columns=columns)
+                cat = fitsio.read(fname_merged, ext=1, lower=True, columns=cat_subset)
                 print(f'Using merged galfile {fname_merged}...')
             else:
                 print('Some zreds are missing. Cannot use merged file...')
@@ -154,7 +155,7 @@ def load_merged_or_in_parallel(cls, indices, tab, filename, zredfile, kw_border,
             # checks nrows in each file match properly (exception can be cryptic...)
             path = os.path.dirname(os.path.abspath(filename))
             with mp.Pool() as p:
-                cat = list(tqdm(p.imap(lambda _: read_cat(path=path, fn=_, columns=columns), tab.filenames[indices])))
+                cat = list(tqdm(p.imap(lambda _: read_cat(path=path, fn=_, columns=cat_subset), tab.filenames[indices])))
             cat = np.concatenate(cat)
 
         assert np.sum(tab.ngals[indices]) == len(cat), 'are you missing galaxies?'
@@ -167,7 +168,7 @@ def load_merged_or_in_parallel(cls, indices, tab, filename, zredfile, kw_border,
             zpath = os.path.dirname(zredfile)
 
             with mp.Pool() as p:
-                zcat = list(tqdm(p.imap(lambda _: read_cat(path=zpath, fn=_), ztab.filenames[indices])))
+                zcat = list(tqdm(p.imap(lambda _: read_cat(path=zpath, fn=_, columns=zcat_subset), ztab.filenames[indices])))
             zcat = np.concatenate(zcat)
 
             assert np.sum(ztab.ngals[indices]) == len(zcat), 'are you missing galaxies?'
@@ -266,8 +267,8 @@ def load_in_loop(cls, tab, indices, path, cat_fields, dtype, columns, use_zred,
             return cls(trim_cat(cat, **kw_border))
 
     @classmethod
-    def from_galfile(cls, filename, zredfile=None, nside=0, hpix=[], border=0.0, truth=False, subset=None,
-                     use_tempfile=False, refmag_range=[-1000.0, 1000.0], chisq_max=1e100):
+    def from_galfile(cls, filename, zredfile=None, nside=0, hpix=[], border=0.0, truth=False, cat_subset=None,
+                     use_tempfile=False, refmag_range=[-1000.0, 1000.0], chisq_max=1e100, zcat_subset=None):
         """
         Generate a GalaxyCatalog from a redmapper "galfile."
 
@@ -401,8 +402,8 @@ def from_galfile(cls, filename, zredfile=None, nside=0, hpix=[], border=0.0, tru
             first_fname = os.path.join(path, tab.filenames[indices[0]])
         elt = fitsio.read(first_fname, ext=1, rows=0, lower=True)
         dtype_in = elt.dtype.descr
-        if subset is not None:
-            dtype_in = [_ for _ in dtype_in if _[0] in subset]
+        if cat_subset is not None:
+            dtype_in = [_ for _ in dtype_in if _[0] in cat_subset]
         if not truth:
             mark = []
             for dt in dtype_in:
@@ -429,7 +430,8 @@ def from_galfile(cls, filename, zredfile=None, nside=0, hpix=[], border=0.0, tru
         if len(_hpix) == 0 or nside == 0:
             # using shortcut from @jacobic
             return cls.load_merged_or_in_parallel(indices=indices, tab=tab, filename=filename, zredfile=zredfile,
-                                              use_zred=use_zred, columns=columns, kw_border=kw_border)
+                                                  use_zred=use_zred, cat_subset=columns, zcat_subset=zcat_subset,
+                                                  kw_border=kw_border)
         else:
             # using memory efficient but slower way from @erykoff
             return cls.load_in_loop(tab=tab, indices=indices, path=path, cat_fields=cat_fields, dtype=dtype, columns=columns,