Make UVFlag.__add__ iterate over _data_params

CHANGELOG.md

@@ -21,6 +21,7 @@ All notable changes to this project will be documented in this file.
 - `utils.uvcalibrate` will error rather than warn if some of the newly added checks do not pass, including if antenna names do not match between the objects, starting in version 2.2. In addition, the `flag_missing` keyword is deprecated and will be removed in version 2.2.
 
 ### Fixed
+- UVFlag.__add__ now properly concatenates all existing data-like parameters of the object, including optional ones.
 - A bug in `UVData.downsample_in_time` where the data were not being properly weighted by the nsample array and the nsample_array was not being properly weighted by the integration times.
 - A bug in `UVData.downsample_in_time` that lead to duplicated data on the final object if a baseline had varying integration times and some integration times were greater than or equal to the requested minimum integration time.
 

pyuvdata/uvflag/tests/test_uvflag.py

@@ -859,6 +859,35 @@ def test_add_frequency():
     assert "Data combined along frequency axis. " in uv3.history
 
 
+def test_add_frequency_with_weights_square():
+    # Same test as above, just checking an optional parameter (also in waterfall mode)
+    uvf1 = UVFlag(test_f_file)
+    uvf1.weights_array = 2 * np.ones_like(uvf1.weights_array)
+    uvf1.to_waterfall(return_weights_square=True)
+    uvf2 = copy.deepcopy(uvf1)
+    uvf2.freq_array += 1e4
+    uvf3 = uvf1.__add__(uvf2, axis="frequency")
+    assert np.array_equal(
+        np.concatenate((uvf1.weights_square_array, uvf2.weights_square_array), axis=1),
+        uvf3.weights_square_array,
+    )
+
+
+def test_add_frequency_mix_weights_square():
+    # Same test as above, checking some error handling
+    uvf1 = UVFlag(test_f_file)
+    uvf1.weights_array = 2 * np.ones_like(uvf1.weights_array)
+    uvf2 = copy.deepcopy(uvf1)
+    uvf1.to_waterfall(return_weights_square=True)
+    uvf2.to_waterfall(return_weights_square=False)
+    uvf2.freq_array += 1e4
+    with pytest.raises(
+        ValueError,
+        match="weights_square_array optional parameter is missing from second UVFlag",
+    ):
+        uvf1.__add__(uvf2, axis="frequency", inplace=True)
+
+
 def test_add_pol():
     uv1 = UVFlag(test_f_file)
     uv2 = copy.deepcopy(uv1)
@@ -1075,6 +1104,10 @@ def test_to_waterfall_bl_ret_wt_sq():
     uvf.to_waterfall(return_weights_square=True)
     assert np.all(uvf.weights_square_array == 4 * Nbls)
 
+    # Switch to flag and check that it is now set to None
+    uvf.to_flag()
+    assert uvf.weights_square_array is None
+
 
 def test_collapse_pol(test_outfile):
     uvf = UVFlag(test_f_file)

pyuvdata/uvflag/uvflag.py

@@ -559,6 +559,9 @@ def _set_mode_flag(self):
         self._flag_array.required = True
         self._metric_array.required = False
         self._weights_array.required = False
+        if self.weights_square_array is not None:
+            self.weights_square_array = None
+
         return
 
     def _set_mode_metric(self):
@@ -570,6 +573,7 @@ def _set_mode_metric(self):
 
         if self.weights_array is None and self.metric_array is not None:
             self.weights_array = np.ones_like(self.metric_array, dtype=float)
+
         return
 
     def _set_type_antenna(self):
@@ -1550,17 +1554,23 @@ def __add__(
             )
             this.Npols = len(this.polarization_array)
 
-        if this.mode == "flag":
-            this.flag_array = np.concatenate(
-                [this.flag_array, other.flag_array], axis=ax
-            )
-        elif this.mode == "metric":
-            this.metric_array = np.concatenate(
-                [this.metric_array, other.metric_array], axis=ax
-            )
-            this.weights_array = np.concatenate(
-                [this.weights_array, other.weights_array], axis=ax
-            )
+        for attr in this._data_params:
+            # Check that 'other' also has the attribute filled
+            if getattr(other, attr) is not None:
+                setattr(
+                    this,
+                    attr,
+                    np.concatenate(
+                        [getattr(this, attr), getattr(other, attr)], axis=ax
+                    ),
+                )
+            # May 21, 2020 - should only happen for weights_square_array attr
+            else:
+                raise ValueError(
+                    f"{attr} optional parameter is missing from second UVFlag"
+                    f" object. To concatenate two {this.mode} objects, they"
+                    " must both contain the same optional parameters set."
+                )
 
         this.history += "Data combined along " + axis + " axis. "
         if not uvutils._check_history_version(this.history, this.pyuvdata_version_str):