HDF5 state format v0.2

phaser/engines/common/noise_models.py

@@ -37,7 +37,7 @@ def calc_loss(
         xp = get_array_module(model_wave, model_intensity, exp_patterns, mask)
         patterns = xp.maximum(exp_patterns, 0.0)
 
-        return ((
+        return (t.cast(numpy.floating,
             2. * xp.sum(mask * (
                 xp.sqrt(patterns + self.offset) - xp.sqrt(model_intensity + self.offset) - self.eps
             )**2) / self.var).astype(exp_patterns.dtype),

phaser/engines/common/output.py

@@ -186,7 +186,7 @@ def _plot_scan(state: ReconsState, out_path: Path, options: SaveOptions):
     ax.set_xlim(left, right)
     ax.set_ylim(bottom, top)
 
-    scan = to_numpy(state.scan)
+    scan = to_numpy(state.scan.data)
     i = numpy.arange(scan[..., 0].size)
     ax.scatter(scan[..., 1].ravel(), scan[..., 0].ravel(), c=i, cmap='plasma', s=0.5, edgecolors='none')
 
@@ -197,7 +197,7 @@ def _plot_scan(state: ReconsState, out_path: Path, options: SaveOptions):
 def _plot_tilt(state: ReconsState, out_path: Path, options: SaveOptions):
     from matplotlib import pyplot
 
-    if state.tilt is None:
+    if state.scan.tilt is None:
         logger = logging.getLogger(__name__)
         logger.warning("Tilt map (`state.tilt`) is missing, skipping `plot_tilt`")
         return
@@ -209,8 +209,8 @@ def _plot_tilt(state: ReconsState, out_path: Path, options: SaveOptions):
     ax.set_xlim(left, right)
     ax.set_ylim(bottom, top)
 
-    scan = to_numpy(state.scan)
-    tilt = to_numpy(state.tilt)
+    scan = to_numpy(state.scan.data)
+    tilt = to_numpy(state.scan.tilt)
     tilt = tilt[..., 1] + tilt[..., 0]*1.j
     max_tilt = max(numpy.max(numpy.abs(tilt)), 1.0)  # at least 1 mrad
     c = colorize_complex(tilt.ravel() / max_tilt, amp=True, rescale=False)

phaser/engines/common/position_correction.py

@@ -41,8 +41,8 @@ def __init__(self, args: None, props: MomentumPositionSolverProps):
         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)
+        xp = get_array_module(sim.scan.data)
+        return xp.zeros_like(sim.scan.data)
 
     def perform_update(
         self,

phaser/engines/common/regularizers.py

@@ -250,7 +250,7 @@ def calc_loss_group(
         xp = get_array_module(sim.object.data)
 
         cost = xp.sum(xp.abs(sim.object.data - 1.0))
-        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
         return (cost * cost_scale * self.cost, state)
 
 
@@ -272,7 +272,7 @@ def calc_loss_group(
 
         cost = xp.sum(abs2(sim.object.data - 1.0))
 
-        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
         return (cost * cost_scale * self.cost, state)  # type: ignore
 
 
@@ -293,7 +293,7 @@ def calc_loss_group(
         xp = get_array_module(sim.object.data)
 
         cost = xp.sum(xp.abs(xp.angle(sim.object.data)))
-        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
         return (cost * cost_scale * self.cost, state)
 
 
@@ -319,7 +319,7 @@ def calc_loss_group(
             xp.abs(fft2(xp.prod(sim.object.data, axis=0)))
         )
         # scale cost by fraction of the total reconstruction in the group
-        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
 
         return (cost * cost_scale * self.cost, state)
 
@@ -351,7 +351,7 @@ def calc_loss_group(
         #)
         # scale cost by fraction of the total reconstruction in the group
         # TODO also scale by # of pixels or similar?
-        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
 
         return (cost * cost_scale * self.cost, state)
 
@@ -377,7 +377,7 @@ def calc_loss_group(
             xp.sum(abs2(xp.diff(sim.object.data, axis=-2)))
         )
         # scale cost by fraction of the total reconstruction in the group
-        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
 
         return (cost * cost_scale * self.cost, state)  # type: ignore
 
@@ -403,7 +403,7 @@ def calc_loss_group(
 
         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 = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
 
         return (cost * cost_scale * self.cost, state)
 
@@ -429,7 +429,7 @@ def calc_loss_group(
 
         cost = xp.sum(abs2(xp.diff(sim.object.data, axis=0)))
         # scale cost by fraction of the total reconstruction in the group
-        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.shape[:-1]), dtype=cost.dtype)
+        cost_scale = xp.array(group.shape[-1] / prod(sim.scan.data.shape[:-1]), dtype=cost.dtype)
 
         return (cost * cost_scale * self.cost, state)  # type: ignore
 
@@ -519,7 +519,7 @@ def __init__(self, args: None, props: UnstructuredGaussianProps):
         self.attr_path = props.attr_path
 
     def init_state(self, sim: ReconsState) -> NDArray[numpy.floating]:
-        xp = get_array_module(sim.scan)
+        xp = get_array_module(sim.scan.data)
         try:
             self.getattr_nested(sim, self.attr_path)
         except AttributeError as e:
@@ -547,8 +547,8 @@ def setattr_nested(self, obj: t.Any, attr_path: str, value: t.Any):
     def apply_iter(self, sim: ReconsState, state: NDArray[numpy.floating]) -> t.Tuple[ReconsState, NDArray[numpy.floating]]:
         from scipy.spatial import KDTree
         obj_samp = sim.object.sampling
-        scan_flat = sim.scan.reshape(-1, 2)
-        scan_ndim = sim.scan.ndim - 1
+        scan_flat = sim.scan.data.reshape(-1, 2)
+        scan_ndim = sim.scan.data.ndim - 1
 
         attr = self.getattr_nested(sim, self.attr_path)
         vals = t.cast(NDArray[numpy.inexact], getattr(attr, 'data', attr))  # Extract raw array

phaser/engines/common/simulation.py

@@ -249,7 +249,7 @@ def cutout_group(
     """Returns (probe, obj) in the cutout region"""
     probes = state.probe.data
 
-    group_scan = state.scan[tuple(group)]
+    group_scan = state.scan.data[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)

phaser/engines/conventional/run.py

@@ -51,7 +51,7 @@ def run_engine(args: EngineArgs, props: ConventionalEnginePlan) -> ReconsState:
 
     solver = props.solver(props)
     sim = solver.init(sim)
-    groups = GroupManager(sim.state.scan, props.grouping, props.compact, seed=seed)
+    groups = GroupManager(sim.state.scan.data, props.grouping, props.compact, seed=seed)
 
     calc_error_mask = mask_fraction_of_groups(len(groups), props.calc_error_fraction)
 
@@ -79,7 +79,7 @@ def run_engine(args: EngineArgs, props: ConventionalEnginePlan) -> ReconsState:
 
     # runs rescaling
     sim = solver.presolve(
-        sim, groups.iter(sim.state.scan),
+        sim, groups.iter(sim.state.scan.data),
         patterns=patterns, pattern_mask=pattern_mask,
         propagators=propagators
     )
@@ -95,7 +95,7 @@ def run_engine(args: EngineArgs, props: ConventionalEnginePlan) -> ReconsState:
         iter_shuffle_groups = shuffle_groups({'state': sim.state, 'niter': props.niter})
 
         sim, pos_update, group_errors = solver.run_iteration(
-            sim, groups.iter(sim.state.scan, i, iter_shuffle_groups),
+            sim, groups.iter(sim.state.scan.data, i, iter_shuffle_groups),
             patterns=patterns, pattern_mask=pattern_mask, propagators=propagators,
             update_object=update_object({'state': sim.state, 'niter': props.niter}),
             update_probe=update_probe({'state': sim.state, 'niter': props.niter}),
@@ -116,16 +116,16 @@ def run_engine(args: EngineArgs, props: ConventionalEnginePlan) -> ReconsState:
 
             # subtract mean position update
             pos_update -= xp.mean(pos_update, tuple(range(pos_update.ndim - 1)))
-            pos_update, position_solver_state = position_solver.perform_update(sim.state.scan, pos_update, position_solver_state)
+            pos_update, position_solver_state = position_solver.perform_update(sim.state.scan.data, pos_update, position_solver_state)
             # subtract mean again (this can change with momentum)
             pos_update -= xp.mean(pos_update, tuple(range(pos_update.ndim - 1)))
             pos_update_rms = float(xp.mean(xp.linalg.norm(pos_update, axis=-1, keepdims=True)))
             logger.info(f"Position update: mean {pos_update_rms}")
-            sim.state.scan += pos_update
-            assert_dtype(sim.state.scan, dtype)
+            sim.state.scan.data += pos_update
+            assert_dtype(sim.state.scan.data, dtype)
 
             # check positions are at least overlapping object
-            sim.state.object.sampling.check_scan(sim.state.scan, sim.state.probe.sampling.extent / 2.)
+            sim.state.object.sampling.check_scan(sim.state.scan.data, sim.state.probe.sampling.extent / 2.)
 
             progress['pos_update_rms'].iters.append(i + start_i)
             progress['pos_update_rms'].values.append(pos_update_rms)
@@ -139,7 +139,7 @@ def run_engine(args: EngineArgs, props: ConventionalEnginePlan) -> ReconsState:
                 progress[k].values.append(error)
 
         sim.state.progress = progress
-        observer.update_iteration(sim.state, i, props.niter, {'total_loss': error})
+        observer.update_iteration(sim.state, i, props.niter, {'total_loss': error} if error is not None else {})
 
     observer.finish_engine(sim.state)
     return sim.state

phaser/engines/conventional/solvers.py

@@ -102,7 +102,7 @@ def run_iteration(
 
         new_obj_mag = xp.zeros_like(self.obj_mag)
         new_probe_mag = xp.zeros_like(self.probe_mag)
-        pos_update = xp.zeros_like(sim.state.scan, dtype=sim.dtype)
+        pos_update = xp.zeros_like(sim.state.scan.data, dtype=sim.dtype)
         iter_errors = []
 
         for (group_i, (group, group_patterns)) in enumerate(self.iter_patterns(groups, patterns, xp)):
@@ -171,7 +171,7 @@ def run_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], st
         return (probe_mag, psi)
 
     props = tilt_propagators(sim.ky, sim.kx, sim.state, props,
-                             sim.state.tilt[tuple(group)] if sim.state.tilt is not None else None)
+                             sim.state.scan.tilt[tuple(group)] if sim.state.scan.tilt is not None else None)
     (probe_mag, psi) = slice_forwards(props, (probe_mag, psi), run_slice)
 
     # modeled and experimental intensity
@@ -238,7 +238,7 @@ def sim_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], st
         return (group_probe_mag, psi)
 
     props = tilt_propagators(sim.ky, sim.kx, sim.state, props,
-                             sim.state.tilt[tuple(group)] if sim.state.tilt is not None else None)
+                             sim.state.scan.tilt[tuple(group)] if sim.state.scan.tilt is not None else None)
     (group_probe_mag, psi) = slice_forwards(props, (group_probe_mag, psi), sim_slice, jit_unroll_slices=jit_unroll_slices)
 
     new_obj_mag += group_obj_mag
@@ -377,7 +377,7 @@ def run_iteration(
         xp = sim.xp
 
         # TODO: ePIE position update
-        pos_update = xp.zeros_like(sim.state.scan)
+        pos_update = xp.zeros_like(sim.state.scan.data)
         iter_errors = []
 
         beta_object = process_schedule(self.plan.beta_object)({'state': sim.state, 'niter': self.engine_plan.niter})
@@ -432,7 +432,7 @@ def run_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], ps
         return psi
 
     props = tilt_propagators(sim.ky, sim.kx, sim.state, props,
-                             sim.state.tilt[tuple(group)] if sim.state.tilt is not None else None)
+                             sim.state.scan.tilt[tuple(group)] if sim.state.scan.tilt is not None else None)
     psi = slice_forwards(props, psi, run_slice)
 
     # modeled and experimental intensity
@@ -479,7 +479,7 @@ def sim_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], ps
         return psi
 
     props = tilt_propagators(sim.ky, sim.kx, sim.state, props,
-                             sim.state.tilt[tuple(group)] if sim.state.tilt is not None else None)
+                             sim.state.scan.tilt[tuple(group)] if sim.state.scan.tilt is not None else None)
     psi = slice_forwards(props, psi, sim_slice, jit_unroll_slices=jit_unroll_slices)
 
     model_wave = fft2(psi[-1] * group_obj[:, -1, None])

phaser/engines/gradient/run.py

@@ -26,6 +26,7 @@
 
 logger = logging.getLogger(__name__)
 _PER_ITER_VARS: t.FrozenSet[ReconsVar] = frozenset({'positions', 'tilt'})
+_PER_ITER_PATHS: t.FrozenSet[str] = frozenset({'initial'}) | _PER_ITER_VARS
 
 
 def process_solvers(
@@ -67,11 +68,12 @@ def process_solvers(
     )
 
 
-_PATH_MAP: t.Dict[t.Tuple[str, ...], ReconsVar] = {
+_PATH_MAP: t.Dict[t.Tuple[str, ...], str] = {
     ('object', 'data'): 'object',
     ('probe', 'data'): 'probe',
-    ('scan',): 'positions',
-    ('tilt',): 'tilt'
+    ('scan', 'data'): 'positions',
+    ('scan', 'tilt'): 'tilt',
+    ('scan', 'initial'): 'initial',  # not a solver variable, but need to apply group indexing
 }
 
 def _normalize_path(path: t.Tuple[tree.GetAttrKey, ...]) -> t.Tuple[str, ...]:
@@ -97,7 +99,7 @@ def f(path: t.Tuple[tree.GetAttrKey, ...], val: t.Any):
         if (var := _PATH_MAP.get(_normalize_path(path))):
             if var in vars:
                 return vars[var]
-            if var in _PER_ITER_VARS and val is not None and group is not None:
+            if var in _PER_ITER_PATHS and val is not None and group is not None:
                 return val[tuple(group)]
         return val
 
@@ -110,13 +112,13 @@ def apply_update(state: ReconsState, update: t.Dict[ReconsVar, numpy.ndarray]) -
     if 'object' in update:
         state.object.data += update['object']
     if 'tilt' in update:
-        state.tilt += update['tilt']
+        state.scan.tilt += update['tilt']
     if 'positions' in update:
         # subtract mean position update
         xp = get_array_module(update['positions'])
         update['positions'] -= xp.mean(update['positions'], tuple(range(update['positions'].ndim - 1)))
 
-        state.scan += update['positions']
+        state.scan.data += update['positions']
 
     return state
 
@@ -177,7 +179,7 @@ def run_engine(args: EngineArgs, props: GradientEnginePlan) -> ReconsState:
     }
     # shuffle_groups defaults to True for sparse groups, False for compact groups
     shuffle_groups = process_flag(props.shuffle_groups or not props.compact)
-    groups = GroupManager(state.scan, props.grouping, props.compact, seed)
+    groups = GroupManager(state.scan.data, props.grouping, props.compact, seed)
 
     observer.init_engine(
         state, recons_name=args['recons_name'],
@@ -210,7 +212,7 @@ def iter_patterns(groups: t.Iterable[NDArray[numpy.int_]]) -> t.Iterable[t.Tuple
 
     # runs rescaling
     rescale_factors = []
-    for (group_i, (group, group_patterns)) in enumerate(iter_patterns(groups.iter(state.scan))):
+    for (group_i, (group, group_patterns)) in enumerate(iter_patterns(groups.iter(state.scan.data))):
         group_rescale_factors = dry_run(
             state, group, propagators, group_patterns,
             xp=xp, dtype=dtype,
@@ -274,7 +276,7 @@ def iter_patterns(groups: t.Iterable[NDArray[numpy.int_]]) -> t.Iterable[t.Tuple
             for (solver, solver_state) in zip(iter_solvers, iter_solver_states)
         ]
 
-        for (group_i, (group, group_patterns)) in enumerate(iter_patterns(groups.iter(state.scan, i, iter_shuffle_groups))):
+        for (group_i, (group, group_patterns)) in enumerate(iter_patterns(groups.iter(state.scan.data, i, iter_shuffle_groups))):
             # prevent the loop running ahead of the GPU stream
             block_until_ready(losses_gpu['total_loss'])
 
@@ -341,8 +343,8 @@ def iter_patterns(groups: t.Iterable[NDArray[numpy.int_]]) -> t.Iterable[t.Tuple
 
         if 'positions' in iter_vars:
             # check positions are at least overlapping object
-            state.object.sampling.check_scan(state.scan, state.probe.sampling.extent / 2.)
-            assert_dtype(state.scan, dtype)
+            state.object.sampling.check_scan(state.scan.data, state.probe.sampling.extent / 2.)
+            assert_dtype(state.scan.data, dtype)
 
         state.progress = progress
         observer.update_iteration(state, i, props.niter, losses)
@@ -434,8 +436,8 @@ def run_model(
 ) -> t.Tuple[Float, t.Tuple[SolverStates, t.Dict[str, Float]]]:
     # apply vars to simulation
     sim = insert_vars(vars, sim, group)
-    group_scan = sim.scan
-    group_tilts = sim.tilt
+    group_scan = sim.scan.data
+    group_tilts = sim.scan.tilt
 
     (ky, kx) = sim.probe.sampling.recip_grid(dtype=dtype, xp=xp)
     xp = get_array_module(sim.probe.data)
@@ -489,7 +491,8 @@ def dry_run(
     dtype: t.Type[numpy.floating],
 ) -> NDArray[numpy.floating]:
     (ky, kx) = sim.probe.sampling.recip_grid(dtype=dtype, xp=xp)
-    group_scan = sim.scan[tuple(group)]
+    group_scan = sim.scan.data[tuple(group)]
+    group_tilt = sim.scan.tilt[tuple(group)] if sim.scan.tilt is not None else None
 
     probes = ifft2shift(sim.probe.data)
     group_obj = ifft2shift(sim.object.sampling.get_view_at_pos(sim.object.data, group_scan, probes.shape[-2:]))
@@ -501,7 +504,7 @@ def sim_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], ps
             return ifft2(fft2(psi * group_obj[:, slice_i, None], shift=False) * prop[:, None], shift=False)
         return psi * group_obj[:, slice_i, None]
 
-    t_props = tilt_propagators(ky, kx, sim, props, sim.tilt[tuple(group)] if sim.tilt is not None else None)
+    t_props = tilt_propagators(ky, kx, sim, props, group_tilt)
     model_wave = fft2(slice_forwards(t_props, probes, sim_slice), shift=False)
     model_intensity = xp.sum(abs2(model_wave), axis=(1, -2, -1))
     exp_intensity = xp.sum(group_patterns, axis=(-2, -1))