Add detector flips & rotation to metadata

Author: hexane360

phaser/hooks/__init__.py

@@ -33,6 +33,7 @@ class LoadEmpadProps(Dataclass):
     diff_step: t.Optional[float] = None
     kv: t.Optional[float] = None
     adu: t.Optional[float] = None
+    det_flips: t.Optional[t.Tuple[bool, bool, bool]] = None
 
 
 class RawDataHook(Hook[None, RawData]):

phaser/hooks/io/empad.py

@@ -21,15 +21,18 @@ def load_empad(args: None, props: LoadEmpadProps) -> RawData:
     if path.suffix.lower() == '.json':  # load as metadata
         meta = EmpadMetadata.from_json(path)
         assert meta.path is not None
-
         path = meta.path / meta.raw_filename
 
+        if meta.empad_version is not None and meta.empad_version > 1:
+            raise ValueError("EMPAD v2 import not currently supported (open a Github issue!)")
+
         voltage = props.kv * 1e3 if props.kv is not None else meta.voltage
         diff_step = props.diff_step or meta.diff_step
         # [x, y] -> [y, x]
         scan_shape = t.cast(t.Tuple[int, int], tuple(reversed(meta.scan_shape)))
         adu = props.adu or meta.adu
         needs_scale = not meta.is_simulated()
+        det_flips = props.det_flips or meta.det_flips
 
         probe_hook = {
             'type': 'focused',
@@ -43,7 +46,7 @@ def load_empad(args: None, props: LoadEmpadProps) -> RawData:
             'shape': scan_shape,
             'step_size': tuple(s*1e10 for s in reversed(meta.scan_step)),  # m to A
             'affine': meta.scan_correction[::-1, ::-1] if meta.scan_correction is not None else None,
-            'rotation': -meta.scan_rotation,
+            'rotation': meta.det_rotation - meta.scan_rotation,
         }
 
         #TODO: add tilt to metafile
@@ -55,6 +58,7 @@ def load_empad(args: None, props: LoadEmpadProps) -> RawData:
         probe_hook = scan_hook = tilt_hook = None
         adu = None
         needs_scale = False
+        det_flips = props.det_flips
 
     if voltage is None:
         raise ValueError("'kv'/'voltage' must be specified by metadata or passed to 'raw_data'")
@@ -66,7 +70,8 @@ def load_empad(args: None, props: LoadEmpadProps) -> RawData:
     if not path.exists():
         raise ValueError(f"Couldn't find raw data at path {path}")
 
-    patterns = numpy.fft.ifftshift(load_4d(path, scan_shape, memmap=True), axes=(-1, -2))
+    logging.info(f"Loading with detector flips: {list(map(int, det_flips or (True, False, False)))} [y, x, transpose]")
+    patterns = numpy.fft.ifftshift(load_4d(path, scan_shape, memmap=True, flips=det_flips), axes=(-1, -2))
 
     if needs_scale:
         if adu is None:

phaser/io/empad.py

@@ -12,6 +12,8 @@
 from typing_extensions import Self
 
 from phaser.types import IsVersion
+from phaser.utils.image import apply_flips
+from phaser.utils.num import to_numpy
 
 
 def _get_dir(f: pane.io.FileOrPath) -> t.Optional[Path]:
@@ -45,9 +47,20 @@ def __post_init__(self):
     version: t.Annotated[str, IsVersion(exactly="2.0")] = "2.0"
     """Metadata version"""
 
+    empad_version: t.Optional[int] = None
+    """Empad version used. Defaults to v1 if not specified."""
+
     raw_filename: str
     """Raw 4DSTEM data filename, relative to metadata location."""
 
+    det_flips: t.Optional[t.Tuple[bool, bool, bool]] = None
+    """
+    Flips to apply to the raw diffraction patterns, (flip_y, flip_x, transpose).
+    Defaults to `(True, False, False)` (appears to be the most common orientation).
+    """
+    det_rotation: float = 0.0
+    """Detector rotation (degrees)."""
+
     orig_path: t.Optional[Path] = None
     """Original path to experimental folder."""
 
@@ -76,7 +89,7 @@ def __post_init__(self):
     diff_step: t.Optional[float] = None
     """Diffraction pixel size (mrad/px)."""
 
-    scan_rotation: float
+    scan_rotation: float = 0.0
     """Scan rotation (degrees)."""
     scan_shape: t.Tuple[int, int]
     """Scan shape (x, y)."""
@@ -112,8 +125,8 @@ def is_simulated(self) -> bool:
         return self.file_type == "pyMultislicer_metadata"
 
 
-def load_4d(path: t.Union[str, Path], scan_shape: t.Optional[t.Tuple[int, int]] = None,
-            memmap: bool = False) -> NDArray[numpy.float32]:
+def load_4d(path: t.Union[str, Path], scan_shape: t.Optional[t.Tuple[int, int]] = None, *,
+            memmap: bool = False, flips: t.Optional[t.Tuple[bool, bool, bool]] = None) -> NDArray[numpy.float32]:
     """
     Load a raw EMPAD dataset into memory.
 
@@ -126,6 +139,8 @@ def load_4d(path: t.Union[str, Path], scan_shape: t.Optional[t.Tuple[int, int]]
      - `path`: Path to file to load
      - `scan_shape`: Scan shape of dataset. Will be inferred from the filename if not specified.
      - `memmap`: If specified, memmap the file as opposed to loading it eagerly.
+     - `flips`: Flips to apply to the diffraction patterns, `(flip_y, flip_x, transpose)`.
+       Defaults to `(True, False, False)` (appears to be the most common orientation).
 
     Returns a numpy array (or `numpy.memmap`)
     """
@@ -148,26 +163,28 @@ def load_4d(path: t.Union[str, Path], scan_shape: t.Optional[t.Tuple[int, int]]
     if not a.size % (130*128) == 0:
         raise ValueError(f"File not divisible by 130x128 (size={a.size}).")
     a.shape = (-1, 130, 128)
-    #a = a[:, :128, :]
 
     if a.shape[0] != n_x * n_y:
         raise ValueError(f"Got {a.shape[0]} probes, expected {n_x}x{n_y} = {n_x * n_y}.")
     a.shape = (n_y, n_x, *a.shape[1:])
-    a = a[..., 127::-1, :]  # flip reciprocal y space, crop junk rows
 
-    return a
+    a = a[..., :128, :]  # crop junk rows
+    return apply_flips(a, flips or (True, False, False))  # defaults to typical EMPAD orientation
 
 
 @t.overload
-def save_4d(arr: NDArray[numpy.float32], *, path: t.Union[str, Path], folder: None = None, name: None = None):
+def save_4d(arr: NDArray[numpy.float32], *, path: t.Union[str, Path], folder: None = None, name: None = None,
+            flips: t.Optional[t.Tuple[bool, bool, bool]] = None):
     ...
 
 @t.overload
-def save_4d(arr: NDArray[numpy.float32], *, path: None = None, folder: t.Union[str, Path], name: t.Optional[str] = None):
+def save_4d(arr: NDArray[numpy.float32], *, path: None = None, folder: t.Union[str, Path], name: t.Optional[str] = None,
+            flips: t.Optional[t.Tuple[bool, bool, bool]] = None):
     ...
 
 def save_4d(arr: NDArray[numpy.float32], *, path: t.Union[str, Path, None] = None,
-            folder: t.Union[str, Path, None] = None, name: t.Optional[str] = None): #):
+            folder: t.Union[str, Path, None] = None, name: t.Optional[str] = None,
+            flips: t.Optional[t.Tuple[bool, bool, bool]] = None):
     """
     Save a raw EMPAD dataset.
 
@@ -183,6 +200,8 @@ def save_4d(arr: NDArray[numpy.float32], *, path: t.Union[str, Path, None] = Non
      - `folder`: Folder to save dataset inside.
      - `name`: When `folder` is specified, format to use to determine filename. Defaults to `"scan_x{x}_y{y}.raw"`.
        Will be formatted using the scan shape `{'x': n_x, 'y': n_y}`.
+     - `flips`: Flips to apply to the diffraction patterns, `(flip_y, flip_x, transpose)`.
+       Defaults to `(True, False, False)` (appears to be the most common orientation).
     """
 
     try:
@@ -206,7 +225,7 @@ def save_4d(arr: NDArray[numpy.float32], *, path: t.Union[str, Path, None] = Non
     out_shape[2] = 130  # dead rows
 
     out = numpy.zeros(out_shape, dtype=numpy.float32)
-    out[..., 127::-1, :] = arr.astype(numpy.float32)
+    out[..., 127::-1, :] = apply_flips(to_numpy(arr.astype(numpy.float32)), flips or (True, False, False))
 
     with open(path, 'wb') as f:
         out.tofile(f)

phaser/utils/image.py

@@ -14,6 +14,34 @@
 NumT = t.TypeVar('NumT', bound=numpy.number)
 
 
+def apply_flips(
+    data: NDArray[NumT], flips: t.Tuple[bool, bool, bool] = (False, False, False)
+) -> NDArray[NumT]:
+    """
+    Applies flips to `data` along the last two axes.
+    If specified, transpose is applied last (after X and Y flips).
+
+    Parameters:
+        data: Input data array.
+        flips: Tuple of three booleans `(flip_y, flip_x, transpose)`.
+
+    Returns: `data` with the specified flips applied.
+    """
+    if not any(flips):
+        return data
+
+    xp = get_array_module(data)
+
+    if flips[0]:
+        data = xp.flip(data, axis=-2)
+    if flips[1]:
+        data = xp.flip(data, axis=-1)
+    if flips[2]:
+        data = xp.moveaxis(data, -2, -1)
+
+    return data
+
+
 @t.overload
 def remove_linear_ramp(  # pyright: ignore[reportOverlappingOverload]
     data: NDArray[NumT], mask: t.Optional[NDArray[numpy.bool_]] = None
@@ -196,6 +224,7 @@ def affine_transform(
 
 
 __all__ = [
+    'apply_flips',
     'remove_linear_ramp', 'colorize_complex', 'scale_to_integral_type',
     'affine_transform', 'to_affine_matrix',
     'scale_matrix', 'rotation_matrix', 'translation_matrix',