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.