[WIP] Add optional near-field corrections to the phasing function

src/pyuvdata/utils/phase_center_catalog.py

@@ -24,6 +24,8 @@ def look_in_catalog(
     cat_pm_dec=None,
     cat_dist=None,
     cat_vrad=None,
+    cat_dist_units="pc",
+    cat_near_field=False,
     ignore_name=False,
     target_cat_id=None,
     phase_dict=None,
@@ -83,6 +85,10 @@ def look_in_catalog(
         Radial velocity of the source, in units of km/s. Only used for sidereal and
         ephem phase centers. Expected to be a float for sidereal and driftscan phase
         centers, and an ndarray of floats of shape (Npts,) for ephem phase centers.
+    cat_dist_units : str
+        Units to interpret the cat_dist parameter in. Defaults to parsecs.
+    cat_near_field : bool
+        Whether or not near-field corrections have been applied. Defaults to False.
     ignore_name : bool
         Nominally, this method will only look at entries where `cat_name`
         matches the name of an entry in the catalog. However, by setting this to
@@ -146,6 +152,8 @@ def look_in_catalog(
             "cat_pm_dec": cat_pm_dec,
             "cat_dist": cat_dist,
             "cat_vrad": cat_vrad,
+            "cat_dist_units": cat_dist_units,
+            "cat_near_field": cat_near_field,
         }
 
     tol_dict = {
@@ -159,6 +167,8 @@ def look_in_catalog(
         "cat_pm_dec": default_tols,
         "cat_dist": default_tols,
         "cat_vrad": default_tols,
+        "cat_dist_units": None,
+        "cat_near_field": None,
     }
 
     if target_cat_id is not None:

src/pyuvdata/utils/phasing.py

@@ -7,12 +7,13 @@
 import erfa
 import numpy as np
 from astropy import units
-from astropy.coordinates import Angle, Distance, EarthLocation, SkyCoord
+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 _get_autocorrelations_mask, _nants_to_nblts, _ntimes_to_nblts
 
 try:
     from lunarsky import MoonLocation, SkyCoord as LunarSkyCoord, Time as LTime
@@ -2560,3 +2561,122 @@ def uvw_track_generator(
         "lst": lst_array,
         "site_loc": site_loc,
     }
+
+
+def _get_focus_xyz(uvd, focus, ra, dec):
+    """
+    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
+        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: ndarray, ndarray, ndarray
+        ENU-frame coordinates of the focal point (meters) (shape (Ntimes,))
+    """
+    # 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)
+
+    # The centre of the ENU frame should be located at the MEDIAN position of the array
+    loc = uvd.telescope.location.itrs.cartesian.xyz.value
+    antpos = uvd.telescope.antenna_positions + loc
+    x, y, z = np.median(antpos, axis=0)
+
+    # 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
+    obj = obj.transform_to(AltAz(obstime=timesteps, location=mwa))
+
+    # Obtain altitude and azimuth
+    theta, phi = obj.alt.to(units.rad), obj.az.to(units.rad)
+
+    # 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 _get_delay(uvd, focus_x, focus_y, focus_z, flipconj):
+    """
+    Calculate near-field phase/delay along the Nblts axis.
+
+    Parameters
+    ----------
+    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 : 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 = uvd.telescope.get_enu_antpos()
+
+    # 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
+    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
+
+    return phi, new_w  # Each of shape (Nblts,)

src/pyuvdata/utils/tools.py

@@ -368,3 +368,100 @@ def _sorted_unique_difference(obj1, obj2=None):
         List containing the difference in unique entries between obj1 and obj2.
     """
     return sorted(set(obj1)) if obj2 is None else sorted(set(obj1).difference(obj2))
+
+
+def _nants_to_nblts(uvd):
+    """
+    Obtain indices to convert (Nants,) to (Nblts,).
+
+    Parameters
+    ----------
+    uvd : UVData object
+
+    Returns
+    -------
+    ind1, ind2 : ndarray, ndarray
+        index pairs to compose (Nblts,) shaped arrays for each
+        baseline from an (Nants,) shaped array
+    """
+    ants = uvd.telescope.antenna_numbers
+
+    ant1 = uvd.ant_1_array
+    ant2 = uvd.ant_2_array
+
+    ind1 = []
+    ind2 = []
+
+    for i in ant1:
+        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):
+    """
+    Obtain indices to convert (Ntimes,) to (Nblts,).
+
+    Parameters
+    ----------
+    uvd : UVData object
+        UVData object
+
+    Returns
+    -------
+    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):
+    """
+    Get a (Nblts,) shaped array that masks autocorrelations.
+
+    Parameters
+    ----------
+    uvd : UVData object
+        UVData object
+
+    Returns
+    -------
+    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.telescope.antenna_numbers:
+        num = uvd.antpair2ind(i, ant2=i)
+
+        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
+    autos = np.asarray(autos).flatten()
+
+    # Initialize mask of ones (1 = not an autocorrelation)
+    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
+        mask[autos] = 0
+
+    return mask