Added a function to unproject near-field corrections and allow dist t…

…o be passed as an astropy.units.Quantity object.

src/pyuvdata/utils/phasing.py

@@ -7,13 +7,13 @@
 import erfa
 import numpy as np
 from astropy import units
-from astropy.coordinates import Angle, Distance, EarthLocation, SkyCoord, AltAz
+from astropy.coordinates import AltAz, Angle, Distance, EarthLocation, SkyCoord
 from astropy.time import Time
 from astropy.utils import iers
 
 from . import _phasing
 from .times import get_lst_for_time
-from .tools import Nants_to_Nblts, Ntimes_to_Nblts, get_autocorrelations_mask
+from .tools import _get_autocorrelations_mask, _nants_to_nblts, _ntimes_to_nblts
 
 try:
     from lunarsky import MoonLocation, SkyCoord as LunarSkyCoord, Time as LTime
@@ -2563,99 +2563,121 @@ def uvw_track_generator(
     }
 
 
-def getFocusXYZ(uvd, focus, ra, dec):
+def _get_focus_xyz(uvd, focus, ra, dec):
     """
-    Returns the x,y,z coordinates of the focal point 
-    (ie location of the NEAR-FIELD object of interest)
-    in the MWA-centred ENU frame at each timestep
-    
+    Return the x,y,z coordinates of the focal point.
+
+    The focal point corresponds to the location of
+    the NEAR-FIELD object of interest in the MWA-centred
+    ENU frame at each timestep.
+
     Parameters
-    -------------
-    uvd: UVData object
-    focus: focal distance of the array (km)
-    ra: right ascension of the focal point ie phase center (deg; shape (Ntimes,))
-    dec: declination of the focal point ie phase center (deg; shape (Ntimes,))
-    
+    ----------
+    uvd : UVData object
+        UVData object
+    focus : float
+        Focal distance of the array (km)
+    ra : float
+        Right ascension of the focal point ie phase center (deg; shape (Ntimes,))
+    dec : float
+        Declination of the focal point ie phase center (deg; shape (Ntimes,))
+
     Returns
-    -------------
-    x, y, z: ENU-frame coordinates of the focal point (meters) (shape (Ntimes,))
+    -------
+    x, y, z: ndarray, ndarray, ndarray
+        ENU-frame coordinates of the focal point (meters) (shape (Ntimes,))
     """
+    # Obtain timesteps
+    timesteps = Time(np.unique(uvd.time_array), format="jd")
 
-    ## Obtain timesteps
-    timesteps = Time(np.unique(uvd.time_array), format='jd')
-
-    ## Initialize sky-based coordinates using right ascension and declination
-    obj = SkyCoord(ra*units.deg, dec*units.deg)
+    # Initialize sky-based coordinates using right ascension and declination
+    obj = SkyCoord(ra * units.deg, dec * units.deg)
 
-    ## The centre of the ENU frame should be located at the MEDIAN position of the array
+    # The centre of the ENU frame should be located at the MEDIAN position of the array
     antpos = uvd.antenna_positions + uvd.telescope_location
-    x,y,z = np.median(antpos, axis=0)
+    x, y, z = np.median(antpos, axis=0)
 
-    ## Initialize EarthLocation object centred on MWA
-    mwa = EarthLocation(x,y,z,unit=units.m)
+    # Initialize EarthLocation object centred on MWA
+    mwa = EarthLocation(x, y, z, unit=units.m)
 
-    ## Convert sky object to an AltAz frame centred on the MWA
+    # Convert sky object to an AltAz frame centred on the MWA
     obj = obj.transform_to(AltAz(obstime=timesteps, location=mwa))
 
-    ## Obtain altitude and azimuth
+    # Obtain altitude and azimuth
     theta, phi = obj.alt.to(units.rad), obj.az.to(units.rad)
 
-    ## Obtain x,y,z ENU coordinates
+    # Obtain x,y,z ENU coordinates
     x = focus * 1e3 * np.cos(theta) * np.sin(phi)
     y = focus * 1e3 * np.cos(theta) * np.cos(phi)
     z = focus * 1e3 * np.sin(theta)
 
     return x, y, z
 
 
-def getPhase(uvd, focus_x, focus_y, focus_z, flipconj):
+def _get_delay(uvd, focus_x, focus_y, focus_z, flipconj):
     """
-    Calculates near-field phase/delay along the Nblts axis
-    
+    Calculate near-field phase/delay along the Nblts axis.
+
     Parameters
-    ------------
-    uvd: UVData object
-    focus_x, focus_y, focus_z: ENU-frame coordinates of focal point (Each of shape (Ntimes,))
-    flipconj: is the uvd conjugation scheme "flipped" compared to what pyuvdata expects?
-    
+    ----------
+    uvd : UVData object
+        UVData object
+    focus_x, focus_y, focus_z : ndarray, ndarray, ndarray
+        ENU-frame coordinates of focal point (Each of shape (Ntimes,))
+    flipconj : bool
+        Is the uvd conjugation scheme "flipped" compared to what pyuvdata expects?
+
     Returns
-    ------------
-    phi: the phase correction to apply to each visibility along the Nblts axis
-    new_w: the calculated near-field delay (or w-term) for each visibility along the Nblts axis
+    -------
+    phi : ndarray
+        The phase correction to apply to each visibility along the Nblts axis
+    new_w : ndarray
+        The calculated near-field delay (or w-term) for each visibility along
+        the Nblts axis
     """
-    ## Get indices to convert between Nants and Nblts
-    ind1, ind2 = Nants_to_Nblts(uvd)
-
-    ## Antenna positions in ENU frame
-    antpos, ants = uvd.get_ENU_antpos(center=True)  # Centred on the MEDIAN of the array; antpos has shape (Nants, 3)
-
-    ## Get tile positions for each baseline
+    # Get indices to convert between Nants and Nblts
+    ind1, ind2 = _nants_to_nblts(uvd)
+
+    # Antenna positions in ENU frame
+    antpos, ants = uvd.get_ENU_antpos(
+        center=True
+    )  # Centred on the MEDIAN of the array; antpos has shape (Nants, 3)
+
+    # Get tile positions for each baseline
     tile1 = antpos[ind1]  # Shape (Nblts, 3)
     tile2 = antpos[ind2]  # Shape (Nblts, 3)
-
-    ## Focus points have shape (Ntimes,); convert to shape (Nblts,)
-    t_inds = Ntimes_to_Nblts(uvd)
-    focus_x, focus_y, focus_z = focus_x[t_inds], focus_y[t_inds], focus_z[t_inds]
-
-    ## Calculate distance from antennas to focal point for each visibility along the Nblts axis
-    r1 = np.sqrt((tile1[:,0] - focus_x)**2 + (tile1[:,1] - focus_y)**2 + (tile1[:,2] - focus_z)**2)
-    r2 = np.sqrt((tile2[:,0] - focus_x)**2 + (tile2[:,1] - focus_y)**2 + (tile2[:,2] - focus_z)**2)
-
-    ## Get the uvw array along the Nblts axis; select only the w's
-    old_w = uvd.uvw_array[:,-1]
-
-    ## Calculate near-field delay
+
+    # Focus points have shape (Ntimes,); convert to shape (Nblts,)
+    t_inds = _ntimes_to_nblts(uvd)
+    (focus_x, focus_y, focus_z) = (focus_x[t_inds], focus_y[t_inds], focus_z[t_inds])
+
+    # Calculate distance from antennas to focal point
+    # for each visibility along the Nblts axis
+    r1 = np.sqrt(
+        (tile1[:, 0] - focus_x) ** 2
+        + (tile1[:, 1] - focus_y) ** 2
+        + (tile1[:, 2] - focus_z) ** 2
+    )
+    r2 = np.sqrt(
+        (tile2[:, 0] - focus_x) ** 2
+        + (tile2[:, 1] - focus_y) ** 2
+        + (tile2[:, 2] - focus_z) ** 2
+    )
+
+    # Get the uvw array along the Nblts axis; select only the w's
+    old_w = uvd.uvw_array[:, -1]
+
+    # Calculate near-field delay
     if flipconj:
         new_w = r2 - r1
     else:
         new_w = r1 - r2
     phi = new_w - old_w
 
-    ## Remove autocorrelations
-    mask = get_autocorrelations_mask(uvd)
-
-    new_w = new_w*mask + old_w*(1-mask)
-    phi = phi*mask
+    # Remove autocorrelations
+    mask = _get_autocorrelations_mask(uvd)
 
-    return phi, new_w  # Each of shape (Nblts,)
+    new_w = new_w * mask + old_w * (1 - mask)
+    phi = phi * mask
 
+    return phi, new_w  # Each of shape (Nblts,)

src/pyuvdata/utils/tools.py

@@ -370,20 +370,22 @@ def _sorted_unique_difference(obj1, obj2=None):
     return sorted(set(obj1)) if obj2 is None else sorted(set(obj1).difference(obj2))
 
 
-def Nants_to_Nblts(uvd):
+def _nants_to_nblts(uvd):
     """
-    Obtain arrays of indices to convert an array of shape (Nants,) to an array of shape (Nblts,)
-    
+    Obtain indices to convert (Nants,) to (Nblts,).
+
     Parameters
-    ----------------
-    uvd: UVData object
-    
+    ----------
+    uvd : UVData object
+
     Returns
-    ----------------
-    ind1, ind2: index pairs to compose Nblts shaped arrays for each baseline from an Nants shaped array
+    -------
+    ind1, ind2 : ndarray, ndarray
+        index pairs to compose (Nblts,) shaped arrays for each
+        baseline from an (Nants,) shaped array
     """
     antpos, ants = uvd.get_ENU_antpos()
-    
+
     ant1 = uvd.ant_1_array
     ant2 = uvd.ant_2_array
 
@@ -394,50 +396,61 @@ def Nants_to_Nblts(uvd):
         ind1.append(np.where(ants == i)[0][0])
     for i in ant2:
         ind2.append(np.where(ants == i)[0][0])
-        
+
     return np.asarray(ind1), np.asarray(ind2)
 
 
-def Ntimes_to_Nblts(uvd):
+def _ntimes_to_nblts(uvd):
     """
-    Obtain array of indices to convert an array of shape (Ntimes,) to an array of shape (Nblts,)
-    
+    Obtain indices to convert (Ntimes,) to (Nblts,).
+
     Parameters
-    -----------------
-    UVData object
-    
+    ----------
+    uvd : UVData object
+        UVData object
+
     Returns
-    -----------------
-    inds: indices that, when applied to an array of shape (Ntimes,), correctly convert it to shape (Nblts,)
+    -------
+    inds : ndarray
+        Indices that, when applied to an array of shape (Ntimes,),
+        correctly convert it to shape (Nblts,)
     """
-
     unique_t = np.unique(uvd.time_array)
     t = uvd.time_array
-    
+
     inds = []
     for i in t:
         inds.append(np.where(unique_t == i)[0][0])
-        
+
     return np.asarray(inds)
 
 
-def get_autocorrelations_mask(uvd):
+def _get_autocorrelations_mask(uvd):
     """
-    Get a (Nblts,) shaped array that masks autocorrelations
-    
+    Get a (Nblts,) shaped array that masks autocorrelations.
+
     Parameters
-    ------------------
-    uvd: UVData object
-    
+    ----------
+    uvd : UVData object
+        UVData object
+
     Returns
-    ------------------
-    mask: (Nblts,) array of 1's and 0's, where 0 indicates an autocorrelation
+    -------
+    mask : ndarray
+        array of shape (Nblts,) of 1's and 0's,
+        where 0 indicates an autocorrelation
     """
     # Get indices along the Nblts axis corresponding to autocorrelations
     autos = []
     for i in uvd.antenna_numbers:
+
         num = uvd.antpair2ind(i, ant2=i)
-        if num is not None:
+
+        if isinstance(num, slice):
+            inds = list(range(num.start, num.stop, num.step))
+            autos.append(inds)
+
+        elif num is not None:
             autos.append(num)
 
     # Flatten it to obtain the 1D array of autocorrelation indices
@@ -447,7 +460,9 @@ def get_autocorrelations_mask(uvd):
     mask = np.ones_like(uvd.baseline_array)
 
     # Populate with zeros (0 = is an autocorrelation)
-    if len(autos) > 0:  # Protect against the case where the uvd is already free of autos
+    if (
+        len(autos) > 0
+    ):  # Protect against the case where the uvd is already free of autos
         mask[autos] = 0
-   
+
     return mask