Merge pull request #316 from carterbox/multi-slice-adjoint

NEW: Implement multislice propagation operator with correct adjoint

src/tike/operators/cupy/convolution.py

@@ -102,9 +102,9 @@ def fwd(self, psi, scan, probe):
 
     def adj(self, nearplane, scan, probe, psi=None, overwrite=False):
         """Combine probe shaped patches into a psi shaped grid by addition."""
-        assert probe.shape[:-4] == scan.shape[:-2]
+        assert probe.shape[:-4] == scan.shape[:-2], (probe.shape, scan.shape)
         assert probe.shape[-4] == 1 or probe.shape[-4] == scan.shape[-2]
-        assert nearplane.shape[:-3] == scan.shape[:-1]
+        assert nearplane.shape[:-3] == scan.shape[:-1], (nearplane.shape, scan.shape)
         if not overwrite:
             nearplane = nearplane.copy()
         nearplane[..., self.pad:self.end, self.pad:self.end] *= probe.conj()

src/tike/operators/cupy/multislice.py

@@ -6,22 +6,24 @@
 import typing
 import numpy.typing as npt
 import numpy as np
-import cupy as cp
 
 from .operator import Operator
 
+from .fresnelspectprop import FresnelSpectProp
 from .propagation import Propagation, ZeroPropagation
 from .convolution import Convolution
 
 
 class Multislice(Operator):
+    """Multiple slice wavefield propgation"""
+
     def __init__(
         self,
         detector_shape: int,
         probe_shape: int,
         nz: int,
         n: int,
-        propagation: typing.Type[Propagation] = ZeroPropagation,
+        propagation: typing.Type[Propagation] = FresnelSpectProp,
         diffraction: typing.Type[Convolution] = Convolution,
         norm: str = "ortho",
         nslices: int = 1,
@@ -46,6 +48,99 @@ def __init__(
         self.n = n
         self.nslices = nslices
 
+    def __enter__(self):
+        self.propagation.__enter__()
+        self.diffraction.__enter__()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.propagation.__exit__(type, value, traceback)
+        self.diffraction.__exit__(type, value, traceback)
+
+    def fwd(
+        self,
+        probe: npt.NDArray[np.csingle],
+        scan: npt.NDArray[np.single],
+        psi: npt.NDArray[np.csingle],
+        **kwargs,
+    ) -> npt.NDArray[np.csingle]:
+        """Please see help(Multislice) for more info."""
+        assert psi.shape[0] == self.nslices and psi.ndim == 3
+        exitwave = self.diffraction.fwd(
+            psi=psi[0],
+            scan=scan,
+            probe=probe,
+        )
+        for s in range(1, self.nslices):
+            exitwave = self.diffraction.fwd(
+                psi=psi[s],
+                scan=scan,
+                probe=self.propagation.fwd(exitwave),
+            )
+        return exitwave
+
+    def adj(
+        self,
+        nearplane: npt.NDArray[np.csingle],
+        probe: npt.NDArray[np.csingle],
+        scan: npt.NDArray[np.single],
+        psi: npt.NDArray[np.csingle],
+        overwrite: bool = False,
+        **kwargs,
+    ) -> npt.NDArray[np.csingle]:
+        """Please see help(Multislice) for more info."""
+        psi_adj = self.xp.zeros_like(psi)
+        probes = [
+            None,
+        ] * len(psi)
+        probes[0] = probe
+        for s in range(1, self.nslices):
+            probes[s] = self.propagation.fwd(
+                self.diffraction.fwd(
+                    psi=psi[s - 1],
+                    scan=scan,
+                    probe=probes[s - 1],
+                )
+            )
+        psi_adj[self.nslices - 1] = self.diffraction.adj(
+            nearplane=nearplane,
+            probe=probes[self.nslices - 1],
+            scan=scan,
+            overwrite=False,
+        )
+        probe_adj = self.diffraction.adj_probe(
+            nearplane=nearplane,
+            scan=scan,
+            psi=psi[self.nslices - 1],
+        )
+        for s in range(self.nslices - 2, -1, -1):
+            probe_adj = self.propagation.adj(probe_adj)
+            psi_adj[s] = self.diffraction.adj(
+                nearplane=probe_adj,
+                probe=probes[s],
+                scan=scan,
+                overwrite=False,
+            )
+            probe_adj = self.diffraction.adj_probe(
+                nearplane=probe_adj,
+                scan=scan,
+                psi=psi[s],
+            )
+        # FIXME: Why does correct adjoint require division by nslices?
+        return psi_adj / self.nslices, probe_adj
+
+    @property
+    def patch(self):
+        return self.diffraction.patch
+
+    @property
+    def pad(self):
+        return self.diffraction.pad
+
+    @property
+    def end(self):
+        return self.diffraction.end
+
 
 class SingleSlice(Multislice):
     """Single slice wavefield propgation"""
@@ -80,15 +175,6 @@ def __init__(
         self.n = n
         self.nslices = 1
 
-    def __enter__(self):
-        self.propagation.__enter__()
-        self.diffraction.__enter__()
-        return self
-
-    def __exit__(self, type, value, traceback):
-        self.propagation.__exit__(type, value, traceback)
-        self.diffraction.__exit__(type, value, traceback)
-
     def fwd(
         self,
         probe: npt.NDArray[np.csingle],
@@ -97,7 +183,7 @@ def fwd(
         **kwargs,
     ) -> npt.NDArray[np.csingle]:
         """Please see help(SingleSlice) for more info."""
-        assert (psi.shape[0] == 1 and psi.ndim == 3)
+        assert psi.shape[0] == 1 and psi.ndim == 3
         return self.diffraction.fwd(
             psi=psi[0],
             scan=scan,
@@ -115,31 +201,16 @@ def adj(
     ) -> npt.NDArray[np.csingle]:
         """Please see help(SingleSlice) for more info."""
         assert psi is None or (psi.shape[0] == 1 and psi.ndim == 3)
-        return self.diffraction.adj(
+        psi_adj = self.diffraction.adj(
             nearplane=nearplane,
             probe=probe,
             scan=scan,
-            overwrite=True,
-            psi=psi[0] if psi is not None else None,
+            overwrite=False,
         )[None, ...]
-
-    def adj_probe(self, nearplane, scan, psi, overwrite=False):
-        assert (psi.shape[0] == 1 and psi.ndim == 3)
-        return self.diffraction.adj_probe(
+        probe_adj = self.diffraction.adj_probe(
             nearplane=nearplane,
             scan=scan,
             psi=psi[0],
-            overwrite=overwrite,
+            overwrite=False,
         )
-
-    @property
-    def patch(self):
-        return self.diffraction.patch
-
-    @property
-    def pad(self):
-        return self.diffraction.pad
-
-    @property
-    def end(self):
-        return self.diffraction.end
+        return psi_adj, probe_adj

src/tike/operators/cupy/ptycho.py

@@ -124,12 +124,12 @@ def adj(
         farplane: npt.NDArray[np.csingle],
         probe: npt.NDArray[np.csingle],
         scan: npt.NDArray[np.single],
-        psi: npt.NDArray[np.csingle] = None,
+        psi: npt.NDArray[np.csingle],
         overwrite: bool = False,
         **kwargs,
     ) -> npt.NDArray[np.csingle]:
         """Please see help(Ptycho) for more info."""
-        return self.diffraction.adj(
+        psi_adj, probe_adj = self.diffraction.adj(
             nearplane=self.propagation.adj(
                 farplane,
                 overwrite=overwrite,
@@ -139,25 +139,7 @@ def adj(
             overwrite=True,
             psi=psi,
         )
-
-    def adj_probe(
-        self,
-        farplane: npt.NDArray[np.csingle],
-        scan: npt.NDArray[np.single],
-        psi: npt.NDArray[np.csingle],
-        overwrite: bool = False,
-        **kwargs,
-    ) -> npt.NDArray[np.csingle]:
-        """Please see help(Ptycho) for more info."""
-        return self.diffraction.adj_probe(
-            psi=psi,
-            scan=scan,
-            nearplane=self.propagation.adj(
-                farplane=farplane,
-                overwrite=overwrite,
-            )[..., 0, :, :, :],
-            overwrite=True,
-        )[..., None, :, :, :]
+        return psi_adj, probe_adj[..., None, :, :, :]
 
     def _compute_intensity(
         self,
@@ -186,64 +168,3 @@ def cost(
         """Please see help(Ptycho) for more info."""
         intensity, _ = self._compute_intensity(data, psi, scan, probe)
         return getattr(objective, model)(data, intensity)
-
-    def grad_psi(
-        self,
-        data: npt.NDArray,
-        psi: npt.NDArray[np.csingle],
-        scan: npt.NDArray[np.single],
-        probe: npt.NDArray[np.csingle],
-        *,
-        model: str,
-    ) -> npt.NDArray[np.csingle]:
-        """Please see help(Ptycho) for more info."""
-        intensity, farplane = self._compute_intensity(data, psi, scan, probe)
-        grad_obj = self.xp.zeros_like(psi)
-        grad_obj = self.adj(
-            farplane=getattr(objective, f'{model}_grad')(
-                data,
-                farplane,
-                intensity,
-            ),
-            probe=probe,
-            scan=scan,
-            psi=grad_obj,
-            overwrite=True,
-        )
-        return grad_obj
-
-    def grad_probe(
-        self,
-        data: npt.NDArray,
-        psi: npt.NDArray[np.csingle],
-        scan: npt.NDArray[np.single],
-        probe: npt.NDArray[np.csingle],
-        mode: typing.List[int] = None,
-        *,
-        model: str,
-    ) -> npt.NDArray[np.csingle]:
-        """Compute the gradient with respect to the probe(s).
-
-        Parameters
-        ----------
-        mode : list(int)
-            Only return the gradient with resepect to these probes.
-
-        """
-        mode = list(range(probe.shape[-3])) if mode is None else mode
-        intensity, farplane = self._compute_intensity(data, psi, scan, probe)
-        # Use the average gradient for all probe positions
-        return self.xp.mean(
-            self.adj_probe(
-                farplane=getattr(objective, f'{model}_grad')(
-                    data,
-                    farplane[..., mode, :, :],
-                    intensity,
-                ),
-                psi=psi,
-                scan=scan,
-                overwrite=True,
-            ),
-            axis=0,
-            keepdims=True,
-        )