Add stateUtils

python/lsst/ts/aos/analysis/__init__.py

@@ -2,3 +2,4 @@
 from .blockUtils import *
 from .efdUtils import *
 from .ofcUtils import *
+from .stateUtils import *

python/lsst/ts/aos/analysis/efdUtils.py

@@ -19,6 +19,7 @@
 ]
 
 
+
 def get_start_end_time_seq_num(day, seq_num):
     butler = Butler(
         "LSSTComCam", instrument="LSSTComCam", collections="LSSTComCam/raw/all"
@@ -209,9 +210,7 @@ async def get_last_m1m3_aos_forces_event(end):
 
 
 def plot_m2_forces(m2_forces, ax, fig):
-    m2_location_path = Path(
-        f'{os.environ["HOME"]}/notebooks/lsst-ts/ts_config_mttcs/MTM2/v2/harrisLUT/cell_geom.yaml'
-    )
+    m2_location_path = Path(f'{os.environ["TS_CONFIG_MTTCS_DIR"]}/MTM2/v2/harrisLUT/cell_geom.yaml')
     with open(m2_location_path, "r") as yaml_file:
         m2_info = yaml.safe_load(yaml_file)
     xy_actuators = np.array(m2_info["locAct_axial"])

python/lsst/ts/aos/analysis/stateUtils.py

@@ -0,0 +1,347 @@
+import numpy as np
+from astropy.table import QTable
+from astropy.time import Time
+import astropy.units as u
+from functools import lru_cache
+from lsst.summit.utils.efdUtils import getEfdData, getMostRecentRowWithDataBefore
+
+__all__ = ['StateFetcher']
+
+
+class StateFetcher:
+    def __init__(self, butler, efdClient):
+        self.butler = butler
+        self.client = efdClient
+
+    @lru_cache
+    def _get_exp_record(self, *, day_obs=None, seq_num=None, exp_id=None, record=None):
+        if record is not None:
+            if any(k is not None for k in [day_obs, seq_num, exp_id]):
+                raise ValueError("Don't specify record if using day_obs/seq_num/exp_id")
+            return record
+        if day_obs is None and seq_num is None:
+            if exp_id is None:
+                raise ValueError("Require (day_obs, seq_num) or exp_id")
+            where = f"exposure.id = {exp_id}"
+        else:
+            if day_obs is None:
+                raise ValueError("Require day_obs when seq_num is given")
+            elif seq_num is None:
+                raise ValueError("Require seq_num when day_obs is given")
+            where = f"exposure.day_obs={day_obs} AND exposure.seq_num={seq_num}"
+        (record,) = self.butler.registry.queryDimensionRecords("exposure", where=where)
+        return record
+
+    @lru_cache
+    def _get_compensated_hexapod_telemetry(self, record):
+        return getEfdData(
+            self.client,
+            "lsst.sal.MTHexapod.application",
+            expRecord=record
+        )
+
+    @lru_cache
+    def _get_uncompensated_hexapod_event(self, record, salIndex):
+        return getMostRecentRowWithDataBefore(
+            self.client,
+            "lsst.sal.MTHexapod.logevent_uncompensatedPosition",
+            timeToLookBefore=record.timespan.end,
+            where=lambda df: df['salIndex'] == salIndex
+        )
+
+    def get_hexapod(
+            self, *,
+            day_obs=None, seq_num=None, exp_id=None, record=None,
+            component=None, compensated=True,
+            doMean=True,
+            out_type='position'
+    ):
+        """Get hexapod state.
+
+        Parameters
+        ----------
+        day_obs : `int`, optional
+            Observation day to query.
+        seq_num : `int`, optional
+            Sequence number to query.
+        exp_id : `int`, optional
+            Exposure id to query.
+        record : `lsst.daf.butler.dimensions.DimensionRecord`, optional
+            The exposure record containing the timespan to query.
+        component : {'M2' | 'Camera'}
+            Which hexapod to query.
+        compensated : `bool`, optional
+            Whether or not to return compensated values
+        doMean : `bool`, optional
+            Return mean value over exposure timespan?  If true, then output is a
+            dictionary with keys indicating hexapod degree-of-freedom and values
+            that are astropy.units.Quantity
+        out_type : {'position', 'error', 'demand', 'raw', 'raw_table'}
+            Adjusts the output.  Allowed values are:
+                position : Returns delivered hexapod positions
+                error : Returns error in hexapod positions
+                demand : Returns demanded hexapod positions
+                raw : Returns raw pandas dataframe with all columns
+                raw_table : Returns astropy.table.Table with all columns
+
+        Raises
+        ------
+        ValueError:
+            If unknown component
+        ValueError:
+            If unknown out_type
+        """
+        if out_type not in ['raw', 'raw_table', 'position', 'error', 'demand']:
+            raise ValueError(f"Unknown out_type {out_type}")
+
+        record = self._get_exp_record(
+            day_obs=day_obs, seq_num=seq_num, exp_id=exp_id, record=record
+        )
+
+        if component == 'Camera':
+            salIndex = 1
+        elif component == 'M2':
+            salIndex = 2
+        else:
+            raise ValueError(f"Unknown component {component}")
+
+        if compensated:
+            df = self._get_compensated_hexapod_telemetry(record)
+            df = df[df['salIndex'] == salIndex]
+        else:
+            df = self._get_uncompensated_hexapod_event(record, salIndex)
+            df = df.to_frame().T
+
+        if out_type == 'raw':
+            return df
+
+        table = QTable.from_pandas(df)
+        if out_type == 'raw_table':
+            return table
+
+        new_names = [component+'_'+name for name in ['x', 'y', 'z', 'rx', 'ry']]
+
+        if compensated:
+            original_names = [f'{out_type}{i}' for i in range(5)]
+        else:
+            original_names = ['x', 'y', 'z', 'u', 'v']
+        table.rename_columns(original_names, new_names)
+
+        time = Time(
+            table['private_efdStamp'].value.astype(np.float64), format='unix'
+        ).tai
+        time.format = 'iso'
+        table['time'] = time
+
+        table = table[['time']+new_names]
+
+        for i, name in enumerate(new_names):
+            table[name] = table[name].astype(np.float64)
+            if i < 3:
+                table[name].unit = u.micron
+            else:
+                table[name].unit = u.degree
+
+        if not doMean:
+            return table
+
+        out = {}
+        for name in ['x', 'y', 'z', 'rx', 'ry']:
+            k = f"{component}_{name}"
+            out[k] = np.mean(table[k])
+
+        return out
+
+    @lru_cache
+    def _get_M2_telemetry(self, record):
+        return getEfdData(
+            self.client,
+            "lsst.sal.MTM2.axialForce",
+            expRecord=record
+        )
+
+    def get_M2_forces(
+        self, *,
+        day_obs=None, seq_num=None, exp_id=None, record=None,
+        doMean=True,
+        topic='appliedAOS'
+    ):
+        """Get M2 actuator forces
+
+        Parameters
+        ----------
+        day_obs : `int`, optional
+            Observation day to query.
+        seq_num : `int`, optional
+            Sequence number to query.
+        exp_id : `int`, optional
+            Exposure id to query.
+        record : `lsst.daf.butler.dimensions.DimensionRecord`, optional
+            The exposure record containing the timespan to query.
+        doMean : `bool`, optional
+            Return mean value over exposure timespan?  If true, then output is a
+            astropy.units.Quantity of size 72 for the number of actuators
+        topic : `str`, optional
+            The kind of force to return.  Must be one of:
+            appliedAOS, hardpointCorrection, applied, lutGravity, lutTemperature
+
+        Returns
+        -------
+        val : `astropy.table.QTable` or `ndarray`
+            Either a table with columns for times and forces, or a numpy array with
+            forces.
+        """
+        record = self._get_exp_record(
+            day_obs=day_obs, seq_num=seq_num, exp_id=exp_id, record=record
+        )
+
+        telemetry = QTable.from_pandas(self._get_M2_telemetry(record))
+        nrow = len(telemetry)
+        if topic == 'appliedAOS':
+            applied = np.empty((nrow, 72), dtype=np.float64)
+            hardpointCorrection = np.empty((nrow, 72), dtype=np.float64)
+            lutGravity = np.empty((nrow, 72), dtype=np.float64)
+            lutTemperature = np.empty((nrow, 72), dtype=np.float64)
+
+            for i in range(72):
+                hardpointCorrection[:, i] = telemetry[f'hardpointCorrection{i}']
+                applied[:, i] = telemetry[f'applied{i}']
+                lutGravity[:, i] = telemetry[f'lutGravity{i}']
+                lutTemperature[:, i] = telemetry[f'lutTemperature{i}']
+            out = applied - hardpointCorrection - lutGravity - lutTemperature
+        else:
+            out = np.empty((nrow, 72), dtype=np.float64)
+            for i in range(72):
+                out[:, i] = telemetry[f'{topic}{i}']
+
+        table = QTable()
+        time = Time(
+            telemetry['private_efdStamp'].value.astype(np.float64), format='unix'
+        ).tai
+        time.format='iso'
+        table['time'] = time
+        table[topic] = out * u.Newton
+
+        if not doMean:
+            return table
+
+        return np.mean(table[topic], axis=0)
+
+    @lru_cache
+    def _get_M1M3_telementry(self, record, topic):
+        # for topics
+        #   '', 'Acceleration', 'Azimuth', 'Balance', 'Elevation', 'Thermal', 'Velocity'
+        return getEfdData(
+            self.client,
+            f"lsst.sal.MTM1M3.applied{topic}Forces",
+            expRecord=record
+        )
+
+    @lru_cache
+    def _get_M1M3_event(self, record, topic):
+        # for topics: 'ActiveOptic', 'Offset', 'Static'
+        return getMostRecentRowWithDataBefore(
+            self.client,
+            f"lsst.sal.MTM1M3.logevent_applied{topic}Forces",
+            timeToLookBefore=record.timespan.end
+        )
+
+    def get_M1M3_forces(
+        self, *,
+        day_obs=None, seq_num=None, exp_id=None, record=None,
+        doMean=True,
+        topic='ActiveOptic'
+    ):
+        """Get M1M3 forces.
+
+        Parameters
+        ----------
+        day_obs : `int`, optional
+            Observation day to query.
+        seq_num : `int`, optional
+            Sequence number to query.
+        exp_id : `int`, optional
+            Exposure id to query.
+        record : `lsst.daf.butler.dimensions.DimensionRecord`, optional
+            The exposure record containing the timespan to query.
+        doMean : `bool`, optional
+            Return mean value over exposure timespan?  If true, then output is a
+            astropy.units.Quantity of size 156 for the number of actuators
+        topic : `str`
+            The kind of force to return.  Must be one of:
+            '', Acceleration, Azimuth, Balance, Elevation, Thermal, Velocity,
+            ActiveOptic, Offset, Static.
+
+        Returns
+        -------
+        val : `astropy.table.QTable` or `ndarray`
+            Either a table with columns for times and forces, or a numpy array with forces.
+
+        """
+        record = self._get_exp_record(
+            day_obs=day_obs, seq_num=seq_num, exp_id=exp_id, record=record
+        )
+
+        if topic in ['ActiveOptic', 'Offset', 'Static']:
+            event = self._get_M1M3_event(record, topic)
+            out = np.empty((156,))
+            for i in range(156):
+                out[i] = event[f'zForces{i}']
+            return out * u.Newton
+
+        telemetry = QTable.from_pandas(self._get_M1M3_telementry(record, topic))
+        nrow = len(telemetry)
+        out = np.empty((nrow, 156), dtype=np.float64)
+        for i in range(156):
+            out[:, i] = telemetry[f'zForces{i}']
+
+        table = QTable()
+        time = Time(telemetry['private_efdStamp'].value.astype(np.float64), format='unix').tai
+        time.format='iso'
+        table['time'] = time
+        table[topic] = out * u.Newton
+
+        if not doMean:
+            return table
+
+        return np.mean(table[topic], axis=0)
+
+    @lru_cache
+    def get_aggregated_state(
+        self, *,
+        day_obs=None, seq_num=None, exp_id=None, record=None,
+    ):
+        """Get the aggregated state of MTAOS
+
+        Parameters
+        ----------
+        day_obs : `int`, optional
+            Observation day to query.
+        seq_num : `int`, optional
+            Sequence number to query.
+        exp_id : `int`, optional
+            Exposure id to query.
+        record : `lsst.daf.butler.dimensions.DimensionRecord`, optional
+            The exposure record containing the timespan to query.
+
+        Returns
+        -------
+        val : `ndarray`
+            The aggregated state inside MTAOS.  The order is:
+                0-4 : M2 hexapod (micron)
+                5-9 : Camera hexapod (degree)
+                10-29 : M1M3 bending modes (micron)
+                30-49 : M2 bending modes (micron)
+        """
+        record = self._get_exp_record(
+            day_obs=day_obs, seq_num=seq_num, exp_id=exp_id, record=record
+        )
+        event = getMostRecentRowWithDataBefore(
+            self.client,
+            f"lsst.sal.MTAOS.logevent_degreeOfFreedom",
+            timeToLookBefore=record.timespan.end
+        )
+        out = np.empty(50,)
+        for i in range(50):
+            out[i] = event[f'aggregatedDoF{i}']
+        return out