Generalize beam shift deflector and stage rotating speed interface

src/instamatic/calibrate/calibrate_beamshift.py

@@ -4,20 +4,20 @@
 import os
 import pickle
 import sys
+from typing import Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
 from skimage.registration import phase_cross_correlation
 from typing_extensions import Self
 
 from instamatic import config
+from instamatic.calibrate.filenames import *
+from instamatic.calibrate.fit import fit_affine_transformation
 from instamatic.image_utils import autoscale, imgscale
 from instamatic.processing.find_holes import find_holes
 from instamatic.tools import find_beam_center, printer
 
-from .filenames import *
-from .fit import fit_affine_transformation
-
 logger = logging.getLogger(__name__)
 
 
@@ -41,11 +41,11 @@ def beamshift_to_pixelcoord(self, beamshift):
         pixelcoord = np.dot(self.reference_shift - beamshift, r_i) + self.reference_pixel
         return pixelcoord
 
-    def pixelcoord_to_beamshift(self, pixelcoord):
+    def pixelcoord_to_beamshift(self, pixelcoord) -> np.ndarray:
         """Converts from pixel coordinates to beamshift x,y."""
         r = self.transform
         beamshift = self.reference_shift - np.dot(pixelcoord - self.reference_pixel, r)
-        return beamshift.astype(int)
+        return beamshift
 
     @classmethod
     def from_data(cls, shifts, beampos, reference_shift, reference_pixel, header=None) -> Self:
@@ -107,13 +107,13 @@ def plot(self, to_file=None, outdir=''):
         else:
             plt.show()
 
-    def center(self, ctrl):
+    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()]
 
         beamshift = self.pixelcoord_to_beamshift(pixel_center)
         if ctrl:
-            ctrl.beamshift.set(*beamshift)
+            ctrl.beamshift.set(*(float(b) for b in beamshift))
         else:
             return beamshift
 
@@ -177,11 +177,11 @@ def calibrate_beamshift_live(
 
     i = 0
     for dx, dy in np.stack([x_grid, y_grid]).reshape(2, -1).T:
-        ctrl.beamshift.set(x=x_cent + dx, y=y_cent + dy)
+        ctrl.beamshift.set(x=float(x_cent + dx), y=float(y_cent + dy))
 
         printer(f'Position: {i + 1}/{tot}: {ctrl.beamshift}')
 
-        outfile = os.path.join(outdir, 'calib_beamshift_{i:04d}') if save_images else None
+        outfile = os.path.join(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(
@@ -204,7 +204,7 @@ def calibrate_beamshift_live(
     print('')
     # print "\nReset to center"
 
-    ctrl.beamshift.set(*beamshift_cent)
+    ctrl.beamshift.set(*(float(_) for _ in beamshift_cent))
 
     # correct for binsize, store in binsize=1
     shifts = np.array(shifts) * binsize / scale

src/instamatic/formats/__init__.py

@@ -51,6 +51,7 @@ def write_tiff(fname: str, data, header: dict = None):
         header = ''
 
     fname = Path(fname).with_suffix('.tiff')
+    fname.parent.mkdir(parents=True, exist_ok=True)
 
     with tifffile.TiffWriter(fname) as f:
         f.write(data=data, software='instamatic', description=header)

src/instamatic/microscope/base.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 from instamatic._typing import float_deg, int_nm
 from instamatic.microscope.utils import StagePositionTuple
 
 
 class MicroscopeBase(ABC):
     @abstractmethod
-    def getBeamShift(self) -> Tuple[int, int]:
+    def getBeamShift(self) -> Tuple[Union[float, int], Union[float, int]]:
         pass
 
     @abstractmethod
@@ -113,7 +113,7 @@ def setBeamBlank(self, mode: bool) -> None:
         pass
 
     @abstractmethod
-    def setBeamShift(self, x: int, y: int) -> None:
+    def setBeamShift(self, x: Union[float, int], y: Union[float, int]) -> None:
         pass
 
     @abstractmethod

src/instamatic/microscope/components/deflectors.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
 from collections import namedtuple
-from typing import Tuple
+from typing import Tuple, Union
 
 from instamatic.microscope.base import MicroscopeBase
 
 DeflectorTuple = namedtuple('DeflectorTuple', ['x', 'y'])
+Number = Union[int, float]
 
 
 class Deflector:
@@ -15,14 +16,14 @@ class Deflector:
     functions.
     """
 
-    def __init__(self, tem: MicroscopeBase):
+    def __init__(self, tem: MicroscopeBase) -> None:
         super().__init__()
         self._tem = tem
         self._getter = None
         self._setter = None
         self.key = 'def'
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         x, y = self.get()
         return f'{self.name}(x={x}, y={y})'
 
@@ -31,45 +32,45 @@ def name(self) -> str:
         """Return name of the deflector."""
         return self.__class__.__name__
 
-    def set(self, x: int, y: int):
+    def set(self, x: Number, y: Number) -> None:
         """Set the X and Y values of the deflector."""
         self._setter(x, y)
 
-    def get(self) -> Tuple[int, int]:
+    def get(self) -> Tuple[Number, Number]:
         """Get X and Y values of the deflector."""
         return DeflectorTuple(*self._getter())
 
     @property
-    def x(self) -> int:
+    def x(self) -> Number:
         """Get/set X value."""
         x, y = self.get()
         return x
 
     @x.setter
-    def x(self, value: int):
+    def x(self, value: Number) -> None:
         self.set(value, self.y)
 
     @property
-    def y(self) -> int:
+    def y(self) -> Number:
         """Get/set Y value."""
         x, y = self.get()
         return y
 
     @y.setter
-    def y(self, value: int):
+    def y(self, value: Number) -> None:
         self.set(self.x, value)
 
     @property
-    def xy(self) -> Tuple[int, int]:
+    def xy(self) -> Tuple[Number, Number]:
         """Get/set x and y values as a tuple."""
         return self.get()
 
     @xy.setter
-    def xy(self, values: Tuple[int, int]):
+    def xy(self, values: Tuple[Number, Number]) -> None:
         x, y = values
         self.set(x=x, y=y)
 
-    def neutral(self):
+    def neutral(self) -> None:
         """Return deflector to stored neutral values."""
         self._tem.setNeutral(self.key)
 

src/instamatic/microscope/components/stage.py

@@ -2,14 +2,16 @@
 
 import time
 from contextlib import contextmanager
-from typing import Optional, Tuple
+from typing import Generator, Optional, Tuple, Union
 
 import numpy as np
 
 from instamatic._typing import float_deg, int_nm
 from instamatic.microscope.base import MicroscopeBase
 from instamatic.microscope.utils import StagePositionTuple
 
+Number = Union[int, float]
+
 
 class Stage:
     """Stage control."""
@@ -74,7 +76,7 @@ def set_with_speed(
             speed=speed,
         )
 
-    def set_rotation_speed(self, speed=1) -> None:
+    def set_rotation_speed(self, speed: Union[float, int] = 1) -> None:
         """Sets the stage (rotation) movement speed on the TEM."""
         self._tem.setRotationSpeed(value=speed)
 
@@ -83,19 +85,19 @@ def set_a_with_speed(self, a: float, speed: int, wait: bool = False):
 
         wait: bool, block until stage movement is complete.
         """
-        with self.rotating_speed(speed):
+        with self.rotation_speed(speed):
             self.set(a=a, wait=False)
         # Do not wait on `set` to return to normal rotation speed quickly
         if wait:
             self.wait()
 
     @contextmanager
-    def rotating_speed(self, speed: int):
+    def rotation_speed(self, speed: Number) -> Generator[None, None, None]:
         """Context manager that sets the rotation speed for the duration of the
-        `with` statement (JEOL only).
+        `with` statement (JEOL, Tecnai only).
 
         Usage:
-            with ctrl.stage.rotating_speed(1):
+            with ctrl.stage.rotation_speed(1):
                 ctrl.stage.a = 40.0
         """
         try:
@@ -145,6 +147,10 @@ def xy(self, values: Tuple[int_nm, int_nm]) -> None:
         x, y = values
         self.set(x=x, y=y, wait=self._wait)
 
+    def get_rotation_speed(self) -> Number:
+        """Gets the stage (rotation) movement speed on the TEM."""
+        return self._tem.getRotationSpeed()
+
     def move_in_projection(self, delta_x: int_nm, delta_y: int_nm) -> None:
         r"""Y and z are always perpendicular to the sample stage. To achieve the
         movement in the projection plane instead, x and y should be broken down
@@ -218,7 +224,7 @@ def wait(self) -> None:
         self._tem.waitForStage()
 
     @contextmanager
-    def no_wait(self):
+    def no_wait(self) -> Generator[None, None, None]:
         """Context manager that prevents blocking stage position calls on
         properties.