Apply tilts to conventional solvers

Author: hexane360

phaser/engines/common/simulation.py

@@ -191,34 +191,33 @@ def make_propagators(state: ReconsState, bwlim_frac: t.Optional[float] = 2/3) ->
 
 
 def tilt_propagators(
+    ky: NDArray[numpy.floating], kx: NDArray[numpy.floating],
     state: ReconsState, 
-    base_props: t.Optional[NDArray[numpy.complexfloating]],  # shape: (Nz-1, Ny, Nx)
-    group_tilts: NDArray[numpy.floating],                    # shape: (batch, 2), in mrad
-    kx: NDArray[numpy.floating],                             # shape: (Ny, Nx)
-    ky: NDArray[numpy.floating],                             # shape: (Ny, Nx)
-    delta_zs: NDArray[numpy.floating],                       # shape: (Nz-1)
+    props: t.Optional[NDArray[numpy.complexfloating]],  # shape: (Nz-1, Ny, Nx)
+    tilts: NDArray[numpy.floating]                      # shape: (..., 2), in mrad
 ) -> t.Optional[NDArray[numpy.complexfloating]]:
     """
-    Applies tilt and slice-dependent propagation phase shifts to base_props.
+    Applies tilt and slice-dependent propagation phase shifts to props.
     -------
     NDArray[complex] or None
-        Tilted propagators of shape (batch, Nz-1, Ny, Nx), or None if no slices.
+        Tilted propagators of shape (n_layers-1, ..., Ny, Nx), or None if no slices.
     """
-    if base_props is None:
+    if props is None:
         return None
 
     xp = get_array_module(state.probe.data)
     dtype = to_real_dtype(state.probe.data.dtype)
     complex_dtype = to_complex_dtype(dtype)
+    delta_zs = state.object.thicknesses[:-1]
 
-    tilt_ramps = xp.exp(  # (batch, Nz-1, Ny, Nx)
-        2.j * xp.pi * delta_zs[:, None, None] * (
-            ufunc_outer(xp.multiply, xp.tan(group_tilts[:, 0] * 1e-3), ky)[:, None, ...] +
-            ufunc_outer(xp.multiply, xp.tan(group_tilts[:, 1] * 1e-3), kx)[:, None, ...]
-        )
+    tilt_ramps = xp.exp(  # (n_layers-1, batch, Ny, Nx)
+        2.j * xp.pi * ufunc_outer(xp.multiply, delta_zs, (
+            ufunc_outer(xp.multiply, xp.tan(tilts[..., 0] * 1e-3), ky) +
+            ufunc_outer(xp.multiply, xp.tan(tilts[..., 1] * 1e-3), kx)
+        ))
     )
 
-    return base_props[None, ...] * tilt_ramps.astype(complex_dtype) 
+    return props[(slice(None), *(None,)*(tilts.ndim - 1), Ellipsis)] * tilt_ramps.astype(complex_dtype)
 
 
 @t.overload
@@ -267,21 +266,21 @@ def slice_forwards(
     if props is None:
         return f(0, None, state)
 
-    n_slices = props.shape[1] + 1  # props shape: (batch, Nz-1, Ny, Nx)
+    # props shape: (N_slices - 1, [batch], Ny, Nx)
+    n_slices = len(props) + 1
 
     if is_jax(props):
         import jax
         def step_fn(carry, slice_i):
-            # props[:, slice_i] has shape (batch, Ny, Nx)
-            new_state = f(slice_i, props[:, slice_i], carry)
+            new_state = f(slice_i, props[slice_i], carry)
             return new_state, None
 
         state, _ = jax.lax.scan(step_fn, state, jax.numpy.arange(n_slices - 1))
         return f(n_slices - 1, None, state)
 
     # fallback numpy mode
     for slice_i in range(n_slices - 1):
-        state = f(slice_i, props[:, slice_i], state)
+        state = f(slice_i, props[slice_i], state)
     return f(n_slices - 1, None, state)
 
 

phaser/engines/conventional/solvers.py

@@ -11,7 +11,7 @@
 from phaser.plan import ConventionalEnginePlan, LSQMLSolverPlan, EPIESolverPlan
 from phaser.execute import Observer
 from phaser.engines.common.simulation import (
-    stream_patterns, SimulationState, cutout_group, slice_forwards, slice_backwards
+    stream_patterns, SimulationState, cutout_group, tilt_propagators, slice_forwards, slice_backwards
 )
 
 
@@ -151,10 +151,11 @@ def run_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], st
         )
 
         if prop is not None:
-            psi = ifft2(fft2(psi * group_obj[:, slice_i, None]) * prop)
+            psi = ifft2(fft2(psi * group_obj[:, slice_i, None]) * prop[:, None])
 
         return (probe_mag, psi)
 
+    props = tilt_propagators(sim.ky, sim.kx, sim.state, props, sim.state.tilt[tuple(group)])
     (probe_mag, psi) = slice_forwards(props, (probe_mag, psi), run_slice)
 
     # modeled and experimental intensity
@@ -214,11 +215,12 @@ def sim_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], st
 
         if prop is not None:
             psi = at(psi, slice_i + 1).set(
-                ifft2(fft2(psi[slice_i] * group_obj[:, slice_i, None]) * prop)
+                ifft2(fft2(psi[slice_i] * group_obj[:, slice_i, None]) * prop[:, None])
             )
 
         return (group_probe_mag, psi)
 
+    props = tilt_propagators(sim.ky, sim.kx, sim.state, props, sim.state.tilt[tuple(group)])
     (group_probe_mag, psi) = slice_forwards(props, (group_probe_mag, psi), sim_slice)
 
     new_obj_mag += group_obj_mag
@@ -253,7 +255,7 @@ def update_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]],
             sim.state.object.data = at(sim.state.object.data, slice_i).add(obj_update)
 
         if prop is not None:
-            chi = ifft2(fft2(delta_P) * prop.conj())
+            chi = ifft2(fft2(delta_P) * prop.conj()[:, None])
         elif update_probe:
             delta_P_avg = ifft2(xp.sum(fft2(delta_P) * subpx_filters.conj(), axis=0))
             delta_P_avg /= (group_obj_mag + illum_reg_probe)
@@ -392,10 +394,11 @@ def epie_dry_run(
 
     def run_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], psi):
         if prop is not None:
-            psi = ifft2(fft2(psi * group_obj[:, slice_i, None]) * prop)
+            psi = ifft2(fft2(psi * group_obj[:, slice_i, None]) * prop[:, None])
 
         return psi
 
+    props = tilt_propagators(sim.ky, sim.kx, sim.state, props, sim.state.tilt[tuple(group)])
     psi = slice_forwards(props, psi, run_slice)
 
     # modeled and experimental intensity
@@ -431,11 +434,12 @@ def epie_run(
     def sim_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], psi):
         if prop is not None:
             psi = at(psi, slice_i + 1).set(
-                ifft2(fft2(psi[slice_i] * group_obj[:, slice_i, None]) * prop)
+                ifft2(fft2(psi[slice_i] * group_obj[:, slice_i, None]) * prop[:, None])
             )
 
         return psi
 
+    props = tilt_propagators(sim.ky, sim.kx, sim.state, props, sim.state.tilt[tuple(group)])
     psi = slice_forwards(props, psi, sim_slice)
 
     model_wave = fft2(psi[-1] * group_obj[:, -1, None])
@@ -468,7 +472,7 @@ def update_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]],
             )
 
         if prop is not None:
-            chi = ifft2(fft2(probe_update) * prop.conj())
+            chi = ifft2(fft2(probe_update) * prop.conj()[:, None])
         elif update_probe:
             # average probe updates in group
             probe_update = ifft2(xp.mean(fft2(probe_update) * subpx_filters.conj(), axis=0))

phaser/engines/gradient/run.py

@@ -400,25 +400,22 @@ def run_model(
     group_tilts = sim.tilt
 
     (ky, kx) = sim.probe.sampling.recip_grid(dtype=dtype, xp=xp)
-    delta_zs = sim.object.thicknesses[:-1]
     xp = get_array_module(sim.probe.data)
     dtype = to_real_dtype(sim.probe.data.dtype)
-    complex_dtype = to_complex_dtype(dtype)
 
     probes = sim.probe.data
     group_obj = sim.object.sampling.get_view_at_pos(sim.object.data, group_scan, probes.shape[-2:])
     group_subpx_filters = fourier_shift_filter(ky, kx, sim.object.sampling.get_subpx_shifts(group_scan, probes.shape[-2:]))[:, None, ...]
     probes = ifft2(fft2(probes) * group_subpx_filters)
-    
-    t_props = tilt_propagators(sim, props, group_tilts, kx, ky, delta_zs,)
 
     def sim_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], psi):
         # psi: (batch, n_probe, Ny, Nx)
         if prop is not None:
             return ifft2(fft2(psi * group_obj[:, slice_i, None]) * prop[:, None])
         return psi * group_obj[:, slice_i, None]
 
-    model_wave = fft2(slice_forwards(t_props, probes, sim_slice))
+    props = tilt_propagators(ky, kx, sim, props, group_tilts)
+    model_wave = fft2(slice_forwards(props, probes, sim_slice))
     model_intensity = xp.sum(abs2(model_wave), axis=1)
 
     (loss, solver_states.noise_model_state) = noise_model.calc_loss(
@@ -447,20 +444,19 @@ def dry_run(
     dtype: t.Type[numpy.floating],
 ) -> NDArray[numpy.floating]:
     (ky, kx) = sim.probe.sampling.recip_grid(dtype=dtype, xp=xp)
-    delta_zs = sim.object.thicknesses[:-1]
 
     probes = sim.probe.data
     group_obj = sim.object.sampling.get_view_at_pos(sim.object.data, sim.scan[tuple(group)], probes.shape[-2:])
     group_subpx_filters = fourier_shift_filter(ky, kx, sim.object.sampling.get_subpx_shifts(sim.scan[tuple(group)], probes.shape[-2:]))[:, None, ...]
     probes = ifft2(fft2(probes) * group_subpx_filters)
-    t_props = tilt_propagators(sim, props, sim.tilt[tuple(group)], kx, ky, delta_zs)
+    props = tilt_propagators(ky, kx, sim, props, sim.tilt[tuple(group)])
 
     def sim_slice(slice_i: int, prop: t.Optional[NDArray[numpy.complexfloating]], psi):
         if prop is not None:
             return ifft2(fft2(psi * group_obj[:, slice_i, None]) * prop[:, None])
         return psi * group_obj[:, slice_i, None]
 
-    model_wave = fft2(slice_forwards(t_props, probes, sim_slice))
+    model_wave = fft2(slice_forwards(props, probes, sim_slice))
     model_intensity = xp.sum(abs2(model_wave), axis=(1, -2, -1))
     exp_intensity = xp.sum(group_patterns, axis=(-2, -1))
 

phaser/utils/_jax_kernels.py

@@ -58,7 +58,11 @@ def get_cutouts(obj: jax.Array, start_idxs: jax.Array, cutout_shape: t.Tuple[int
 
 @partial(jax.jit, static_argnums=0)
 def outer(ufunc: t.Any, x: jax.Array, y: jax.Array) -> jax.Array:
-    return jax.vmap(jax.vmap(ufunc, (None, 0)), (0, None))(x, y)
+    if x.ndim == 0 or y.ndim == 0:
+        return ufunc(x, y)
+
+    out_shape = (*x.shape, *y.shape)
+    return jax.vmap(jax.vmap(ufunc, (None, 0)), (0, None))(x.ravel(), y.ravel()).reshape(out_shape)
 
 
 @partial(jax.jit, static_argnames=('output_shape', 'order', 'mode', 'cval'))

tests/test_num.py

@@ -1,6 +1,6 @@
 
 import numpy
-from numpy.testing import assert_array_almost_equal
+from numpy.testing import assert_array_almost_equal, assert_array_equal
 import pytest
 
 from .utils import with_backends, get_backend_module, get_backend_scipy, mock_importerror
@@ -9,7 +9,8 @@
     get_array_module, get_scipy_module,
     to_real_dtype, to_complex_dtype,
     fft2, ifft2, abs2,
-    to_numpy, as_array
+    to_numpy, as_array,
+    ufunc_outer
 )
 
 
@@ -197,4 +198,16 @@ def test_to_array(backend: str):
     assert_array_almost_equal(
         arr,
         numpy.array([1., 2., 3., 4.])
-    )
+    )
+
+
+@with_backends('cpu', 'jax', 'cuda')
+def test_ufunc_outer(backend: str):
+    xp = get_backend_module(backend)
+
+    xs = numpy.arange(12).reshape(4, 3)
+    ys = numpy.arange(30).reshape(5, 6)
+
+    expected = numpy.multiply.outer(xs, ys)
+    actual = to_numpy(ufunc_outer(xp.multiply, xp.array(xs), xp.array(ys)))
+    assert_array_equal(expected, actual)