Test raster scan planner and z estimate from raster and snake

tests/unit_tests/test_scan_planners.py

@@ -443,3 +443,229 @@ def test_snake_scan_single_column():
         (0, 10),
         (0, 15),
     ]
+
+
+def test_snake_scan_z_propagation():
+    """Test that snake planner selects correct previous focus height."""
+    initial_position = (0, 0)
+    planner_settings = {
+        "dx": 50,
+        "dy": 50,
+        "x_count": 5,
+        "y_count": 5,
+        "style": "snake",
+    }
+
+    planner = scan_planners.RegularGridPlanner(
+        initial_position=initial_position,
+        planner_settings=planner_settings,
+    )
+
+    expected_z = 1
+
+    while not planner.scan_complete:
+        xy_pos, z_est = planner.get_next_location_and_z_estimate()
+        print(xy_pos, z_est)
+
+        # First position should have no estimate
+        if xy_pos == (0, 0):
+            assert z_est is None
+        else:
+            # Should estimate from previous focused point
+            assert z_est == expected_z - 1
+
+        xyz_pos = (xy_pos[0], xy_pos[1], expected_z)
+
+        planner.mark_location_visited(
+            xyz_pos,
+            imaged=True,
+            focused=True,
+        )
+
+        expected_z += 1
+
+
+def test_raster_planner_basic_grid():
+    """Check that raster scan planner generates a single point for a 1x1 scan."""
+    initial_position = (100, 50)
+    planner_settings = {
+        "dx": 100,
+        "dy": 100,
+        "x_count": 1,
+        "y_count": 1,
+        "style": "raster",
+    }
+
+    planner = scan_planners.RegularGridPlanner(
+        initial_position=initial_position,
+        planner_settings=planner_settings,
+    )
+
+    assert not planner.scan_complete
+    # When we start it should want to stay in the initial pos and have
+    # no z_estimate
+    xy_pos, z_pos = planner.get_next_location_and_z_estimate()
+    assert xy_pos == initial_position
+    assert z_pos is None
+
+    # Try to mark location as imaged with only xy_position
+    with pytest.raises(ValueError, match="3 value tuple expected"):
+        planner.mark_location_visited(xy_pos, imaged=False, focused=False)
+    # scan still not complete
+    assert not planner.scan_complete
+    # if we mark this position as visited but not imaged
+    planner.mark_location_visited(
+        (xy_pos[0], xy_pos[1], 10), imaged=False, focused=False
+    )
+    # scan is now complete
+    assert planner.scan_complete
+
+    # if scan is complete, asking for the next location returns an error
+    with pytest.raises(RuntimeError):
+        planner.get_next_location_and_z_estimate()
+
+
+def test_raster_scan_basic_length():
+    """Raster scan planner should generate the correct number of locations."""
+    initial_position = (100, 50)
+    planner_settings = {
+        "dx": 100,
+        "dy": 100,
+        "x_count": 3,
+        "y_count": 4,
+        "style": "raster",
+    }
+
+    planner = scan_planners.RegularGridPlanner(
+        initial_position=initial_position,
+        planner_settings=planner_settings,
+    )
+
+    coords = planner.remaining_locations
+
+    assert len(coords) == 3 * 4
+
+
+def test_raster_scan_ordering():
+    """Test that raster scan returns a path in the right order."""
+    initial_position = (0, 0)
+    planner_settings = {
+        "dx": 10,
+        "dy": 10,
+        "x_count": 4,
+        "y_count": 3,
+        "style": "raster",
+    }
+
+    planner = scan_planners.RegularGridPlanner(
+        initial_position=initial_position,
+        planner_settings=planner_settings,
+    )
+
+    coords = planner.remaining_locations
+
+    expected = [
+        (0, 0),
+        (10, 0),
+        (20, 0),
+        (30, 0),
+        (0, 10),
+        (10, 10),
+        (20, 10),
+        (30, 10),
+        (0, 20),
+        (10, 20),
+        (20, 20),
+        (30, 20),
+    ]
+
+    assert coords == expected
+
+
+def test_raster_scan_single_row():
+    """Test edge case of a single row scan."""
+    initial_position = (0, 0)
+    planner_settings = {"dx": 5, "dy": 5, "x_count": 4, "y_count": 1, "style": "raster"}
+
+    planner = scan_planners.RegularGridPlanner(
+        initial_position=initial_position,
+        planner_settings=planner_settings,
+    )
+
+    assert planner.remaining_locations == [(0, 0), (5, 0), (10, 0), (15, 0)]
+
+
+def test_raster_scan_single_column():
+    """Test edge case of a single column scan."""
+    initial_position = (0, 0)
+    planner_settings = {"dx": 5, "dy": 5, "x_count": 1, "y_count": 4, "style": "raster"}
+
+    planner = scan_planners.RegularGridPlanner(
+        initial_position=initial_position,
+        planner_settings=planner_settings,
+    )
+
+    assert planner.remaining_locations == [
+        (0, 0),
+        (0, 5),
+        (0, 10),
+        (0, 15),
+    ]
+
+
+def test_raster_z_propagation():
+    """Test that snake planner selects correct previous focus height.
+
+    Constructs a 5x5 grid in a raster pattern, and test that for each movement,
+    the chosen next z position is either
+        - None, for the first point
+        - the start of the previous row, for the first point in a row
+        - the previous site otherwise
+    """
+    initial_position = (0, 0)
+    x_count = 5
+    y_count = 5
+    dx = 50
+    dy = 50
+
+    planner = scan_planners.RegularGridPlanner(
+        initial_position=initial_position,
+        planner_settings={
+            "dx": dx,
+            "dy": dy,
+            "x_count": x_count,
+            "y_count": y_count,
+            "style": "raster",
+        },
+    )
+
+    visited_positions = []
+    current_z = 1
+
+    while not planner.scan_complete:
+        visited_count = len(visited_positions)
+        xy_pos, z_est = planner.get_next_location_and_z_estimate()
+        print(visited_count)
+
+        if visited_count == 0:
+            assert z_est is None
+        else:
+            # check whether this site is at the start of a new row
+            if visited_count % x_count == 0:
+                # if so, get the z position from the start of the previous row
+                expected_z = visited_positions[-x_count][2]
+            else:
+                expected_z = visited_positions[-1][2]
+
+            assert z_est == expected_z
+
+        xyz_pos = (xy_pos[0], xy_pos[1], current_z)
+
+        planner.mark_location_visited(
+            xyz_pos,
+            imaged=True,
+            focused=True,
+        )
+
+        visited_positions.append(xyz_pos)
+        current_z += 1