Make UVPSpec labels optional; better checking for optional labels

HERA-Team · Apr 28, 2018 · 528f8db · 528f8db
1 parent 0c0fd5f
commit 528f8db
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 6 deletions.
diff --git a/hera_pspec/pspecdata.py b/hera_pspec/pspecdata.py
@@ -1366,8 +1366,10 @@ def pspec(self, bls1, bls2, dsets, input_data_weight='identity', norm='I',
             uvp.filename1 = dset1.extra_keywords['filename']
         if hasattr(dset2.extra_keywords, 'filename'): 
             uvp.filename2 = dset2.extra_keywords['filename']
-        uvp.label1 = self.labels[self.dset_idx(dsets[0])]
-        uvp.label2 = self.labels[self.dset_idx(dsets[1])]
+        lbl1 = self.labels[self.dset_idx(dsets[0])]
+        lbl2 = self.labels[self.dset_idx(dsets[1])]
+        if lbl1 is not None: uvp.label1 = lbl1
+        if lbl2 is not None: uvp.label2 = lbl2
 
         # fill data arrays
         uvp.data_array = data_array

diff --git a/hera_pspec/tests/test_uvpspec.py b/hera_pspec/tests/test_uvpspec.py
@@ -55,6 +55,8 @@ def build_example_uvpspec():
     norm = "I"
     git_hash = "random"
     scalar_array = np.ones((Nspws, Npols), np.float)
+    label1 = 'red'
+    #label2 = 'blue' # Leave commented out to make sure non-named UVPSpecs work!
 
     telescope_location = np.array([5109325.85521063, 
                                    2005235.09142983, 
@@ -77,7 +79,7 @@ def build_example_uvpspec():
               'integration_array', 'bl_array', 'bl_vecs', 'telescope_location', 
               'units', 'channel_width', 'weighting', 'history', 'taper', 'norm', 
               'git_hash', 'nsample_array', 'time_avg_array', 'lst_avg_array', 
-              'cosmo', 'scalar_array']
+              'cosmo', 'scalar_array', 'label1']
 
     # Set all parameters
     for p in params:

diff --git a/hera_pspec/uvpspec.py b/hera_pspec/uvpspec.py
@@ -79,7 +79,7 @@ def __init__(self):
                             "data_array", "wgt_array", "integration_array", "spw_array", "freq_array", "dly_array",
                             "pol_array", "lst_1_array", "lst_2_array", "time_1_array", "time_2_array", "blpair_array",
                             "Nbls", "bl_vecs", "bl_array", "channel_width", "telescope_location", "weighting", "units",
-                            "taper", "norm", "git_hash", "nsample_array", 'lst_avg_array', 'time_avg_array', 'folded', "label1", "label2"]
+                            "taper", "norm", "git_hash", "nsample_array", 'lst_avg_array', 'time_avg_array', 'folded']
 
         self._immutables = ["Ntimes", "Nblpairts", "Nblpairs", "Nspwdlys", "Nspws", "Ndlys", "Npols", "Nfreqs", "history",
                             "Nbls", "channel_width", "weighting", "units", "filename1", "filename2", "label1", "label2",
@@ -1220,8 +1220,8 @@ def average_spectra(self, blpair_groups=None, time_avg=False, inplace=True):
         uvp.integration_array = ints_array
         uvp.wgt_array = wgts_array
         uvp.nsample_array = nsmp_array
-        uvp.label1 = self.label1
-        uvp.label2 = self.label2
+        if hasattr(self, 'label1'): uvp.label1 = self.label1
+        if hasattr(self, 'label2'): uvp.label2 = self.label2
 
         uvp.check()