432 lines
16 KiB
Python
432 lines
16 KiB
Python
"""Functionality for planning scan routes.
|
|
|
|
A scan route can be planned by a ScanPlanner class currently there
|
|
is only one type the SmartSpiral. More can be added using by
|
|
subclassing the ScanPlanner
|
|
"""
|
|
|
|
from typing import TypeAlias, Optional
|
|
import logging
|
|
from copy import copy
|
|
|
|
import numpy as np
|
|
|
|
LOGGER = logging.getLogger(__name__)
|
|
|
|
XYPos: TypeAlias = tuple[int, int]
|
|
XYZPos: TypeAlias = tuple[int, int, int]
|
|
XYPosList: TypeAlias = list[XYPos]
|
|
XYZPosList: TypeAlias = list[XYZPos]
|
|
|
|
|
|
# how many times the minimum distance between images to include as a "nearby" image
|
|
# default 1.4 includes images offset in x or y, but not diagonally.
|
|
# This wis based of a 4:3 aspect ratio. So x moves are 1.33 times larger than y
|
|
NEIGHBOUR_CUTOFF = 1.4
|
|
|
|
|
|
def enforce_xy_tuple(value: XYPos) -> XYPos:
|
|
"""Check input is a tuple and is of length 2.
|
|
|
|
If possible it will coerce the value to a tuple.
|
|
|
|
:raises ValueError: if the input cannot be coerced to a tuple of length 2.
|
|
"""
|
|
if not isinstance(value, (list, tuple)):
|
|
raise ValueError("2 value tuple expected")
|
|
if not len(value) == 2:
|
|
raise ValueError("2 value tuple expected")
|
|
if isinstance(value, list):
|
|
return tuple(value)
|
|
return value
|
|
|
|
|
|
def enforce_xyz_tuple(value: XYZPos) -> XYZPos:
|
|
"""Check input is a tuple and is of length 3.
|
|
|
|
If possible it will coerce the value to a tuple.
|
|
|
|
:raises ValueError: if the input cannot be coerced to a tuple of length 3.
|
|
"""
|
|
if not isinstance(value, (list, tuple)):
|
|
raise ValueError("3 value tuple expected")
|
|
if not len(value) == 3:
|
|
raise ValueError("3 value tuple expected")
|
|
if isinstance(value, list):
|
|
return tuple(value)
|
|
return value
|
|
|
|
|
|
class ScanPlanner:
|
|
"""A base class for a scan planner.
|
|
|
|
This should never be used directly for a scan, it should be subclassed.
|
|
|
|
Each subclass should implement at least the methods with NotImplementedError
|
|
set:
|
|
|
|
* ``_parse()`` - to parse the planner_settings dictionary, saving values to class
|
|
variables
|
|
* ``_initial_location_list()`` - Sets the list of locations for the scan to follow
|
|
|
|
For a simple scan pattern this should be sufficient. For more complex ones that
|
|
dynamically adjust the path it is suggested to override ``mark_location_visited()``
|
|
calling ``super().mark_location_visited()`` at the start of the method so that all
|
|
locations are adjusted.
|
|
|
|
When subclassing be sure to use ``enforce_xy_tuple`` and ``enforce_xyz_tuple`` on
|
|
any user data before running.
|
|
"""
|
|
|
|
def __init__(
|
|
self, initial_position: XYPos, planner_settings: Optional[dict] = None
|
|
):
|
|
"""Set up lists for the path planning, and scan history."""
|
|
self._initial_position = enforce_xy_tuple(initial_position)
|
|
self._parse(planner_settings)
|
|
|
|
# The remaining (x,y) locations to scan
|
|
# (This was `path` before refactoring from the long `sample_scan` code)
|
|
self._remaining_locations: XYPosList = self._initial_location_list()
|
|
|
|
# This holds a list of all (x,y,z) locations where images were taken
|
|
# this may not be equivalent to the x,y positions ins self._path_history
|
|
# if background detect is used
|
|
# (This was not used in the `sample_scan` code)
|
|
self._imaged_locations: XYZPosList = []
|
|
|
|
# This holds a list of all (x,y,z) locations where autofocus was successful
|
|
# (This was `focused_path` before refactoring from the long `sample_scan` code)
|
|
self._focused_locations: XYZPosList = []
|
|
|
|
# This holds a list of all x,y locations visited in order since the start
|
|
# (This was `true_path` before refactoring from the long `sample_scan` code
|
|
# previously it had z set, but if we don't take an image, not z is needed and
|
|
# it slows other checks)
|
|
self._path_history: XYPosList = []
|
|
|
|
@property
|
|
def scan_complete(self) -> bool:
|
|
"""Return True if there are no locations left to scan."""
|
|
return not self._remaining_locations
|
|
|
|
@property
|
|
def remaining_locations(self) -> XYPosList:
|
|
"""Property to access a copy of the remaining_locations."""
|
|
return copy(self._remaining_locations)
|
|
|
|
@property
|
|
def imaged_locations(self) -> XYZPosList:
|
|
"""Property to access a copy of the imaged_locations."""
|
|
return copy(self._imaged_locations)
|
|
|
|
@property
|
|
def focused_locations(self) -> XYZPosList:
|
|
"""Property to access a copy of the focused_locations."""
|
|
return copy(self._focused_locations)
|
|
|
|
@property
|
|
def path_history(self) -> XYPosList:
|
|
"""Property to access a copy of the path_history."""
|
|
return copy(self._path_history)
|
|
|
|
def _parse(self, planner_settings: Optional[dict] = None) -> None:
|
|
"""Parse any settings sent to this planner and store them if needed."""
|
|
raise NotImplementedError("Did you call the ScanPlanner base class?")
|
|
|
|
def _initial_location_list(self) -> XYPosList:
|
|
"""Set the initial list of locations for this scan planner.
|
|
|
|
This is called on initialisation.
|
|
|
|
For a simple grid scan/snake scan this would be all locations to move to.
|
|
|
|
Note for implementation that this _must_ contain (x,y) tuples, not [x, y]
|
|
lists or matching errors could occur.
|
|
"""
|
|
raise NotImplementedError("Did you call the ScanPlanner base class?")
|
|
|
|
def position_visited(self, position: XYPos) -> bool:
|
|
"""Return True if input xy position has been visited before."""
|
|
# Ensure tuple for correct matching!
|
|
return tuple(position) in self._path_history
|
|
|
|
def position_planned(self, position: XYPos) -> bool:
|
|
"""Return True if input xy position is planned."""
|
|
# Ensure tuple for correct matching!
|
|
return tuple(position) in self._remaining_locations
|
|
|
|
def get_next_location_and_z_estimate(self) -> tuple[XYPos, Optional[int]]:
|
|
"""Return the next location to scan and its estimated z-position.
|
|
|
|
Note z-position may be None! This indicates that the current z, position
|
|
should be used.
|
|
"""
|
|
if self.scan_complete:
|
|
raise RuntimeError("Can't get next position, scan is complete")
|
|
|
|
next_location = self._remaining_locations[0]
|
|
|
|
# If focussed locations exist return closest location, favouring most recent
|
|
closest_pos = self.closest_focus_site(next_location)
|
|
if closest_pos is None:
|
|
z = None
|
|
else:
|
|
z = closest_pos[2]
|
|
|
|
return next_location, z
|
|
|
|
def closest_focus_site(self, xy_pos: XYPos) -> Optional[XYZPos]:
|
|
"""Return the xyz position of the closest site where focus was achieved.
|
|
|
|
The most recently taken image is returned in the case of a tie.
|
|
|
|
:param xy_pos: The xy_position which the returned position should be closest
|
|
to.
|
|
|
|
Returns None if there if no focussed locations are present
|
|
"""
|
|
if not self._focused_locations:
|
|
return None
|
|
|
|
# must be float64 (double precision) to deal with the huge numbers involved!
|
|
current_pos = np.array(xy_pos, dtype="float64")
|
|
path_pos = np.array(self._focused_locations, dtype="float64")[:, :2]
|
|
|
|
# Use linalg.norm to calculate the direct distance bweween the points
|
|
# Note linalg.norm always uses float64
|
|
dists = np.linalg.norm((path_pos - current_pos), axis=1)
|
|
|
|
# Get indices of all minima.
|
|
# Note np.where always returns a tuple of arrays, hence the trailing [0]
|
|
indices = np.where(dists == np.min(dists))[0]
|
|
|
|
# The last index is most recent
|
|
return self._focused_locations[indices[-1]]
|
|
|
|
def mark_location_visited(
|
|
self, xyz_pos: XYZPos, imaged: bool, focused: bool
|
|
) -> None:
|
|
"""Mark the location as visited.
|
|
|
|
Args:
|
|
xyz_pos: the x_y_z position
|
|
imaged: true if an image was taken, false if not (due to background detect)
|
|
focused: true if autofocus completed successfully
|
|
|
|
"""
|
|
# ensure is tuple!
|
|
xyz_pos = enforce_xyz_tuple(xyz_pos)
|
|
xy_pos = xyz_pos[:2]
|
|
|
|
# Remove the expected position from the remaining locations list
|
|
# and check it's correct
|
|
expected_pos = tuple(self._remaining_locations.pop(0))
|
|
if xy_pos != expected_pos:
|
|
raise RuntimeError("Wrong scan location visited!")
|
|
|
|
# Append xy position for path_history
|
|
self._path_history.append(xy_pos)
|
|
# And full x,y,z for imaged and foucsed if appropriate
|
|
if imaged:
|
|
self._imaged_locations.append(xyz_pos)
|
|
if focused:
|
|
self._focused_locations.append(xyz_pos)
|
|
|
|
|
|
class SmartSpiral(ScanPlanner):
|
|
"""A scan planner that spirals outward from the centre, prioritising short moves.
|
|
|
|
This planner spirals out from the centre, but prioritises short moves over rigidly
|
|
sticking to minimising radius from the centre of the scan.
|
|
|
|
Each time and image is taken the four neighbouring images are added
|
|
to the list of positions to image (unless they are already listed or
|
|
tried). However, if a location is not imaged due no sample being detected
|
|
then neibouring positions are not imaged.
|
|
|
|
The next image taken is the closes to the centre (considering the largest
|
|
of vertical or horizontal distance), ties are broken by the distance from
|
|
the current position.
|
|
"""
|
|
|
|
_max_dist: int = 0
|
|
_dx: int = 0
|
|
_dy: int = 0
|
|
|
|
def _parse(self, planner_settings: Optional[dict] = None) -> None:
|
|
"""Parse SmartSpiral Settings dictionary.
|
|
|
|
* ``dx`` - the movement size in x
|
|
* ``dy`` - the movement size in y
|
|
* ``max_dist`` - The maximum distance to a location can be from the centre.
|
|
"""
|
|
expected_keys = ["max_dist", "dx", "dy"]
|
|
invalid_msg = "SmartSpiral requires a planner_settings dictionary with keys: "
|
|
if not planner_settings:
|
|
raise ValueError(invalid_msg + ",".join(expected_keys))
|
|
if not all(keys in planner_settings for keys in expected_keys):
|
|
raise KeyError(invalid_msg + ",".join(expected_keys))
|
|
|
|
self._dx = int(planner_settings["dx"])
|
|
self._dy = int(planner_settings["dy"])
|
|
self._max_dist = int(planner_settings["max_dist"])
|
|
|
|
def _initial_location_list(self) -> XYPosList:
|
|
"""Set the initial list of locations for this scan planner.
|
|
|
|
This is salled on initialisation.
|
|
|
|
For smart spiral this is just the first point
|
|
"""
|
|
return [self._initial_position]
|
|
|
|
def mark_location_visited(
|
|
self, xyz_pos: XYZPos, imaged: bool = True, focused: bool = True
|
|
) -> None:
|
|
"""Mark the location as visited. Adjust extra positions accordingly.
|
|
|
|
Args:
|
|
xyz_pos: the x_y_z position
|
|
imaged: true if an image was taken, false if not (due to background detect)
|
|
focused: true if autofocus completed successfully
|
|
|
|
"""
|
|
# First call the base class to update the positions
|
|
super().mark_location_visited(xyz_pos, imaged, focused)
|
|
|
|
xy_pos = enforce_xy_tuple(xyz_pos[:2])
|
|
if imaged:
|
|
self._add_surrounding_positions(xy_pos)
|
|
self._re_sort_remaining_locations(xy_pos)
|
|
|
|
def _add_surrounding_positions(self, xy_pos: XYPos) -> None:
|
|
"""Add the 4 surrounding positions to the list of remaining locations to visit.
|
|
|
|
This adds the surrounding positions (with 4 point connectivity) to the
|
|
remaining locations list if they are not:
|
|
|
|
* too far away
|
|
* already planned
|
|
* already visited
|
|
"""
|
|
new_positions = [
|
|
(xy_pos[0] - self._dx, xy_pos[1]),
|
|
(xy_pos[0] + self._dx, xy_pos[1]),
|
|
(xy_pos[0], xy_pos[1] - self._dy),
|
|
(xy_pos[0], xy_pos[1] + self._dy),
|
|
]
|
|
|
|
for new_pos in new_positions:
|
|
# Skip position if already planned or visited
|
|
if self.position_planned(new_pos) or self.position_visited(new_pos):
|
|
continue
|
|
|
|
dist = distance_between(new_pos, self._initial_position)
|
|
if dist > self._max_dist:
|
|
LOGGER.debug("Rejected moving to %s as it is out of range", new_pos)
|
|
continue
|
|
self._remaining_locations.append(new_pos)
|
|
|
|
def _re_sort_remaining_locations(self, current_pos: XYPos) -> None:
|
|
"""Sort the remaining positions based on the current location."""
|
|
|
|
# Defined rather than use a lambda for readability
|
|
def sort_key(pos):
|
|
return (
|
|
self.moves_between(current_pos, pos),
|
|
self.moves_between(self._initial_position, pos),
|
|
distance_between(current_pos, pos),
|
|
)
|
|
|
|
self._remaining_locations.sort(key=sort_key)
|
|
|
|
def get_next_location_and_z_estimate(self) -> tuple[XYPos, Optional[int]]:
|
|
"""Return the next location to scan and its estimated z-position.
|
|
|
|
This overrides the default behaviour of ScanPlanner to take the lowest value of
|
|
nearest neighbours as this works best for smart stack.
|
|
|
|
Note z-position may be None! This indicates that the current z position
|
|
should be used.
|
|
"""
|
|
if self.scan_complete:
|
|
raise RuntimeError("Can't get next position, scan is complete")
|
|
|
|
next_location = self._remaining_locations[0]
|
|
|
|
# If focused locations exist, return the neighbour with the lowest z position
|
|
closest_pos = self.select_nearby_focus_site(next_location)
|
|
if closest_pos is None:
|
|
z = None
|
|
else:
|
|
z = closest_pos[2]
|
|
|
|
return next_location, z
|
|
|
|
def select_nearby_focus_site(self, xy_pos: XYPos) -> Optional[XYZPos]:
|
|
"""Return the xyz position of the nearby site with the lowest z position.
|
|
|
|
Lowest position is best, as starting too high causes smart stacking to
|
|
autofocus and restart. Starting too low just requires extra movements in +z.
|
|
Nearby is defined as within NEIGHBOUR_CUTOFF times the distance to the closest
|
|
neighbour.
|
|
|
|
Returns None if there if no focused locations are present
|
|
"""
|
|
if not self._focused_locations:
|
|
return None
|
|
|
|
# must be float64 (double precision) to deal with the huge numbers involved!
|
|
current_pos = np.array(xy_pos, dtype="float64")
|
|
path_pos = np.array(self._focused_locations, dtype="float64")[:, :2]
|
|
|
|
# Use linalg.norm to calculate the direct distance between the points
|
|
# Note linalg.norm always uses float64
|
|
dists = np.linalg.norm((path_pos - current_pos), axis=1)
|
|
|
|
# Get indices of all focused sites within NEIGHBOUR_CUTOFF the minimum distance.
|
|
# Note np.where always returns a tuple of arrays, hence the trailing [0]
|
|
indices = np.where(dists <= NEIGHBOUR_CUTOFF * np.min(dists))[0]
|
|
|
|
# Turning into an array allows slicing based on a list
|
|
focused_locations_array = np.array(self._focused_locations)
|
|
|
|
# Choose the lowest (smallest z) of the neighbouring sites. Smart stack works best
|
|
# if started too low, so the lowest z will perform best
|
|
chosen_focused_site = min(focused_locations_array[indices], key=lambda x: x[-1])
|
|
|
|
# Convert back into list so values are of type int instead of np.int32
|
|
return tuple(chosen_focused_site.tolist())
|
|
|
|
def moves_between(
|
|
self,
|
|
starting_pos: XYPos | np.ndarray,
|
|
ending_pos: XYPos | np.ndarray,
|
|
) -> float:
|
|
"""Return the larger of x moves or y moves between two xy positions.
|
|
|
|
:param starting_pos: the position to measure from
|
|
:param ending_pos: the position to measure to
|
|
|
|
"""
|
|
move_size = np.array([self._dx, self._dy])
|
|
|
|
starting_pos = np.array(starting_pos, dtype="float64")
|
|
ending_pos = np.array(ending_pos, dtype="float64")
|
|
|
|
displacement_in_moves = (ending_pos - starting_pos) / move_size
|
|
|
|
return np.max(np.abs(displacement_in_moves))
|
|
|
|
|
|
def distance_between(
|
|
current_pos: XYPos | np.ndarray, next_pos: XYPos | np.ndarray
|
|
) -> float:
|
|
"""Calculate the distance between the two xy positions.
|
|
|
|
This was previously called ``distance_to_site``
|
|
"""
|
|
next_pos = np.array(next_pos, dtype="float64")
|
|
current_pos = np.array(current_pos, dtype="float64")
|
|
return float(np.linalg.norm(next_pos - current_pos))
|