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

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
@@ -2572,3 +2573,124 @@
         "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() - 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
+    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 @@
         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

src/pyuvdata/uvdata/uvdata.py

@@ -24,6 +24,7 @@
 from ..telescopes import known_telescopes
 from ..utils import phasing as phs_utils
 from ..utils.io import hdf5 as hdf5_utils
+from ..utils.phasing import _get_delay, _get_focus_xyz
 from ..uvbase import UVBase
 from .initializers import new_uvdata
 
@@ -4625,11 +4626,68 @@
             elif (key == "cat_id") and (phase_dict[key] is not None):
                 # If this is the cat_id, make it an int
                 phase_dict[key] = int(phase_dict[key])
-            elif not ((phase_dict[key] is None) or isinstance(phase_dict[key], str)):
-                phase_dict[key] = float(phase_dict[key])
-
         return phase_dict
 
+    def _apply_near_field_corrections(self, focus, ra, dec, flipconj=False):
+        """
+        Apply near-field corrections by focusing the array to the specified focal point.
+
+        Parameters
+        ----------
+        focus : astropy.units.Quantity object
+            Focal point of the array
+        ra : ndarray
+            Right ascension of the focal point ie phase center (rad; shape (Ntimes,))
+        dec : ndarray
+            Declination of the focal point ie phase center (rad; shape (Ntimes,))
+        flipconj : bool
+            Is the uvd conjugation scheme "flipped" compared to
+            what pyuvdata expects? (default False)
+
+        Returns
+        -------
+        None (performs operations inplace)
+        """
+        # ------- Parameters that are independent of frequency --------
+
+        # Obtain focal distance in km
+        focus = focus.to(units.km).value
+
+        # Convert ra, dec from radians to degrees
+        ra, dec = np.degrees(ra), np.degrees(dec)
+
+        # Calculate the x, y, z coordinates of the focal point
+        # in ENU frame for each vis along Nblts axis
+        focus_x, focus_y, focus_z = _get_focus_xyz(self, focus, ra, dec)
+
+        # Calculate near-field correction at the specified timestep
+        # for each vis along Nblts axis
+        phi, new_w = _get_delay(self, focus_x, focus_y, focus_z, flipconj)
+
+        # Update old w with new w
+        self.uvw_array[:, -1] = new_w
+
+        # ---------------- Frequency-dependent calculations ---------------------
+
+        # Calculate wavelength associate with each frequency
+        wavelengths = 299792458 / self.freq_array
+
+        # Reshape the phi and wavelengths arrays in order to
+        # be able to broadcast them together
+        phi = np.reshape(phi, (phi.size, 1))  # (Nblts, 1)
+        wavelengths = np.reshape(wavelengths, (1, wavelengths.size))  # (1, Nfreqs)
+
+        # Calculate phase corrections at all frequencies
+        # -- produces shape (Nblts, Nfreqs)
+        phase_corrections = np.exp(-2j * np.pi * phi / wavelengths)
+
+        # Set data at each polarization (Npols = 4 usually)
+        for pol in self.polarization_array:
+            prev = np.reshape(self.get_data(pol), (self.Nblts, self.Nfreqs, 1))
+            corr = np.reshape(phase_corrections, (self.Nblts, self.Nfreqs, 1))
+
+            self.set_data(corr * prev, pol)
+
     def phase(
         self,
         *,
@@ -4690,7 +4748,9 @@
         cat_type : str
             Type of phase center to be added. Must be one of:
             "sidereal" (fixed RA/Dec), "ephem" (RA/Dec that moves with time),
-            "driftscan" (fixed az/el position). Default is "sidereal".
+            "driftscan" (fixed az/el position), "near_field" (applies near-field
+            corrections to the specified dist, assuming sidereal phase center).
+            Default is "sidereal".
         ephem_times : ndarray of float
             Only used when `cat_type="ephem"`. Describes the time for which the values
             of `cat_lon` and `cat_lat` are caclulated, in units of JD. Shape is (Npts,).
@@ -4700,10 +4760,14 @@
         pm_dec : float
             Proper motion in Dec, in units of mas/year. Only used for sidereal phase
             centers.
-        dist : float or ndarray of float
-            Distance of the source, in units of pc. Only used for sidereal and ephem
-            phase centers. Expected to be a float for sidereal phase
-            centers, and an ndarray of floats of shape (Npts,) for ephem phase centers.
+        dist : float or ndarray of float or astropy.units.Quantity object.
+            Distance to the source. Used for sidereal and ephem phase centers,
+            and for applying near-field corrections. If passed either as a float
+            (for sidereal phase centers) or as an ndarray of floats of shape (Npts,)
+            (for ephem phase centers), will be interpreted in units of parsec for all
+            cat_types except near_field; in the latter case it will be interpreted
+            in meters. Alternatively, an astropy.units.Quantity object may be passed
+            instead, in which case the units will be infered automatically.
         vrad : float or ndarray of float
             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 phase
@@ -4747,6 +4811,25 @@
         # Before moving forward with the heavy calculations, we need to do some
         # basic housekeeping to make sure that we've got the coordinate data that
         # we need in order to proceed.
+        if dist is not None:
+            if isinstance(dist, units.Quantity):
+                dist_qt = copy.deepcopy(dist)
+            else:
+                if cat_type == "near_field":
+                    dist_qt = dist * units.m
+                else:
+                    dist_qt = dist * units.parsec
+
+            dist = dist_qt.to(
+                units.parsec
+            ).value  # phase_dict internally stores in parsecs
+
+        if cat_type == "near_field":
+            near_field = True
+            cat_type = "sidereal"  # apply far-field phasing BEFORE near-field
+        else:
+            near_field = False
+
         phase_dict = self._phase_dict_helper(
             lon=lon,
             lat=lat,
@@ -4901,6 +4984,12 @@
         if cleanup_old_sources:
             self._clear_unused_phase_centers()
 
+        # Lastly, apply near-field corrections if specified
+        if near_field:
+            self._apply_near_field_corrections(
+                focus=dist_qt, ra=phase_dict["cat_lon"], dec=phase_dict["cat_lat"]
+            )
+
     def phase_to_time(
         self, time, *, phase_frame="icrs", use_ant_pos=True, select_mask=None
     ):