Repository: RadioAstronomySoftwareGroup/pyuvdata (python)

Issue:
baselines_to_antums is not vectorized
Seems you can't give a list of baseline numbers and get a list of antenna pair tuples.  That would be useful.

Pull request:
Vectorize baseline_to_antnums
<!--- Provide a general summary of your changes in the Title above -->
allows array_like inputs to `baseline_to_antnums`
## Description
<!--- Describe your changes in detail -->
casts input of `baseline_to_antnums` to `np.asarray` to perform vectorized operations on baseline numbers given as lists. 

updated docstring to numpy style
## Motivation and Context
<!--- Why is this change required? What problem does it solve? -->
<!--- If it fixes an open issue, please link to the issue here. If this PR closes an issue, put the word 'closes' before the issue link to auto-close the issue when the PR is merged. -->
allows for more consistent behavior in `baseline_to_antnums` between numpy arrays and other array_like objects.
closes #619
## Types of changes
<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [ ] Documentation change (documentation changes only)
- [ ] Version change
- [ ] Build or continuous integration change


## Checklist:
<!--- You may remove the checklists that don't apply to your change type(s) or just leave them empty -->
<!--- Go over all the following points, and replace the space with an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
- [x] I have read the [contribution guide](https://github.com/RadioAstronomySoftwareGroup/pyuvdata/blob/master/.github/CONTRIBUTING.md).
- [x] My code follows the code style of this project.

Bug fix checklist:
- [x] My fix includes a new test that breaks as a result of the bug (if possible).
- [ ] All new and existing tests pass in both python 2 and python 3.
- [x] I have updated the [CHANGELOG](https://github.com/RadioAstronomySoftwareGroup/pyuvdata/blob/master/CHANGELOG.md).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f581386f..0d42d581 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,11 +2,14 @@
 All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
+- `utils.uvcalibrate` flag propagation bug fix
+- `UVCal.ant2ind` indexing bug fix
 - `UVCal.get_*` methods for accessing data arrays with antenna-polarization keys
 - `utils.uvcalibrate` for automated calibration of `UVData` by `UVCal`
 
 ### Fixed
 - Fixed a bug in select that caused bls and antenna_names/numbers to be or'ed rather than and'ed together.
+- Fixed a bug where `baseline_to_antnums` could accept a numpy array as input but not other array_like objects.
 
 ## [1.4.0] - 2019-05-23
 
diff --git a/pyuvdata/utils.py b/pyuvdata/utils.py
index 3a669b19..55394e44 100644
--- a/pyuvdata/utils.py
+++ b/pyuvdata/utils.py
@@ -464,12 +464,11 @@ def uvcalibrate(uvdata, uvcal, inplace=True, prop_flags=True, flag_missing=True,
                 gain[mask] = 1.0
                 uvdata.flag_array[blt_inds, 0, :, pol_ind] += mask
 
-            else:
-                # apply to data
-                if uvcal.gain_convention == 'multiply':
-                    uvdata.data_array[blt_inds, 0, :, pol_ind] *= gain
-                elif uvcal.gain_convention == 'divide':
-                    uvdata.data_array[blt_inds, 0, :, pol_ind] /= gain
+            # apply to data
+            if uvcal.gain_convention == 'multiply':
+                uvdata.data_array[blt_inds, 0, :, pol_ind] *= gain
+            elif uvcal.gain_convention == 'divide':
+                uvdata.data_array[blt_inds, 0, :, pol_ind] /= gain
 
     # update history
     uvdata.history += "\nCalibrated with pyuvdata.utils.uvcalibrate."
@@ -988,16 +987,25 @@ def baseline_to_antnums(baseline, Nants_telescope):
     """
     Get the antenna numbers corresponding to a given baseline number.
 
-    Args:
-        baseline: integer baseline number
-        Nant_telescope: integer number of antennas
+    Parameters
+    ----------
+    baseline : int or array_like of ints
+        baseline number
+    Nant_telescope : int
+        number of antennas
 
-    Returns:
-        tuple with the two antenna numbers corresponding to the baseline.
+    Returns
+    -------
+    int or array_like of int
+        first antenna number(s)
+    int or array_like of int
+        second antenna number(s)
     """
     if Nants_telescope > 2048:
         raise Exception('error Nants={Nants}>2048 not '
                         'supported'.format(Nants=Nants_telescope))
+
+    baseline = np.asarray(baseline, dtype=np.int64)
     if np.min(baseline) > 2**16:
         ant2 = (baseline - 2**16) % 2048 - 1
         ant1 = (baseline - 2**16 - (ant2 + 1)) / 2048 - 1
@@ -1011,16 +1019,23 @@ def antnums_to_baseline(ant1, ant2, Nants_telescope, attempt256=False):
     """
     Get the baseline number corresponding to two given antenna numbers.
 
-    Args:
-        ant1: first antenna number (integer)
-        ant2: second antenna number (integer)
-        Nant_telescope: integer number of antennas
-        attempt256: Option to try to use the older 256 standard used in
-            many uvfits files (will use 2048 standard if there are more
-            than 256 antennas). Default is False.
+    Parameters
+    ----------
+    ant1 : int or array_like of int
+        first antenna number
+    ant2 : int or array_like of int
+        second antenna number
+    Nant_telescope : int
+        number of antennas
+    attempt256 : bool
+        Option to try to use the older 256 standard used in
+        many uvfits files (will use 2048 standard if there are more
+        than 256 antennas). Default is False.
 
-    Returns:
-        integer baseline number corresponding to the two antenna numbers.
+    Returns
+    -------
+    int or array of int
+        baseline number corresponding to the two antenna numbers.
     """
     ant1, ant2 = np.int64((ant1, ant2))
     if Nants_telescope is not None and Nants_telescope > 2048:
diff --git a/pyuvdata/uvcal.py b/pyuvdata/uvcal.py
index 41f51b17..6f86fa83 100644
--- a/pyuvdata/uvcal.py
+++ b/pyuvdata/uvcal.py
@@ -738,7 +738,7 @@ class UVCal(UVBase):
         if not self._has_key(antnum=antnum):
             raise ValueError("{} not found in ant_array".format(antnum))
 
-        return np.argmin(self.ant_array - antnum)
+        return np.argmin(np.abs(self.ant_array - antnum))
 
     def jpol2ind(self, jpol):
         """
diff --git a/pyuvdata/uvdata.py b/pyuvdata/uvdata.py
index 934c71eb..7c2369e8 100644
--- a/pyuvdata/uvdata.py
+++ b/pyuvdata/uvdata.py
@@ -577,15 +577,15 @@ class UVData(UVBase):
 
         Parameters
         ----------
-        baseline : int
+        baseline : int or array_like of int
             baseline number
 
         Returns
         -------
-        int
-            first antenna number
-        int
-            second antenna number
+        int or array_like of int
+            first antenna number(s)
+        int or array_like of int
+            second antenna number(s)
         """
         return uvutils.baseline_to_antnums(baseline, self.Nants_telescope)
 
@@ -595,9 +595,9 @@ class UVData(UVBase):
 
         Parameters
         ----------
-        ant1 : int
+        ant1 : int or array_like of int
             first antenna number
-        ant2 : int
+        ant2 : int or array_like of int
             second antenna number
         attempt256 : bool
             Option to try to use the older 256 standard used in many uvfits files
@@ -605,7 +605,7 @@ class UVData(UVBase):
 
         Returns
         -------
-        int
+        int or array of int
             baseline number corresponding to the two antenna numbers.
         """
         return uvutils.antnums_to_baseline(ant1, ant2, self.Nants_telescope, attempt256=attempt256)


Test patch:
diff --git a/pyuvdata/tests/test_utils.py b/pyuvdata/tests/test_utils.py
index dd405843..01ce1e4e 100644
--- a/pyuvdata/tests/test_utils.py
+++ b/pyuvdata/tests/test_utils.py
@@ -948,7 +948,7 @@ def test_uvcalibrate_apply_gains():
 
     # division calibrate
     uvc.gain_convention = 'divide'
-    uvdcal = uvutils.uvcalibrate(uvd, uvc, prop_flags=False, flag_missing=False, inplace=False)
+    uvdcal = uvutils.uvcalibrate(uvd, uvc, prop_flags=True, flag_missing=False, inplace=False)
     np.testing.assert_array_almost_equal(uvdcal.get_data(key), uvd.get_data(key) / (uvc.get_gains(ant1) * uvc.get_gains(ant2).conj()).T)
 
     # multiplication calibrate
diff --git a/pyuvdata/tests/test_uvcal.py b/pyuvdata/tests/test_uvcal.py
index c4392f22..5f8b14bc 100644
--- a/pyuvdata/tests/test_uvcal.py
+++ b/pyuvdata/tests/test_uvcal.py
@@ -1442,17 +1442,19 @@ def test_uvcal_get_methods():
     uvc = UVCal()
     uvc.read_calfits(os.path.join(DATA_PATH, 'zen.2457698.40355.xx.gain.calfits'))
 
-    # test get methods
-    key = (9, 'Jxx')
+    # test get methods: add in a known value and make sure it is returned
+    key = (10, 'Jxx')
+    uvc.gain_array[1] = 0.0
     d, f, q = uvc.get_gains(key), uvc.get_flags(key), uvc.get_quality(key)
 
     # test shapes
+    assert np.all(np.isclose(d, 0.0))
     assert d.shape == (uvc.Nfreqs, uvc.Ntimes)
     assert f.shape == (uvc.Nfreqs, uvc.Ntimes)
     assert q.shape == (uvc.Nfreqs, uvc.Ntimes)
 
     # test against by-hand indexing
-    np.testing.assert_array_almost_equal(d, uvc.gain_array[uvc.ant_array.tolist().index(9), 0, :, :, uvc.jones_array.tolist().index(-5)])
+    np.testing.assert_array_almost_equal(d, uvc.gain_array[uvc.ant_array.tolist().index(10), 0, :, :, uvc.jones_array.tolist().index(-5)])
 
     # test variable key input
     d2 = uvc.get_gains(*key)
@@ -1461,23 +1463,23 @@ def test_uvcal_get_methods():
     np.testing.assert_array_almost_equal(d, d2)
     d2 = uvc.get_gains(key[:1])
     np.testing.assert_array_almost_equal(d, d2)
-    d2 = uvc.get_gains(9, -5)
+    d2 = uvc.get_gains(10, -5)
     np.testing.assert_array_almost_equal(d, d2)
-    d2 = uvc.get_gains(9, 'x')
+    d2 = uvc.get_gains(10, 'x')
     np.testing.assert_array_almost_equal(d, d2)
 
     # check has_key
-    assert uvc._has_key(antnum=9)
+    assert uvc._has_key(antnum=10)
     assert uvc._has_key(jpol='Jxx')
-    assert uvc._has_key(antnum=9, jpol='Jxx')
-    assert not uvc._has_key(antnum=9, jpol='Jyy')
+    assert uvc._has_key(antnum=10, jpol='Jxx')
+    assert not uvc._has_key(antnum=10, jpol='Jyy')
     assert not uvc._has_key(antnum=101, jpol='Jxx')
 
     # test exceptions
     pytest.raises(ValueError, uvc.get_gains, 1)
-    pytest.raises(ValueError, uvc.get_gains, (9, 'Jyy'))
+    pytest.raises(ValueError, uvc.get_gains, (10, 'Jyy'))
     uvc.cal_type = 'delay'
-    pytest.raises(ValueError, uvc.get_gains, 9)
+    pytest.raises(ValueError, uvc.get_gains, 10)
 
 
 def test_deprecated_x_orientation():
diff --git a/pyuvdata/tests/test_uvdata.py b/pyuvdata/tests/test_uvdata.py
index 77b63b9e..1ef3a22b 100644
--- a/pyuvdata/tests/test_uvdata.py
+++ b/pyuvdata/tests/test_uvdata.py
@@ -297,7 +297,10 @@ def uvdata_baseline():
 def test_baseline_to_antnums(uvdata_baseline):
     """Test baseline to antnum conversion for 256 & larger conventions."""
     assert uvdata_baseline.uv_object.baseline_to_antnums(67585) == (0, 0)
-    pytest.raises(Exception, uvdata_baseline.uv_object2.baseline_to_antnums, 67585)
+    with pytest.raises(Exception) as cm:
+        uvdata_baseline.uv_object2.baseline_to_antnums(67585)
+    assert str(cm.value).startswith('error Nants={Nants}>2048'
+                                    ' not supported'.format(Nants=uvdata_baseline.uv_object2.Nants_telescope))
 
     ant_pairs = [(10, 20), (280, 310)]
     for pair in ant_pairs:
@@ -313,6 +316,18 @@ def test_baseline_to_antnums(uvdata_baseline):
         assert pair == ant_pair_out
 
 
+def test_baseline_to_antnums_vectorized(uvdata_baseline):
+    """Test vectorized antnum to baseline conversion."""
+    ant_1 = [10, 280]
+    ant_2 = [20, 310]
+    baseline_array = uvdata_baseline.uv_object.antnums_to_baseline(ant_1, ant_2)
+    assert np.array_equal(baseline_array, [88085, 641335])
+    ant_1_out, ant_2_out = uvdata_baseline.uv_object.baseline_to_antnums(baseline_array.tolist())
+    print('out:', ant_1_out, ant_2_out)
+    assert np.array_equal(ant_1, ant_1_out)
+    assert np.array_equal(ant_2, ant_2_out)
+
+
 def test_antnums_to_baselines(uvdata_baseline):
     """Test antums to baseline conversion for 256 & larger conventions."""
     assert uvdata_baseline.uv_object.antnums_to_baseline(0, 0) == 67585