Refactor scan shape normalization

Author: hexane360

phaser/execute.py

@@ -1,6 +1,7 @@
 import dataclasses
 import itertools
 import logging
+import math
 import typing as t
 
 import numpy
@@ -122,6 +123,40 @@ def _normalize_observers(
     return ObserverSet(obs)
 
 
+def _normalize_scan_shape(
+    patterns: Patterns, state: ReconsState
+) -> t.Tuple[Patterns, ReconsState]:
+    """
+    Normalizes 'patterns' and 'state' to share a common scan shape.
+
+    Requires that there are an equal number of patterns and scan positions.
+    Reshapes 'state.scan' and 'patterns' to match shape, choosing the highest
+    dimensional shape of the two. 'state.tilt' is reshaped as well.
+    """
+    patterns_shape = patterns.patterns.shape[:-2]
+    scan_shape = state.scan.shape[:-1]
+
+    n_patterns = math.prod(patterns_shape)
+    n_scan = math.prod(scan_shape)
+    if n_scan != n_patterns:
+        raise ValueError(f"# of scan positions {n_scan} doesn't match # of patterns {n_patterns}")
+
+    # choose the highest dimensional shape
+    new_shape = scan_shape if len(scan_shape) > len(patterns_shape) else patterns_shape
+
+    patterns.patterns = patterns.patterns.reshape((*new_shape, *patterns.patterns.shape[-2:]))
+    state.scan = state.scan.reshape((*new_shape, 2))
+
+    if state.tilt is not None:
+        n_tilt = math.prod(state.tilt.shape[:-1])
+        if n_tilt != n_patterns:
+            raise ValueError(f"# of tilt positions {n_scan} doesn't match # of patterns {n_patterns}")
+
+        state.tilt = state.tilt.reshape((*new_shape, 2))
+
+    return patterns, state
+
+
 def load_raw_data(
     plan: ReconsPlan, xp: t.Any, seed: t.Any = None,
     init_state: t.Union[ReconsState, PartialReconsState, None] = None
@@ -295,6 +330,7 @@ def initialize_reconstruction(
         progress=ProgressState(iters=numpy.array([]), detector_errors=numpy.array([])),
         wavelength=wavelength
     )
+    data, state = _normalize_scan_shape(data, state)
 
     # process post_init hooks
     for p in plan.post_init:
@@ -305,15 +341,6 @@ def initialize_reconstruction(
 
     # perform some checks on preprocessed data
 
-    if state.scan.shape[:-1] != data.patterns.shape[:-2]:
-        n_pos = int(numpy.prod(state.scan.shape[:-1]))
-        n_pat = int(numpy.prod(data.patterns.shape[:-2]))
-        if n_pos != n_pat:
-            raise ValueError(f"# of scan positions {n_pos} doesn't match # of patterns {n_pat}")
-
-        # reshape patterns to match scan
-        data.patterns = data.patterns.reshape((*state.scan.shape[:-1], *data.patterns.shape[-2:]))
-
     avg_pattern_intensity = float(numpy.nanmean(numpy.nansum(data.patterns, axis=(-1, -2))))
 
     if avg_pattern_intensity < 5.0:

tests/test_initialization.py

@@ -16,6 +16,21 @@
 
 
 def load_empty(args, props) -> RawData:
+    scan_shape = props['scan_shape']
+    det_shape = props['det_shape']
+
+    return {
+        'patterns': numpy.zeros((*scan_shape, *det_shape), dtype=numpy.float32),
+        'mask': numpy.ones(det_shape, dtype=numpy.float32),
+        'sampling': Sampling(det_shape, sampling=(1.0, 1.0)),
+        'wavelength': 1.0,
+        'scan_hook': None,
+        'probe_hook': None,
+        'seed': None,
+    }
+
+
+def load_no_probe(args, props) -> RawData:
     return {
         'patterns': numpy.zeros((32, 32, 64, 64), dtype=numpy.float32),
         'mask': numpy.ones((64, 64), dtype=numpy.float32),
@@ -34,7 +49,7 @@ def load_empty(args, props) -> RawData:
 def test_load_raw_data_missing():
     plan = ReconsPlan.from_data({
         'name': 'test',
-        'raw_data': 'tests.test_initialization:load_empty',
+        'raw_data': 'tests.test_initialization:load_no_probe',
         'engines': [],
     })
     xp = numpy
@@ -46,7 +61,7 @@ def test_load_raw_data_missing():
 def test_load_raw_data_override():
     plan = {
         'name': 'test',
-        'raw_data': 'tests.test_initialization:load_empty',
+        'raw_data': 'tests.test_initialization:load_no_probe',
         'engines': [],
         'init': {
             'probe': {
@@ -90,7 +105,7 @@ def test_load_raw_data_override():
 def test_load_raw_data_prev_state(caplog):
     plan = {
         'name': 'test',
-        'raw_data': 'tests.test_initialization:load_empty',
+        'raw_data': 'tests.test_initialization:load_no_probe',
         'engines': [],
     }
 
@@ -135,4 +150,35 @@ def test_load_raw_data_prev_state(caplog):
 
     # both should be modeled
     assert recons.state.probe is not probe_state
-    assert recons.state.scan is not scan_state
+    assert recons.state.scan is not scan_state
+
+
+def test_load_3d_raw_data():
+    scan_shape = (64, 64)
+    det_shape = (128, 128)
+
+    plan = ReconsPlan.from_data({
+        'name': 'test',
+        'raw_data': {
+            'type': 'tests.test_initialization:load_empty',
+            'scan_shape': (4096,),
+            'det_shape': det_shape,
+        },
+        'init': {
+            'scan': {
+                'type': 'raster',
+                'shape': scan_shape,
+                'step_size': (1.0, 1.0),
+            },
+            'probe': {
+                'type': 'focused',
+                'conv_angle': 20.0,
+                'defocus': 300.0,
+            }
+        },
+        'engines': [],
+    })
+    recons = initialize_reconstruction(plan)
+
+    assert recons.state.scan.shape == (*scan_shape, 2)
+    assert recons.patterns.patterns.shape == (*scan_shape, *det_shape)