phaserEM 0.1__py3-none-any.whl

phaser/io/empad.py

@@ -0,0 +1,212 @@
+
+from pathlib import Path
+import re
+import typing as t
+
+import numpy
+import pane
+import pane.io
+from numpy.typing import NDArray
+from pane.annotations import shape
+from pane.convert import IntoConverterHandlers
+from typing_extensions import Self
+
+from phaser.types import IsVersion
+
+
+def _get_dir(f: pane.io.FileOrPath) -> t.Optional[Path]:
+    if isinstance(f, (str, Path)):
+        return Path(f).parent
+
+    name = getattr(f, 'name', None)
+    if name in (None, '<stdout>', '<stderr>'):
+        return None
+    path = Path(name)
+    return path.parent if path.exists() else None
+
+
+class EmpadMetadata(pane.PaneBase, frozen=False, kw_only=True, allow_extra=True):
+    file_type: t.Literal['pyMultislicer_metadata', 'empad_metadata'] = 'empad_metadata'
+
+    @classmethod
+    def from_json(cls, f: pane.io.FileOrPath, *,
+                  custom: t.Optional[IntoConverterHandlers] = None) -> Self:
+        path = _get_dir(f)
+        self = pane.io.from_json(f, cls, custom=custom)
+        object.__setattr__(self, 'path', path)
+        return self
+
+    def __post_init__(self):
+        object.__setattr__(self, 'path', None)
+
+    name: str
+    """Experiment name"""
+
+    version: t.Annotated[str, IsVersion(exactly="2.0")] = "2.0"
+    """Metadata version"""
+
+    raw_filename: str
+    """Raw 4DSTEM data filename, relative to metadata location."""
+
+    orig_path: t.Optional[Path] = None
+    """Original path to experimental folder."""
+
+    path: t.Optional[Path] = pane.field(init=False, exclude=True)
+    """Current path to experimental folder (based on metadata loading)"""
+
+    author: t.Optional[str] = None
+    """Author of dataset"""
+    time: t.Optional[str] = None
+    """Image acquisition time (RFC 2822 format)"""
+    time_unix: t.Optional[float] = None
+    """Image acquisition time (seconds since Unix epoch)"""
+    bg_unix: t.Optional[float] = None
+    """Background image acquisition time (seconds since Unix epoch)"""
+    has_bg: t.Optional[bool] = None
+    """Whether background image is valid"""
+
+    voltage: float
+    """Accelerating voltage (V)."""
+    conv_angle: t.Optional[float] = None
+    """Convergence angle (mrad)."""
+    defocus: t.Optional[float] = None
+    """Defocus (m). Positive is overfocus."""
+    camera_length: t.Optional[float] = None
+    """Camera length (m)."""
+    diff_step: t.Optional[float] = None
+    """Diffraction pixel size (mrad/px)."""
+
+    scan_rotation: float
+    """Scan rotation (degrees)."""
+    scan_shape: t.Tuple[int, int]
+    """Scan shape (x, y)."""
+    scan_fov: t.Tuple[float, float]
+    """Scan field of view (m)."""
+    scan_step: t.Tuple[float, float]
+    """Scan step (m/px)."""
+
+    exposure_time: t.Optional[float] = None
+    """Pixel exposure time (s)."""
+    post_exposure_time: t.Optional[float] = None
+    """Pixel post-exposure time (s)."""
+    beam_current: t.Optional[float] = None
+    """Approx. beam current (A)."""
+    adu: t.Optional[float] = None
+    """Single-electron intensity (data units)."""
+
+    scan_correction: t.Optional[t.Annotated[NDArray[numpy.floating], shape((2, 2))]] = None
+    """Scan correction matrix, [x', y'] = scan_correction @ [x, y]"""
+
+    scan_positions: t.Optional[t.List[t.Tuple[float, float]]] = None
+    """
+    Scan position override (m).
+    Should be specified as a 1d list of (x, y) positions, in scan order. `scan_correction` is applied to these positions (if present).
+    """
+
+    notes: t.Optional[str] = None
+
+    crop: t.Optional[t.Tuple[int, int, int, int]] = None
+    """Region scan is valid within, (min_y, max_y, min_x, max_x). Python-style slicing."""
+
+    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]:
+    """
+    Load a raw EMPAD dataset into memory.
+
+    The file is loaded so the dimensions are: (scan_y, scan_x, k_y, k_x), with y decreasing downwards.
+
+    Patterns are not fftshifted or normalized upon loading.
+
+    # Parameters
+
+     - `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.
+
+    Returns a numpy array (or `numpy.memmap`)
+    """
+    path = Path(path)
+
+    if scan_shape is None:
+        match = re.search(r"x(\d+)_y(\d+)", path.name)
+        if match:
+            n_x, n_y = map(int, (match[1], match[2]))
+        else:
+            raise ValueError(f"Unable to infer probe dimensions from name {path.name}")
+    else:
+        n_y, n_x = scan_shape
+
+    if memmap:
+        a = numpy.memmap(path, dtype=numpy.float32, mode='r')
+    else:
+        a = numpy.fromfile(path, dtype=numpy.float32)
+
+    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
+
+
+@t.overload
+def save_4d(arr: NDArray[numpy.float32], *, path: t.Union[str, Path], folder: None = None, name: None = 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: t.Union[str, Path, None] = None,
+            folder: t.Union[str, Path, None] = None, name: t.Optional[str] = None): #):
+    """
+    Save a raw EMPAD dataset.
+
+    Either `path` or `folder` can be specified. If `folder` is specified,
+    `name` will be used as a format string to determine the filename.
+    `path` and `folder` cannot be specified simultaneously.
+
+    Patterns are not fftshifted or normalized upon saving.
+
+    Parameters:
+     - `arr`: Array to save
+     - `path`: Path to save dataset to.
+     - `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}`.
+    """
+
+    try:
+        assert len(arr.shape) == 4
+        assert arr.shape[2:] == (128, 128)
+    except AssertionError as e:
+        raise ValueError("Invalid data format") from e
+
+    if folder is not None:
+        if path is not None:
+            raise ValueError("Cannot specify both 'path' and 'folder'")
+
+        n_y, n_x = arr.shape[:2]
+        path = Path(folder) / (name or "scan_x{x}_y{y}.raw").format(x=n_x, y=n_y)
+    elif path is not None:
+        path = Path(path)
+    else:
+        raise ValueError("Must specify either 'path' or 'folder'")
+
+    out_shape = list(arr.shape)
+    out_shape[2] = 130  # dead rows
+
+    out = numpy.zeros(out_shape, dtype=numpy.float32)
+    out[..., 127::-1, :] = arr.astype(numpy.float32)
+
+    with open(path, 'wb') as f:
+        out.tofile(f)

phaser/main.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+import sys
+import typing as t
+
+import click
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command('run')
+@click.argument('path', type=click.Path(exists=True, dir_okay=False))
+def run(path: t.Union[str, Path]):
+    from .plan import ReconsPlan
+    from .execute import execute_plan
+    plans = ReconsPlan.from_yaml_all(path)
+
+    for plan in plans:
+        execute_plan(plan)
+
+
+@cli.command('serve')
+@click.option('--host', type=str, default='localhost')
+@click.option('--port', type=int)
+@click.option('-v', '--verbose', count=True)
+def serve(host: str = 'localhost', port: t.Optional[int] = None, verbose: int = 0):
+    from phaser.web.server import server
+
+    if ':' in host:
+        (host, port_from_host) = host.rsplit(':', maxsplit=1)
+        try:
+            port_from_host = int(port_from_host)
+        except ValueError:
+            print(f"Invalid host '{host}:{port_from_host}'", file=sys.stderr)
+            sys.exit(1)
+
+        port = port or port_from_host
+
+    server.run(hostname=host, port=port, verbosity=verbose)
+
+
+@cli.command('validate')
+@click.argument('path', type=click.Path(allow_dash=True), default='-')
+@click.option('--json/--no-json', default=False)
+def validate(path: t.Union[str, Path], json: bool = False):
+    from contextlib import nullcontext
+    from .plan import ReconsPlan
+
+    try:
+        if path == '-':
+            file = nullcontext(sys.stdin)
+        else:
+            file = open(Path(path).expanduser(), 'r')
+
+        with file as file:
+            plans = ReconsPlan.from_yaml_all(file)
+    except Exception as e:
+        print(f"Validation failed:\n{e}", file=sys.stderr)
+
+        if json:
+            from json import dump
+            dump({'result': 'error', 'error': str(e)}, sys.stdout)
+            print()
+
+        sys.exit(1)
+
+    if len(plans) == 1:
+        print("Validation of plan successful!", file=sys.stderr)
+    else:
+        print(f"Validation of {len(plans)} plans successful!", file=sys.stderr)
+
+    if json:
+        from json import dump
+        dump({
+            'result': 'success',
+            'plans': [plan.into_data() for plan in plans],
+        }, sys.stdout)
+        print()
+
+
+@cli.command('worker')
+@click.argument('url', type=str, required=True)
+@click.option('--quiet/--loud', default=False)
+def worker(url: str, quiet: bool = False):
+    from phaser.web.worker import run_worker
+
+    run_worker(url, quiet=quiet)
+
+
+if __name__ == '__main__':
+    cli()

phaser/plan.py

@@ -0,0 +1,184 @@
+from pathlib import Path
+import typing as t
+
+from .types import Dataclass, Slices, BackendName, Flag, ReconsVars, IsVersion, EmptyDict
+from .hooks import RawDataHook, ProbeHook, ObjectHook, ScanHook, EngineHook, PostInitHook, PostLoadHook
+from .hooks.solver import NoiseModelHook, ConventionalSolverHook, PositionSolverHook, GradientSolverHook
+from .hooks.schedule import FlagLike, ScheduleLike
+from .hooks.regularization import IterConstraintHook, GroupConstraintHook, CostRegularizerHook
+
+
+SaveType: t.TypeAlias = t.Literal[
+    'probe', 'probe_mag', 'probe_recip', 'probe_recip_mag',
+    'object_phase_stack', 'object_phase_sum',
+    'object_mag_stack', 'object_mag_sum',
+]
+
+
+class InitPlan(Dataclass, kw_only=True):
+    state: t.Optional[Path] = None
+
+    scan: t.Union[EmptyDict, ScanHook, None] = None
+    probe: t.Union[EmptyDict, ProbeHook, None] = None
+    object: t.Optional[ObjectHook] = None  # ObjectHook('random')
+
+
+class SaveOptions(Dataclass, kw_only=True):
+    images: t.Tuple[SaveType, ...] = ('probe', 'object_phase_stack')
+    crop_roi: bool = True
+    unwrap_phase: bool = True
+    img_dtype: t.Literal['float', '8bit', '16bit', '32bit'] = '16bit'
+
+    out_dir: str = "{name}"
+    img_fmt: str = "{type}_iter{iter.total_iter}.tiff"
+    hdf5_fmt: str = "iter{iter.total_iter}.h5"
+
+
+class EnginePlan(Dataclass, kw_only=True):
+    sim_shape: t.Optional[t.Tuple[int, int]] = None
+    resize_method: t.Literal['pad_crop', 'resample'] = 'pad_crop'
+
+    probe_modes: int = 1
+    base_mode_power: float = 0.7
+    """Intensity to assign to the base mode when creating incoherent probe modes."""
+
+    bwlim_frac: t.Optional[float] = 2/3
+    obj_pad_px: float = 5.0
+
+    slices: t.Optional[Slices] = None
+
+    niter: int = 10
+    grouping: t.Optional[int] = None
+    compact: bool = False
+    shuffle_groups: t.Optional[FlagLike] = None
+    buffer_n_groups: int = 2
+
+    update_probe: FlagLike = True
+    update_object: FlagLike = True
+    update_positions: FlagLike = False
+
+    calc_error: FlagLike = Flag(every=1)
+    calc_error_fraction: float = 0.1
+
+    save: FlagLike = False
+    save_images: FlagLike = False
+    save_options: SaveOptions = SaveOptions()
+
+    send_every_group: bool = False
+
+
+class AmplitudeNoisePlan(Dataclass, kw_only=True):
+    gaussian_variance: float = 0.1
+    eps: float = 1.0e-3
+    offset: float = 0.0
+
+
+class AnscombeNoisePlan(AmplitudeNoisePlan, kw_only=True):
+    offset: float = 0.375
+
+
+class PoissonNoisePlan(AmplitudeNoisePlan, kw_only=True):
+    eps: float = 1.0e-3
+
+
+NoiseModelHook.known['amplitude'] = ('phaser.engines.common.noise_models:AmplitudeNoiseModel', AmplitudeNoisePlan)
+NoiseModelHook.known['anscombe'] = ('phaser.engines.common.noise_models:AnscombeNoiseModel', AnscombeNoisePlan)
+NoiseModelHook.known['poisson'] = ('phaser.engines.common.noise_models:PoissonNoiseModel', PoissonNoisePlan)
+
+
+class LSQMLSolverPlan(Dataclass, kw_only=True):
+    stochastic: bool = True
+
+    beta_object: ScheduleLike = 1.0
+    beta_probe: ScheduleLike = 1.0
+
+    illum_reg_object: ScheduleLike = 1e-2
+    illum_reg_probe: ScheduleLike = 1e-2
+
+    gamma: ScheduleLike = 1e-4
+
+
+class EPIESolverPlan(Dataclass, kw_only=True):
+    beta_object: ScheduleLike = 1.0
+    beta_probe: ScheduleLike = 1.0
+
+
+ConventionalSolverHook.known['lsqml'] = ('phaser.engines.conventional.solvers:LSQMLSolver', LSQMLSolverPlan)
+ConventionalSolverHook.known['epie'] = ('phaser.engines.conventional.solvers:EPIESolver', EPIESolverPlan)
+
+
+class ConventionalEnginePlan(EnginePlan, kw_only=True):
+    noise_model: NoiseModelHook
+    solver: ConventionalSolverHook
+    position_solver: t.Optional[PositionSolverHook] = None
+
+    group_constraints: t.List[GroupConstraintHook]
+    iter_constraints: t.List[IterConstraintHook]
+
+
+class GradientEnginePlan(EnginePlan):
+    noise_model: NoiseModelHook
+    solvers: t.Dict[ReconsVars, GradientSolverHook]
+
+    regularizers: t.List[CostRegularizerHook]
+    group_constraints: t.List[GroupConstraintHook]
+    iter_constraints: t.List[IterConstraintHook]
+
+
+class SGDSolverPlan(Dataclass, kw_only=True):
+    learning_rate: ScheduleLike
+    momentum: t.Optional[ScheduleLike] = None
+    nesterov: bool = True
+
+
+class AdamSolverPlan(Dataclass, kw_only=True):
+    learning_rate: ScheduleLike
+
+    b1: float = 0.9
+    b2: float = 0.999
+    eps: float = 1.0e-8
+    eps_root: float = 0.0
+
+    nesterov: bool = False
+
+
+class PolyakSGDSolverPlan(Dataclass, kw_only=True):
+    max_learning_rate: ScheduleLike
+    f_min: float
+    scaling: ScheduleLike = 1.0
+    eps: float = 0.0
+
+
+GradientSolverHook.known['sgd'] = ('phaser.engines.gradient.solvers:SGDSolver', SGDSolverPlan)
+GradientSolverHook.known['adam'] = ('phaser.engines.gradient.solvers:AdamSolver', AdamSolverPlan)
+GradientSolverHook.known['polyak_sgd'] = ('phaser.engines.gradient.solvers:PolyakSGDSolver', PolyakSGDSolverPlan)
+
+EngineHook.known['conventional'] = ('phaser.engines.conventional.run:run_engine', ConventionalEnginePlan)
+EngineHook.known['gradient'] = ('phaser.engines.gradient.run:run_engine', GradientEnginePlan)
+
+
+class ReconsPlan(Dataclass, kw_only=True):
+    file_type: t.Literal['phaser_plan'] = 'phaser_plan'
+    version: t.Annotated[str, IsVersion(exactly="1.0")] = "1.0"
+
+    name: str
+
+    backend: t.Optional[BackendName] = None
+    dtype: t.Literal['float32', 'float64'] = 'float32'
+
+    wavelength: t.Optional[float] = None
+
+    raw_data: RawDataHook
+
+    post_load: t.Sequence[PostLoadHook] = ()
+
+    init: InitPlan = InitPlan()
+
+    post_init: t.Sequence[PostInitHook] = ()
+
+    slices: t.Optional[Slices] = None
+
+    engines: t.List[EngineHook]
+    #engines: t.List[t.Annotated[t.Union[ConventionalEngine, GradientEngine], Tagged('type')]]
+
+