rahix/openflexure-microscope-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
openflexure-microscope-server/tests/utilities/scan_test_helpers.py

270 lines
8.7 KiB
Python
Executable file
Raw Blame History

#! /usr/bin/env python3
"""Utility functions for testing scan planners.
These including fake sample creation, scan path visualisation, and
persistent storage of expected scan paths for samples.
"""
import argparse
import os
import pickle
import numpy as np
from scipy import interpolate
from matplotlib import pyplot as plt
from matplotlib.path import Path as MatPath
from matplotlib.patches import PathPatch
from matplotlib.figure import Figure
from openflexure_microscope_server import scan_planners
THIS_DIR = os.path.dirname(os.path.realpath(__file__))
ALL_SAMPLE_NAMES = ("lobed", "regular", "core")
class FakeSample:
"""A fake sample to test scan algorithms.
The sample is able to return whether a given position is sample, there is no image
associated with the sample
"""
def __init__(self, xy_points: list[tuple[int, int]]):
"""Create the sample from a spline interpolation around the given points."""
self._sample_perimeter = interp_closed_path(xy_points, 500)
def is_sample(self, pos: tuple[int, int], im_size: tuple[int, int]) -> bool:
"""Return True if an image at a given location is on the sample.
The image size is specified to check if it overlaps the sample. It doesn't
check the entire image field as this is designed to be used where the fake
sample is much larger than the image and has smooth edges. It just checks the
4 corners.
"""
img_corners = [
(pos[0] + im_size[0], pos[1] + im_size[1]),
(pos[0] + im_size[0], pos[1] - im_size[1]),
(pos[0] - im_size[0], pos[1] + im_size[1]),
(pos[0] - im_size[0], pos[1] - im_size[1]),
]
return any(
self._sample_perimeter.contains_point(corner) for corner in img_corners
)
@property
def patch(self) -> PathPatch:
"""The sample as a matplotlib patch for plotting."""
patch = PathPatch(self._sample_perimeter)
patch.set(color=(1.0, 0.8, 1.0, 1.0))
return patch
def visualise_scan(sample: FakeSample, planner: scan_planners.ScanPlanner) -> Figure:
"""For a given sample and scanner object return a matplotlib figure of the scan."""
fig, ax = plt.subplots(figsize=(8, 8))
ax.add_artist(sample.patch)
xh, yh = zip(*planner.path_history, strict=True)
xi, yi, _zi = zip(*planner.imaged_locations, strict=True)
if planner.secondary_locations:
xs, ys, _zs = zip(*planner.secondary_locations, strict=True)
else:
xs, ys, _zs = [], [], []
# convert history to numpy array so can calculate quiver arrows
xh = np.array(xh)
yh = np.array(yh)
plt.quiver(
xh[:-1],
yh[:-1],
xh[1:] - xh[:-1],
yh[1:] - yh[:-1],
scale_units="xy",
angles="xy",
scale=1,
)
plt.plot(xh, yh, "r.")
plt.plot(xi, yi, "g*")
plt.plot(xs, ys, "o", mfc="none", mec="blue")
ax.axis("equal")
return fig
def interp_closed_path(xy_points: list[tuple[int, int]], n_points: int) -> MatPath:
"""Interpolate an n_point closed curve from a lists of xy_points.
This can be used to create an arbitrary sample shape for testing a scan planner.
Modified from:
https://stackoverflow.com/questions/33962717/interpolating-a-closed-curve-using-scipy
"""
# Use zip to separate x and y points into tuples
x, y = zip(*xy_points, strict=True)
# Append first point and convert to array
x = np.array(x + (x[0],))
y = np.array(y + (y[0],))
# fit splines to x=f(u) and y=g(u), treating both as periodic. also note that s=0
# is needed in order to force the spline fit to pass through all the input points.
spline_data, *_unused = interpolate.splprep([x, y], s=0, per=True)
# evaluate the spline
xi, yi = interpolate.splev(np.linspace(0, 1, n_points), spline_data)
# Convert to a matplotlib closed path
path_points = [[xp, yp] for xp, yp in (zip(xi, yi, strict=True))]
return MatPath(path_points, closed=True)
def example_smart_spiral(
sample_name: str = "lobed",
) -> tuple[FakeSample, scan_planners.ScanPlanner]:
"""Run an example scan.
:returns: The sample scanned and the planner object after scan is complete.
"""
xy_sample_points = load_sample_points(sample_name)
sample = FakeSample(xy_sample_points)
img_size = (1000, 1000)
initial_position = (0, 0)
planner_settings = {"dx": 1200, "dy": 800, "max_dist": 100000}
planner = scan_planners.SmartSpiral(
initial_position=initial_position, planner_settings=planner_settings
)
while not planner.scan_complete:
xy_pos, _ = planner.get_next_location_and_z_estimate()
xyz_pos = (xy_pos[0], xy_pos[1], 0)
imaged = sample.is_sample(xy_pos, img_size)
planner.mark_location_visited(xyz_pos, imaged=imaged, focused=imaged)
return sample, planner
def profile_example_smart_spiral():
"""Profile running an example scan and print the cumulative profile stats."""
import pstats
import cProfile
profiler = cProfile.Profile()
profiler.runcall(example_smart_spiral)
stats_fname = os.path.join(THIS_DIR, "scan_example_stats.pstats")
profiler.dump_stats(stats_fname)
run_stats = pstats.Stats(stats_fname)
run_stats.strip_dirs()
run_stats.sort_stats("cumulative")
run_stats.print_stats("scan_planners.py")
def plot_all_examples():
"""Plot all examples as png files."""
for sample_name in ALL_SAMPLE_NAMES:
sample, planner = example_smart_spiral(sample_name)
png_fname = os.path.join(THIS_DIR, f"scan_example_plot_{sample_name}.png")
fig = visualise_scan(sample, planner)
fig.savefig(png_fname, dpi=200)
def update_example_smart_spiral_pickles():
"""Pickle the ScanPlanner for the example_smart_spiral().
This is done so the history can be compared by testing to check
the algorithm is unchanged.
If the algorithm is purposefully changed then this will need to be
run to update the pickle for the test to pass.
Takes sample, the sample type we have generated, so we can make a
pickle for each sample type
"""
for sample_name in ALL_SAMPLE_NAMES:
pkl_fname = os.path.join(THIS_DIR, f"example_smart_spiral_{sample_name}.pkl")
_, planner = example_smart_spiral(sample_name)
with open(pkl_fname, "wb") as pkl_file_obj:
pickle.dump(planner, pkl_file_obj, pickle.HIGHEST_PROTOCOL)
def get_expected_result_for_example_smart_spiral(
sample_name: str,
) -> scan_planners.ScanPlanner:
"""Return the expected ScanPlanner object for the example_smart_spiral().
This is loaded from a pickle so that the object can be committed to the repo.
"""
pkl_fname = os.path.join(THIS_DIR, f"example_smart_spiral_{sample_name}.pkl")
with open(pkl_fname, "rb") as pkl_file_obj:
return pickle.load(pkl_file_obj)
def load_sample_points(sample_name: str):
"""Return the points to generate the FakeSample corresponding to the given input name.
Options are "lobed", "regular", and "core".
"""
sample_options = {
"lobed": [
(-5000, -5000),
(-2000, 16000),
(1000, 2000),
(6000, 7000),
(9000, 2000),
],
"regular": [
(-5000, -5000),
(-5000, 5000),
(5000, 5000),
(5000, -5000),
],
"core": [
(-12000, 2000),
(-12000, 1000),
(0, -2000),
(10000, -2000),
(10000, -1000),
(1000, 0),
],
}
if sample_name not in sample_options:
all_samples = ", ".join(sample_options.keys())
raise ValueError(
f"{sample_name} is not a valid sample name. Valid names : {all_samples}"
)
return sample_options[sample_name]
def main():
"""Run the profiler, the plotting, or update the pickles based on command line input.
This only runs if run as a command line script.
"""
parser = argparse.ArgumentParser(description="Simulated scan-planning utility")
subparsers = parser.add_subparsers(dest="command", required=True)
# profile
subparsers.add_parser("profile", help="Run simulation under cProfile")
# plot
subparsers.add_parser("plot", help="Run simulation and produce plots")
# update
subparsers.add_parser("update", help="Update stored reference pickles")
args = parser.parse_args()
if args.command == "profile":
profile_example_smart_spiral()
elif args.command == "plot":
plot_all_examples()
elif args.command == "update":
update_example_smart_spiral_pickles()
else:
parser.error(f"Unknown command: {args.command}")
if __name__ == "__main__":
main()
Reference in a new issue View git blame Copy permalink