phaserEM 0.1__py3-none-any.whl

phaser/hooks/io/empad.py

@@ -0,0 +1,88 @@
+
+from pathlib import Path
+import warnings
+import logging
+import typing as t
+
+import numpy
+
+from phaser.utils.num import Sampling
+from phaser.utils.physics import Electron
+from phaser.io.empad import load_4d, EmpadMetadata
+from phaser.types import cast_length
+from .. import LoadEmpadProps, RawData
+
+
+def load_empad(args: None, props: LoadEmpadProps) -> RawData:
+    logger = logging.getLogger(__name__)
+
+    path = Path(props.path).expanduser()
+
+    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
+
+        voltage = props.kv * 1e3 if props.kv is not None else meta.voltage
+        diff_step = props.diff_step or meta.diff_step
+        scan_shape = meta.scan_shape
+        adu = props.adu or meta.adu
+        needs_scale = not meta.is_simulated()
+
+        probe_hook = {
+            'type': 'focused',
+            'conv_angle': meta.conv_angle,
+            'defocus': meta.defocus * 1e10 if meta.defocus is not None else None,
+        }
+        # TODO: handle explicit scan_positions here
+        scan_hook = {
+            'type': 'raster',
+            # [x, y] -> [y, x]
+            'shape': tuple(reversed(meta.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,
+        }
+
+    else:
+        voltage = props.kv * 1e3 if props.kv is not None else None
+        diff_step = props.diff_step
+        scan_shape = None
+        probe_hook = scan_hook = None
+        adu = None
+        needs_scale = False
+
+    if voltage is None:
+        raise ValueError("'kv'/'voltage' must be specified by metadata or passed to 'raw_data'")
+    if diff_step is None:
+        raise ValueError("'diff_step' must be specified by metadata or passed to 'raw_data'")
+
+    wavelength = Electron(voltage).wavelength
+
+    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))
+
+    if needs_scale:
+        if adu is None:
+            warnings.warn("ADU not supplied for experimental dataset. This is not recommended.")
+        else:
+            logger.info(f"Scaling patterns by ADU ({adu:.1f})")
+            patterns /= adu
+
+    a = wavelength / (diff_step * 1e-3)  # recip. pixel size -> 1 / real space extent
+    sampling = Sampling(cast_length(patterns.shape[-2:], 2), extent=(a, a))
+
+    mask = numpy.zeros_like(patterns, shape=patterns.shape[-2:])
+    mask[2:-2, 2:-2] = 1.
+
+    return {
+        'patterns': patterns,
+        'mask': numpy.fft.ifftshift(mask, axes=(-1, -2)),
+        'sampling': sampling,
+        'wavelength': wavelength,
+        'probe_hook': probe_hook,
+        'scan_hook': scan_hook,
+        'seed': None,
+    }

phaser/hooks/object.py

@@ -0,0 +1,25 @@
+
+import numpy
+
+from phaser.utils.num import cast_array_module, to_complex_dtype
+from phaser.utils.object import random_phase_object
+from ..state import ObjectState
+from . import ObjectHookArgs, RandomObjectProps
+
+
+def random_object(args: ObjectHookArgs, props: RandomObjectProps) -> ObjectState:
+    sampling = args['sampling']
+
+    if args['slices'] is not None:
+        thicknesses = numpy.array(args['slices'].thicknesses, dtype=args['dtype'])
+        shape = (len(thicknesses), *sampling.shape)
+    else:
+        thicknesses = numpy.array([], dtype=args['dtype'])
+        shape = sampling.shape
+
+    obj = random_phase_object(
+        shape, props.sigma,
+        dtype=to_complex_dtype(args['dtype']),
+        xp=cast_array_module(args['xp'])
+    )
+    return ObjectState(sampling, obj, thicknesses)

phaser/hooks/preprocessing.py

@@ -0,0 +1,133 @@
+import logging
+import math
+import typing as t
+
+import numpy
+from numpy.typing import NDArray
+
+from phaser.types import cast_length
+from phaser.utils.num import get_array_module, cast_array_module, to_numpy, fft2, ifft2, at, Sampling
+from phaser.utils.misc import create_rng, create_sparse_groupings
+from phaser.utils.optics import fourier_shift_filter
+from phaser.utils.image import affine_transform
+from phaser.state import Patterns, ReconsState
+from . import RawData, PostInitArgs, PoissonProps, ScaleProps, DropNanProps, CropDataProps
+
+logger = logging.getLogger(__name__)
+
+
+def crop_data(raw_data: RawData, props: CropDataProps) -> RawData:
+    if raw_data['patterns'].ndim != 4:
+        raise ValueError(f"'crop_data' expects a 4D array of patterns, got shape {raw_data['patterns'].shape} instead")
+
+    (y_i, y_f, x_i, x_f) = props.crop
+    logging.info(f"Cropping raw data to {0 if y_i is None else y_i}:{raw_data['patterns'].shape[0] if y_f is None else y_f},"
+                 f" {0 if x_i is None else x_i}:{raw_data['patterns'].shape[1] if x_f is None else x_f}")
+    raw_data['patterns'] = raw_data['patterns'][slice(y_i, y_f), slice(x_i, x_f)]
+
+    if raw_data['scan_hook'] is not None:
+        if raw_data['scan_hook']['type'] == 'raster':
+            raw_data['scan_hook'] = {
+                **raw_data['scan_hook'],
+                'shape': raw_data['patterns'].shape[:2],
+            }
+
+    return raw_data
+
+
+def scale_patterns(raw_data: RawData, props: ScaleProps) -> RawData:
+    raw_data['patterns'] *= props.scale
+    return raw_data
+
+
+def add_poisson_noise(raw_data: RawData, props: PoissonProps) -> RawData:
+    xp = get_array_module(raw_data['patterns'])
+    dtype = raw_data['patterns'].dtype
+
+    if props.scale is not None:
+        logger.info(f"Adding poisson noise to raw patterns, after scaling by {props.scale:.2e}")
+        raw_data['patterns'] *= props.scale
+    else:
+        logger.info("Adding poisson noise to raw patterns")
+
+    rng = create_rng(raw_data['seed'], 'poisson_noise')
+
+    # TODO do this in batches?
+    patterns = rng.poisson(to_numpy(raw_data['patterns'])).astype(dtype)
+
+    if props.gaussian is not None:
+        patterns += rng.normal(scale=props.gaussian, size=patterns.shape)
+
+    logger.info(f"Mean pattern intensity: {numpy.nanmean(numpy.nansum(patterns, axis=(-1, -2)))}")
+
+    raw_data['patterns'] = xp.array(patterns)
+    return raw_data
+
+
+def drop_nan_patterns(args: PostInitArgs, props: DropNanProps) -> t.Tuple[Patterns, ReconsState]:
+    xp = get_array_module(args['data'].patterns)
+
+    # flatten scan and patterns
+    scan = args['state'].scan.reshape(-1, 2)
+    patterns = args['data'].patterns.reshape(-1, *args['data'].patterns.shape[-2:])
+
+    fraction_nan = xp.sum(xp.isnan(args['data'].patterns), axis=(-1, -2)) / xp.prod(patterns.shape[-2:])
+
+    mask = fraction_nan > props.threshold
+
+    if (n := int(xp.sum(mask))):
+        logger.info(f"Dropping {n}/{patterns.shape[0]} patterns which are at least {props.threshold:.1%} NaN values")
+        patterns = patterns[mask]
+
+        if scan.shape[0] == mask.size:
+            # apply mask to scan as well
+            scan = scan[mask]
+        elif scan.shape[0] != patterns.shape[0]:
+            raise ValueError(f"# of scan positions {scan.shape[0]} doesn't match # of patterns"
+                             f" before ({mask.size}) or after ({patterns.shape[0]}) filtering")
+        # otherwise, we assume the mask has already been applied to the scan
+
+    args['state'].scan = scan
+    args['data'].patterns = patterns
+
+    return (args['data'], args['state'])
+
+
+def diffraction_align(args: PostInitArgs, props: t.Any = None) -> t.Tuple[Patterns, ReconsState]:
+    patterns, state = args['data'], args['state']
+
+    xp = cast_array_module(args['xp'])
+    grouping = 128
+    groups = create_sparse_groupings(patterns.patterns.shape[:-2], grouping)
+
+    sum_pattern = xp.zeros(patterns.patterns.shape[-2:], dtype=patterns.patterns.dtype)
+
+    for group in groups:
+        pats = xp.array(patterns.patterns[tuple(group)]) * xp.array(patterns.pattern_mask)
+        sum_pattern += t.cast(NDArray[numpy.floating], xp.nansum(pats, axis=tuple(range(pats.ndim - 2))))
+
+    mean_pattern = sum_pattern / math.prod(patterns.patterns.shape[:-2])
+
+    ky, kx = Sampling(
+        cast_length(mean_pattern.shape, 2), extent=(1.0, 1.0)
+    ).recip_grid(dtype=patterns.patterns.dtype, xp=xp)
+
+    shift = xp.array([
+        xp.nansum(ky * mean_pattern), xp.nansum(kx * mean_pattern)
+    ]) / xp.nansum(mean_pattern)
+
+    logging.info(f"Shifting diffraction patterns by ({shift[1]}, {shift[0]}) px")
+
+    def bilinear_shift(arr: NDArray[numpy.floating]) -> NDArray[numpy.floating]:
+        return to_numpy(xp.fft.ifftshift(affine_transform(
+            xp.fft.fftshift(xp.array(arr), axes=(-2, -1)), [1., 1.], shift,
+            output_shape=arr.shape[-2:], order=1
+        ), axes=(-2, -1)))
+
+    for group in groups:
+        patterns.patterns[tuple(group)] = bilinear_shift(patterns.patterns[tuple(group)])
+
+    # fftshift mask as well
+    patterns.pattern_mask = bilinear_shift(patterns.pattern_mask)
+
+    return (patterns, state)

phaser/hooks/probe.py

@@ -0,0 +1,24 @@
+import logging
+
+from phaser.utils.optics import make_focused_probe
+from ..state import ProbeState
+from . import ProbeHookArgs, FocusedProbeProps
+
+
+def focused_probe(args: ProbeHookArgs, props: FocusedProbeProps) -> ProbeState:
+    logger = logging.getLogger(__name__)
+
+    if props.conv_angle is None:
+        raise ValueError("Probe 'conv_angle' must be specified by metadata or manually")
+    if props.defocus is None:
+        raise ValueError("Probe 'defocus' must be specified by metadata or manually")
+
+    logger.info(f"Making probe, conv_angle {props.conv_angle} mrad, defocus {props.defocus} A")
+
+    sampling = args['sampling']
+    ky, kx = sampling.recip_grid(dtype=args['dtype'], xp=args['xp'])
+    probe = make_focused_probe(
+        ky, kx, args['wavelength'],
+        props.conv_angle, defocus=props.defocus
+    )
+    return ProbeState(sampling, probe)

phaser/hooks/regularization.py

@@ -0,0 +1,97 @@
+import abc
+import typing as t
+
+import numpy
+from numpy.typing import NDArray
+
+from phaser.types import Dataclass
+from phaser.utils.num import Float
+from . import Hook
+
+from .solver import HasState, StateT
+
+if t.TYPE_CHECKING:
+    from phaser.state import ReconsState
+
+
+@t.runtime_checkable
+class GroupConstraint(HasState[StateT], t.Protocol[StateT]):
+    def apply_group(self, group: NDArray[numpy.integer], sim: 'ReconsState', state: StateT) -> t.Tuple['ReconsState', StateT]:
+        ...
+
+
+@t.runtime_checkable
+class IterConstraint(HasState[StateT], t.Protocol[StateT]):
+    def apply_iter(self, sim: 'ReconsState', state: StateT) -> t.Tuple['ReconsState', StateT]:
+        ...
+
+
+@t.runtime_checkable
+class CostRegularizer(HasState[StateT], t.Protocol[StateT]):
+    def calc_loss_group(self, group: NDArray[numpy.integer], sim: 'ReconsState', state: StateT) -> t.Tuple[Float, StateT]:
+        ...
+
+
+class ClampObjectAmplitudeProps(Dataclass):
+    amplitude: float = 1.1
+
+
+class LimitProbeSupportProps(Dataclass):
+    max_angle: float
+
+
+class RegularizeLayersProps(Dataclass):
+    weight: float = 0.9  # weight of regularization to apply
+    sigma: float = 50.0  # standard deviation of gaussian filter
+
+
+class ObjLowPassProps(Dataclass):
+    max_freq: float = 0.4  # 1/px (nyquist = 0.5)
+
+
+class IterConstraintHook(Hook[None, IterConstraint]):
+    known = {
+        'clamp_object_amplitude': ('phaser.engines.common.regularizers:ClampObjectAmplitude', ClampObjectAmplitudeProps),
+        'limit_probe_support': ('phaser.engines.common.regularizers:LimitProbeSupport', LimitProbeSupportProps),
+        'layers': ('phaser.engines.common.regularizers:RegularizeLayers', RegularizeLayersProps),
+        'obj_low_pass': ('phaser.engines.common.regularizers:ObjLowPass', ObjLowPassProps),
+        'remove_phase_ramp': ('phaser.engines.common.regularizers:RemovePhaseRamp', t.Dict[str, t.Any]),
+    }
+
+
+class GroupConstraintHook(Hook[None, GroupConstraint]):
+    known = {
+        'clamp_object_amplitude': ('phaser.engines.common.regularizers:ClampObjectAmplitude', ClampObjectAmplitudeProps),
+        'limit_probe_support': ('phaser.engines.common.regularizers:LimitProbeSupport', LimitProbeSupportProps),
+        'obj_low_pass': ('phaser.engines.common.regularizers:ObjLowPass', ObjLowPassProps),
+        'remove_phase_ramp': ('phaser.engines.common.regularizers:RemovePhaseRamp', t.Dict[str, t.Any]),
+    }
+
+
+class CostRegularizerProps(Dataclass):
+    cost: float
+
+
+class TVRegularizerProps(Dataclass):
+    cost: float
+    eps: float = 1.0e-8
+
+
+class CostRegularizerHook(Hook[None, CostRegularizer]):
+    known = {
+        'obj_l1': ('phaser.engines.common.regularizers:ObjL1', CostRegularizerProps),
+        'obj_l2': ('phaser.engines.common.regularizers:ObjL2', CostRegularizerProps),
+        'obj_phase_l1': ('phaser.engines.common.regularizers:ObjPhaseL1', CostRegularizerProps),
+        'obj_recip_l1': ('phaser.engines.common.regularizers:ObjRecipL1', CostRegularizerProps),
+        'obj_tv': ('phaser.engines.common.regularizers:ObjTotalVariation', TVRegularizerProps),
+        'obj_tikh': ('phaser.engines.common.regularizers:ObjTikhonov', CostRegularizerProps),
+        'obj_tikhonov': ('phaser.engines.common.regularizers:ObjTikhonov', CostRegularizerProps),
+        'layers_tv': ('phaser.engines.common.regularizers:LayersTotalVariation', CostRegularizerProps),
+        'layers_tikh': ('phaser.engines.common.regularizers:LayersTikhonov', CostRegularizerProps),
+        'layers_tikhonov': ('phaser.engines.common.regularizers:LayersTikhonov', CostRegularizerProps),
+        'probe_phase_tikh': ('phaser.engines.common.regularizers:ProbePhaseTikhonov', CostRegularizerProps),
+        'probe_phase_tikhonov': ('phaser.engines.common.regularizers:ProbePhaseTikhonov', CostRegularizerProps),
+        'probe_recip_tv': ('phaser.engines.common.regularizers:ProbeRecipTotalVariation', TVRegularizerProps),
+        'probe_recip_tikh': ('phaser.engines.common.regularizers:ProbeRecipTikhonov', CostRegularizerProps),
+        'probe_recip_tikhonov': ('phaser.engines.common.regularizers:ProbeRecipTikhonov', CostRegularizerProps),
+    }

phaser/hooks/scan.py

@@ -0,0 +1,27 @@
+
+import numpy
+from numpy.typing import NDArray
+
+from phaser.utils.num import cast_array_module
+from phaser.utils.scan import make_raster_scan
+from . import ScanHookArgs, RasterScanProps
+
+
+def raster_scan(args: ScanHookArgs, props: RasterScanProps) -> NDArray[numpy.floating]:
+    xp = cast_array_module(args['xp'])
+
+    if props.shape is None:
+        raise ValueError("scan 'shape' must be specified by metadata or manually")
+    if props.step_size is None:
+        raise ValueError("scan 'step_size' must be specified by metadata or manually")
+
+    scan = make_raster_scan(
+        props.shape, props.step_size, props.rotation or 0.0,
+        dtype=args['dtype'], xp=xp,
+    )
+
+    if props.affine is not None:
+        affine = xp.array(props.affine, dtype=scan.dtype)
+        scan = scan @ affine
+
+    return scan

phaser/hooks/schedule.py

@@ -0,0 +1,75 @@
+import logging
+import typing as t
+
+import numpy
+
+from ..types import Dataclass, Flag, process_schedule
+from .hook import Hook
+
+if t.TYPE_CHECKING:
+    from ..state import ReconsState
+
+
+class FlagArgs(t.TypedDict):
+    state: 'ReconsState'
+    niter: int
+
+
+class FlagHook(Hook[FlagArgs, bool]):
+    known = {}
+
+
+class ScheduleHook(Hook[FlagArgs, float]):
+    known = {}
+
+
+FlagLike: t.TypeAlias = t.Union[bool, Flag, FlagHook]
+ScheduleLike: t.TypeAlias = t.Union[float, ScheduleHook]
+
+
+# TODO: make these hooks two-step
+
+
+class ConstantScheduleProps(Dataclass):
+    value: float
+
+
+def constant_schedule(args: FlagArgs, props: ConstantScheduleProps) -> float:
+    return props.value
+
+
+class PiecewiseScheduleProps(Dataclass):
+    init: ScheduleLike
+    steps: t.Dict[int, ScheduleLike]
+
+
+def piecewise_schedule(args: FlagArgs, props: PiecewiseScheduleProps) -> float:
+    i = args['state'].iter.engine_iter
+
+    for thresh, val in sorted(props.steps.items(), key=lambda t: t[0], reverse=True):
+        if i >= thresh:
+            return process_schedule(val)(args)
+
+    return process_schedule(props.init)(args)
+
+
+class ExprScheduleProps(Dataclass):
+    expr: str
+
+
+def expr_schedule(args: FlagArgs, props: ExprScheduleProps) -> float:
+    val = float(eval(props.expr, {
+        'i': args['state'].iter.engine_iter,
+        'iter': args['state'].iter,
+        'state': args['state'],
+        'niter': args['niter'],
+        'np': numpy,
+    }))
+    logging.getLogger(__name__).debug(f"expr_schedule expr: {props.expr} value: {val}")
+
+    return val
+
+
+ScheduleHook.known['constant'] = ('phaser.hooks.schedule:constant_schedule', ConstantScheduleProps)
+ScheduleHook.known['piecewise'] = ('phaser.hooks.schedule:piecewise_schedule', PiecewiseScheduleProps)
+ScheduleHook.known['expr'] = ('phaser.hooks.schedule:expr_schedule', ExprScheduleProps)

phaser/hooks/solver.py

@@ -0,0 +1,169 @@
+import abc
+import typing as t
+
+import numpy
+from numpy.typing import NDArray
+
+from phaser.types import Dataclass, ReconsVar
+from phaser.utils.num import Float
+from . import Hook
+
+StateT = t.TypeVar('StateT')
+
+if t.TYPE_CHECKING:
+    from phaser.engines.common.simulation import SimulationState
+    from phaser.execute import Observer
+    from phaser.plan import ConventionalEnginePlan, GradientEnginePlan
+    from phaser.state import ReconsState
+
+
+class HasState(t.Protocol[StateT]):  # type: ignore
+    def init_state(self, sim: 'ReconsState') -> StateT:
+        ...
+
+
+class NoiseModel(HasState[StateT], t.Protocol[StateT]):
+    @classmethod
+    def name(cls) -> str:
+        ...
+
+    def calc_loss(
+        self,
+        model_wave: NDArray[numpy.complexfloating],
+        model_intensity: NDArray[numpy.floating],
+        exp_patterns: NDArray[numpy.floating],
+        mask: NDArray[numpy.floating], 
+        state: StateT,
+    ) -> t.Tuple[Float, StateT]:
+        """
+        Return the calculated loss, summed across the detector and averaged across the scan.
+
+        May be called in a JAX jit context, so must have no side effects.
+        """
+        ...
+
+    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: StateT,
+    ) -> t.Tuple[NDArray[numpy.complexfloating], StateT]:
+        """
+        Return the calculated wave update `chi` in reciprocal space.
+
+        May be called in a JAX jit context, so must have no side effects.
+        """
+        ...
+
+
+class NoiseModelHook(Hook[None, NoiseModel]):
+    known = {}
+
+
+class PositionSolver(HasState[StateT], t.Protocol[StateT]):
+    def perform_update(
+        self,
+        positions: NDArray[numpy.floating],
+        gradients: NDArray[numpy.floating],
+        state: StateT
+    ) -> t.Tuple[NDArray[numpy.floating], StateT]:
+        """
+        Return the calculated position updates
+        """
+        ...
+
+
+class SteepestDescentPositionSolverProps(Dataclass):
+    # fraction of optimal step to take
+    step_size: float = 1e-2
+    # maximum step size (in angstroms)
+    max_step_size: t.Optional[float] = None
+
+
+class MomentumPositionSolverProps(Dataclass):
+    # fraction of optimal step to take
+    step_size: float = 1e-2
+    # maximum step size (in angstroms)
+    max_step_size: t.Optional[float] = None
+    # momentum decay rate
+    momentum: float = 0.9
+
+
+class PositionSolverHook(Hook[None, PositionSolver]):
+    known = {
+        'steepest_descent': ('phaser.engines.common.position_correction:SteepestDescentPositionSolver', SteepestDescentPositionSolverProps),
+        'momentum': ('phaser.engines.common.position_correction:MomentumPositionSolver', MomentumPositionSolverProps),
+    }
+
+
+class ConventionalSolver(abc.ABC):
+    @classmethod
+    @abc.abstractmethod
+    def name(cls) -> str:
+        ...
+
+    @abc.abstractmethod
+    def init(self, sim: 'SimulationState') -> 'SimulationState':
+        ...
+
+    @abc.abstractmethod
+    def presolve(
+        self,
+        sim: 'SimulationState',
+        groups: t.Iterator[NDArray[numpy.int_]], *,
+        patterns: NDArray[numpy.floating],
+        pattern_mask: NDArray[numpy.floating],
+        propagators: t.Optional[NDArray[numpy.complexfloating]],
+    ) -> 'SimulationState':
+        ...
+
+    @abc.abstractmethod
+    def run_iteration(
+        self,
+        sim: 'SimulationState',
+        groups: t.Iterator[NDArray[numpy.int_]], *,
+        patterns: NDArray[numpy.floating],
+        pattern_mask: NDArray[numpy.floating],
+        propagators: t.Optional[NDArray[numpy.complexfloating]],
+        update_object: bool,
+        update_probe: bool,
+        update_positions: bool,
+        calc_error: bool,
+        calc_error_mask: NDArray[numpy.bool_],
+        observer: 'Observer',
+    ) -> t.Tuple['SimulationState', NDArray[numpy.floating], t.List[NDArray[numpy.floating]]]:
+        """
+        Run an iteration of the reconstruction. Return a tuple `(sim, pos_update, iter_errors)`.
+        """
+        ...
+
+
+class ConventionalSolverHook(Hook['ConventionalEnginePlan', ConventionalSolver]):
+    known = {}
+
+
+class GradientSolver(HasState[StateT], t.Protocol[StateT]):
+    name: str
+    params: t.FrozenSet[ReconsVar]
+
+    def init_state(self, sim: 'ReconsState') -> StateT:
+        ...
+
+    def update_for_iter(self, sim: 'ReconsState', state: StateT, niter: int) -> StateT:
+        ...
+
+    def update(
+        self, sim: 'ReconsState', state: StateT, grad: t.Dict[ReconsVar, numpy.ndarray], loss: float,
+    ) -> t.Tuple[t.Dict[ReconsVar, numpy.ndarray], StateT]:
+        ...
+
+
+class GradientSolverArgs(t.TypedDict):
+    plan: 'GradientEnginePlan'
+    params: t.Iterable[ReconsVar]
+
+
+class GradientSolverHook(Hook['GradientSolverArgs', GradientSolver]):
+    known = {}