phaserEM 0.1__py3-none-any.whl

phaser/engines/common/regularizers.py

@@ -0,0 +1,403 @@
+from functools import partial
+import logging
+import typing as t
+
+import numpy
+from numpy.typing import NDArray
+
+from phaser.utils.num import (
+    get_array_module, get_scipy_module, Float,
+    jit, fft2, ifft2, abs2, xp_is_jax, to_real_dtype
+)
+from phaser.state import ReconsState
+from phaser.hooks.regularization import (
+    ClampObjectAmplitudeProps, LimitProbeSupportProps,
+    RegularizeLayersProps, ObjLowPassProps,
+    CostRegularizerProps, TVRegularizerProps
+)
+
+
+class ClampObjectAmplitude:
+    def __init__(self, args: None, props: ClampObjectAmplitudeProps):
+        self.amplitude = props.amplitude
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def apply_group(self, group: NDArray[numpy.integer], sim: ReconsState, state: None) -> t.Tuple[ReconsState, None]:
+        return self.apply_iter(sim, state)
+
+    def apply_iter(self, sim: ReconsState, state: None) -> t.Tuple[ReconsState, None]:
+        amp = to_real_dtype(sim.object.data.dtype)(self.amplitude)
+        sim.object.data = clamp_amplitude(sim.object.data, amp)
+        return (sim, None)
+
+
+@partial(jit, donate_argnames=('obj',), cupy_fuse=True)
+def clamp_amplitude(obj: NDArray[numpy.complexfloating], amplitude: t.Union[float, numpy.floating]) -> NDArray[numpy.complexfloating]:
+    xp = get_array_module(obj)
+
+    obj_amp = xp.abs(obj)
+    scale = xp.minimum(obj_amp, amplitude) / obj_amp
+    return obj * scale
+
+
+class LimitProbeSupport:
+    def __init__(self, args: None, props: LimitProbeSupportProps):
+        self.max_angle = props.max_angle
+
+    def init_state(self, sim: ReconsState) -> NDArray[numpy.bool_]:
+        xp = get_array_module(sim.probe.data)
+        (ky, kx) = sim.probe.sampling.recip_grid(xp=xp)
+        mask = kx**2 + ky**2 <= (self.max_angle*1e-3 / sim.wavelength)**2
+        return mask
+
+    def apply_group(self, group: NDArray[numpy.integer], sim: ReconsState, state: NDArray[numpy.bool_]) -> t.Tuple[ReconsState, NDArray[numpy.bool_]]:
+        return self.apply_iter(sim, state)
+
+    def apply_iter(self, sim: ReconsState, state: NDArray[numpy.bool_]) -> t.Tuple[ReconsState, NDArray[numpy.bool_]]:
+        mask = state
+        #xp = get_array_module(sim.state.probe.data)
+        #print(f"intensity before: {xp.sum(abs2(sim.state.probe.data))}")
+        sim.probe.data = ifft2(fft2(sim.probe.data) * mask)
+        #print(f"intensity after: {xp.sum(abs2(sim.state.probe.data))}")
+        return (sim, mask)
+
+
+class RemovePhaseRamp:
+    def __init__(self, args: None, props: t.Any):
+        ...
+
+    def init_state(self, sim: ReconsState) -> NDArray[numpy.bool_]:
+        xp = get_array_module(sim.object.data)
+        return sim.object.sampling.get_region_mask(xp=xp)
+
+    def apply_group(self, group: NDArray[numpy.integer], sim: ReconsState, state: NDArray[numpy.bool_]) -> t.Tuple[ReconsState, NDArray[numpy.bool_]]:
+        return self.apply_iter(sim, state)
+
+    def apply_iter(self, sim: ReconsState, state: NDArray[numpy.bool_]) -> t.Tuple[ReconsState, NDArray[numpy.bool_]]:
+        from phaser.utils.image import remove_linear_ramp
+        xp = get_array_module(sim.object.data)
+        phase = remove_linear_ramp(xp.angle(sim.object.data), state)
+        sim.object.data = t.cast(NDArray[numpy.complexfloating], xp.abs(sim.object.data) * xp.exp(1.j * phase))
+        return (sim, state)
+
+
+class RegularizeLayers:
+    def __init__(self, args: None, props: RegularizeLayersProps):
+        self.weight = props.weight
+        self.sigma = props.sigma
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def apply_iter(self, sim: ReconsState, state: None) -> t.Tuple[ReconsState, None]:
+        xp = get_array_module(sim.object.data)
+        scipy = get_scipy_module(sim.object.data)
+        dtype = to_real_dtype(sim.object.data)
+
+        if len(sim.object.thicknesses) < 2:
+            return (sim, None)
+
+        # approximate layers as equally spaced
+        layer_spacing = numpy.mean(sim.object.thicknesses)
+        # calculate size of filter (go to ~sigma in each direction)
+        r = int(numpy.ceil(2. * self.sigma / layer_spacing))
+        n = 2*r + 1
+
+        # make Gaussian filter
+        zs = ((xp.arange(0, n) - (n-1)//2) * layer_spacing).astype(dtype)
+        kernel = xp.exp(-(zs / self.sigma)**2 / 2.)
+        kernel /= xp.sum(kernel)
+
+        # we convolve the log of object, because the transmission
+        # function is multiplicative, not additive
+
+        if xp_is_jax(xp):
+            new_obj = xp.exp(scipy.signal.convolve(
+                xp.pad(xp.log(sim.object.data), ((r, r), (0, 0), (0, 0)), mode='edge'),
+                kernel[:, None, None],
+                mode="valid"
+            ))
+        else:
+            new_obj = xp.exp(scipy.ndimage.convolve1d(xp.log(
+                sim.object.data
+            ), kernel, axis=0, mode='nearest'))
+
+        assert new_obj.shape == sim.object.data.shape
+        assert new_obj.dtype == sim.object.data.dtype
+        sim.object.data = (
+            self.weight * new_obj + (1 - self.weight) * sim.object.data
+        )
+        return (sim, None)
+
+
+class ObjLowPass:
+    def __init__(self, args: None, props: ObjLowPassProps):
+        self.logger = logging.getLogger(__name__)
+        self.max_freq = props.max_freq
+
+    def init_state(self, sim: ReconsState) -> NDArray[numpy.bool_]:
+        samp = sim.object.sampling
+        xp = get_array_module(sim.object.data)
+
+        ky = xp.fft.fftfreq(samp.shape[0], 1.0)
+        kx = xp.fft.fftfreq(samp.shape[1], 1.0)
+        (ky, kx) = xp.meshgrid(ky, kx, indexing='ij')
+        k2 = ky**2 + kx**2
+
+        return k2 <= self.max_freq**2
+
+    def apply_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: NDArray[numpy.bool_]
+    ) -> t.Tuple[ReconsState, NDArray[numpy.bool_]]:
+        return self.apply_iter(sim, state)
+
+    def apply_iter(
+        self, sim: ReconsState, state: NDArray[numpy.bool_]
+    ) -> t.Tuple[ReconsState, NDArray[numpy.bool_]]:
+        # TODO: should this be done in-place?
+        sim.object.data = ifft2(state * fft2(sim.object.data))
+        return (sim, state)
+
+
+class ObjL1:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        cost = xp.sum(xp.abs(sim.object.data - 1.0))
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+        return (cost * cost_scale * self.cost, state)
+
+
+class ObjL2:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: Float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        cost = xp.sum(abs2(sim.object.data - 1.0))
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+        return (cost * cost_scale * self.cost, state)
+
+
+class ObjPhaseL1:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        cost = xp.sum(xp.abs(xp.angle(sim.object.data)))
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+        return (cost * cost_scale * self.cost, state)
+
+
+class ObjRecipL1:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        # l1 norm of diff. pattern amplitude
+        # TODO log object before this?
+        cost = xp.sum(
+            xp.abs(fft2(xp.prod(sim.object.data, axis=0)))
+        )
+        # scale cost by fraction of the total reconstruction in the group
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+
+        return (cost * cost_scale * self.cost, state)
+
+
+class ObjTotalVariation:
+    def __init__(self, args: None, props: TVRegularizerProps):
+        self.cost: float = props.cost
+        self.eps: float = props.eps
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        # isotropic total variation
+        g_y, g_x = img_grad(sim.object.data)
+        cost = xp.sum(xp.sqrt(abs2(g_y) + abs2(g_x) + self.eps))
+        # anisotropic total variation
+        #cost = (
+        #    xp.sum(xp.abs(xp.diff(sim.object.data, axis=-1))) +
+        #    xp.sum(xp.abs(xp.diff(sim.object.data, axis=-2)))
+        #)
+        # scale cost by fraction of the total reconstruction in the group
+        # TODO also scale by # of pixels or similar?
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+
+        return (cost * cost_scale * self.cost, state)
+
+
+class ObjTikhonov:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        cost = (
+            xp.sum(abs2(xp.diff(sim.object.data, axis=-1))) +
+            xp.sum(abs2(xp.diff(sim.object.data, axis=-2)))
+        )
+        # scale cost by fraction of the total reconstruction in the group
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+
+        return (cost * cost_scale * self.cost, state)
+
+
+class LayersTotalVariation:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        if sim.object.data.shape[0] < 2:
+            return (0.0, state)
+
+        cost = xp.sum(xp.abs(xp.diff(sim.object.data, axis=0)))
+        # scale cost by fraction of the total reconstruction in the group
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+
+        return (cost * cost_scale * self.cost, state)
+
+
+class LayersTikhonov:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.object.data)
+
+        if sim.object.data.shape[0] < 2:
+            return (0.0, state)
+
+        cost = xp.sum(abs2(xp.diff(sim.object.data, axis=0)))
+        # scale cost by fraction of the total reconstruction in the group
+        cost_scale = (group.shape[-1] / numpy.prod(sim.scan.shape[:-1])).astype(cost.dtype)
+
+        return (cost * cost_scale * self.cost, state)
+
+
+class ProbePhaseTikhonov:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.probe.data)
+
+        phase = xp.angle(fft2(sim.probe.data))
+
+        cost = (
+            xp.sum(abs2(xp.diff(phase, axis=-1))) +
+            xp.sum(abs2(xp.diff(phase, axis=-2)))
+        )
+        cost_scale = 1.0
+
+        return (cost * cost_scale * self.cost, state)
+
+
+class ProbeRecipTikhonov:
+    def __init__(self, args: None, props: CostRegularizerProps):
+        self.cost: float = props.cost
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.probe.data)
+        probe_recip = xp.fft.fftshift(fft2(sim.probe.data), axes=(-1, -2))
+
+        cost = (
+            xp.sum(abs2(xp.diff(probe_recip, axis=-1))) +
+            xp.sum(abs2(xp.diff(probe_recip, axis=-2)))
+        )
+        cost_scale = 1.0
+
+        return (cost * cost_scale * self.cost, state)
+
+
+class ProbeRecipTotalVariation:
+    def __init__(self, args: None, props: TVRegularizerProps):
+        self.cost: float = props.cost
+        self.eps: float = props.eps
+
+    def init_state(self, sim: ReconsState) -> None:
+        return None
+
+    def calc_loss_group(
+        self, group: NDArray[numpy.integer], sim: ReconsState, state: None
+    ) -> t.Tuple[Float, None]:
+        xp = get_array_module(sim.probe.data)
+        probe_recip = xp.fft.fftshift(fft2(sim.probe.data), axes=(-1, -2))
+
+        g_y, g_x = img_grad(probe_recip)
+        cost = xp.sum(xp.sqrt(abs2(g_y) + abs2(g_x) + self.eps))
+        cost_scale = 1.0
+
+        return (cost * cost_scale * self.cost, state)
+
+
+def img_grad(img: numpy.ndarray) -> t.Tuple[numpy.ndarray, numpy.ndarray]:
+    xp = get_array_module(img)
+    return (
+        xp.diff(img, axis=-2, append=img[..., -1:, :]),
+        xp.diff(img, axis=-1, append=img[..., :, -1:]),
+    )

phaser/engines/common/simulation.py

@@ -0,0 +1,270 @@
+import collections
+import logging
+import typing as t
+
+import numpy
+from numpy.typing import NDArray, DTypeLike
+from typing_extensions import Self
+
+from phaser.utils.num import (
+    get_array_module, to_real_dtype, to_complex_dtype,
+    fft2, ifft2, is_jax, to_numpy, block_until_ready,
+)
+from phaser.utils.misc import FloatKey, jax_dataclass, create_compact_groupings, create_sparse_groupings, shuffled
+from phaser.utils.optics import fresnel_propagator, fourier_shift_filter
+from phaser.state import ReconsState
+from phaser.hooks.solver import NoiseModel
+from phaser.hooks.regularization import GroupConstraint, IterConstraint, StateT
+
+logger = logging.getLogger(__name__)
+
+
+class GroupManager:
+    def __init__(
+        self,
+        scan: NDArray[numpy.floating],
+        grouping: t.Optional[int] = None,
+        compact: bool = False,
+        seed: t.Any = None,
+    ):
+        self.grouping = grouping or 64
+        self.compact = compact
+        self.seed = seed
+        self.groups: t.Optional[t.List[NDArray[numpy.int64]]] = None
+        self.n_groups: int = int(numpy.ceil(numpy.prod(scan.shape[:-1]) / self.grouping))
+
+    def _make(self, scan: NDArray[numpy.floating], i: int = 0) -> t.List[NDArray[numpy.int64]]:
+        if self.compact:
+            return create_compact_groupings(scan, self.grouping, seed=self.seed, i=i)
+        else:
+            return create_sparse_groupings(scan, self.grouping, seed=self.seed, i=i)
+
+    def iter(
+        self, scan: NDArray[numpy.floating],
+        i: int = 0, shuffle_groups: bool = False,
+    ) -> t.Iterator[NDArray[numpy.int64]]:
+        if shuffle_groups or self.groups is None:
+            self.groups = self._make(scan, i)
+            return iter(self.groups)
+        # shuffle order of groups (though not the groups themselves)
+        return shuffled(self.groups, seed=self.seed, i=i)
+
+    def __len__(self) -> int:
+        return self.n_groups
+
+
+def stream_patterns(
+    groups: t.Iterable[NDArray[numpy.int64]], patterns: NDArray[numpy.floating],
+    xp: t.Any, buf_n: int = 1
+) -> t.Iterator[t.Tuple[NDArray[numpy.int64], NDArray[numpy.floating]]]:
+    if buf_n == 0:
+        for group in groups:
+            group_patterns = xp.asarray(patterns[tuple(group)])
+            yield group, block_until_ready(group_patterns)
+        return
+
+    buf = collections.deque()
+    it = iter(groups)
+
+    for group in it:
+        buf.append((group, xp.asarray(patterns[tuple(group)])))
+        if len(buf) >= buf_n:
+            break
+
+    while len(buf) > 0:
+        (group, group_patterns) = buf.popleft()
+        yield group, block_until_ready(group_patterns)
+
+        # attempt to feed queue
+        try:
+            group = next(it)
+            buf.append((group, xp.asarray(patterns[tuple(group)])))
+        except StopIteration:
+            continue
+
+
+@jax_dataclass(init=False, static_fields=('xp', 'dtype', 'noise_model', 'group_constraints', 'iter_constraints'), drop_fields=('ky', 'kx'))
+class SimulationState:
+    state: ReconsState
+
+    ky: NDArray[numpy.floating]
+    kx: NDArray[numpy.floating]
+
+    noise_model: NoiseModel
+    group_constraints: t.Tuple[GroupConstraint[t.Any], ...]
+    iter_constraints: t.Tuple[IterConstraint[t.Any], ...]
+
+    noise_model_state: t.Any
+    group_constraint_states: t.Tuple[t.Any, ...]
+    iter_constraint_states: t.Tuple[t.Any, ...]
+
+    xp: t.Any
+    dtype: DTypeLike
+    start_iter: int
+
+    def __init__(
+        self, *,
+        state: ReconsState,
+        noise_model: NoiseModel[t.Any],
+        group_constraints: t.Tuple[GroupConstraint[t.Any], ...],
+        iter_constraints: t.Tuple[IterConstraint[t.Any], ...],
+        xp: t.Any,
+        dtype: DTypeLike,
+        noise_model_state: t.Optional[t.Any] = None,
+        group_constraint_states: t.Optional[t.Tuple[t.Any, ...]] = None,
+        iter_constraint_states: t.Optional[t.Tuple[t.Any, ...]] = None,
+        start_iter: t.Optional[int] = None,
+    ):
+        self.xp = xp
+        self.dtype = dtype
+        self.state = state
+
+        self.noise_model = noise_model
+        self.group_constraints = group_constraints
+        self.iter_constraints = iter_constraints
+
+        self.noise_model_state = noise_model_state or noise_model.init_state(self.state)
+        self.group_constraint_states = group_constraint_states if group_constraint_states is not None else tuple(
+            reg.init_state(self.state) for reg in group_constraints
+        )
+        self.iter_constraint_states = iter_constraint_states if iter_constraint_states is not None else tuple(
+            reg.init_state(self.state) for reg in iter_constraints
+        )
+
+        self.start_iter = start_iter if start_iter is not None else self.state.iter.total_iter
+        (self.ky, self.kx) = state.probe.sampling.recip_grid(dtype=dtype, xp=xp)
+
+    def apply_group_constraints(self, group: NDArray[numpy.integer]) -> Self:
+        def apply_reg(reg: GroupConstraint[t.Any], state: t.Any):
+            (self.state, state) = reg.apply_group(group, self.state, state)
+            return state
+
+        self.group_constraint_states = tuple(
+            apply_reg(reg, state) for (reg, state) in zip(self.group_constraints, self.group_constraint_states)
+        )
+
+        return self
+
+    def apply_iter_constraints(self) -> Self:
+        def apply_reg(reg: IterConstraint[t.Any], state: t.Any):
+            (self.state, state) = reg.apply_iter(self.state, state)
+            return state
+
+        self.iter_constraint_states = tuple(
+            apply_reg(reg, state) for (reg, state) in zip(self.iter_constraints, self.iter_constraint_states)
+        )
+
+        return self
+
+
+def make_propagators(state: ReconsState, bwlim_frac: t.Optional[float] = 2/3) -> t.Optional[NDArray[numpy.complexfloating]]:
+    xp = get_array_module(state.probe.data)
+    dtype = to_real_dtype(state.probe.data.dtype)
+    complex_dtype = to_complex_dtype(dtype)
+
+    (ky, kx) = state.probe.sampling.recip_grid(xp=xp, dtype=dtype)
+
+    # ignore last slice; we don't need it
+    delta_zs = to_numpy(state.object.thicknesses)[:-1]
+    if len(delta_zs) == 0:
+        return None
+
+    unique_zs = set(map(FloatKey, delta_zs))
+
+    if bwlim_frac is not None:
+        bwlim = numpy.min(state.probe.sampling.k_max) * bwlim_frac
+        k2 = ky**2 + kx**2
+        bwlim_mask = k2 <= bwlim**2
+        logger.info(f"Bandwidth limit: {bwlim * state.wavelength * 1e3:6.2f} mrad")
+    else:
+        bwlim_mask = xp.ones(ky.shape, dtype=numpy.bool_)
+
+    props = {
+        z: fresnel_propagator(ky, kx, state.wavelength, z).astype(complex_dtype) * bwlim_mask
+        for z in unique_zs
+    }
+
+    return xp.stack(
+        [props[FloatKey(z)] for z in delta_zs],
+        axis = 0
+    )
+
+
+@t.overload
+def cutout_group(
+    ky: NDArray[numpy.floating], kx: NDArray[numpy.floating],
+    state: ReconsState, group: NDArray[numpy.integer],
+    return_filters: t.Literal[False] = False
+) -> t.Tuple[NDArray[numpy.complexfloating], NDArray[numpy.complexfloating], NDArray[numpy.floating]]:
+    ...
+
+@t.overload
+def cutout_group(
+    ky: NDArray[numpy.floating], kx: NDArray[numpy.floating],
+    state: ReconsState, group: NDArray[numpy.integer],
+    return_filters: t.Literal[True]
+) -> t.Tuple[NDArray[numpy.complexfloating], NDArray[numpy.complexfloating], NDArray[numpy.floating], NDArray[numpy.complexfloating]]:
+    ...
+
+def cutout_group(
+    ky: NDArray[numpy.floating], kx: NDArray[numpy.floating],
+    state: ReconsState, group: NDArray[numpy.integer],
+    return_filters: bool = False
+):
+    """Returns (probe, obj) in the cutout region"""
+    probes = state.probe.data
+
+    group_scan = state.scan[tuple(group)]
+    group_obj = state.object.sampling.get_view_at_pos(state.object.data, group_scan, probes.shape[-2:])
+    # group probes in real space
+    # shape (len(group), 1, Ny, Nx)
+    group_subpx_filters = fourier_shift_filter(ky, kx, state.object.sampling.get_subpx_shifts(group_scan, probes.shape[-2:]))[:, None, ...]
+    # shape (len(group), probe modes, Ny, Nx)
+    shifted_probes = ifft2(fft2(probes) * group_subpx_filters)
+
+    if return_filters:
+        return (shifted_probes, group_obj, group_scan, group_subpx_filters)
+
+    return (shifted_probes, group_obj, group_scan)
+
+
+def slice_forwards(
+    props: t.Optional[NDArray[numpy.complexfloating]],
+    state: StateT,
+    f: t.Callable[[int, t.Optional[NDArray[numpy.complexfloating]], StateT], StateT]
+) -> StateT:
+    if props is None:
+        return f(0, None, state)
+
+    n_slices = len(props) + 1
+
+    if is_jax(props):
+        import jax
+        state = jax.lax.fori_loop(0, n_slices - 1, lambda slice_i, state: f(slice_i, props[slice_i], state), state, unroll=False)
+        return f(n_slices - 1, None, state)
+
+    for slice_i in range(n_slices - 1):
+        state = f(slice_i, props[slice_i], state)
+
+    return f(n_slices - 1, None, state)
+
+
+def slice_backwards(
+    props: t.Optional[NDArray[numpy.complexfloating]],
+    state: StateT,
+    f: t.Callable[[int, t.Optional[NDArray[numpy.complexfloating]], StateT], StateT]
+) -> StateT:
+    if props is None:
+        return f(0, None, state)
+
+    n_slices = len(props) + 1
+
+    if is_jax(props):
+        import jax
+        state = jax.lax.fori_loop(1, n_slices, lambda i, state: f(n_slices - i, props[n_slices - i - 1], state), state, unroll=False)
+        return f(0, None, state)
+
+    for slice_i in range(n_slices - 1, 0, -1):
+        state = f(slice_i, props[slice_i - 1], state)
+
+    return f(0, None, state)