diff --git a/src/pyuvdata/utils/phasing.py b/src/pyuvdata/utils/phasing.py index 5d9321e98..a9fbb5866 100644 --- a/src/pyuvdata/utils/phasing.py +++ b/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,44 +2563,50 @@ 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) @@ -2608,54 +2614,70 @@ def getFocusXYZ(uvd, focus, ra, dec): 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,) diff --git a/src/pyuvdata/utils/tools.py b/src/pyuvdata/utils/tools.py index 9617fba29..481bb9501 100644 --- a/src/pyuvdata/utils/tools.py +++ b/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 diff --git a/src/pyuvdata/uvdata/uvdata.py b/src/pyuvdata/uvdata/uvdata.py index c54e0e089..c203b7990 100644 --- a/src/pyuvdata/uvdata/uvdata.py +++ b/src/pyuvdata/uvdata/uvdata.py @@ -24,9 +24,9 @@ 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 -from ..utils.phasing import getFocusXYZ, getPhase __all__ = ["UVData"] import contextlib @@ -4344,6 +4344,9 @@ def unproject_phase( ValueError If the object is alread unprojected. """ + # Start by undoing the near-field phasing + self.unproject_near_field() + # select_mask_use is length Nblts, True means should be unprojected # only select blts that are actually phased. if select_mask is not None: @@ -4414,6 +4417,8 @@ def _phase_dict_helper( cat_name, lookup_name, time_array, + dist_units="pc", + near_field=False, ): """ Supplies a dictionary with parameters for the phase method to use. @@ -4500,6 +4505,8 @@ def _phase_dict_helper( cat_pm_dec = pm_dec cat_dist = dist cat_vrad = vrad + cat_dist_units = dist_units + cat_near_field = near_field else: temp_dict = self.phase_center_catalog[cat_id] cat_type = temp_dict["cat_type"] @@ -4514,6 +4521,8 @@ def _phase_dict_helper( cat_pm_dec = temp_dict.get("cat_pm_dec") cat_dist = temp_dict.get("cat_dist") cat_vrad = temp_dict.get("cat_vrad") + cat_dist_units = temp_dict.get("cat_dist_units") + cat_near_field = temp_dict.get("cat_near_field") else: # The name of the source is unique! cat_type = "sidereal" if cat_type is None else cat_type @@ -4526,6 +4535,8 @@ def _phase_dict_helper( cat_pm_dec = pm_dec cat_dist = dist cat_vrad = vrad + cat_dist_units = dist_units + cat_near_field = near_field if (cat_epoch is None) and (cat_type != "unprojected"): cat_epoch = 1950.0 if (cat_frame in ["fk4", "fk4noeterms"]) else 2000.0 @@ -4583,6 +4594,8 @@ def _phase_dict_helper( "cat_vrad": cat_vrad, "info_source": info_source, "cat_id": cat_id, + "cat_dist_units": cat_dist_units, + "cat_near_field": cat_near_field, } # Finally, make sure everything is a float or an ndarray of floats @@ -4593,85 +4606,101 @@ def _phase_dict_helper( # 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]) + # cat_dist_units should be str and cat_near_field should be bool + if (key != "cat_dist_units") and (key != "cat_near_field"): + phase_dict[key] = float(phase_dict[key]) return phase_dict - def _apply_near_field_corrections( - self, - focus, - ra, - dec, - f_units="km", - flipconj=False, - ): + 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 : float - Focal point of the array (km or pc) - ra : float + ---------- + 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 : float + dec : ndarray Declination of the focal point ie phase center (rad; shape (Ntimes,)) - f_units : str - Units for the focus parameter. Accepts either "km" or "pc". Defaults to "km". flipconj : bool - Is the uvd conjugation scheme "flipped" compared to what pyuvdata expects? (default False) + Is the uvd conjugation scheme "flipped" compared to + what pyuvdata expects? (default False) - Raises - --------------- - ValueError if f_units not 'km' or 'pc'. - Returns - --------------- + ------- None (performs operations inplace) """ - # ---------------- Parameters that are independent of frequency ------------------------------- + # ------- Parameters that are independent of frequency -------- - ## Make sure focus is in km - if f_units != "km": - if f_units == "pc": - focus = focus * 30856775814913.673 - else: - raise ValueError(f"f_units parameter accepts one of 'km' or 'pc'. Entered value: {f_units}.") + # 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 = getFocusXYZ(self, focus, ra, dec) - ## Calculate near-field correction at the specified timestep for each vis along Nblts axis - phi, new_w = getPhase(self, focus_x, focus_y, focus_z, flipconj) + # 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) - ## Update old w with new w - self.uvw_array[:,-1] = new_w + # 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 ------------------------------------------- + # ---------------- 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) + # 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) + # 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: # For some reason the data need an extra axis? prev = np.reshape(self.get_data(pol), (self.Nblts, self.Nfreqs, 1)) - corr = np.reshape(phase_corrections, (self.Nblts, self.Nfreqs, 1)) + corr = np.reshape(phase_corrections, (self.Nblts, self.Nfreqs, 1)) + + self.set_data(corr * prev, pol) + + def unproject_near_field(self, flipconj=False): + """ + Undo near-field phasing. + + Works by near-field projecting to a distance of infinity. - self.set_data(corr*prev, pol) + Parameters + ---------- + flipconj : bool + Is the conjugation scheme "flipped" compared to + what pyuvdata expects? (Default False) + Returns + ------- + None (performs operations inplace) + """ + # Keep the far-field Ra and Dec + key = list(self.phase_center_catalog.keys())[-1] + ra = self.phase_center_catalog[key]["cat_lon"] + dec = self.phase_center_catalog[key]["cat_lat"] + + self._apply_near_field_corrections( + focus=1e19 * units.km, # TODO: Is there a better way to do this? + ra=ra, + dec=dec, + flipconj=flipconj, + ) def phase( self, @@ -4744,10 +4773,13 @@ def phase( 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. + 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 @@ -4766,8 +4798,9 @@ def phase( Optional mask for selecting which data to operate on along the blt-axis. Shape is (Nblts,). Ignored if `use_old_proj` is True. near_field : bool - Option to apply near-field corrections to the provided ra, dec, and distance. - Defaults to False (apply far-field corrections only). + Option to apply near-field corrections to the provided + ra, dec, and distance. Defaults to False + (apply far-field corrections only). Raises ------ @@ -4775,6 +4808,10 @@ def phase( If the `cat_name` is None. """ + key = list(self.phase_center_catalog.keys())[-1] + if self.phase_center_catalog[key]["cat_name"] != "": + self.unproject_near_field() + if cat_type != "unprojected": if lon is None: if ra is None: @@ -4794,6 +4831,23 @@ def phase( # 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. + dist_units = "pc" + if dist is not None: + if isinstance(dist, units.Quantity): + dist_units = dist.unit + dist_qt = copy.deepcopy(dist) + dist = dist.value + else: + dist_qt = dist * units.parsec + + if near_field: + warnings.warn( + "near_field is set to True, but dist is not an " + "astropy.units.Quantity object. " + "Distance will be interpreted in units of parsec.", + UserWarning, + ) + phase_dict = self._phase_dict_helper( lon=lon, lat=lat, @@ -4808,6 +4862,8 @@ def phase( cat_name=cat_name, lookup_name=lookup_name, time_array=self.time_array, + dist_units=dist_units, + near_field=near_field, ) if phase_dict["cat_type"] not in ["ephem", "unprojected"]: @@ -4951,10 +5007,7 @@ def phase( # Lastly, apply near-field corrections if specified if near_field: self._apply_near_field_corrections( - focus=phase_dict["cat_dist"], - ra=phase_dict["cat_lon"], - dec=phase_dict["cat_lat"], - f_units="pc", + focus=dist_qt, ra=phase_dict["cat_lon"], dec=phase_dict["cat_lat"] ) def phase_to_time(