phaserEM 0.1__py3-none-any.whl

phaser/__main__.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python3
+
+if __name__ == '__main__':
+    from phaser.main import cli
+    cli()

phaser/engines/common/noise_models.py

@@ -0,0 +1,113 @@
+
+import typing as t
+
+import numpy
+from numpy.typing import NDArray
+
+from phaser.hooks.solver import NoiseModel
+from phaser.plan import AmplitudeNoisePlan, AnscombeNoisePlan, PoissonNoisePlan
+from phaser.utils.num import get_array_module, Float
+from phaser.state import ReconsState
+
+
+class AmplitudeNoiseModel(NoiseModel[None]):
+    @classmethod
+    def name(cls) -> str:
+        return "amplitude"
+
+    def __init__(self, args: None, props: AmplitudeNoisePlan):
+        self.offset: float = props.offset
+        self.gaussian_variance: float = props.gaussian_variance
+
+        self.var: float = 1 + self.gaussian_variance
+
+        self.eps: float = props.eps
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss(
+        self,
+        model_wave: NDArray[numpy.complexfloating],
+        model_intensity: NDArray[numpy.floating],
+        exp_patterns: NDArray[numpy.floating],
+        mask: NDArray[numpy.floating], 
+        state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(model_wave, model_intensity, exp_patterns, mask)
+        patterns = xp.maximum(exp_patterns, 0.0)
+
+        return ((
+            2. * xp.sum(mask * (
+                xp.sqrt(patterns + self.offset) - xp.sqrt(model_intensity + self.offset) - self.eps
+            )**2) / self.var).astype(exp_patterns.dtype),
+            state
+        )
+
+    def calc_wave_update(
+        self,
+        model_wave: NDArray[numpy.complexfloating],
+        model_intensity: NDArray[numpy.floating],
+        exp_patterns: NDArray[numpy.floating],
+        mask: NDArray[numpy.floating], 
+        state: None
+    ) -> t.Tuple[NDArray[numpy.complexfloating], None]:
+        xp = get_array_module(model_wave, model_intensity, exp_patterns, mask)
+        patterns = xp.maximum(exp_patterns, 0.0)
+
+        update = xp.sqrt(patterns + self.offset) / (xp.sqrt(model_intensity + self.offset) + self.eps) - 1.0
+        update *= mask # / self.var
+        #print(f"min update: {xp.min(update).get()}")
+        #print(f"max update: {xp.max(update).get()}")
+
+        return (update * model_wave, state)
+
+
+class AnscombeNoiseModel(AmplitudeNoiseModel):
+    @classmethod
+    def name(cls) -> str:
+        return "anscombe"
+
+    def __init__(self, args: None, props: AnscombeNoisePlan):
+        super().__init__(args, props)
+
+
+class PoissonNoiseModel(NoiseModel[None]):
+    @classmethod
+    def name(cls) -> str:
+        return "poisson"
+
+    def __init__(self, args: None, props: PoissonNoisePlan):
+        self.eps: float = props.eps
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss(
+        self,
+        model_wave: NDArray[numpy.complexfloating],
+        model_intensity: NDArray[numpy.floating],
+        exp_patterns: NDArray[numpy.floating],
+        mask: NDArray[numpy.floating], 
+        state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(model_wave, model_intensity, exp_patterns, mask)
+        patterns = xp.maximum(exp_patterns, 0.0)
+        #intensity - patterns * xp.log(intensity + self.offset)
+
+        loss = xp.sum(mask * (
+            #model_intensity - patterns * xp.log(model_intensity + self.eps)
+            model_intensity + patterns * (xp.log(patterns + self.eps) - xp.log(model_intensity + self.eps) - 1.0)
+            #patterns - (model_intensity + self.offset) * xp.log(patterns)
+        )).astype(exp_patterns.dtype)
+        return (loss, state)
+
+    def calc_wave_update(
+        self,
+        model_wave: NDArray[numpy.complexfloating],
+        model_intensity: NDArray[numpy.floating],
+        exp_patterns: NDArray[numpy.floating],
+        mask: NDArray[numpy.floating], 
+        state: None
+    ) -> t.Tuple[NDArray[numpy.complexfloating], None]:
+        raise NotImplementedError()

phaser/engines/common/output.py

@@ -0,0 +1,189 @@
+from functools import partial
+from pathlib import Path
+import typing as t
+
+import numpy
+from numpy.typing import NDArray
+import tifffile
+
+from phaser.utils.num import to_numpy, abs2, fft2, get_array_module
+from phaser.utils.image import remove_linear_ramp, colorize_complex, scale_to_integral_type
+from phaser.utils.io import tiff_write_opts, tiff_write_opts_recip
+from phaser.state import ReconsState
+from phaser.plan import SaveOptions
+
+
+def output_images(state: ReconsState, out_dir: Path, options: SaveOptions):
+    for ty in options.images:
+        if ty not in _SAVE_FUNCS:
+            raise ValueError(f"Unknown image type '{ty}'")
+
+        try:
+            out_name = options.img_fmt.format(
+                type=ty, iter=state.iter,
+            )
+            out_path = out_dir / out_name
+        except KeyError as e:
+            raise ValueError(f"Invalid format string in 'img_fmt' (unknown key {e})") from None
+        except Exception as e:
+            raise ValueError("Invalid format string in 'img_fmt'") from e
+
+        _SAVE_FUNCS[ty](state, out_path, options)
+
+
+def output_state(state: ReconsState, out_dir: Path, options: SaveOptions):
+    try:
+        out_name = options.hdf5_fmt.format(
+            iter=state.iter
+        )
+    except KeyError as e:
+        raise ValueError(f"Invalid format string in 'hdf5_fmt' (unknown key {e})") from None
+    except Exception as e:
+        raise ValueError("Invalid format string in 'hdf5_fmt'") from e
+
+    state.write_hdf5(out_dir / out_name)
+
+
+def _save_probe(state: ReconsState, out_path: Path, options: SaveOptions):
+    probe = to_numpy(state.probe.data)
+    write_opts = tiff_write_opts(state.probe.sampling, n_slices=probe.shape[0])
+
+    if options.img_dtype == 'float':
+        # save complex image
+        with tifffile.TiffWriter(out_path, ome=True) as w:
+            write_opts['metadata']['axes'] = 'CYX'
+            w.write(probe, **write_opts)
+        return
+
+    img = scale_to_integral_type(
+        colorize_complex(probe), options.img_dtype
+    )
+    with tifffile.TiffWriter(out_path, ome=True) as w:
+        write_opts['metadata']['axes'] = 'CYXS'
+        w.write(img, photometric='rgb', **write_opts)
+
+
+def _save_probe_mag(state: ReconsState, out_path: Path, options: SaveOptions):
+    probe_mag = abs2(state.probe.data)
+    write_opts = tiff_write_opts(state.probe.sampling, n_slices=probe_mag.shape[0])
+
+    if options.img_dtype != 'float':
+        probe_mag = scale_to_integral_type(to_numpy(probe_mag), options.img_dtype, min_range=0.2)
+
+    with tifffile.TiffWriter(out_path, ome=True) as w:
+        write_opts['metadata']['axes'] = 'CYX'
+        w.write(to_numpy(probe_mag), **write_opts)
+
+
+def _save_probe_recip(state: ReconsState, out_path: Path, options: SaveOptions):
+    xp = get_array_module(state.probe.data)
+    probe = to_numpy(xp.fft.fftshift(fft2(state.probe.data), axes=(-1, -2)))
+    write_opts = tiff_write_opts_recip(state.probe.sampling, n_slices=probe.shape[0])
+
+    if options.img_dtype == 'float':
+        # save complex image
+        with tifffile.TiffWriter(out_path, ome=True) as w:
+            write_opts['metadata']['axes'] = 'CYX'
+            w.write(probe, **write_opts)
+        return
+
+    img = scale_to_integral_type(
+        colorize_complex(probe), options.img_dtype
+    )
+    with tifffile.TiffWriter(out_path, ome=True) as w:
+        write_opts['metadata']['axes'] = 'CYXS'
+        w.write(img, photometric='rgb', **write_opts)
+
+
+def _save_probe_recip_mag(state: ReconsState, out_path: Path, options: SaveOptions):
+    xp = get_array_module(state.probe.data)
+    probe_mag = to_numpy(abs2(xp.fft.fftshift(fft2(state.probe.data), axes=(-1, -2))))
+    write_opts = tiff_write_opts_recip(state.probe.sampling, n_slices=probe_mag.shape[0])
+
+    if options.img_dtype != 'float':
+        probe_mag = scale_to_integral_type(probe_mag, options.img_dtype, min_range=0.2)
+
+    with tifffile.TiffWriter(out_path, ome=True) as w:
+        write_opts['metadata']['axes'] = 'CYX'
+        w.write(probe_mag, **write_opts)
+
+
+def _save_object_phase(state: ReconsState, out_path: Path, options: SaveOptions, stack: bool = False):
+    crop = options.crop_roi
+
+    xp = get_array_module(state.object.data)
+    obj_phase = xp.angle(state.object.data)
+
+    obj_sampling = state.object.sampling
+    write_opts = tiff_write_opts(
+        obj_sampling,
+        corner=obj_sampling.region_min if crop else None,
+        zs=state.object.zs() if stack else None,
+    )
+
+    if crop:
+        obj_phase = obj_phase[..., *state.object.sampling.get_region_crop()]
+        mask = xp.ones(obj_phase.shape[-2:], dtype=numpy.bool_)
+    else:
+        # include whole image, but only scale based on ROI
+        mask = state.object.sampling.get_region_mask(xp=xp)
+
+    if options.unwrap_phase:
+        obj_phase = xp.unwrap(xp.unwrap(obj_phase, axis=-1), axis=-2)
+
+    if not stack:
+        obj_phase = xp.sum(obj_phase, axis=0)
+
+    mask = to_numpy(mask)
+    obj_phase = remove_linear_ramp(to_numpy(obj_phase), mask)
+
+    if options.img_dtype != 'float':
+        obj_phase = scale_to_integral_type(obj_phase, options.img_dtype, mask)
+
+    with tifffile.TiffWriter(out_path, ome=True) as w:
+        write_opts['metadata']['axes'] = 'ZYX' if stack else 'YX'
+        w.write(obj_phase, **write_opts)
+
+
+def _save_object_mag(state: ReconsState, out_path: Path, options: SaveOptions, stack: bool = False):
+    crop = options.crop_roi
+
+    obj_sampling = state.object.sampling
+    write_opts = tiff_write_opts(
+        obj_sampling,
+        corner=obj_sampling.region_min if crop else None,
+        zs=state.object.zs() if stack else None,
+    )
+
+    xp = get_array_module(state.object.data)
+    obj_mag = abs2(state.object.data)
+    if crop:
+        obj_mag = obj_mag[..., *state.object.sampling.get_region_crop()]
+        mask = numpy.ones(obj_mag.shape[-2:], dtype=numpy.bool_)
+    else:
+        # include whole image, but only scale based on ROI
+        mask = state.object.sampling.get_region_mask(xp=numpy)
+
+    if not stack:
+        obj_mag = xp.prod(obj_mag, axis=0)
+
+    obj_mag = to_numpy(obj_mag)
+    if options.img_dtype != 'float':
+        obj_mag = scale_to_integral_type(obj_mag, options.img_dtype, mask, min_range=0.2)
+
+    with tifffile.TiffWriter(out_path, ome=True) as w:
+        write_opts['metadata']['axes'] = 'ZYX' if stack else 'YX'
+        w.write(obj_mag, **write_opts)
+
+
+
+_SAVE_FUNCS: t.Dict[str, t.Callable[[ReconsState, Path, SaveOptions], t.Any]] = {
+    'probe': _save_probe,
+    'probe_mag': _save_probe_mag,
+    'probe_recip': _save_probe_recip,
+    'probe_recip_mag': _save_probe_recip_mag,
+    'object_phase_stack': partial(_save_object_phase, stack=True),
+    'object_phase_sum': partial(_save_object_phase, stack=False),
+    'object_mag_stack': partial(_save_object_mag, stack=True),
+    'object_mag_sum': partial(_save_object_mag, stack=False),
+}

phaser/engines/common/position_correction.py

@@ -0,0 +1,62 @@
+import typing as t
+
+import numpy
+from numpy.typing import NDArray
+
+from phaser.utils.num import get_array_module
+from phaser.hooks.solver import (
+    PositionSolver, SteepestDescentPositionSolverProps, MomentumPositionSolverProps
+)
+from phaser.state import ReconsState
+
+
+class SteepestDescentPositionSolver(PositionSolver[None]):
+    def __init__(self, args: None, props: SteepestDescentPositionSolverProps):
+        self.step_size = props.step_size
+        self.max_step_size = props.step_size
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def perform_update(
+        self,
+        positions: NDArray[numpy.floating],
+        gradients: NDArray[numpy.floating],
+        state: None
+    ) -> t.Tuple[NDArray[numpy.floating], None]:
+        xp = get_array_module(positions, gradients)
+        update = self.step_size * gradients
+
+        if self.max_step_size is not None:
+            update_mag = xp.linalg.norm(update, axis=-1, keepdims=True)
+            update *= xp.minimum(update_mag, self.max_step_size) / update_mag
+
+        return (update, state)
+
+
+class MomentumPositionSolver(PositionSolver[NDArray[numpy.floating]]):
+    def __init__(self, args: None, props: MomentumPositionSolverProps):
+        self.step_size = props.step_size
+        self.max_step_size = props.max_step_size
+        self.momentum = props.momentum
+
+    def init_state(self, sim: ReconsState) -> NDArray[numpy.floating]:
+        xp = get_array_module(sim.scan)
+        return xp.zeros_like(sim.scan)
+
+    def perform_update(
+        self,
+        positions: NDArray[numpy.floating],
+        gradients: NDArray[numpy.floating],
+        state: NDArray[numpy.floating]
+    ) -> t.Tuple[NDArray[numpy.floating], NDArray[numpy.floating]]:
+        xp = get_array_module(positions, gradients, state)
+
+        update = self.step_size * gradients + self.momentum * state
+
+        if self.max_step_size is not None:
+            update_mag = xp.linalg.norm(update, axis=-1, keepdims=True)
+            update *= xp.minimum(update_mag, self.max_step_size) / update_mag
+
+        # state is just previous update step
+        return (update, update)