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

src/pyuvdata/utils/phase_center_catalog.py

@@ -7,7 +7,7 @@
 
 from . import RADIAN_TOL
 
-allowed_cat_types = ["sidereal", "ephem", "unprojected", "driftscan"]
+allowed_cat_types = ["sidereal", "ephem", "unprojected", "driftscan", "near_field"]
 
 
 def look_in_catalog(
@@ -610,6 +610,8 @@ def generate_phase_center_cat_entry(
             "driftscan" (fixed az/el position),
             "unprojected" (no w-projection, equivalent to the old
             `phase_type` == "drift").
+            "near-field" (equivalent to sidereal with the addition
+            of near-field corrections)
     cat_lon : float or ndarray
         Value of the longitudinal coordinate (e.g., RA, Az, l) in radians of the
         phase center. No default unless `cat_type="unprojected"`, in which case the
@@ -684,7 +686,7 @@ def generate_phase_center_cat_entry(
     if not isinstance(cat_name, str):
         raise ValueError("cat_name must be a string.")
 
-    # We currently only have 4 supported types -- make sure the user supplied
+    # We currently only have 5 supported types -- make sure the user supplied
     # one of those
     if cat_type not in allowed_cat_types:
         raise ValueError(f"cat_type must be one of {allowed_cat_types}.")

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
@@ -2085,6 +2086,8 @@ def calc_app_coords(
             "ephem" (RA/Dec that moves with time),
             "driftscan" (fixed az/el position),
             "unprojected" (alias for "driftscan" with (Az, Alt) = (0 deg, 90 deg)).
+            "near_field" (equivalent to sidereal, with the addition of
+                near-field corrections)
     time_array : float or ndarray of float or Time object
         Times for which the apparent coordinates are to be calculated, in UTC JD.
         If more than a single element, must be the same shape as lon_coord and
@@ -2175,7 +2178,7 @@ def calc_app_coords(
         else:
             unique_lst = lst_array[unique_mask]
 
-    if coord_type == "sidereal":
+    if coord_type == "sidereal" or coord_type == "near_field":
         # If the coordinates are not in the ICRS frame, go ahead and transform them now
         if coord_frame != "icrs":
             icrs_ra, icrs_dec = transform_sidereal_coords(
@@ -2571,3 +2574,120 @@ 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 ENU
+    frame centered on the median position of the
+    antennas.
+
+    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 center 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 the telescope
+    telescope = EarthLocation(x, y, z, unit=units.m)
+
+    # Convert sky object to an AltAz frame centered on the telescope
+    obj = obj.transform_to(AltAz(obstime=timesteps, location=telescope))
+
+    # 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):
+    """
+    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,))
+
+    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() - np.median(
+        uvd.telescope.get_enu_antpos(), axis=0
+    )
+
+    # 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
+    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,98 @@ 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):
+            step = num.step if num.step is not None else 1
+            inds = list(range(num.start, num.stop, step))
+            autos.append(inds)
+
+    # 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

src/pyuvdata/uvdata/miriad.py

@@ -1643,6 +1643,14 @@ def write_miriad(
                 "Only ITRS telescope locations are supported in Miriad files."
             )
 
+        if any(
+            entry.get("cat_type") == "near_field"
+            for entry in self.phase_center_catalog.values()
+        ):
+            raise NotImplementedError(
+                "Writing near-field phased data to miriad format is not yet supported."
+            )
+
         # change time_array and lst_array to mark beginning of integration,
         # per Miriad format
         miriad_time_array = self.time_array - self.integration_time / (24 * 3600.0) / 2

src/pyuvdata/uvdata/ms.py

@@ -107,6 +107,15 @@ def write_ms(
         if not casa_present:  # pragma: no cover
             raise ImportError(no_casa_message) from casa_error
 
+        if any(
+            entry.get("cat_type") == "near_field"
+            for entry in self.phase_center_catalog.values()
+        ):
+            raise NotImplementedError(
+                "Writing near-field phased data to Measurement Set format "
+                + "is not yet supported."
+            )
+
         if run_check:
             self.check(
                 check_extra=check_extra,