Merge pull request #138 from simonsobs/koopman/hwp-module

Add HWP module

docs/api.rst

@@ -27,6 +27,14 @@ sorunlib.commands
     :undoc-members:
     :show-inheritance:
 
+sorunlib.hwp
+------------
+
+.. automodule:: sorunlib.hwp
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
 sorunlib.seq
 ------------
 

src/sorunlib/__init__.py

@@ -1,4 +1,4 @@
-from . import acu, seq, smurf, wiregrid
+from . import acu, hwp, seq, smurf, wiregrid
 
 from .commands import wait_until
 from .util import create_clients
@@ -18,7 +18,13 @@ def initialize(test_mode=False):
     CLIENTS = create_clients(test_mode=test_mode)
 
 
-__all__ = ["acu", "seq", "smurf", "wiregrid", "wait_until", "initialize"]
+__all__ = ["acu",
+           "hwp",
+           "seq",
+           "smurf",
+           "wiregrid",
+           "wait_until",
+           "initialize"]
 
 from . import _version
 __version__ = _version.get_versions()['version']

src/sorunlib/hwp.py

@@ -0,0 +1,44 @@
+import sorunlib as run
+from sorunlib._internal import check_response
+
+
+# Public API
+def set_freq(freq):
+    """Set the rotational frequency of the HWP.
+
+    Args:
+        freq (float): Target frequency to rotate the HWP in Hz. This is a
+            *signed float*, the meaning of which depends on the OCS site
+            configuration. For details see the `documentation for the HWP
+            Supervisor Agent <docs_>`_.
+
+    .. _docs: https://socs.readthedocs.io/en/main/agents/hwp_supervisor_agent.html
+
+    """
+    hwp = run.CLIENTS['hwp']
+    resp = hwp.pid_to_freq(target_freq=freq)
+    check_response(hwp, resp)
+
+
+def stop(active=True, brake_voltage=None):
+    """Stop the HWP.
+
+    Args:
+        active (bool, optional): If True, actively try to stop the HWP by
+            applying the brake. If False, simply turn off the PMX power and let
+            it spin down on its own. Defaults to True.
+        brake_voltage (float, optional): Voltage used when actively stopping
+            the HWP. Only considered when active is True.
+
+    """
+    hwp = run.CLIENTS['hwp']
+
+    if active:
+        if brake_voltage is None:
+            resp = hwp.brake()
+        else:
+            resp = hwp.brake(brake_voltage=brake_voltage)
+        check_response(hwp, resp)
+    else:
+        resp = hwp.pmx_off()
+        check_response(hwp, resp)

src/sorunlib/util.py

@@ -171,6 +171,7 @@ def create_clients(config=None, sorunlib_config=None, test_mode=False):
         in the format::
 
             clients = {'acu': acu_client,
+                       'hwp': hwp_supervisor_client,
                        'smurf': [smurf_client1, smurf_client2, smurf_client3],
                        'wiregrid': {'actuator': actuator_client,
                                     'encoder': encoder_client,
@@ -186,11 +187,15 @@ def create_clients(config=None, sorunlib_config=None, test_mode=False):
         smurf_agent_class = 'PysmurfController'
 
     acu_id = _find_active_instances('ACUAgent')
+    hwp_id = _find_active_instances('HWPSupervisor')
     smurf_ids = _find_active_instances(smurf_agent_class)
 
     if acu_id:
         acu_client = _try_client(acu_id)
         clients['acu'] = acu_client
+    if hwp_id:
+        hwp_client = _try_client(hwp_id)
+        clients['hwp'] = hwp_client
 
     # Always create smurf client list, even if empty
     smurf_clients = [_try_client(x) for x in smurf_ids]

tests/test_hwp.py

@@ -0,0 +1,31 @@
+import os
+os.environ["OCS_CONFIG_DIR"] = "./test_util/"
+
+import pytest
+
+from sorunlib import hwp
+
+from util import create_patch_clients
+
+
+patch_clients_satp = create_patch_clients('satp')
+
+
+@pytest.mark.parametrize("active", [True, False])
+def test_stop(patch_clients_satp, active):
+    hwp.stop(active=active)
+    if active:
+        hwp.run.CLIENTS['hwp'].brake.assert_called_with()
+    else:
+        hwp.run.CLIENTS['hwp'].pmx_off.assert_called()
+
+
+def test_stop_brake_voltage(patch_clients_satp):
+    VOLTAGE = 5.0
+    hwp.stop(active=True, brake_voltage=VOLTAGE)
+    hwp.run.CLIENTS['hwp'].brake.assert_called_with(brake_voltage=VOLTAGE)
+
+
+def test_set_freq(patch_clients_satp):
+    hwp.set_freq(freq=2.0)
+    hwp.run.CLIENTS['hwp'].pid_to_freq.assert_called_with(target_freq=2.0)

tests/test_util.py

@@ -73,6 +73,22 @@ def mock_client(*args, **kwargs):
                            'delay_task': 1},
                        'agent_class': 'FakeDataAgent',
                        'agent_address': 'observatory.fake-data-1'},
+                   'observatory.hwp-supervisor': {
+                       'expired': False,
+                       'time_expired': None,
+                       'last_updated': 1710168924.336035,
+                       'op_codes': {
+                           'monitor': 3,
+                           'spin_control': 3,
+                           'pid_to_freq': 1,
+                           'set_const_voltage': 1,
+                           'brake': 1,
+                           'pmx_off': 1,
+                           'abort_action': 7
+                       },
+                       'agent_class': 'HWPSupervisor',
+                       'agent_address': 'observatory.hwp-supervisor'
+                   },
                    'observatory.acu-sat1': {
                        'expired': False,
                        'time_expired': None,

tests/util.py

@@ -77,6 +77,7 @@ def mocked_clients(**kwargs):
     smurfs = [_mock_smurf_client(id_) for id_ in smurf_ids]
 
     clients = {'acu': _mock_acu_client(platform_type),
+               'hwp': MagicMock(),
                'smurf': smurfs,
                'wiregrid': {'actuator': MagicMock(),
                             'encoder': MagicMock(),