diff --git a/src/openflexure_microscope_server/things/camera_stage_mapping.py b/src/openflexure_microscope_server/things/camera_stage_mapping.py index 75ce1968..409aa800 100644 --- a/src/openflexure_microscope_server/things/camera_stage_mapping.py +++ b/src/openflexure_microscope_server/things/camera_stage_mapping.py @@ -31,8 +31,8 @@ from camera_stage_mapping.exceptions import MappingError import labthings_fastapi as lt from labthings_fastapi.types.numpy import DenumpifyingDict - from camera_stage_mapping.camera_stage_tracker import Tracker + from .camera import CameraDependency as CameraClient from .stage import StageDependency as Stage @@ -218,8 +218,8 @@ class CameraStageMapper(lt.Thing): .. code-block:: python stage_disp = np.dot( - np.array(image_to_stage_displacement_matrix), np.array([dy,dx]), + np.array(image_to_stage_displacement_matrix), ) """ @@ -277,14 +277,19 @@ class CameraStageMapper(lt.Thing): @lt.thing_action def convert_image_to_stage_coordinates( - self, x: float, y: float + self, x: float, y: float, **_kwargs: float ) -> Mapping[str, int]: - """Convert image coordinates to stage coordinates.""" + """Convert image coordinates to stage coordinates. Only x and y are returned.""" self.assert_calibrated() - relative_move: np.ndarray = np.dot( - np.array([y, x]), np.array(self.image_to_stage_displacement_matrix) - ) - return {"x": int(relative_move[0]), "y": int(relative_move[1])} + return csm_img_to_stage(self.image_to_stage_displacement_matrix, x=x, y=y) + + @lt.thing_action + def convert_stage_to_image_coordinates( + self, x: int, y: int, **_kwargs: int + ) -> Mapping[str, float]: + """Convert stage coordinates to image coordinates. Only x and y are returned.""" + self.assert_calibrated() + return csm_stage_to_img(self.image_to_stage_displacement_matrix, x=x, y=y) @lt.thing_property def thing_state(self) -> Mapping[str, Any]: @@ -293,3 +298,54 @@ class CameraStageMapper(lt.Thing): k: getattr(self, k) for k in ["image_to_stage_displacement_matrix", "image_resolution"] } + + +def csm_img_to_stage( + matrix: np.ndarray | list[list[float]], + *, + x: float | int, + y: float | int, + **_kwargs: int, +) -> Mapping[str, int]: + """Apply any CSM matrix to image coordinates. + + x and y must be kwargs and extra kwargs are ignored, allowing: + ``csm_img_to_stage(matrix, **position)`` to run for a mapping position. + + Note that x and y are the actual (x, y) of the image, not the (m, n) indices used + by numpy + + :param matrix: The matrix to use in the calculation + :param x: the x image coordinate (keyword only) + :param y: the y image coordinate (keyword only) + :return: The resulting stage coordinates as a mapping. + """ + # Note this is (y,x) not (x,y) to put it in numpy image indices + relative_move: np.ndarray = np.dot(np.array([y, x]), np.array(matrix)) + return {"x": round(relative_move[0]), "y": round(relative_move[1])} + + +def csm_stage_to_img( + matrix: np.ndarray | list[list[float]], + *, + x: float | int, + y: float | int, + **_kwargs: int, +) -> Mapping[str, float]: + """Apply any CSM matrix to stage coordinates to get image coordinates. + + x and y must be kwargs and extra kwargs are ignored, allowing: + ``csm_img_to_stage(matrix, **position)`` to run for a mapping position. + + Note that x and y are the actual (x, y) of the image, not the (m, n) indices used + numpy + + :param matrix: The matrix to use in the calculation + :param x: the x stage coordinate (keyword only) + :param y: the y stage coordinate (keyword only) + :return: The resulting img coordinates as a mapping. + """ + inverse_matrix = np.linalg.inv(np.array(matrix)) + relative_move = np.dot(np.array([x, y]), inverse_matrix) + # Note that the relative move is (y, x) as it is from numpy and is in matrix coords. + return {"x": float(relative_move[1]), "y": float(relative_move[0])} diff --git a/src/openflexure_microscope_server/things/stage_measure.py b/src/openflexure_microscope_server/things/stage_measure.py index 45175c2c..6b1837be 100644 --- a/src/openflexure_microscope_server/things/stage_measure.py +++ b/src/openflexure_microscope_server/things/stage_measure.py @@ -18,14 +18,11 @@ import time from dataclasses import dataclass from threading import Lock -from scipy.optimize import curve_fit import numpy as np from camera_stage_mapping import fft_image_tracking import labthings_fastapi as lt -from openflexure_microscope_server.utilities import quadratic - # Things from .autofocus import AutofocusThing from .camera_stage_mapping import CameraStageMapper @@ -72,6 +69,18 @@ class RomDataTracker: """The last stage coordinate recorded.""" return self.stage_coords[-1].copy() + def fit_axis(self, axis: Literal["x", "y"]) -> np.poly1d: + """Quadratic fit stage z against x or y and return poly1d of fit. + + Note that this considers only one of x and y, and ignores the other coordinate. If + this function is used on moves where both x and y are changing, it will give misleading + results. + """ + lateral_positions = [i[axis] for i in self.stage_coords] + z_positions = [i["z"] for i in self.stage_coords] + fit_params = np.polyfit(lateral_positions, z_positions, 2) + return np.poly1d(fit_params) + def predict_z_displacement( self, axis: Literal["x", "y"], @@ -85,14 +94,22 @@ class RomDataTracker: :param stage_position: The current stage position in stage coordinates. :return: The predicted relative z displacement needed to stay in focus. """ - # x or y positions - lateral_positions = [i[axis] for i in self.stage_coords] - z_positions = [i["z"] for i in self.stage_coords] - fit_params, *_others = curve_fit(quadratic, lateral_positions, z_positions) - z_dest = quadratic(stage_position[axis] + stage_movement[axis], *fit_params) + fit_func = self.fit_axis(axis) + z_dest = fit_func(stage_position[axis] + stage_movement[axis]) return int(z_dest - stage_position["z"]) + def find_turning_point(self, axis: Literal["x", "y"]) -> dict[str, int]: + """Find the turing point from the recorded coordinates.""" + fit_func = self.fit_axis(axis) + turning_loc = fit_func.deriv().roots[0] + turning_z = fit_func(turning_loc) + + # As we only move in 1 direction, pull the other axis from the coords. + other_axis = "x" if axis == "y" else "y" + other_coord = self.stage_coords[-1][other_axis] + return {axis: int(turning_loc), other_axis: other_coord, "z": int(turning_z)} + @dataclass class RomDeps: @@ -197,6 +214,51 @@ class RangeofMotionThing(lt.Thing): "Step Range": step_range, } + @lt.thing_action + def perform_recentre( + self, + autofocus: AutofocusDep, + stage: StageDep, + cam: CamDep, + csm: CSMDep, + logger: lt.deps.InvocationLogger, + ) -> None: + """Measures the curvature of the motion to centre x and y axes. + + :param autofocus: A raw_thing_client dependency for autofocus. + :param stage: A raw_thing_client depeendency for the microscope stage. + :param cam: A raw_thing_client depeendency for the camera. + :param csm: A raw_thing_client depeendency for camera stage mapping. + :param logger: A raw_thing_client depeendency for the logger. + """ + got_lock = self._lock.acquire(blocking=False) + if not got_lock: + raise RuntimeError("Trying to run recentre when a test is already running.") + + try: + rom_deps = RomDeps( + autofocus=autofocus, stage=stage, csm=csm, cam=cam, logger=logger + ) + logger.info("Recentring the stage.") + + self._set_stream_resolution(cam) + + # Strictly typed definitions to iterate over for MyPys sake. + axes: tuple[Literal["x"], Literal["y"]] = ("x", "y") + + for axis in axes: + self._recentre_axis(axis, rom_deps) + + centre = rom_deps.stage.position + centre_str = f"({centre['x']}, {centre['y']})" + rom_deps.logger.info(f"Centre is estimated at {centre_str}.") + # Set the central position to (0,0,0) + rom_deps.stage.set_zero_position() + rom_deps.logger.info("Position reset to (0, 0, 0).") + + finally: + self._lock.release() + def _set_stream_resolution(self, cam: CamDep) -> None: """Set the self._stream_resolution attribute by reading camera. @@ -236,7 +298,7 @@ class RangeofMotionThing(lt.Thing): ) rom_deps.logger.info("Moving the stage in 5 medium sized steps.") - self._initial_moves_for_z_prediction( + self._moves_for_z_prediction( axis=axis, direction=direction, rom_deps=rom_deps, @@ -276,13 +338,87 @@ class RangeofMotionThing(lt.Thing): finally: rom_deps.stage.move_absolute(**starting_position, block_cancellation=True) + def _recentre_axis(self, axis: Literal["x", "y"], rom_deps: RomDeps) -> None: + """Recentre a single axis. + + :param axis: The axis to recentre + :param rom_deps: All dependencies that were passed to the calling Action. + """ + rom_deps.logger.info(f"Finding centre in {axis}.") + # A new tracker for this axis. + self._rom_data = RomDataTracker() + # Direction is in image coords, initial assumption is that centre is at (0,0). + direction = self._img_dir_from_stage_coords( + {"x": 0, "y": 0, "z": 0}, axis, rom_deps + ) + + i = 0 + while True: + i += 1 + # Find z then make a number of moves (autofocussing and logging position) + # 5 moves initially (2 subsequently). + rom_deps.autofocus.looping_autofocus(dz=1000) + self._rom_data.stage_coords.append(rom_deps.stage.position) + self._moves_for_z_prediction( + axis=axis, + direction=direction, + rom_deps=rom_deps, + n_moves=5 if i == 1 else 2, + ) + + centred, direction = self._recentre_decision(axis, rom_deps) + if centred: + break + + # If still going at iteration 9 exit + if i > 9: + raise RuntimeError(f"Couldn't find centre of {axis}-axis") + + # Make a big z-corrected move towards estimate of centre. + self._big_z_corrected_movement(axis, direction, rom_deps) + + def _recentre_decision( + self, axis: Literal["x", "y"], rom_deps: RomDeps + ) -> tuple[bool, Literal[1, -1]]: + """Decide what to do next during recentreing. + + The algorithm here: + + * Estimate turning point location + * Check if distance to point is further than a "BIG_STEP" + * If smaller, move to estimated centre, and return that it is centred + * If larger, return that it is not centred, and the direction to move in. + + :param axis: The axis to recentre + :param rom_deps: All dependencies that were passed to the calling Action. + :return: A tuple. The first value is True for "the stage is now centred" and + False for "the stage is not centred". The second value is the direction to + move in, this only has meaning if the first value is False. + """ + estimate = self._rom_data.find_turning_point(axis=axis) + img_perc = self._distance_in_img_percentage(estimate, axis, rom_deps) + + # if the distance is less than 1 big step away then move to it and exit + if abs(img_perc) < BIG_STEP: + rom_deps.logger.info(f"Estimated centre of {axis}-axis is {estimate[axis]}") + rom_deps.stage.move_absolute(**estimate) + # Note the second return, the direction, is meaningless here. + return True, 1 + rom_deps.logger.info( + f"Estimated centre {abs(img_perc):.0f}% of a field of view away, " + "that is too far to move in one move." + ) + # Else calculate the desired direction in image coordinates. + direction = self._img_dir_from_stage_coords(estimate, axis, rom_deps) + return False, direction + def _img_percentage_to_img_coords( self, fov_perc: int, axis: Literal["x", "y"] ) -> float: """For a given image percentage and axis return the distance in img coords. :param fov_perc: The percentage of field of view the stage should move by. - :param axis: The resolution of the stream from the camera. + :param axis: The axis which is being measured. This must be 'x' or 'y'. :return: Distance in image coordinates (pixels) """ if self._stream_resolution is None: @@ -293,6 +429,45 @@ class RangeofMotionThing(lt.Thing): img_index = 0 if axis == "x" else 1 return (fov_perc / 100) * self._stream_resolution[img_index] + def _img_dir_from_stage_coords( + self, target: dict[str, int], axis: Literal["x", "y"], rom_deps: RomDeps + ) -> Literal[1, -1]: + """For a target location in stage coords, return the direction in image coordinates. + + :param target: The target poisiton in stage coordinates + :param axis: The axis which is being measured. This must be 'x' or 'y'. + :param rom_deps: All dependencies that were passed to the calling Action. + :return: Direction to move in image coordinates. + """ + target_im_coords = rom_deps.csm.convert_stage_to_image_coordinates(**target) + current_loc = rom_deps.stage.position + current_loc_im_coords = rom_deps.csm.convert_stage_to_image_coordinates( + **current_loc + ) + return -1 if target_im_coords[axis] < current_loc_im_coords[axis] else 1 + + def _distance_in_img_percentage( + self, target: dict[str, int], axis: Literal["x", "y"], rom_deps: RomDeps + ) -> float: + """For a target location in stage coords return the distance in percentage of FOV. + + :param target: The target poisiton in stage coordinates + :param axis: The axis which is being measured. This must be 'x' or 'y'. + :param rom_deps: All dependencies that were passed to the calling Action. + :return: Percentage of field of view the stage should move by. + """ + if self._stream_resolution is None: + raise RuntimeError( + "Stream resolution must be set before converting coords to percentage" + ) + + current_loc = rom_deps.stage.position + move_stage = {key: target[key] - current_loc[key] for key in target} + move_img = rom_deps.csm.convert_stage_to_image_coordinates(**move_stage) + + img_index = 0 if axis == "x" else 1 + return (move_img[axis] / self._stream_resolution[img_index]) * 100 + def _movement_in_img_coords( self, fov_perc: int, @@ -314,36 +489,37 @@ class RangeofMotionThing(lt.Thing): return {"x": distance * direction, "y": 0} return {"x": 0, "y": distance * direction} - def _initial_moves_for_z_prediction( + def _moves_for_z_prediction( self, axis: Literal["x", "y"], direction: Literal[1, -1], rom_deps: RomDeps, + n_moves: int = 5, ) -> None: - """Perform 5 medium sized moves with autofocus for z feed-forward. + """Perform medium sized moves with autofocus for z feed-forward. z-feed forward allows prediction of the z-position as the stage moves. For the feed forward calculation to work an initial number of measurements must be - taken. This method performs these initial measurements. + taken. :param direction: The direction the stage moves. :param axis: The axis which is being measured. This must be 'x' or 'y'. :param rom_deps: All dependencies that were passed to the calling Action + :param n_moves: Number of moves to make. Default is 5 which is enough for an + initial z estimate. """ movement = self._movement_in_img_coords( fov_perc=MEDIUM_STEP, axis=axis, direction=direction ) - for _loop in range(5): + for _loop in range(n_moves): offset = self._move_and_measure(movement=movement, rom_deps=rom_deps) - rom_deps.logger.info(f"Offset measured as {offset[axis]}") self._rom_data.record_movement(rom_deps.stage.position, offset) if _parasitic_motion_detected(movement, offset): raise ParasiticMotionError( - "Parasitic motion detected during initial images to calculate " - "z-curvature. This may indicate you have started at the end of the " - "range of travel, or that camera stage mapping is poorly " - "calibrated." + "Parasitic motion detected during images to calculate z-curvature. " + "This may indicate you have started at the end of the range of " + "travel, or that camera stage mapping is poorly calibrated." ) def _big_z_corrected_movement( @@ -361,11 +537,11 @@ class RangeofMotionThing(lt.Thing): fov_perc=BIG_STEP, axis=axis, direction=direction ) # Convert to stage coordinates - stage_movemenet = rom_deps.csm.convert_image_to_stage_coordinates(**movement) + stage_movement = rom_deps.csm.convert_image_to_stage_coordinates(**movement) z_disp = self._rom_data.predict_z_displacement( axis=axis, - stage_movement=stage_movemenet, + stage_movement=stage_movement, stage_position=rom_deps.stage.position, ) rom_deps.stage.move_relative(z=z_disp) diff --git a/src/openflexure_microscope_server/utilities.py b/src/openflexure_microscope_server/utilities.py index 68cde53b..77c2ab06 100644 --- a/src/openflexure_microscope_server/utilities.py +++ b/src/openflexure_microscope_server/utilities.py @@ -23,7 +23,6 @@ from functools import wraps import json from pydantic import BaseModel -import numpy as np T = TypeVar("T") P = ParamSpec("P") @@ -345,28 +344,6 @@ def _get_version_from_toml(toml_path: str) -> str: return "Undefined" -# Let MyPy know that the output type matches x. -@overload -def quadratic(x: np.ndarray, a: float, b: float, c: float) -> np.ndarray: ... -@overload -def quadratic(x: float, a: float, b: float, c: float) -> float: ... - - -def quadratic( - x: float | np.ndarray, a: float, b: float, c: float -) -> float | np.ndarray: - """Quadratic function. Used for predicting z. - - :param x: The point or points at which to evaluate the quadratic. This can be a - float or a numpy array. The return will be the same type. - :param a: The coefficient of x^2. - :param b: The coefficient of x. - :param c: The constant coefficient. - :return: The quadratic, evaluated at each point in ``x`` - """ - return a * x**2 + b * x + c - - # Use overload to clarify to MyPy that if enforce_dict is true, inputs and outputs are # dictionaries @overload diff --git a/tests/test_stage_measure.py b/tests/test_stage_measure.py index 9f4e3a4e..5e0462d3 100644 --- a/tests/test_stage_measure.py +++ b/tests/test_stage_measure.py @@ -12,6 +12,10 @@ import pytest import labthings_fastapi as lt from openflexure_microscope_server.things import stage_measure +from openflexure_microscope_server.things.camera_stage_mapping import ( + csm_img_to_stage, + csm_stage_to_img, +) LOGGER = logging.getLogger("mock-invocation_logger") @@ -56,10 +60,10 @@ def example_rom_data(): mock_positions = [ {"x": 0, "y": 0, "z": 42}, {"x": 727, "y": 2, "z": 154}, - {"x": 1454, "y": 4, "z": 228}, + {"x": 1454, "y": 4, "z": 248}, {"x": 2181, "y": 6, "z": 351}, - {"x": 2908, "y": 8, "z": 509}, - {"x": 3635, "y": 10, "z": 617}, + {"x": 2908, "y": 8, "z": 430}, + {"x": 3635, "y": 10, "z": 490}, ] rom_data = stage_measure.RomDataTracker() @@ -75,12 +79,19 @@ def test_predict_z(example_rom_data): mock_z_diff = example_rom_data.predict_z_displacement( axis="x", stage_movement={"x": 5243, "y": 0}, - stage_position={"x": 3635, "y": 10, "z": 617}, + stage_position={"x": 3635, "y": 10, "z": 490}, ) - expected_z_diff = 1343 + expected_z_diff = 153 assert mock_z_diff == expected_z_diff +def test_find_tuning_point(example_rom_data): + """Check that the prediction for the tuning point and z position.""" + mock_turn = example_rom_data.find_turning_point(axis="x") + expected_turning = {"x": 7580, "y": 10, "z": 661} + assert mock_turn == expected_turning + + @pytest.mark.parametrize( ("movement", "axis", "other_axis"), [ @@ -156,10 +167,13 @@ def rom_thing(example_rom_data) -> stage_measure.RangeofMotionThing: def mock_rom_deps(csm_matrix, mocker) -> stage_measure.RomDeps: """Return a RomDeps object full of mocks, except the logger which is LOGGER.""" - def apply_csm(x: float, y: float) -> dict[str, int]: + def apply_csm(x: float, y: float, **_kwargs: float) -> dict[str, int]: """Convert image coordinates to stage coordinates.""" - vec = np.dot(np.array([y, x]), np.array(csm_matrix)) - return {"x": int(vec[0]), "y": int(vec[1])} + return csm_img_to_stage(csm_matrix, x=x, y=y) + + def un_apply_csm(x: float, y: float, **_kwargs: float) -> dict[str, float]: + """Convert stage coordinates to image coordinates.""" + return csm_stage_to_img(csm_matrix, x=x, y=y) mock_cam = mocker.Mock() mock_cam.image_is_sample.return_value = (True, "Mocked not measured.") @@ -168,6 +182,7 @@ def mock_rom_deps(csm_matrix, mocker) -> stage_measure.RomDeps: # Set up mock csm to return a CSM matrix mock_csm.image_to_stage_displacement_matrix = csm_matrix mock_csm.convert_image_to_stage_coordinates.side_effect = apply_csm + mock_csm.convert_stage_to_image_coordinates.side_effect = un_apply_csm return stage_measure.RomDeps( autofocus=mocker.Mock(), @@ -178,6 +193,53 @@ def mock_rom_deps(csm_matrix, mocker) -> stage_measure.RomDeps: ) +@pytest.mark.parametrize( + ("pos", "target_pos", "axis", "expected_dir"), + [ + ({"x": 5000, "y": 0, "z": 0}, {"x": 0, "y": 0, "z": 0}, "x", -1), + ({"x": -5000, "y": 0, "z": 0}, {"x": 0, "y": 0, "z": 0}, "x", 1), + ({"x": 0, "y": 0, "z": 0}, {"x": 5000, "y": 0, "z": 0}, "x", 1), + # Y is flipped due to CSM sign + ({"x": 0, "y": 5000, "z": 0}, {"x": 0, "y": 0, "z": 0}, "y", 1), + ], +) +def test_img_dir_from_stage_coords( + pos, target_pos, axis, expected_dir, rom_thing, mock_rom_deps +): + """Check image direction is correctly generated.""" + mock_rom_deps.stage.position = pos + direction = rom_thing._img_dir_from_stage_coords(target_pos, axis, mock_rom_deps) + assert direction == expected_dir + + +def test_distance_in_img_percentage_err(rom_thing, mock_rom_deps): + """Check that _distance_in_img_percentage throws error if stream res is not set.""" + rom_thing._stream_resolution = None + with pytest.raises(RuntimeError): + rom_thing._distance_in_img_percentage( + {"x": 0, "y": 0, "z": 0}, "x", mock_rom_deps + ) + + +@pytest.mark.parametrize( + ("pos", "target_pos", "axis", "expected_perc"), + [ + ({"x": 5000, "y": 0, "z": 0}, {"x": 0, "y": 0, "z": 0}, "x", -352.44), + ({"x": -5000, "y": 0, "z": 0}, {"x": 0, "y": 0, "z": 0}, "x", 352.44), + ({"x": 0, "y": 0, "z": 0}, {"x": 5000, "y": 0, "z": 0}, "x", 352.44), + # Y is flipped due to CSM sign + ({"x": 0, "y": 5000, "z": 0}, {"x": 0, "y": 0, "z": 0}, "y", 462.131), + ], +) +def test_distance_in_img_percentage( + pos, target_pos, axis, expected_perc, rom_thing, mock_rom_deps +): + """Check _distance_in_img_percentage calculates expected values based on mock CSM.""" + mock_rom_deps.stage.position = pos + img_perc = rom_thing._distance_in_img_percentage(target_pos, axis, mock_rom_deps) + assert round(img_perc, 3) == expected_perc + + def test_offset_from(rom_thing, mock_rom_deps, mocker): """Check the calls and returns for RangeofMotionThing._offset_from.""" # Set up mock for the FFT displacement @@ -402,7 +464,7 @@ def test_big_z_corrected_movement(rom_thing, mock_rom_deps): assert "x" not in move_kwargs assert "y" not in move_kwargs assert "z" in move_kwargs - assert move_kwargs["z"] == 1148 + assert move_kwargs["z"] == 160 # And one move in image coordinates assert mock_rom_deps.csm.move_in_image_coordinates.call_count == 1 @@ -410,7 +472,7 @@ def test_big_z_corrected_movement(rom_thing, mock_rom_deps): assert lat_mov_kwargs == expected_movement -def test_initial_moves_for_z_prediction(rom_thing, mock_rom_deps, mocker): +def test_moves_for_z_prediction(rom_thing, mock_rom_deps, mocker): """Check the initial moves are of the correct size and are recorded.""" # Mock the _offset_from and stage.position to return generated dictionaries that # increment each time they are called. (All values 0 the first time, all values 1 @@ -428,7 +490,7 @@ def test_initial_moves_for_z_prediction(rom_thing, mock_rom_deps, mocker): rom_thing._rom_data = stage_measure.RomDataTracker() # Run it! - rom_thing._initial_moves_for_z_prediction("x", direction=-1, rom_deps=mock_rom_deps) + rom_thing._moves_for_z_prediction("x", direction=-1, rom_deps=mock_rom_deps) # Check that _rom_data now contains the 5 mocked returns in order. assert rom_thing._rom_data.offsets == [{"x": i, "y": i} for i in range(5)] @@ -450,7 +512,7 @@ def test_move_until_edge_error(rom_thing, mock_rom_deps, mocker): return_value=mock_position_dict ) mocker.patch.object( - rom_thing, "_initial_moves_for_z_prediction", side_effect=RuntimeError("Mock") + rom_thing, "_moves_for_z_prediction", side_effect=RuntimeError("Mock") ) # Remove the mock RomData before starting @@ -493,7 +555,7 @@ def test_move_until_edge(rom_thing, mock_rom_deps, mocker): # Mock the main movement functions mock_init_moves = mocker.patch.object( rom_thing, - "_initial_moves_for_z_prediction", + "_moves_for_z_prediction", side_effect=add_fake_initial_positions, ) mock_big_moves = mocker.patch.object( @@ -544,13 +606,16 @@ def test_move_until_edge(rom_thing, mock_rom_deps, mocker): ] -def test_perform_rom_test_locked(rom_thing, mock_rom_deps): - """Check the error if the Rom test is locked.""" +def test_perform_rom_actions_locked(rom_thing, mock_rom_deps): + """Check the error if running one of the calibration actions while locked.""" # Not an RLock so no need to thread. rom_thing._lock.acquire() err_msg = "Trying to run ROM test when a test is already running." with pytest.raises(RuntimeError, match=err_msg): rom_thing.perform_rom_test(**dataclasses.asdict(mock_rom_deps)) + err_msg = "Trying to run recentre when a test is already running." + with pytest.raises(RuntimeError, match=err_msg): + rom_thing.perform_recentre(**dataclasses.asdict(mock_rom_deps)) def test_perform_rom_test_(rom_thing, mock_rom_deps, mocker): @@ -598,3 +663,111 @@ def test_perform_rom_test_(rom_thing, mock_rom_deps, mocker): assert mock_move_until_edge.call_args_list[3].kwargs["axis"] == "y" assert mock_move_until_edge.call_args_list[3].kwargs["direction"] == -1 + + +def test_perform_recenter(rom_thing, mock_rom_deps, mocker, caplog): + """Check that performing the recentre runs through expected operations.""" + + def check_lock(*_args, **_kwargs): + """Check the thing is locked.""" + assert not rom_thing._lock.acquire(blocking=False) + + mock_set_stream_res = mocker.patch.object( + rom_thing, "_set_stream_resolution", side_effect=check_lock + ) + mock_recentre_axis = mocker.patch.object(rom_thing, "_recentre_axis") + + # Set a mock stage position, as we patch _recentre_axis this should be logged + # as the centre. + mock_rom_deps.stage.position = {"x": 4321, "y": 1234, "z": 0} + + # Unpack the dict and save because it the dictionary is a copy so we can't check + # mocks from `mock_rom_deps` + mock_deps_dict = dataclasses.asdict(mock_rom_deps) + + with caplog.at_level(logging.INFO): + rom_thing.perform_recentre(**mock_deps_dict) + + assert len(caplog.messages) == 3 + assert caplog.messages[0] == "Recentring the stage." + assert caplog.messages[1] == "Centre is estimated at (4321, 1234)." + assert caplog.messages[2] == "Position reset to (0, 0, 0)." + + # Check the lock was freed + assert rom_thing._lock.acquire(blocking=False) + rom_thing._lock.release() + + assert mock_set_stream_res.call_count == 1 + + # Recentre called first in x then in y + assert mock_recentre_axis.call_count == 2 + assert mock_recentre_axis.call_args_list[0].args[0] == "x" + assert mock_recentre_axis.call_args_list[1].args[0] == "y" + + # check that the position set to zero once + assert mock_deps_dict["stage"].set_zero_position.call_count == 1 + + +@pytest.mark.parametrize("true_on", [1, 4, 10, 11]) +def test_recentre_axis(true_on, rom_thing, mock_rom_deps, mocker): + """Test the high level algorithm of recentring an axis. + + This doesn't include deciding if we are centred or choosing the direction to move. + """ + ## Start such that movement starts negative + mock_rom_deps.stage.position = {"x": 10000, "y": 10000, "z": 0} + mock_moves = mocker.patch.object(rom_thing, "_moves_for_z_prediction") + + # Mock _recentre_decision so it returns (False, 1) a number of times then finally + # (True, 1). + decisions = [(False, 1)] * (true_on - 1) + [(True, 1)] + mock_recentre_decision = mocker.patch.object( + rom_thing, "_recentre_decision", side_effect=decisions + ) + + if true_on > 10: + with pytest.raises(RuntimeError, match="Couldn't find centre"): + rom_thing._recentre_axis("x", mock_rom_deps) + else: + rom_thing._recentre_axis("x", mock_rom_deps) + + # _recentre_decision is called until True or exits after 10 attempts + assert mock_recentre_decision.call_count == min(true_on, 10) + + # The _moves_for_z_prediction is called once before any decision, and not called + # after a decision of centred. The max call count is 10 + assert mock_moves.call_count == min(true_on, 10) + for i, arg_list in enumerate(mock_moves.call_args_list): + # First direction is -ve due to starting pos, then it changes direction + # as the mock always returns 1 + expected_dir = -1 if i < 1 else 1 + assert arg_list.kwargs["direction"] == expected_dir + + +@pytest.mark.parametrize( + ("dist", "direction", "expected_decision"), + [ + (500, 1, (False, 1)), # Big step is 200% this is over, movement is positive. + (-500, -1, (False, -1)), + (200, 1, (False, 1)), + (199, 1, (True, 1)), # Less than 200% FOV movement return centred=True + (-199, -1, (True, 1)), # Always return 1 when c + ], +) +def test_recentre_decision( + dist, direction, expected_decision, rom_thing, mock_rom_deps, mocker +): + """What the algorithm decides to do in different situations.""" + # Mock find turning point just because there is no _rom_data + mocker.patch.object(rom_thing._rom_data, "find_turning_point") + + # As _distance_in_img_percentage and _img_dir_from_stage_coords are tested + # this test mocks their values and checks the return is as expected + mocker.patch.object(rom_thing, "_distance_in_img_percentage", return_value=dist) + mocker.patch.object(rom_thing, "_img_dir_from_stage_coords", return_value=direction) + + assert rom_thing._recentre_decision("x", mock_rom_deps) == expected_decision + + # Check that we make an absolute move (to the centre) before returning centred. + centred = expected_decision[0] + mock_rom_deps.stage.move_absolute.call_count == 1 if centred else 0