Improve the readability of beam shift calibration code & results (#144)

* Streamline the `calibrate_beamshift_live` function

* Streamline `CalibBeamShift.plot`

* Improvements to `CalibBeamShift` readability (WIP)

* Add option to calibrate beamshift with vsp

* Fix errors, add beam center

* Switch calibration output format from pickle to yaml

* Allow delay during calibrate beamshift

* Add necessary reflections to fix plotting

* Final tweaks

* Make the yaml produced by CalibBeamShift human-readable

* Update src/instamatic/calibrate/calibrate_beamshift.py

Co-authored-by: Stef Smeets <stefsmeets@users.noreply.github.com>

* Update src/instamatic/calibrate/calibrate_beamshift.py

Co-authored-by: Stef Smeets <stefsmeets@users.noreply.github.com>

* Minor post-review type-hint improvements + ruff

---------

Co-authored-by: Stef Smeets <stefsmeets@users.noreply.github.com>

Author: Baharis

docs/config.md

@@ -243,9 +243,9 @@ This file holds the specifications of the camera. This file is must be located t
 : Set the correction ratio for the cross pixels in the Timepix detector, default: 3.
 
 **calib_beamshift**
-: Set up the grid and stepsize for the calibration of the beam shift in SerialED. The calibration will run a grid of `stepsize` by `stepsize` points, with steps of `stepsize`. The stepsize must be given corresponding to 2500x, and instamatic will then adjust the stepsize depending on the actual magnification, if needed. For example:
+: Set up the grid and stepsize for the calibration of the beam shift in SerialED. The calibration will run a grid of `stepsize` by `stepsize` points, with steps of `stepsize`. The stepsize must be given corresponding to 2500x, and instamatic will then adjust the stepsize depending on the actual magnification, if needed. If the beam moves too slow, a `delay` between setting beam shift and getting image can be introduced. For example:
 ```yaml
-{gridsize: 5, stepsize: 500}
+{gridsize: 5, stepsize: 500, delay: 0.5}
 ```
 
 **calib_directbeam**

src/instamatic/calibrate/calibrate_beamshift.py

@@ -1,112 +1,149 @@
 from __future__ import annotations
 
 import logging
-import os
-import pickle
 import sys
-from typing import Optional
+import time
+from contextlib import contextmanager, nullcontext
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Optional, Sequence
 
 import matplotlib.pyplot as plt
 import numpy as np
+import yaml
 from skimage.registration import phase_cross_correlation
+from tqdm import tqdm
 from typing_extensions import Self
 
 from instamatic import config
+from instamatic._typing import AnyPath
 from instamatic.calibrate.filenames import *
 from instamatic.calibrate.fit import fit_affine_transformation
+from instamatic.formats import read_tiff
 from instamatic.image_utils import autoscale, imgscale
 from instamatic.processing.find_holes import find_holes
-from instamatic.tools import find_beam_center, printer
+from instamatic.tools import find_beam_center
+from instamatic.utils.yaml import Numpy2DDumper
+
+if TYPE_CHECKING:
+    from instamatic.gui.videostream_processor import DeferredImageDraw, VideoStreamProcessor
+
 
 logger = logging.getLogger(__name__)
 
 
+Vector2 = np.ndarray  # numpy array with two float (or int) elements
+VectorNx2 = np.ndarray  # numpy array with N Vector2-s
+Matrix2x2 = np.ndarray  # numpy array of shape (2, 2) with float elements
+
+
+@dataclass
 class CalibBeamShift:
     """Simple class to hold the methods to perform transformations from one
-    setting to another based on calibration results."""
+    setting to another based on calibration results.
+
+    Throughout this class, the following two terms are used consistently:
+    - pixel: the (x, y) beam position in pixels as determined from camera image
+    - shift: the unitless (x, y) value pair reported by the BeamShift deflector
+    """
 
-    def __init__(self, transform, reference_shift, reference_pixel):
-        super().__init__()
-        self.transform = transform
-        self.reference_shift = reference_shift
-        self.reference_pixel = reference_pixel
-        self.has_data = False
+    transform: Matrix2x2
+    reference_pixel: Vector2
+    reference_shift: Vector2
+    pixels: Optional[VectorNx2] = field(default=None, repr=False)
+    shifts: Optional[VectorNx2] = field(default=None, repr=False)
+    images: Optional[list[np.ndarray]] = field(default=None, repr=False)
 
     def __repr__(self):
         return f'CalibBeamShift(transform=\n{self.transform},\n   reference_shift=\n{self.reference_shift},\n   reference_pixel=\n{self.reference_pixel})'
 
-    def beamshift_to_pixelcoord(self, beamshift):
+    def beamshift_to_pixelcoord(self, beamshift: Sequence[float, float]) -> Vector2:
         """Converts from beamshift x,y to pixel coordinates."""
         r_i = np.linalg.inv(self.transform)
-        pixelcoord = np.dot(self.reference_shift - beamshift, r_i) + self.reference_pixel
-        return pixelcoord
+        return np.dot(self.reference_shift - np.array(beamshift), r_i) + self.reference_pixel
 
-    def pixelcoord_to_beamshift(self, pixelcoord) -> np.ndarray:
+    def pixelcoord_to_beamshift(self, pixelcoord: Sequence[float, float]) -> Vector2:
         """Converts from pixel coordinates to beamshift x,y."""
-        r = self.transform
-        beamshift = self.reference_shift - np.dot(pixelcoord - self.reference_pixel, r)
-        return beamshift
+        pc = np.array(pixelcoord)
+        return self.reference_shift - np.dot(pc - self.reference_pixel, self.transform)
 
     @classmethod
-    def from_data(cls, shifts, beampos, reference_shift, reference_pixel, header=None) -> Self:
-        fit_result = fit_affine_transformation(shifts, beampos)
-        r = fit_result.r
-        t = fit_result.t
-
-        c = cls(transform=r, reference_shift=reference_shift, reference_pixel=reference_pixel)
-        c.data_shifts = shifts
-        c.data_beampos = beampos
-        c.has_data = True
-        c.header = header
-
-        return c
+    def from_data(
+        cls,
+        pixels: VectorNx2,
+        shifts: VectorNx2,
+        reference_pixel: Vector2,
+        reference_shift: Vector2,
+        images: Optional[list[np.ndarray]] = None,
+    ) -> Self:
+        return cls(
+            transform=fit_affine_transformation(pixels, shifts).r,
+            reference_pixel=reference_pixel,
+            reference_shift=reference_shift,
+            pixels=pixels,
+            shifts=shifts,
+            images=images,
+        )
 
     @classmethod
-    def from_file(cls, fn=CALIB_BEAMSHIFT) -> Self:
+    def from_file(cls, fn: AnyPath = CALIB_BEAMSHIFT) -> Self:
         """Read calibration from file."""
-        import pickle
-
         try:
-            return pickle.load(open(fn, 'rb'))
+            with open(Path(fn), 'r') as yaml_file:
+                return cls(**{k: np.array(v) for k, v in yaml.safe_load(yaml_file).items()})
         except OSError as e:
             prog = 'instamatic.calibrate_beamshift'
             raise OSError(f'{e.strerror}: {fn}. Please run {prog} first.')
 
     @classmethod
-    def live(cls, ctrl, outdir='.') -> Self:
+    def live(
+        cls, ctrl, outdir: AnyPath = '.', vsp: Optional[VideoStreamProcessor] = None
+    ) -> Self:
         while True:
             c = calibrate_beamshift(ctrl=ctrl, save_images=True, outdir=outdir)
-            if input(' >> Accept? [y/n] ') == 'y':
-                return c
+            with c.annotate_videostream(vsp) if vsp else nullcontext():
+                if input(' >> Accept? [y/n] ') == 'y':
+                    return c
 
-    def to_file(self, fn=CALIB_BEAMSHIFT, outdir='.'):
+    def to_file(self, fn: AnyPath = CALIB_BEAMSHIFT, outdir: AnyPath = '.') -> None:
         """Save calibration to file."""
-        fout = os.path.join(outdir, fn)
-        pickle.dump(self, open(fout, 'wb'))
-
-    def plot(self, to_file=None, outdir=''):
-        if not self.has_data:
-            return
-
-        if to_file:
-            to_file = f'calib_{beamshift}.png'
-
-        beampos = self.data_beampos
-        shifts = self.data_shifts
-
-        r_i = np.linalg.inv(self.transform)
-        beampos_ = np.dot(beampos, r_i)
-
-        plt.scatter(*shifts.T, marker='>', label='Observed pixel shifts')
-        plt.scatter(*beampos_.T, marker='<', label='Positions in pixel coords')
+        yaml_path = Path(outdir) / fn
+        yaml_dict = asdict(self)  # type: ignore[arg-type]
+        yaml_dict = {k: v.tolist() for k, v in yaml_dict.items() if k != 'images'}
+        with open(yaml_path, 'w') as yaml_file:
+            yaml.dump(yaml_dict, yaml_file, Dumper=Numpy2DDumper, default_flow_style=None)
+
+    def plot(self, to_file: Optional[AnyPath] = None):
+        """Assuming the data is present, plot the data."""
+        shifts = np.dot(self.shifts, np.linalg.inv(self.transform))
+        plt.scatter(*self.pixels.T, marker='>', label='Observed pixel shifts')
+        plt.scatter(*shifts.T, marker='<', label='Reconstructed pixel shifts')
         plt.legend()
         plt.title('BeamShift vs. Direct beam position (Imaging)')
         if to_file:
-            plt.savefig(os.path.join(outdir, to_file))
+            plt.savefig(Path(to_file) / 'calib_beamshift.png')
             plt.close()
         else:
             plt.show()
 
+    @contextmanager
+    def annotate_videostream(self, vsp: Optional[VideoStreamProcessor] = None) -> None:
+        shifts = np.dot(self.shifts, np.linalg.inv(self.transform))
+        ins: list[DeferredImageDraw.Instruction] = []
+
+        vsp.temporary_frame = np.max(self.images, axis=0)
+        print('Determined (blue) vs calibrated (orange) beam positions:')
+        for p, s in zip(self.pixels, shifts):
+            p = (p + self.reference_pixel)[::-1]  # xy coords inverted for plot
+            s = (s + self.reference_pixel)[::-1]  # xy coords inverted for plot
+            ins.append(vsp.draw.circle(p, radius=3, fill='blue'))
+            ins.append(vsp.draw.circle(s, radius=3, fill='orange'))
+        ins.append(vsp.draw.circle(self.reference_pixel[::-1], radius=3, fill='black'))
+        yield
+        vsp.temporary_frame = None
+        for i in ins:
+            vsp.draw.instructions.remove(i)
+
     def center(self, ctrl) -> Optional[np.ndarray]:
         """Return beamshift values to center the beam in the frame."""
         pixel_center = [val / 2.0 for val in ctrl.cam.get_image_dimensions()]
@@ -119,8 +156,13 @@ def center(self, ctrl) -> Optional[np.ndarray]:
 
 
 def calibrate_beamshift_live(
-    ctrl, gridsize=None, stepsize=None, save_images=False, outdir='.', **kwargs
-):
+    ctrl,
+    gridsize: Optional[int] = None,
+    stepsize: Optional[float] = None,
+    save_images: bool = False,
+    outdir: AnyPath = '.',
+    **kwargs,
+) -> CalibBeamShift:
     """Calibrate pixel->beamshift coordinates live on the microscope.
 
     ctrl: instance of `TEMController`
@@ -141,87 +183,54 @@ def calibrate_beamshift_live(
     """
     exposure = kwargs.get('exposure', ctrl.cam.default_exposure)
     binsize = kwargs.get('binsize', ctrl.cam.default_binsize)
+    gridsize = gridsize or config.camera.calib_beamshift.get('gridsize', 5)
+    stepsize = stepsize or config.camera.calib_beamshift.get('stepsize', 250)
+    outfile = Path(outdir) / 'calib_beamshift_center' if save_images else None
+    kwargs = {'exposure': exposure, 'binsize': binsize, 'out': outfile}
 
-    if not gridsize:
-        gridsize = config.camera.calib_beamshift.get('gridsize', 5)
-    if not stepsize:
-        stepsize = config.camera.calib_beamshift.get('stepsize', 250)
-
-    img_cent, h_cent = ctrl.get_image(
-        exposure=exposure, binsize=binsize, comment='Beam in center of image'
-    )
+    comment = 'Beam in the center of the image'
+    img_cent, h_cent = ctrl.get_image(comment=comment, **kwargs)
     x_cent, y_cent = beamshift_cent = np.array(h_cent['BeamShift'])
 
-    magnification = h_cent['Magnification']
-    stepsize = 2500.0 / magnification * stepsize
+    stepsize = 2500.0 / h_cent['Magnification'] * stepsize
 
     print(f'Gridsize: {gridsize} | Stepsize: {stepsize:.2f}')
 
     img_cent, scale = autoscale(img_cent)
-
-    outfile = os.path.join(outdir, 'calib_beamcenter') if save_images else None
-
     pixel_cent = find_beam_center(img_cent) * binsize / scale
 
     print('Beamshift: x={} | y={}'.format(*beamshift_cent))
     print('Pixel: x={} | y={}'.format(*pixel_cent))
 
-    shifts = []
-    beampos = []
-
-    n = int((gridsize - 1) / 2)  # number of points = n*(n+1)
-    x_grid, y_grid = np.meshgrid(
-        np.arange(-n, n + 1) * stepsize, np.arange(-n, n + 1) * stepsize
-    )
-    tot = gridsize * gridsize
+    images, pixels, shifts = [], [], []
+    dx_dy = ((np.indices((gridsize, gridsize)) - gridsize // 2) * stepsize).reshape(2, -1).T
 
-    i = 0
-    for dx, dy in np.stack([x_grid, y_grid]).reshape(2, -1).T:
+    progress_bar = tqdm(dx_dy, desc='Beamshift calibration')
+    for i, (dx, dy) in enumerate(progress_bar):
         ctrl.beamshift.set(x=float(x_cent + dx), y=float(y_cent + dy))
+        progress_bar.set_postfix_str(ctrl.beamshift)
+        time.sleep(config.camera.calib_beamshift.get('delay', 0.0))
 
-        printer(f'Position: {i + 1}/{tot}: {ctrl.beamshift}')
-
-        outfile = os.path.join(outdir, f'calib_beamshift_{i:04d}') if save_images else None
-
+        kwargs['out'] = Path(outdir) / f'calib_beamshift_{i:04d}' if save_images else None
         comment = f'Calib image {i}: dx={dx} - dy={dy}'
-        img, h = ctrl.get_image(
-            exposure=exposure,
-            binsize=binsize,
-            out=outfile,
-            comment=comment,
-            header_keys=('BeamShift',),
-        )
+        img, h = ctrl.get_image(comment=comment, header_keys=('BeamShift',), **kwargs)
         img = imgscale(img, scale)
 
-        shift, error, phasediff = phase_cross_correlation(img_cent, img, upsample_factor=10)
-
-        beamshift = np.array(h['BeamShift'])
-        beampos.append(beamshift)
-        shifts.append(shift)
-
-        i += 1
+        images.append(img)
+        pixels.append(phase_cross_correlation(img_cent, img, upsample_factor=10)[0])
+        shifts.append(np.array(h['BeamShift']))
 
     print('')
-    # print "\nReset to center"
-
     ctrl.beamshift.set(*(float(_) for _ in beamshift_cent))
 
-    # correct for binsize, store in binsize=1
-    shifts = np.array(shifts) * binsize / scale
-    beampos = np.array(beampos) - np.array(beamshift_cent)
-
+    # normalize to binsize = 1 and 512-pixel image scale before initializing
     c = CalibBeamShift.from_data(
-        shifts,
-        beampos,
-        reference_shift=beamshift_cent,
+        np.array(pixels) * binsize / scale,
+        np.array(shifts) - beamshift_cent,
         reference_pixel=pixel_cent,
-        header=h_cent,
+        reference_shift=beamshift_cent,
+        images=images,
     )
-
-    # Calling c.plot with videostream crashes program
-    # if not hasattr(ctrl.cam, "VideoLoop"):
-    #     c.plot()
-
     return c
 
 
@@ -239,7 +248,7 @@ def calibrate_beamshift_from_image_fn(center_fn, other_fn):
     print()
     print('Center:', center_fn)
 
-    img_cent, h_cent = load_img(center_fn)
+    img_cent, h_cent = read_tiff(center_fn)
     beamshift_cent = np.array(h_cent['BeamShift'])
 
     img_cent, scale = autoscale(img_cent, maxdim=512)
@@ -252,11 +261,12 @@ def calibrate_beamshift_from_image_fn(center_fn, other_fn):
     print('Beamshift: x={} | y={}'.format(*beamshift_cent))
     print('Pixel: x={:.2f} | y={:.2f}'.format(*pixel_cent))
 
+    images = []
     shifts = []
     beampos = []
 
     for fn in other_fn:
-        img, h = load_img(fn)
+        img, h = read_tiff(fn)
         img = imgscale(img, scale)
 
         beamshift = np.array(h['BeamShift'])
@@ -266,6 +276,7 @@ def calibrate_beamshift_from_image_fn(center_fn, other_fn):
 
         shift, error, phasediff = phase_cross_correlation(img_cent, img, upsample_factor=10)
 
+        images.append(img)
         beampos.append(beamshift)
         shifts.append(shift)
 
@@ -276,9 +287,9 @@ def calibrate_beamshift_from_image_fn(center_fn, other_fn):
     c = CalibBeamShift.from_data(
         shifts,
         beampos,
-        reference_shift=beamshift_cent,
         reference_pixel=pixel_cent,
-        header=h_cent,
+        reference_shift=beamshift_cent,
+        images=images,
     )
     c.plot()
 

src/instamatic/calibrate/filenames.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 CALIB_STAGE_LOWMAG = 'calib_stage_lowmag.pickle'
-CALIB_BEAMSHIFT = 'calib_beamshift.pickle'
+CALIB_BEAMSHIFT = 'calib_beamshift.yaml'
 CALIB_BRIGHTNESS = 'calib_brightness.pickle'
 CALIB_DIFFSHIFT = 'calib_diffshift.pickle'
 CALIB_DIRECTBEAM = 'calib_directbeam.pickle'