ratt-ru · sjperkins · Jan 13, 2020 · Jan 13, 2020 · Jan 15, 2020
diff --git a/daskms/observation.py b/daskms/observation.py
@@ -0,0 +1,167 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from pyrap.measures import measures
+from pyrap.quanta import quantity as q
+
+VLA_ANTENNA_POSITIONS = np.asarray([
+         [-1.60171e+06, -5.04201e+06, 3.5546e+06],
+         [-1.60115e+06, -5.042e+06, 3.55486e+06],
+         [-1.60072e+06, -5.04227e+06, 3.55467e+06],
+         [-1.60119e+06, -5.042e+06, 3.55484e+06],
+         [-1.60161e+06, -5.042e+06, 3.55465e+06],
+         [-1.60116e+06, -5.04183e+06, 3.5551e+06],
+         [-1.60101e+06, -5.04209e+06, 3.5548e+06],
+         [-1.60119e+06, -5.04198e+06, 3.55488e+06],
+         [-1.60095e+06, -5.04213e+06, 3.55477e+06],
+         [-1.60118e+06, -5.04193e+06, 3.55495e+06],
+         [-1.60107e+06, -5.04205e+06, 3.55482e+06],
+         [-1.6008e+06, -5.04222e+06, 3.55471e+06],
+         [-1.60116e+06, -5.04178e+06, 3.55516e+06],
+         [-1.60145e+06, -5.04199e+06, 3.55474e+06],
+         [-1.60123e+06, -5.04198e+06, 3.55486e+06],
+         [-1.60153e+06, -5.042e+06, 3.5547e+06],
+         [-1.60114e+06, -5.04168e+06, 3.55532e+06],
+         [-1.60132e+06, -5.04199e+06, 3.55481e+06],
+         [-1.60117e+06, -5.04187e+06, 3.55504e+06],
+         [-1.60119e+06, -5.04202e+06, 3.55481e+06],
+         [-1.60117e+06, -5.0419e+06, 3.55499e+06],
+         [-1.60088e+06, -5.04217e+06, 3.55474e+06],
+         [-1.60138e+06, -5.04199e+06, 3.55478e+06],
+         [-1.60118e+06, -5.04195e+06, 3.55492e+06],
+         [-1.60127e+06, -5.04198e+06, 3.55483e+06],
+         [-1.60111e+06, -5.04202e+06, 3.55484e+06],
+         [-1.60115e+06, -5.04173e+06, 3.55524e+06]])
+
+
+def _julian_day(year, month, day):
+    """
+    Given a Anno Dominei date, computes the Julian Date in days.
+
+    Parameters
+    ----------
+    year : int
+    month : int
+    day : float
+
+    Returns
+    -------
+    float
+        Julian Date
+    """
+
+    # Formula below from
+    # http://scienceworld.wolfram.com/astronomy/JulianDate.html
+    # Also agrees with https://gist.github.com/jiffyclub/1294443
+    return (367*year - int(7*(year + int((month+9)/12))/4)
+            - int((3*(int(year + (month - 9)/7)/100)+1)/4)
+            + int(275*month/9) + day + 1721028.5)
+
+
+def _modified_julian_date(year, month, day):
+    """
+    Given a Anno Dominei date, computes the Modified Julian Date in days.
+
+    Parameters
+    ----------
+    year : int
+    month : int
+    day : int
+
+    Returns
+    -------
+    float
+        Modified Julian Date
+    """
+
+    return _julian_day(year, month, day) - 2400000.5
+
+
+def time(year, month, day, hours, minutes, seconds, intervals):
+    """
+    Parameters
+    ----------
+
+    .. math::
+
+        t_i = mjd +  \frac{interval_i}{2} + \frac{interval_{i + 1}}{2}}
+
+    year : int
+    month : int
+    day : int
+    hours : int
+    minutes : int
+    seconds : int
+    interval : :class:`numpy.ndarray`
+        A sequence of intervals **around** the time
+
+    Returns
+    -------
+    time : :class:`numpy.ndarray`
+        Array of time values
+    """
+    day += (hours / 24.) + (minutes / 1440.) + (seconds / 86400.)
+    start = _modified_julian_date(year, month, day) * 86400.
+    half = intervals / 2
+
+    half[1:] += half[:-1]
+    half[0] = 0
+
+    return start + np.cumsum(half)
+
+
+def synthesize_uvw(antenna_positions, time, phase_dir,
+                   auto_correlations=True):
+    """
+    Synthesizes new UVW coordinates based on time according to
+    NRAO CASA convention (same as in fixvis)
+    User should check these UVW coordinates carefully:
+    if time centroid was used to compute
+    original uvw coordinates the centroids
+    of these new coordinates may be wrong, depending on whether
+    data timesteps were heavily flagged.
+    """
+
+    dm = measures()
+    epoch = dm.epoch("UT1", q(time[0], "s"))
+    ref_dir = dm.direction("j2000",
+                           q(phase_dir[0], "rad"),
+                           q(phase_dir[1], "rad"))
+    ox, oy, oz = antenna_positions[0]
+    obs = dm.position("ITRF", q(ox, "m"), q(oy, "m"), q(oz, "m"))
+
+    # Setup local horizon coordinate frame with antenna 0 as reference position
+    dm.do_frame(obs)
+    dm.do_frame(ref_dir)
+    dm.do_frame(epoch)
+
+    ant1, ant2 = np.triu_indices(antenna_positions.shape[0],
+                                 0 if auto_correlations else 1)
+
+    ntime = time.shape[0]
+    nbl = ant1.shape[0]
+    rows = ntime * nbl
+    uvw = np.empty((rows, 3), dtype=np.float64)
+
+    # For each timestep
+    for ti, t in enumerate(time):
+        epoch = dm.epoch("UT1", q(t, "s"))
+        dm.do_frame(epoch)
+
+        ant_uvw = np.zeros_like(antenna_positions)
+
+        # Calculate antenna UVW positions
+        for ai, (x, y, z) in enumerate(antenna_positions):
+            bl = dm.baseline("ITRF",
+                             q([x, ox], "m"),
+                             q([y, oy], "m"),
+                             q([z, oz], "m"))
+
+            ant_uvw[ai] = dm.to_uvw(bl)["xyz"].get_value()[0:3]
+
+        # Now calculate baseline UVW positions
+        # noting that ant1 - ant2 is the CASA convention
+        base = ti*nbl
+        uvw[base:base + nbl, :] = ant_uvw[ant1] - ant_uvw[ant2]
+
+    return uvw
diff --git a/daskms/tests/test_observation.py b/daskms/tests/test_observation.py
@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+
+import astropy.units as u
+import numpy as np
+from numpy.testing import assert_almost_equal
+import pytest
+
+from daskms.observation import time as sim_time
+
+
+@pytest.mark.parametrize("date", [(2020, 1, 1)])
+def test_time(date):
+    year, month, day = date
+    intervals = np.array([14, 15, 16, 17])
+    res = sim_time(year, month, day, 0, 0, 0, intervals)
+    start = 58849.0 * 86400.
+
+    expected = start + np.array([0, 14.5, 30, 46.5])
+    assert_almost_equal(res, expected)
+
+
+def xyz_to_uvw_rot_mat(hour_angle, declination):
+    har = hour_angle.to(u.rad).value  # HA-radians
+    decr = declination.to(u.rad).value  # Dec-radians
+
+    ntime = hour_angle.shape[0]
+
+    sin_har = np.sin(har)
+    cos_har = np.cos(har)
+    sin_decr = np.full(ntime, np.sin(decr))
+    cos_decr = np.full(ntime, np.cos(decr))
+    time_zeros = np.zeros(sin_har.shape, dtype=har.dtype)
+
+    rotation = np.array(
+        [[sin_har, cos_har, time_zeros],
+         [-sin_decr * cos_har, sin_decr * sin_har, cos_decr],
+         [cos_decr * cos_har, -cos_decr * sin_har, sin_decr]])
+
+    return rotation
+
+
+def z_rot_mat(rotation_angle):
+    ar = rotation_angle.to(u.rad).value  # Angle in radians
+
+    rotation = np.array(
+        [[np.cos(ar), np.sin(ar), 0],
+         [-np.sin(ar), np.cos(ar), 0],
+         [0, 0, 1], ],
+        dtype=np.float_)
+
+    return rotation
+
+
+@pytest.mark.parametrize("auto_corrs", [True])
+def test_synthesize_uvw(wsrt_antenna_positions, auto_corrs):
+    from daskms.observation import synthesize_uvw
+    intervals = np.full(4, 15.0, dtype=np.float64)
+    t = sim_time(2020, 1, 1, 0, 0, 0, intervals)
+    phase_dir = np.deg2rad([30, 60])
+    uvw = synthesize_uvw(wsrt_antenna_positions, t, phase_dir,
+                         auto_correlations=auto_corrs)
+
+    # print(uvw.shape)
+    # print(wsrt_antenna_positions)
+
+    from astropy.time import Time
+    from astropy.coordinates import EarthLocation
+    import astropy.units as u
+
+    # mean_posn = np.mean(wsrt_antenna_positions, axis=0)
+    mean_posn = wsrt_antenna_positions[0]
+    centre = EarthLocation.from_geocentric(mean_posn[0],
+                                           mean_posn[1],
+                                           mean_posn[2],
+                                           unit=u.m)
+
+    lon, lat, height = centre.to_geodetic()
+
+    mean_subbed_itrf = wsrt_antenna_positions - mean_posn
+    obs_times = Time(t / 86400.00, format='mjd', scale='utc')
+
+    rotation = z_rot_mat(lon)
+    ant_local_xyz = np.dot(rotation, mean_subbed_itrf.T).T
+
+    ant1, ant2 = np.triu_indices(wsrt_antenna_positions.shape[0],
+                                 0 if auto_corrs else 1)
+
+    baseline_local_xyz = ant_local_xyz[ant1] - ant_local_xyz[ant2]
+    baseline_local_xyz *= u.m
+
+    ntime = len(obs_times)
+    nbl = len(baseline_local_xyz)
+    uvw_array = np.zeros((ntime*nbl, 3), dtype=np.float_)
+    uvw_array *= baseline_local_xyz.unit
+
+    from astropy.coordinates import SkyCoord
+    phase_dir = SkyCoord(ra=phase_dir[0], dec=phase_dir[1], unit=u.rad)
+
+    lha = obs_times.sidereal_time('apparent', longitude=lon) - phase_dir.ra
+
+    rotation = xyz_to_uvw_rot_mat(lha, phase_dir.dec)
+    res = np.einsum("ijt,bj->tbi", rotation, baseline_local_xyz)
+    res = res.reshape(-1, 3)
+
+    uvw *= u.m
+    v = res - uvw
+    print(v)
+
+    rmsd = np.sqrt(np.sum((res - uvw)**2, axis=0)/res.shape[0])
+    print("RMSD", rmsd)
+    print(uvw)
+    print(res)
+    print("max", (res - uvw).max(axis=0))
+    assert_almost_equal(res, uvw)