Add computed MSID to return Quat for quaternion MSID sets

Ska/engarchive/derived/comps.py

@@ -324,6 +324,75 @@ def get_stats_attrs(self, tstart, tstop, msid, match_args, interval):
 ############################################################################
 
 
+class Comp_Quat(ComputedMsid):
+    """Computed MSID for returning the quaternion telemetry as a Quat object
+
+    This defines the following computed MSIDs based on the corresponding
+    TDB MSIDs:
+
+    - quat_aoattqt = AOATTQT[1-4]
+    - quat_aoatupq = AOATUPQ[1-3]
+    - quat_aocmdqt = AOCMDQT[1-3]
+    - quat_aotarqt = AOTARQT[1-3]
+
+    Example::
+
+      >>> from cheta import fetch
+      >>> qatt = fetch.Msid('quat_aoattqt', '2022:001:00:00:00', '2022:001:00:00:04')
+      >>> qatt.vals
+      Quat(array([[-0.07434856, -0.55918674, -0.80432653,  0.18665828],
+                  [-0.07434854, -0.55918679, -0.8043265 ,  0.18665825],
+                  [-0.07434849, -0.55918674, -0.80432653,  0.18665829],
+                  [-0.07434849, -0.55918667, -0.80432653,  0.18665852]]))
+      >>> qatt.vals.equatorial
+      array([[193.28905806,  19.16894296,  67.36207683],
+             [193.28905485,  19.1689407 ,  67.36208471],
+             [193.28906329,  19.16893787,  67.36207699],
+             [193.28908839,  19.16895134,  67.36206404]])
+    """
+    msid_match = r'quat_(aoattqt|aoatupq|aocmdqt|aotarqt)'
+
+    def get_msid_attrs(self, tstart, tstop, msid, msid_args):
+        from Quaternion import Quat, normalize
+
+        if 'maude' in self.fetch_sys.data_source.sources():
+            raise ValueError(
+                f'{msid} is not available from MAUDE due to issues aligning telemetry')
+
+        msid_root = msid_args[0]
+        n_comp = 4 if msid_root == 'aoattqt' else 3
+        msids = [f'{msid_root}{ii}' for ii in range(1, n_comp + 1)]
+        dat = self.fetch_sys.MSIDset(msids, tstart, tstop)
+
+        # Interpolate to a common time base, leaving in flagged bad data and
+        # marking data bad if any of the set at each time are bad. See:
+        # https://sot.github.io/eng_archive/fetch_tutorial.html#filtering-and-bad-values
+        dat.interpolate(
+            times=dat[msids[0]].times,
+            filter_bad=False,
+            bad_union=True)
+
+        q1 = dat[msids[0]].vals.astype(np.float64)
+        q2 = dat[msids[1]].vals.astype(np.float64)
+        q3 = dat[msids[2]].vals.astype(np.float64)
+        if n_comp == 4:
+            q4 = dat[msids[3]].vals.astype(np.float64)
+        else:
+            q4 = np.sqrt((1.0 - q1**2 - q2**2 - q3**2).clip(0.0))
+
+        q = np.array([q1, q2, q3, q4]).transpose()
+        quat = Quat(q=normalize(q))
+        bads = np.zeros_like(q1, dtype=bool)
+        for msid in msids:
+            bads |= dat[msid].bads
+
+        out = {'vals': quat,
+               'bads': bads,
+               'times': dat.times,
+               'unit': None}
+        return out
+
+
 class Comp_MUPS_Valve_Temp_Clean(ComputedMsid):
     """Computed MSID for cleaned MUPS valve temps PM2THV1T, PM1THV2T
 

Ska/engarchive/tests/test_comps.py

@@ -3,6 +3,7 @@
 """Test that computed MSIDs work as expected."""
 
 import numpy as np
+from Quaternion import Quat
 import pytest
 
 from .. import fetch_eng, fetch_sci, fetch as fetch_cxc
@@ -202,3 +203,27 @@ def test_stats(stat):
             assert np.allclose(val, valc)
         else:
             assert np.all(val == valc)
+
+
+@pytest.mark.parametrize('msid', ['aoattqt', 'aoatupq', 'aocmdqt', 'aotarqt'])
+def test_quat_comp(msid):
+    start, stop = '2022:002', '2022:003'
+    datq = fetch_eng.MSID(f'quat_{msid}', start, stop)
+    dats = fetch_eng.MSIDset([f'{msid}*'], start, stop)
+    n_comp = len(dats)
+    for ii in range(n_comp):
+        vq = datq.vals.q[:, ii]
+        vn = dats[f'{msid}{ii + 1}'].vals
+        # Code handles q1**2 + q2**2 + q3**2 > 1 by clipping the norm and renormalizing
+        # the quaternion. This generates a significant difference for the cmd quat.
+        ok = np.isclose(vq, vn, rtol=0, atol=(1e-4 if msid == 'aocmdqt' else 1e-8))
+        assert np.all(ok)
+    assert isinstance(datq.vals, Quat)
+
+
+@pytest.mark.parametrize('msid', ['aoattqt', 'aoatupq', 'aocmdqt', 'aotarqt'])
+def test_quat_comp_exception(msid):
+    start, stop = '2022:002', '2022:003'
+    with pytest.raises(ValueError, match=f'quat_{msid} is not available from MAUDE'):
+        with fetch_eng.data_source('maude'):
+            fetch_eng.MSID(f'quat_{msid}', start, stop)

pytest.ini

@@ -4,3 +4,5 @@ filterwarnings =
     ignore:numpy.ufunc size changed:RuntimeWarning
     ignore:the imp module is deprecated in favour of importlib:DeprecationWarning
     ignore:parse functions are required to provide a named argument:PendingDeprecationWarning
+    ignore:'soft_unicode' has been renamed to 'soft_str':DeprecationWarning
+    ignore:`np.object` is a deprecated alias for the builtin `object`:DeprecationWarning