Merge pull request #104 from sot/transforms

Migrate transforms here and make them work for N-d inputs

chandra_aca/tests/test_all.py

@@ -14,7 +14,8 @@
 
 import chandra_aca
 from chandra_aca.transform import (snr_mag_for_t_ccd, radec_to_yagzag,
-                                   yagzag_to_radec, pixels_to_yagzag, yagzag_to_pixels)
+                                   yagzag_to_radec, pixels_to_yagzag, yagzag_to_pixels,
+                                   eci_to_radec, radec_to_eci)
 from chandra_aca import drift
 
 dirname = os.path.dirname(__file__)
@@ -302,3 +303,102 @@ def test_snr_mag():
     assert np.allclose(arr, [9.0, 8.8112, 8.1818], atol=0.0001, rtol=0)
     arr = snr_mag_for_t_ccd(np.array([-11.5, -10, -5]), ref_mag=9.0, ref_t_ccd=-9.5, scale_4c=1.59)
     assert np.allclose(arr, [9.2517, 9.0630, 8.4336], atol=0.0001, rtol=0)
+
+
+def test_eci_to_radec():
+    eci = np.array([0.92541658, 0.16317591, 0.34202014])
+    ra, dec = eci_to_radec(eci)
+    assert np.allclose(ra, 9.9999999129952908)
+    assert np.allclose(dec, 19.999999794004037)
+    assert isinstance(ra, np.float64)
+    assert isinstance(dec, np.float64)
+
+
+def test_vectorized_eci_to_radec():
+    eci = np.array([[0.92541658, 0.16317591, 0.34202014],
+                    [0.9248273, -0.16307201, 0.34365969]])
+    ra, dec = eci_to_radec(eci)
+    assert np.allclose(ra[0], 9.9999999129952908)
+    assert np.allclose(ra[1], 349.9999997287627)
+    assert np.allclose(dec[0], 19.999999794004037)
+    assert np.allclose(dec[1], 20.099999743270516)
+    assert isinstance(ra, np.ndarray)
+    assert isinstance(dec, np.ndarray)
+
+
+def test_radec_to_eci():
+    eci = radec_to_eci(10, 20)
+    assert np.allclose(eci[0], 0.92541658)
+    assert np.allclose(eci[1], 0.16317591)
+    assert np.allclose(eci[2], 0.34202014)
+    assert isinstance(eci, np.ndarray)
+
+
+@pytest.mark.parametrize('shape', [(24,), (6, 4), (3, 4, 2)])
+def test_radec_eci_multidim(shape):
+    """Test radec_to_eci and eci_to_radec for multidimensional inputs"""
+    ras = np.linspace(0., 359., 24)
+    decs = np.linspace(-90., 90., 24)
+    ras_nd = ras.reshape(shape)
+    decs_nd = decs.reshape(shape)
+
+    # First do everything as scalars
+    ecis_list = [radec_to_eci(ra, dec) for ra, dec in zip(ras.tolist(), decs.tolist())]
+    ecis_nd_from_list = np.array(ecis_list).reshape(shape + (3,))
+
+    ecis = radec_to_eci(ras_nd, decs_nd)
+    assert ecis.shape == shape + (3,)
+    assert np.allclose(ecis, ecis_nd_from_list)
+
+    ras_rt, decs_rt = eci_to_radec(ecis)
+    assert ras_rt.shape == shape
+    assert decs_rt.shape == shape
+    assert np.allclose(ras_rt, ras_nd)
+    assert np.allclose(decs_rt, decs_nd)
+
+
+@pytest.mark.parametrize('shape', [(24,), (6, 4), (3, 4, 2)])
+def test_radec_yagzag_multidim(shape):
+    """Test radec_to_yagzag and yagzag_to_radec for multidimensional inputs"""
+    dec = 0.5
+    ras = np.linspace(0., 1., 24)
+    eqs = np.empty(shape=(24, 3))
+    eqs[..., 0] = 0.1
+    eqs[..., 1] = np.linspace(-0.1, 0.1, 24)
+    eqs[..., 2] = np.linspace(-1, 1, 24)
+
+    ras_nd = ras.reshape(shape)
+    qs_nd = Quat(equatorial=eqs.reshape(shape + (3,)))
+
+    # First do everything as scalars
+    yagzags_list = [radec_to_yagzag(ra, dec, Quat(equatorial=eq))
+                    for ra, eq in zip(ras.tolist(), eqs.tolist())]
+    yagzags_nd_from_list = np.array(yagzags_list).reshape(shape + (2,))
+
+    # Test broadcasting by providing a scalar for `dec`
+    yags, zags = radec_to_yagzag(ras_nd, dec, qs_nd)
+    assert yags.shape == shape
+    assert zags.shape == shape
+    assert np.allclose(yags, yagzags_nd_from_list[..., 0])
+    assert np.allclose(zags, yagzags_nd_from_list[..., 1])
+
+    ras_rt, decs_rt = yagzag_to_radec(yags, zags, qs_nd)
+    assert ras_rt.shape == shape
+    assert decs_rt.shape == shape
+    assert np.allclose(ras_rt, ras_nd)
+    assert np.allclose(decs_rt, dec)
+
+
+def test_radec_yagzag_quat_init():
+    att = [1.0, 2.0, 3.0]
+    q_att = Quat(att)
+
+    yag0, zag0 = radec_to_yagzag(1.1, 2.1, q_att)
+    yag1, zag1 = radec_to_yagzag(1.1, 2.1, att)
+    assert yag0 == yag1
+    assert zag0 == zag1
+
+    ra0, dec0 = yagzag_to_radec(100.1, 200.1, q_att)
+    ra1, dec1 = yagzag_to_radec(100.1, 200.1, att)
+    assert dec0 == dec1
+    assert ra0 == ra1

chandra_aca/transform.py

@@ -7,8 +7,6 @@
 - Science target coordinate to ACA frame conversions
 """
 
-from __future__ import division
-
 import numpy as np
 
 from chandra_aca import dark_model
@@ -138,8 +136,8 @@ def pixels_to_yagzag(row, col, allow_bad=False, flight=False, t_aca=20,
     :param pix_zero_loc: row/col coords are integral at 'edge' or 'center'
     :rtype: (yang, zang) each vector of the same length as row/col
     """
-    row = np.array(row)
-    col = np.array(col)
+    row = np.asarray(row)
+    col = np.asarray(col)
 
     if pix_zero_loc == 'center':
         # Transform row/col values from 'center' convention to 'edge'
@@ -180,8 +178,8 @@ def yagzag_to_pixels(yang, zang, allow_bad=False, pix_zero_loc='edge'):
     :param pix_zero_loc: row/col coords are integral at 'edge' or 'center'
     :rtype: (row, col) each vector of the same length as row/col
     """
-    yang = np.array(yang)
-    zang = np.array(zang)
+    yang = np.asarray(yang)
+    zang = np.asarray(zang)
     row, col = _poly_convert(yang, zang, ACA2PIX_coeff)
 
     # Row/col are in edge coordinates at this point, check if they are on
@@ -240,46 +238,57 @@ def _poly_convert(y, z, coeffs, t_aca=None):
     return newy, newz
 
 
-def radec_to_yagzag(ra, dec, q_att):
+def radec_to_eci(ra, dec):
     """
-    Given RA, Dec, and pointing quaternion, determine ACA Y-ang, Z-ang.  The
-    input ``ra`` and ``dec`` values can be 1-d arrays in which case the output
-    ``yag`` and ``zag`` will be corresponding arrays of the same length.
-
-    This is a wrapper around Ska.quatutil.radec2yagzag but uses arcsec instead
-    of deg for yag, zag.
+    Convert from RA,Dec to ECI.  The input ``ra`` and ``dec`` values can be 1-d
+    arrays of length N in which case the output ``ECI`` will be an array with
+    shape (N, 3). The N dimension can actually be any multidimensional shape.
 
     :param ra: Right Ascension (degrees)
     :param dec: Declination (degrees)
-    :param q_att: ACA pointing quaternion
-
-    :returns:  yag, zag (arcsec)
+    :returns: numpy array ECI (3-vector or N x 3 array)
     """
-    from Ska.quatutil import radec2yagzag
-    yag, zag = radec2yagzag(ra, dec, q_att)
-    yag *= 3600
-    zag *= 3600
+    r = np.deg2rad(ra)
+    d = np.deg2rad(dec)
+    out = np.broadcast(r, d)
 
-    return yag, zag
+    out = np.empty(out.shape + (3,), dtype=np.float64)
+    out[..., 0] = np.cos(r) * np.cos(d)
+    out[..., 1] = np.sin(r) * np.cos(d)
+    out[..., 2] = np.sin(d)
+    return out
 
 
-def yagzag_to_radec(yag, zag, q_att):
+def eci_to_radec(eci):
     """
-    Given ACA Y-ang, Z-ang and pointing quaternion determine RA, Dec. The
-    input ``yag`` and ``zag`` values can be 1-d arrays in which case the output
-    ``ra`` and ``dec`` will be corresponding arrays of the same length.
-
-    This is a wrapper around Ska.quatutil.yagzag2radec but uses arcsec instead
-    of deg for yag, zag.
-
-    :param yag: ACA Y angle (arcsec)
-    :param zag: ACA Z angle (arcsec)
-    :param q_att: ACA pointing quaternion
+    Convert from ECI vector(s) to RA, Dec.  The input ``eci`` value
+    can be an array of 3-vectors having shape (N, 3) in which case
+    the output RA, Dec will be arrays of shape N. The N dimension can
+    actually be any multidimensional shape.
 
-    :returns: ra, dec (arcsec)
+    :param eci: ECI as 3-vector or (N, 3) array
+    :rtype: scalar or array ra, dec (degrees)
     """
-    from Ska.quatutil import yagzag2radec
-    ra, dec = yagzag2radec(yag / 3600, zag / 3600, q_att)
+    eci = np.asarray(eci)
+    if eci.shape[-1] != 3:
+        raise ValueError('final dimension of `eci` must be 3')
+
+    ra = np.degrees(np.arctan2(eci[..., 1], eci[..., 0]))
+    dec = np.degrees(np.arctan2(eci[..., 2], np.sqrt(eci[..., 1]**2 + eci[..., 0]**2)))
+    # Enforce strictly 0 <= RA < 360. Note weird corner case that one can get
+    # ra being negative and very small, e.g. -7.7e-18, which then has 360 added
+    # and turns into exactly 360.0 because of floating point precision.
+    bad = ra < 0
+    if eci.ndim > 1:
+        ra[bad] += 360
+    elif bad:
+        ra += 360
+
+    bad = ra >= 360
+    if eci.ndim > 1:
+        ra[bad] -= 360
+    elif bad:
+        ra -= 360
 
     return ra, dec
 
@@ -392,3 +401,65 @@ def calc_targ_from_aca(aca, y_off, z_off, si_align=None):
     q_targ = q_aca * q_si_align * q_off
 
     return q_targ
+
+
+def radec_to_yagzag(ra, dec, q_att):
+    """
+    Given RA, Dec, and pointing quaternion, determine ACA Y-ang, Z-ang.  The
+    input ``ra`` and ``dec`` values can be 1-d arrays in which case the output
+    ``yag`` and ``zag`` will be corresponding arrays of the same length.
+
+    :param ra: Right Ascension (degrees)
+    :param dec: Declination (degrees)
+    :param q_att: ACA pointing quaternion (Quat or Quat-compatible input)
+
+    :returns:  yag, zag (arcsec)
+    """
+    if not isinstance(q_att, Quat):
+        q_att = Quat(q_att)
+    eci = radec_to_eci(ra, dec)  # N x 3
+    qt = q_att.transform.swapaxes(-2, -1)  # Transpose, allowing for leading dimensions
+    d_aca = np.einsum('...jk,...k->...j', qt, eci)
+    yag = np.rad2deg(np.arctan2(d_aca[..., 1], d_aca[..., 0]))
+    zag = np.rad2deg(np.arctan2(d_aca[..., 2], d_aca[..., 0]))
+    return yag * 3600, zag * 3600
+
+
+def yagzag_to_radec(yag, zag, q_att):
+    """
+    Given ACA Y-ang, Z-ang and pointing quaternion determine RA, Dec. The
+    input ``yag`` and ``zag`` values can be 1-d arrays in which case the output
+    ``ra`` and ``dec`` will be corresponding arrays of the same length.
+
+    :param yag: ACA Y angle (arcsec)
+    :param zag: ACA Z angle (arcsec)
+    :param q_att: ACA pointing quaternion (Quat or Quat-compatible input)
+
+    :returns: ra, dec (degrees)
+    """
+    if not isinstance(q_att, Quat):
+        q_att = Quat(q_att)
+    yag = np.asarray(yag)
+    zag = np.asarray(zag)
+    out = np.broadcast(yag, zag)  # Object with the right broadcasted shape
+    d_aca = np.empty(out.shape + (3,), dtype=np.float64)
+    d_aca[..., 0] = np.ones(out.shape)
+    d_aca[..., 1] = np.tan(np.deg2rad(yag / 3600))
+    d_aca[..., 2] = np.tan(np.deg2rad(zag / 3600))
+    d_aca = normalize_vector(d_aca)
+    eci = np.einsum('...jk,...k->...j', q_att.transform, d_aca)
+    return eci_to_radec(eci)
+
+
+def normalize_vector(vec, ord=None):
+    """Normalize ``vec`` over the last dimension.
+
+    For an L x M x N input array, this normalizes over the N dimension
+    using ``np.linalg.norm``.
+
+    :param vec: input vector or array of vectors
+    :param ord: ord parameter for np.linalg.norm (default=None => 2-norm)
+    :returns: normed array of vectors
+    """
+    norms = np.linalg.norm(vec, axis=-1, keepdims=True, ord=ord)
+    return vec / norms