Repository: instamatic-dev/instamatic
Pull request #124: Generalize beam shift deflector and stage rotating speed interface
Author: login: Baharis
Labels: name: bug
Created: 2025-05-05T15:01:41Z | Merged: 2025-05-09T12:34:36Z

### Context
Below I would like to suggest some small changes to the interface that may warrant an adequately short discussion:

Firstly, the `Deflector` class, parent for `BeamShift`, `BeamTilt`, `DiffShift` etc. is type-hinted to accept and return only instances of `int`. However, FEI microscopes do not store these as integers, but rather floats between -1 and 1 exclusive. This carries major implications because some functions, e.g. `CalibBeamShift.pixelcoord_to_beamshift` explicitly map its result to `int` using `beamshift.astype(int)`. rendering use on FEI machine impossible. So here I suggest modifying `microscope/component/deflectors:Deflector` base class hints to suggest that these values can be of either numeric type.

Secondly, although following the above `CalibBeamShift` should theoretically return shifts as `int`, it does not. The return type of `pixelcoord_to_beamshift` and the setting methods use as a container `np.ndarray` full of `np.int`. Consequently, if your server does not have numpy installed, `np.int` is an unknown type. This can be easily fixed by mapping all set values to float: `ctrl.beamshift.set(*(float(b) for b in beamshift))`. This completely fixes the issue for FEI, and should be benign for Jeol since its interface converts it back to `int` anyway, and the precision of converting int to float and back is likely good enough:
```python
>>> a = 123456789123456789123456789
>>> abs(a - int(float(a))) / a
1.793587896398062e-17
```

Thirdly, the `Stage` class currently has two (unused) methods to work with stage speed: the `set_rotation_speed` setter and the `rotating_speed` context. I suggest adding method `get_rotation_speed` - otherwise, one can get it only via a camel case parameter of a private attribute. I also suggest renaming `rotating_speed` to `rotation_speed` so that it matches the other two. I could not find this context used anywhere outside my own code, so I doubt anyone will mind.

Finally, I have introduced here some QOL improvements: made imports absolute (otherwise calibration could not be run as script, plus relative imports are good only in dynamically developed packages, [link](https://stackoverflow.com/questions/4209641/absolute-vs-explicit-relative-import-of-python-module)), added or improved type hinting, made it so that you don't need to manually make a folder to save a tiff file during calibration.

### Minor change

- `CalibBeamShift`: allow saving tiff in a new directory so that calibration can be run without manually creating the output directory for collected tiff files each time.

### Bugfix

- `CalibBeamShift`: some methods return or set `float` instead of `np.int` type, so that you can use this calibration of FEI;
- `calibrate_beamshift_live`: can now be run as a script thanks to making imports absolute;
- `calibrate_beamshift_live`: fix typo so that calibration files are not all saved under the same name `calib_beamshift_{i.tiff`.

### Code maintanance
- `Deflector` child classes are now properly type hinted to handle either `int` or `float` type values.



Diff:
diff --git a/src/instamatic/calibrate/calibrate_beamshift.py b/src/instamatic/calibrate/calibrate_beamshift.py
index 6e54f151..60b6dcb5 100644
--- a/src/instamatic/calibrate/calibrate_beamshift.py
+++ b/src/instamatic/calibrate/calibrate_beamshift.py
@@ -4,6 +4,7 @@
 import os
 import pickle
 import sys
+from typing import Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -11,13 +12,12 @@
 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
diff --git a/src/instamatic/formats/__init__.py b/src/instamatic/formats/__init__.py
index a66ee7cc..c3f930b8 100644
--- a/src/instamatic/formats/__init__.py
+++ b/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)
diff --git a/src/instamatic/microscope/base.py b/src/instamatic/microscope/base.py
index 1f826d2d..5926be86 100644
--- a/src/instamatic/microscope/base.py
+++ b/src/instamatic/microscope/base.py
@@ -1,7 +1,7 @@
 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
@@ -9,7 +9,7 @@
 
 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
diff --git a/src/instamatic/microscope/components/deflectors.py b/src/instamatic/microscope/components/deflectors.py
index b194ac1b..cf956d41 100644
--- a/src/instamatic/microscope/components/deflectors.py
+++ b/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)
 
diff --git a/src/instamatic/microscope/components/stage.py b/src/instamatic/microscope/components/stage.py
index 5c011cb0..72a5e554 100644
--- a/src/instamatic/microscope/components/stage.py
+++ b/src/instamatic/microscope/components/stage.py
@@ -2,7 +2,7 @@
 
 import time
 from contextlib import contextmanager
-from typing import Optional, Tuple
+from typing import Generator, Optional, Tuple, Union
 
 import numpy as np
 
@@ -10,6 +10,8 @@
 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.
 


Reviews:
user: login: stefsmeets
body: Looks good to me, I left a minor comment.
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-08T07:46:21Z
user: login: Baharis
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-08T12:47:06Z

Review comments:
diff_hunk: @@ -83,19 +83,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: Union[float, int]) -> Generator[None, None, None]:
path: src/instamatic/microscope/components/stage.py
user: login: stefsmeets
body: Use `Number` here as well?

```suggestion
    def rotation_speed(self, speed: Number) -> Generator[None, None, None]:
```
created_at: 2025-05-08T07:45:13Z
author_association: MEMBER
diff_hunk: @@ -145,6 +145,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) -> Union[float, int]:
path: src/instamatic/microscope/components/stage.py
user: login: stefsmeets
body: ```suggestion
    def get_rotation_speed(self) -> Number:
```
created_at: 2025-05-08T07:45:51Z
author_association: MEMBER
diff_hunk: @@ -145,6 +145,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) -> Union[float, int]:
path: src/instamatic/microscope/components/stage.py
user: login: Baharis
body: Can do though it will require adding declaration there as well. In the other place it was more needed since this patter repeated multiple times. 👍 
created_at: 2025-05-08T12:47:05Z
author_association: MEMBER

PR comments:
user: login: Baharis
created_at: 2025-05-05T15:05:26Z
body: @viljarjf Tagging you because you are the author of the `MicroscopeBase` that propagates the information that `getBeamShift` and `setBeamShift` operate on `int`s; these use `float` on FEI machines. I modified only hints for these two there, maybe I should actually generalize all of these. Also AFAIK you are an active Jeol machine user.
author_association: MEMBER
user: login: Baharis
created_at: 2025-05-08T14:16:19Z
body: Alright, done. If that is the case and there are no further comments, I will merge this tomorrow.
author_association: MEMBER