diff --git a/src/sorunlib/hwp.py b/src/sorunlib/hwp.py index 2fda9cec..8b696615 100644 --- a/src/sorunlib/hwp.py +++ b/src/sorunlib/hwp.py @@ -2,6 +2,33 @@ from sorunlib._internal import check_response +def _get_direction(): + """Get the rotational direction ('cw' or 'ccw') of the HWP. The direction + is determined in part by the configuration of the HWP supervisor agent. For + details see the `HWP Supervisor agent docs `_. + + * 'cw' is clockwise as seen from the sky to window. + * 'ccw' is counter-clockwise as seen from the sky to window. + + Returns: + str: The direction of the HWP, either 'cw' or 'ccw'. + + Raises: + RuntimeError: If the direction is not either 'cw' or 'ccw'. + + .. _docs: https://socs.readthedocs.io/en/main/agents/hwp_supervisor_agent.html + + """ + hwp = run.CLIENTS['hwp'] + resp = hwp.monitor.status() + direction = resp.session['data']['hwp_state']['direction'] + + if direction not in ['cw', 'ccw']: + raise RuntimeError("The HWP direction is unknown. Aborting...") + + return direction + + # Public API def set_freq(freq): """Set the rotational frequency of the HWP. diff --git a/src/sorunlib/wiregrid.py b/src/sorunlib/wiregrid.py index b1a25555..f8b993c6 100644 --- a/src/sorunlib/wiregrid.py +++ b/src/sorunlib/wiregrid.py @@ -198,6 +198,24 @@ def _check_temperature_sensors(): _verify_temp_response(resp, 'AIN2C', 0) +def _check_wiregrid_position(): + """Check the wiregrid position. + + Returns: + str: The wiregrid position, either 'inside' or 'outside'. + + Raises: + RuntimeError: When the wiregrid position is unknown. + + """ + actuator = run.CLIENTS['wiregrid']['actuator'] + resp = actuator.acq.status() + position = resp.session['data']['fields']['position'] + if position not in ['inside', 'outside']: + raise RuntimeError("The wiregrid position is unknown. Aborting...") + return position + + # Public API def insert(): """Insert the wiregrid.""" @@ -305,3 +323,144 @@ def calibrate(continuous=False, elevation_check=True, boresight_check=True, finally: # Stop SMuRF streams run.smurf.stream('off') + + +def time_constant(num_repeats=1): + """ + Run a wiregrid time constant measurement. + + Args: + num_repeats (int): Number of repeats. Default is 1. + If this is odd, the HWP direction will be changed to the opposite + of the initial direction. If this is even, the HWP direction will be + the same as the initial direction. + + """ + # Check the number of repeats + if num_repeats < 1 or not isinstance(num_repeats, int): + error = "The ``num_repeats`` should be int and larger than 0." + raise RuntimeError(error) + + _check_agents_online() + _check_motor_on() + _check_telescope_position(elevation_check=True, boresight_check=False) + _check_wiregrid_position() + if _check_wiregrid_position() == 'inside': + error = "The wiregrid is already inserted before the wiregrid time " + \ + "constant measurement. Please inspect wiregrid and HWP " + \ + "before continuing observations." + raise RuntimeError(error) + + if _check_zenith(): + el_tag = ', wg_el90' + else: + el_tag = '' + + # Check the current HWP direction + try: + current_hwp_direction = run.hwp._get_direction() # 'cw' or 'ccw' + except RuntimeError as e: + error = "Wiregrid time constant measurment has failed " + \ + "due to a failure in getting the HWP direction.\n" + str(e) + raise RuntimeError(error) + + # Rotate to get encoder reference before insertion + rotate(continuous=True, duration=10) + + # Bias step (the wire grid is off the window) + bs_tag = 'wiregrid, wg_time_constant, wg_ejected, ' + \ + f'hwp_{current_hwp_direction}' + el_tag + run.smurf.bias_step(tag=bs_tag, concurrent=True) + time.sleep(5) + + # Insert the wiregrid while streaming + try: + stream_tag = 'wiregrid, wg_time_constant, wg_inserting, ' + \ + f'hwp_{current_hwp_direction}' + el_tag + run.smurf.stream('on', tag=stream_tag, subtype='cal') + insert() + time.sleep(5) + finally: + run.smurf.stream('off') + + for i in range(num_repeats): + if current_hwp_direction == 'ccw': + target_hwp_direction = 'cw' + elif current_hwp_direction == 'cw': + target_hwp_direction = 'ccw' + + # Bias step (the wire grid is on the window) + # Before stopping the HWP + bs_tag = 'wiregrid, wg_time_constant, wg_inserted, ' + \ + f'hwp_{current_hwp_direction}' + el_tag + run.smurf.bias_step(tag=bs_tag, concurrent=True) + + # Run stepwise rotation before stopping the HWP + try: + stream_tag = 'wiregrid, wg_time_constant, ' + \ + f'wg_stepwise, hwp_{current_hwp_direction}' + \ + el_tag + run.smurf.stream('on', tag=stream_tag, subtype='cal') + # Run stepwise rotation + rotate(continuous=False) + finally: + run.smurf.stream('off') + + # Stop the HWP while streaming + try: + stream_tag = 'wiregrid, wg_time_constant, ' + \ + f'hwp_change_{current_hwp_direction}_to_stop' + el_tag + run.smurf.stream('on', tag=stream_tag, subtype='cal') + run.hwp.stop(active=True) + finally: + run.smurf.stream('off') + + # Reverse the HWP while streaming + try: + stream_tag = 'wiregrid, wg_time_constant, ' + \ + f'hwp_change_stop_to_{target_hwp_direction}' + el_tag + run.smurf.stream('on', tag=stream_tag, subtype='cal') + # Note: This is hardcoding the correspondance between direction and + # the sign of the frequency, which is subject to change depending + # on the hardware/agent configuration. This should be removed in + # the future, if possible. + if target_hwp_direction == 'ccw': + run.hwp.set_freq(freq=2.0) + elif target_hwp_direction == 'cw': + run.hwp.set_freq(freq=-2.0) + current_hwp_direction = target_hwp_direction + finally: + run.smurf.stream('off') + + # Run stepwise rotation after changing the HWP rotation + try: + stream_tag = 'wiregrid, wg_time_constant, ' + \ + f'wg_stepwise, hwp_{current_hwp_direction}' + \ + el_tag + run.smurf.stream('on', tag=stream_tag, subtype='cal') + # Run stepwise rotation + rotate(continuous=False) + finally: + run.smurf.stream('off') + + # Bias step (the wire grid is on the window) + # After changing the HWP rotation + bs_tag = 'wiregrid, wg_time_constant, wg_inserted, ' + \ + f'hwp_{current_hwp_direction}' + el_tag + run.smurf.bias_step(tag=bs_tag, concurrent=True) + time.sleep(5) + + # Eject the wiregrid while streaming + try: + stream_tag = 'wiregrid, wg_time_constant, wg_ejecting, ' + \ + f'hwp_{current_hwp_direction}' + el_tag + run.smurf.stream('on', tag=stream_tag, subtype='cal') + eject() + time.sleep(5) + finally: + run.smurf.stream('off') + + # Bias step (the wire grid is off the window) + bs_tag = 'wiregrid, wg_time_constant, wg_ejected, ' + \ + f'hwp_{current_hwp_direction}' + el_tag + run.smurf.bias_step(tag=bs_tag, concurrent=True) diff --git a/tests/test_hwp.py b/tests/test_hwp.py index 4e5ccd24..c9e01231 100644 --- a/tests/test_hwp.py +++ b/tests/test_hwp.py @@ -2,15 +2,59 @@ os.environ["OCS_CONFIG_DIR"] = "./test_util/" import pytest +from unittest.mock import MagicMock +import time +import ocs +from ocs.ocs_client import OCSReply from sorunlib import hwp -from util import create_patch_clients - +from util import create_patch_clients, create_session patch_clients_satp = create_patch_clients('satp') +def create_hwp_client(direction): + """Create a HWP client with mock acq Process session.data. + + Args: + direction (str): direction of the HWP. 'ccw' (counter-clockwise) or 'cw' (clockwise). + + """ + client = MagicMock() + session = create_session('acq') + session.data = { + 'hwp_state': { + 'direction': direction, + }, + 'timestamp': time.time(), + } + reply = OCSReply(ocs.OK, 'msg', session.encoded()) + client.monitor.status = MagicMock(return_value=reply) + + return client + + +@pytest.mark.parametrize('direction', ['ccw', 'cw']) +def test__get_direction(patch_clients_satp, direction): + hwp.run.CLIENTS['hwp'] = create_hwp_client(direction) + ret = hwp._get_direction() + if direction == 'ccw': + assert ret == 'ccw' + elif direction == 'cw': + assert ret == 'cw' + hwp.run.CLIENTS['hwp'].monitor.status.assert_called_once() + + +@pytest.mark.parametrize('direction', [None, '']) +def test__get_direction_invalid(patch_clients_satp, direction): + hwp.run.CLIENTS['hwp'] = create_hwp_client(direction) + with pytest.raises(RuntimeError) as e: + hwp._get_direction() + assert str(e.value) == "The HWP direction is unknown. Aborting..." + hwp.run.CLIENTS['hwp'].monitor.status.assert_called_once() + + @pytest.mark.parametrize("active", [True, False]) def test_stop(patch_clients_satp, active): hwp.stop(active=active) diff --git a/tests/test_wiregrid.py b/tests/test_wiregrid.py index 3c5500d7..ece8abab 100644 --- a/tests/test_wiregrid.py +++ b/tests/test_wiregrid.py @@ -1,10 +1,11 @@ import os import time os.environ["OCS_CONFIG_DIR"] = "./test_util/" +os.environ["SORUNLIB_CONFIG"] = "./data/example_config.yaml" import pytest -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, patch, call import ocs from ocs.ocs_client import OCSReply @@ -55,17 +56,23 @@ def create_labjack_client(): return client -def create_actuator_client(motor): +def create_actuator_client(motor, position): """Create an actuator client with mock acq Process session.data. Args: motor (int): Motor state, 0 is off, 1 is on. + position (str): Position of the wiregrid, either 'inside' or 'outside'. """ client = MagicMock() session = create_session('acq') - session.data = {'fields': {'motor': motor}, - 'timestamp': time.time()} + session.data = { + 'fields': { + 'motor': motor, + 'position': position + }, + 'timestamp': time.time(), + } session.set_status('running') reply = OCSReply(ocs.OK, 'msg', session.encoded()) client.acq.status = MagicMock(return_value=reply) @@ -201,7 +208,8 @@ def test__check_telescope_position_invalid(el, boresight): @pytest.mark.parametrize('motor', [(0), (1)]) @patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) def test__check_motor_on(motor): - wiregrid.run.CLIENTS['wiregrid']['actuator'] = create_actuator_client(motor=motor) + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=motor, position='inside') wiregrid._check_motor_on() if motor == 1: wiregrid.run.CLIENTS['wiregrid']['actuator'].acq.status.assert_called_once() @@ -213,7 +221,8 @@ def test__check_motor_on(motor): @patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) def test__check_motor_on_invalid_state(): - wiregrid.run.CLIENTS['wiregrid']['actuator'] = create_actuator_client(motor=3) + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=3, position='inside') with pytest.raises(RuntimeError): wiregrid._check_motor_on() wiregrid.run.CLIENTS['wiregrid']['actuator'].acq.status.assert_called_once() @@ -221,7 +230,8 @@ def test__check_motor_on_invalid_state(): @patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) def test__check_agents_online(): - wiregrid.run.CLIENTS['wiregrid']['actuator'] = create_actuator_client(motor=1) + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=1, position='inside') wiregrid.run.CLIENTS['wiregrid']['kikusui'] = create_kikusui_client() wiregrid.run.CLIENTS['wiregrid']['encoder'] = create_encoder_client() wiregrid.run.CLIENTS['wiregrid']['labjack'] = create_labjack_client() @@ -239,7 +249,8 @@ def test__check_agents_online(): def test_calibrate_stepwise(patch_clients, continuous, el, tag): # Setup all mock clients wiregrid.run.CLIENTS['acu'] = create_acu_client(180, el, 0) - wiregrid.run.CLIENTS['wiregrid']['actuator'] = create_actuator_client(motor=1) + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=1, position='inside') wiregrid.run.CLIENTS['wiregrid']['kikusui'] = create_kikusui_client() wiregrid.run.CLIENTS['wiregrid']['encoder'] = create_encoder_client() wiregrid.run.CLIENTS['wiregrid']['labjack'] = create_labjack_client() @@ -262,3 +273,211 @@ def test__check_process_data_stale_data(): with pytest.raises(RuntimeError): stale_time = time.time() - wiregrid.AGENT_TIMEDIFF_THRESHOLD wiregrid._check_process_data('test process', stale_time) + + +@pytest.mark.parametrize('position', [('inside'), ('outside')]) +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +def test__check_wiregrid_position(position): + wiregrid.run.CLIENTS['wiregrid']['actuator'] = create_actuator_client(1, position) + return_position = wiregrid._check_wiregrid_position() + wiregrid.run.CLIENTS['wiregrid']['actuator'].acq.status.assert_called_once() + assert return_position == position + + +@pytest.mark.parametrize('position', [('unknown'), ('')]) +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +def test__check_wiregrid_position_invalid(position): + wiregrid.run.CLIENTS['wiregrid']['actuator'] = create_actuator_client(1, position) + with pytest.raises(RuntimeError): + wiregrid._check_wiregrid_position() + wiregrid.run.CLIENTS['wiregrid']['actuator'].acq.status.assert_called_once() + + +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +@patch('sorunlib.wiregrid.time.sleep', MagicMock()) +def test_time_constant_cw(): + # Setup all mock clients + wiregrid.run.CLIENTS['acu'] = create_acu_client(180, 50, 0) + wiregrid.run.hwp._get_direction = MagicMock(return_value='cw') # cw + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=1, position='outside') + wiregrid.run.CLIENTS['wiregrid']['kikusui'] = create_kikusui_client() + wiregrid.run.CLIENTS['wiregrid']['encoder'] = create_encoder_client() + wiregrid.run.CLIENTS['wiregrid']['labjack'] = create_labjack_client() + wiregrid.run.wiregrid.rotate = MagicMock() + + wiregrid.time_constant(num_repeats=1) + + # just make sure bias_steps and streams because other functions are already + # tested separately. + expected_calls_of_bias_steps = [ + call(tag='wiregrid, wg_time_constant, wg_ejected, hwp_cw'), + call(tag='wiregrid, wg_time_constant, wg_inserted, hwp_cw'), + call(tag='wiregrid, wg_time_constant, wg_inserted, hwp_ccw'), + call(tag='wiregrid, wg_time_constant, wg_ejected, hwp_ccw') + ] + + common_kwargs_of_streams = { + "downsample_factor": None, + "filter_disable": False + } + expected_tags_of_streams = [ + 'wiregrid, wg_time_constant, wg_inserting, hwp_cw', + 'wiregrid, wg_time_constant, wg_stepwise, hwp_cw', + 'wiregrid, wg_time_constant, hwp_change_cw_to_stop', + 'wiregrid, wg_time_constant, hwp_change_stop_to_ccw', + 'wiregrid, wg_time_constant, wg_stepwise, hwp_ccw', + 'wiregrid, wg_time_constant, wg_ejecting, hwp_ccw' + ] + expected_calls_of_streams = [ + call(tag=stream_tag, subtype='cal', kwargs=common_kwargs_of_streams) + for stream_tag in expected_tags_of_streams + ] + + for client in wiregrid.run.CLIENTS['smurf']: + assert client.take_bias_steps.start.call_args_list == expected_calls_of_bias_steps + assert client.stream.start.call_args_list == expected_calls_of_streams + assert client.stream.stop.call_count == 6 + + assert wiregrid.run.wiregrid.rotate.call_count == 3 + + +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +@patch('sorunlib.wiregrid.time.sleep', MagicMock()) +def test_time_constant_ccw_el90(): + # Setup all mock clients + wiregrid.run.CLIENTS['acu'] = create_acu_client(180, 90, 0) + wiregrid.run.hwp._get_direction = MagicMock(return_value='ccw') # ccw + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=1, position='outside') + wiregrid.run.CLIENTS['wiregrid']['kikusui'] = create_kikusui_client() + wiregrid.run.CLIENTS['wiregrid']['encoder'] = create_encoder_client() + wiregrid.run.CLIENTS['wiregrid']['labjack'] = create_labjack_client() + wiregrid.run.wiregrid.rotate = MagicMock() + + wiregrid.time_constant(num_repeats=1) + + # just make sure bias_steps and streams because other functions are already + # tested separately. + expected_calls_of_bias_steps = [ + call(tag='wiregrid, wg_time_constant, wg_ejected, hwp_ccw, wg_el90'), + call(tag='wiregrid, wg_time_constant, wg_inserted, hwp_ccw, wg_el90'), + call(tag='wiregrid, wg_time_constant, wg_inserted, hwp_cw, wg_el90'), + call(tag='wiregrid, wg_time_constant, wg_ejected, hwp_cw, wg_el90') + ] + + common_kwargs_of_streams = { + "downsample_factor": None, + "filter_disable": False + } + expected_tags_of_streams = [ + 'wiregrid, wg_time_constant, wg_inserting, hwp_ccw, wg_el90', + 'wiregrid, wg_time_constant, wg_stepwise, hwp_ccw, wg_el90', + 'wiregrid, wg_time_constant, hwp_change_ccw_to_stop, wg_el90', + 'wiregrid, wg_time_constant, hwp_change_stop_to_cw, wg_el90', + 'wiregrid, wg_time_constant, wg_stepwise, hwp_cw, wg_el90', + 'wiregrid, wg_time_constant, wg_ejecting, hwp_cw, wg_el90' + ] + expected_calls_of_streams = [ + call(tag=stream_tag, subtype='cal', kwargs=common_kwargs_of_streams) + for stream_tag in expected_tags_of_streams + ] + + for client in wiregrid.run.CLIENTS['smurf']: + assert client.take_bias_steps.start.call_args_list == expected_calls_of_bias_steps + assert client.stream.start.call_args_list == expected_calls_of_streams + assert client.stream.stop.call_count == 6 + + assert wiregrid.run.wiregrid.rotate.call_count == 3 + + +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +@patch('sorunlib.wiregrid.time.sleep', MagicMock()) +def test_time_constant_repeats(): + # Setup all mock clients + wiregrid.run.CLIENTS['acu'] = create_acu_client(180, 50, 0) + wiregrid.run.hwp._get_direction = MagicMock(return_value='cw') # cw + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=1, position='outside') + wiregrid.run.CLIENTS['wiregrid']['kikusui'] = create_kikusui_client() + wiregrid.run.CLIENTS['wiregrid']['encoder'] = create_encoder_client() + wiregrid.run.CLIENTS['wiregrid']['labjack'] = create_labjack_client() + wiregrid.run.wiregrid.rotate = MagicMock() + + wiregrid.time_constant(num_repeats=2) + + # just make sure bias_steps and streams because other functions are already + # tested separately. + expected_calls_of_bias_steps = [ + call(tag='wiregrid, wg_time_constant, wg_ejected, hwp_cw'), + call(tag='wiregrid, wg_time_constant, wg_inserted, hwp_cw'), + call(tag='wiregrid, wg_time_constant, wg_inserted, hwp_ccw'), + call(tag='wiregrid, wg_time_constant, wg_inserted, hwp_cw'), + call(tag='wiregrid, wg_time_constant, wg_ejected, hwp_cw') + ] + + common_kwargs_of_streams = { + "downsample_factor": None, + "filter_disable": False + } + expected_tags_of_streams = [ + 'wiregrid, wg_time_constant, wg_inserting, hwp_cw', + 'wiregrid, wg_time_constant, wg_stepwise, hwp_cw', + 'wiregrid, wg_time_constant, hwp_change_cw_to_stop', + 'wiregrid, wg_time_constant, hwp_change_stop_to_ccw', + 'wiregrid, wg_time_constant, wg_stepwise, hwp_ccw', + 'wiregrid, wg_time_constant, hwp_change_ccw_to_stop', + 'wiregrid, wg_time_constant, hwp_change_stop_to_cw', + 'wiregrid, wg_time_constant, wg_stepwise, hwp_cw', + 'wiregrid, wg_time_constant, wg_ejecting, hwp_cw' + ] + expected_calls_of_streams = [ + call(tag=stream_tag, subtype='cal', kwargs=common_kwargs_of_streams) + for stream_tag in expected_tags_of_streams + ] + + for client in wiregrid.run.CLIENTS['smurf']: + assert client.take_bias_steps.start.call_args_list == expected_calls_of_bias_steps + assert client.stream.start.call_args_list == expected_calls_of_streams + assert client.stream.stop.call_count == 9 + + assert wiregrid.run.wiregrid.rotate.call_count == 4 + + +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +@patch('sorunlib.wiregrid.time.sleep', MagicMock()) +def test_time_constant_num_repeats_failed(): + with pytest.raises(RuntimeError): + wiregrid.time_constant(num_repeats=-2) + + +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +@patch('sorunlib.wiregrid.time.sleep', MagicMock()) +def test_time_constant_wiregrid_position_failed(): + wiregrid.run.CLIENTS['acu'] = create_acu_client(180, 50, 0) + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=1, position='inside') + wiregrid.run.CLIENTS['wiregrid']['kikusui'] = create_kikusui_client() + wiregrid.run.CLIENTS['wiregrid']['encoder'] = create_encoder_client() + wiregrid.run.CLIENTS['wiregrid']['labjack'] = create_labjack_client() + + with pytest.raises(RuntimeError): + wiregrid.time_constant(num_repeats=1) + + +@patch('sorunlib.wiregrid.run.CLIENTS', mocked_clients()) +@patch('sorunlib.wiregrid.time.sleep', MagicMock()) +def test_time_constant_hwp_direction_failed(): + wiregrid.run.CLIENTS['acu'] = create_acu_client(180, 50, 0) + wiregrid.run.CLIENTS['wiregrid']['actuator'] = \ + create_actuator_client(motor=1, position='outside') + wiregrid.run.CLIENTS['wiregrid']['kikusui'] = create_kikusui_client() + wiregrid.run.CLIENTS['wiregrid']['encoder'] = create_encoder_client() + wiregrid.run.CLIENTS['wiregrid']['labjack'] = create_labjack_client() + + # Set up expected raise from invalid direction + wiregrid.run.hwp._get_direction = MagicMock(return_value=None) + wiregrid.run.hwp._get_direction.side_effect = RuntimeError("The HWP direction is unknown. Aborting...") + + with pytest.raises(RuntimeError): + wiregrid.time_constant(num_repeats=1)