Repository: RadioAstronomySoftwareGroup/pyuvdata

Issue:
Resolve init logic in uvflag
There is an outstanding issue in hera_qm to figure out the init logic in uvflag, and be sure the weights array ends up the way one would expect. Since we have moved uvflag over to pyuvdata, we need to move the issue too.

HERA-Team/hera_qm#264

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602?src=pr&el=h1) Report
> Merging [#602](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602?src=pr&el=desc) into [master](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/commit/ea29fd13d0a8eefb2c4f839e74e11f3636e09d4d?src=pr&el=desc) will **decrease** coverage by `<.01%`.
> The diff coverage is `100%`.

[![Impacted file tree graph](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602/graphs/tree.svg?width=650&token=4QykcCeB1H&height=150&src=pr)](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #602      +/-   ##
==========================================
- Coverage   98.86%   98.86%   -0.01%     
==========================================
  Files          20       20              
  Lines        8012     8009       -3     
==========================================
- Hits         7921     7918       -3     
  Misses         91       91
```


| [Impacted Files](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [pyuvdata/uvflag.py](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602/diff?src=pr&el=tree#diff-cHl1dmRhdGEvdXZmbGFnLnB5) | `100% <100%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602?src=pr&el=footer). Last update [ea29fd1...e1e023c](https://codecov.io/gh/RadioAstronomySoftwareGroup/pyuvdata/pull/602?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

bhazelton: @adampbeardsley if you rebase this branch on master the tests should pass.
adampbeardsley: @mkolopanis I knocked off your comments. Ready for re-review.
adampbeardsley: Good call @mkolopanis - this was the whole reason I went down this rabbit hole last week, and I totally forgot about it by now.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e44f2896..f581386f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -19,6 +19,7 @@ All notable changes to this project will be documented in this file.
 - `lst_array` is now saved to UVFITS files (even though it's not a standard parameter) so that it doesn't have to be recalculated
 
 ### Fixed
+- Fixed init logic in UVFlag.
 - Fixed a bug in how FHD uvw vectors were oriented (and visibilities were conjugated)
 - Fixed a bug in `UVData.inflate_by_redundancy` when Nblts is not equal to Nbls * Ntimes.
 - Fixed a bug in UVData when reading in an FHD file with a single time integration.
diff --git a/pyuvdata/uvflag.py b/pyuvdata/uvflag.py
index c22cd05b..38b523d4 100644
--- a/pyuvdata/uvflag.py
+++ b/pyuvdata/uvflag.py
@@ -16,26 +16,36 @@ from six.moves import map
 
 
 class UVFlag(object):
-    ''' Object to handle flag arrays and waterfalls. Supports reading/writing,
+    """ Object to handle flag arrays and waterfalls. Supports reading/writing,
     and stores all relevant information to combine flags and apply to data.
-    '''
+    """
 
     def __init__(self, input, mode='metric', copy_flags=False, waterfall=False, history='',
                  label=''):
-        '''Initialize UVFlag object.
-        Args:
-            input: UVData object, UVCal object, or path to previously saved UVFlag object.
-                   Can also be a list of any compatible combination of the above options.
-                   UVData and UVCal objects cannot be directly combined, unless waterfall is True.
-            mode: "metric" (default) or "flag" to initialize UVFlag in given mode.
-                  The mode determines whether the object has a floating point metric_array
-                  or a boolean flag_array.
-            copy_flags: Whether to copy flags from input to new UVFlag object. Default is False.
-            waterfall: Whether to immediately initialize as a waterfall object,
-                       with flag/metric axes: time, frequency, polarization. Default is False.
-            history: History string to attach to object.
-            label: string used for labeling the object (e.g. 'FM')
-        '''
+        """Initialize UVFlag object. Metadata is copied from input object. If input
+        is subclass of UVData or UVCal, the weights_array will be set to all ones.
+        Input lists or tuples are iterated through, treating each entry with an
+        individual UVFlag init.
+
+        Parameters
+        ----------
+        input : UVData or UVCal or str or list of compatible combination of options
+            Input to initialize UVFlag object. If str, assumed to be path to previously
+            saved UVFlag object. UVData and UVCal objects cannot be directly combined,
+            unless waterfall is True.
+        mode : {"metric", "flag"}, optional
+            The mode determines whether the object has a floating point metric_array
+            or a boolean flag_array. Default is "metric".
+        copy_flags : bool, optional
+            Whether to copy flags from input to new UVFlag object. Default is False.
+        waterfall : bool, optional
+            Whether to immediately initialize as a waterfall object, with flag/metric
+            axes: time, frequency, polarization. Default is False.
+        history : str, optional
+            History string to attach to object. Default is empty string.
+        label: str, optional
+            String used for labeling the object (e.g. 'FM'). Default is empty string.
+        """
 
         self.mode = mode.lower()  # Gets overwritten if reading file
         self.history = ''  # Added to at the end
@@ -45,13 +55,16 @@ class UVFlag(object):
             self.pyuvdata_version_str += ('  Git origin: ' + uvversion.git_origin
                                           + '.  Git hash: ' + uvversion.git_hash
                                           + '.  Git branch: ' + uvversion.git_branch
-                                          + '.  Git description: ' + uvversion.git_description + '.')
+                                          + '.  Git description: '
+                                          + uvversion.git_description + '.')
         self.label = ''  # Added to at the end
         if isinstance(input, (list, tuple)):
-            self.__init__(input[0], mode=mode, copy_flags=copy_flags, waterfall=waterfall, history=history)
+            self.__init__(input[0], mode=mode, copy_flags=copy_flags,
+                          waterfall=waterfall, history=history)
             if len(input) > 1:
                 for i in input[1:]:
-                    fobj = UVFlag(i, mode=mode, copy_flags=copy_flags, waterfall=waterfall, history=history)
+                    fobj = UVFlag(i, mode=mode, copy_flags=copy_flags,
+                                  waterfall=waterfall, history=history)
                     self += fobj
                 del(fobj)
 
@@ -60,7 +73,8 @@ class UVFlag(object):
             self.read(input, history)
         elif waterfall and issubclass(input.__class__, (UVData, UVCal)):
             self.type = 'waterfall'
-            self.history += 'Flag object with type "waterfall" created by ' + self.pyuvdata_version_str
+            self.history += ('Flag object with type "waterfall" created by '
+                             + self.pyuvdata_version_str)
             self.time_array, ri = np.unique(input.time_array, return_index=True)
             self.freq_array = input.freq_array[0, :]
             if issubclass(input.__class__, UVData):
@@ -70,11 +84,8 @@ class UVFlag(object):
                 self.polarization_array = input.jones_array
                 self.lst_array = lst_from_uv(input)[ri]
             if copy_flags:
-                self.metric_array = flags2waterfall(input, keep_pol=True)
-                self.history += ' Waterfall generated from ' + str(input.__class__) + ' object.'
-                if self.mode == 'flag':
-                    warnings.warn('Copying flags into waterfall results in mode=="metric".')
-                    self.mode = 'metric'
+                raise NotImplementedError('Cannot copy flags when initializing '
+                                          ' waterfall UVFlag from UVData or UVCal.')
             else:
                 if self.mode == 'flag':
                     self.flag_array = np.zeros((len(self.time_array),
@@ -87,7 +98,8 @@ class UVFlag(object):
 
         elif issubclass(input.__class__, UVData):
             self.type = 'baseline'
-            self.history += 'Flag object with type "baseline" created by ' + self.pyuvdata_version_str
+            self.history += ('Flag object with type "baseline" created by '
+                             + self.pyuvdata_version_str)
             self.baseline_array = input.baseline_array
             self.ant_1_array = input.ant_1_array
             self.ant_2_array = input.ant_2_array
@@ -110,7 +122,8 @@ class UVFlag(object):
 
         elif issubclass(input.__class__, UVCal):
             self.type = 'antenna'
-            self.history += 'Flag object with type "antenna" created by ' + self.pyuvdata_version_str
+            self.history += ('Flag object with type "antenna" created by '
+                             + self.pyuvdata_version_str)
             self.ant_array = input.ant_array
             self.time_array = input.time_array
             self.lst_array = lst_from_uv(input)
@@ -127,8 +140,11 @@ class UVFlag(object):
                     self.flag_array = np.zeros_like(input.flag_array)
                 elif self.mode == 'metric':
                     self.metric_array = np.zeros_like(input.flag_array).astype(np.float)
+        else:
+            raise ValueError('input to UVFlag.__init__ must be one of: list, tuple, '
+                             'string, UVData, or UVCal.')
 
-        if not isinstance(input, str):
+        if issubclass(input.__class__, (UVData, UVCal)):
             if self.mode == 'flag':
                 self.weights_array = np.ones(self.flag_array.shape)
             else:
@@ -203,7 +219,8 @@ class UVFlag(object):
                 self.time_array = header['time_array'][()]
                 self.lst_array = header['lst_array'][()]
                 self.freq_array = header['freq_array'][()]
-                self.history = uvutils._bytes_to_str(header['history'][()]) + ' Read by ' + self.pyuvdata_version_str
+                self.history = (uvutils._bytes_to_str(header['history'][()])
+                                + ' Read by ' + self.pyuvdata_version_str)
                 self.history += history
                 if 'label' in header.keys():
                     self.label = uvutils._bytes_to_str(header['label'][()])
@@ -259,7 +276,8 @@ class UVFlag(object):
             header['lst_array'] = self.lst_array
             header['freq_array'] = self.freq_array
             header['polarization_array'] = self.polarization_array
-            header['history'] = uvutils._str_to_bytes(self.history + 'Written by ' + self.pyuvdata_version_str)
+            header['history'] = uvutils._str_to_bytes(self.history + 'Written by '
+                                                      + self.pyuvdata_version_str)
             header['label'] = uvutils._str_to_bytes(self.label)
 
             if self.type == 'baseline':
@@ -284,12 +302,12 @@ class UVFlag(object):
                                            compression=data_compression)
 
     def __add__(self, other, inplace=False, axis='time'):
-        '''Add two UVFlag objects together along a given axis.
+        """Add two UVFlag objects together along a given axis.
         Args:
             other: second UVFlag object to concatenate with self.
             inplace: Whether to concatenate to self, or create a new UVFlag object. Default is False.
             axis: Axis along which to combine UVFlag objects. Default is time.
-        '''
+        """
 
         # Handle in place
         if inplace:
@@ -367,11 +385,11 @@ class UVFlag(object):
         return self
 
     def __or__(self, other, inplace=False):
-        '''Combine two UVFlag objects in "flag" mode by "OR"-ing their flags.
+        """Combine two UVFlag objects in "flag" mode by "OR"-ing their flags.
         Args:
             other: second UVFlag object to combine with self.
             inplace: Whether to combine to self, or create a new UVFlag object. Default is False.
-        '''
+        """
         if (self.mode != 'flag') or (other.mode != 'flag'):
             raise ValueError('UVFlag object must be in "flag" mode to use "or" function.')
 
@@ -388,10 +406,10 @@ class UVFlag(object):
             return this
 
     def __ior__(self, other):
-        '''In place or
+        """In place or
         Args:
             other: second UVFlag object to combine with self.
-        '''
+        """
         self.__or__(other, inplace=True)
         return self
 
@@ -415,7 +433,7 @@ class UVFlag(object):
             del(self.flag_array)
 
     def copy(self):
-        ''' Simply return a copy of this object '''
+        """ Simply return a copy of this object """
         return copy.deepcopy(self)
 
     def combine_metrics(self, others, method='quadmean', inplace=True):
@@ -473,7 +491,8 @@ class UVFlag(object):
             darr = self.metric_array
         if len(self.polarization_array) > 1:
             # Collapse pol dimension. But note we retain a polarization axis.
-            d, w = uvutils.collapse(darr, method, axis=-1, weights=self.weights_array, return_weights=True)
+            d, w = uvutils.collapse(darr, method, axis=-1, weights=self.weights_array,
+                                    return_weights=True)
             darr = np.expand_dims(d, axis=d.ndim)
             self.weights_array = np.expand_dims(w, axis=w.ndim)
             self.polarization_array = np.array([','.join(map(str, self.polarization_array))],
@@ -546,7 +565,7 @@ class UVFlag(object):
         self.clear_unused_attributes()
 
     def to_baseline(self, uv, force_pol=False):
-        '''Convert a UVFlag object of type "waterfall" to type "baseline".
+        """Convert a UVFlag object of type "waterfall" to type "baseline".
         Broadcasts the flag array to all baselines.
         This function does NOT apply flags to uv.
         Args:
@@ -555,7 +574,7 @@ class UVFlag(object):
                        Otherwise, will require polarizations match.
                        For example, this keyword is useful if one flags on all
                        pols combined, and wants to broadcast back to individual pols.
-        '''
+        """
         if self.type == 'baseline':
             return
         if not (issubclass(uv.__class__, UVData) or (isinstance(uv, UVFlag) and uv.type == 'baseline')):
@@ -608,7 +627,7 @@ class UVFlag(object):
         self.history += 'Broadcast to type "baseline" with ' + self.pyuvdata_version_str
 
     def to_antenna(self, uv, force_pol=False):
-        '''Convert a UVFlag object of type "waterfall" to type "antenna".
+        """Convert a UVFlag object of type "waterfall" to type "antenna".
         Broadcasts the flag array to all antennas.
         This function does NOT apply flags to uv.
         Args:
@@ -617,7 +636,7 @@ class UVFlag(object):
                        Otherwise, will require polarizations match.
                        For example, this keyword is useful if one flags on all
                        pols combined, and wants to broadcast back to individual pols.
-        '''
+        """
         if self.type == 'antenna':
             return
         if not (issubclass(uv.__class__, UVCal) or (isinstance(uv, UVFlag) and uv.type == 'antenna')):
@@ -663,13 +682,13 @@ class UVFlag(object):
         self.history += 'Broadcast to type "antenna" with ' + self.pyuvdata_version_str
 
     def to_flag(self, threshold=np.inf):
-        '''Convert to flag mode. NOT SMART. Removes metric_array and creates a
+        """Convert to flag mode. NOT SMART. Removes metric_array and creates a
         flag_array from a simple threshold on the metric values.
 
         Args:
             threshold (float): Metric value over which the corresponding flag is
                 set to True. Default is np.inf, which results in flags of all False.
-        '''
+        """
         if self.mode == 'flag':
             return
         elif self.mode == 'metric':
@@ -683,7 +702,7 @@ class UVFlag(object):
         self.clear_unused_attributes()
 
     def to_metric(self, convert_wgts=False):
-        '''Convert to metric mode. NOT SMART. Simply recasts flag_array as float
+        """Convert to metric mode. NOT SMART. Simply recasts flag_array as float
         and uses this as the metric array.
 
         Args:
@@ -693,7 +712,7 @@ class UVFlag(object):
                 flags as metrics to calculate flag fraction. Zero weighting
                 completely flagged rows/columns prevents those from counting
                 against a threshold along the other dimension.
-        '''
+        """
         if self.mode == 'metric':
             return
         elif self.mode == 'flag':
@@ -816,13 +835,13 @@ def and_rows_cols(waterfall):
 
 
 def lst_from_uv(uv):
-    ''' Calculate the lst_array for a UVData or UVCal object.
+    """ Calculate the lst_array for a UVData or UVCal object.
     Args:
         uv: a UVData or UVCal object.
     Returns:
         lst_array: lst_array corresponding to time_array and at telescope location.
                    Units are radian.
-    '''
+    """
     if not isinstance(uv, (UVCal, UVData)):
         raise ValueError('Function lst_from_uv can only operate on '
                          'UVCal or UVData object.')


Test patch:
diff --git a/pyuvdata/tests/test_uvflag.py b/pyuvdata/tests/test_uvflag.py
index c5baee21..4418cf39 100644
--- a/pyuvdata/tests/test_uvflag.py
+++ b/pyuvdata/tests/test_uvflag.py
@@ -169,18 +169,33 @@ def test_init_waterfall_flag():
 def test_init_waterfall_copy_flags():
     uv = UVCal()
     uv.read_calfits(test_c_file)
-    uvf = uvtest.checkWarnings(UVFlag, [uv], {'copy_flags': True, 'mode': 'flag', 'waterfall': True},
-                               nwarnings=1, message='Copying flags into waterfall')
-    assert not hasattr(uvf, 'flag_array')  # Should be metric due to copy flags
-    assert uvf.metric_array.shape == (uv.Ntimes, uv.Nfreqs, uv.Njones)
-    assert uvf.weights_array.shape == (uv.Ntimes, uv.Nfreqs, uv.Njones)
-    assert uvf.type == 'waterfall'
-    assert uvf.mode == 'metric'
-    assert np.all(uvf.time_array == np.unique(uv.time_array))
-    assert np.all(uvf.freq_array == uv.freq_array[0])
-    assert np.all(uvf.polarization_array == uv.jones_array)
-    assert 'Flag object with type "waterfall"' in uvf.history
-    assert pyuvdata_version_str in uvf.history
+    pytest.raises(NotImplementedError, UVFlag, uv, copy_flags=True, mode='flag',
+                  waterfall=True)
+
+
+def test_init_invalid_input():
+    # input is not UVData, UVCal, path, or list/tuple
+    with pytest.raises(ValueError) as cm:
+        foo = UVFlag(14)
+    assert str(cm.value).startswith('input to UVFlag.__init__ must be one of:')
+
+
+@uvtest.skipIf_no_h5py
+def test_init_list_files_weights(tmpdir):
+    # Test that weights are preserved when reading list of files
+    tmp_path = tmpdir.strpath
+    # Create two files to read
+    uvf = UVFlag(test_f_file)
+    np.random.seed(0)
+    wts1 = np.random.rand(*uvf.weights_array.shape)
+    uvf.weights_array = wts1.copy()
+    uvf.write(os.path.join(tmp_path, 'test1.h5'))
+    wts2 = np.random.rand(*uvf.weights_array.shape)
+    uvf.weights_array = wts2.copy()
+    uvf.write(os.path.join(tmp_path, 'test2.h5'))
+    uvf2 = UVFlag([os.path.join(tmp_path, 'test1.h5'),
+                   os.path.join(tmp_path, 'test2.h5')])
+    assert np.all(uvf2.weights_array == np.concatenate([wts1, wts2], axis=0))
 
 
 @uvtest.skipIf_no_h5py