Add wavefield preconditioner

ptypy/engines/ML.py

@@ -106,6 +106,22 @@ class ML(PositionCorrectionEngine):
     help = How many coefficients to be used in the the linesearch
     doc = choose between the 'quadratic' approximation (default) or 'all'
 
+    [wavefield_precond]
+    default = False
+    type = bool
+    help = Whether to use the wavefield preconditioner
+    doc = This parameter can give faster convergence.
+
+    [wavefield_delta_object]
+    default = 0.1
+    type = float
+    help = Wavefield preconditioner damping constant for the object.
+
+    [wavefield_delta_probe]
+    default = 0.1
+    type = float
+    help = Wavefield preconditioner damping constant for the probe.
+
     """
 
     SUPPORTED_MODELS = [Full, Vanilla, Bragg3dModel, BlockVanilla, BlockFull, GradFull, BlockGradFull]
@@ -137,6 +153,10 @@ def __init__(self, ptycho_parent, pars=None):
         # Probe gradient
         self.pr_grad_new = None
 
+        # Object and probe fluence maps
+        if self.p.wavefield_precond:
+            self.ob_fln = None
+            self.pr_fln = None
 
         # Other
         self.tmin = None
@@ -172,6 +192,11 @@ def engine_initialize(self):
         self.pr_grad_new = self.pr.copy(self.pr.ID + '_grad_new', fill=0.)
         self.pr_h = self.pr.copy(self.pr.ID + '_h', fill=0.)
 
+        # Object and probe fluence maps
+        if self.p.wavefield_precond:
+            self.ob_fln = self.ob.copy(self.ob.ID + '_fln', fill=0., dtype='real')
+            self.pr_fln = self.pr.copy(self.pr.ID + '_fln', fill=0., dtype='real')
+
         self.tmin = 1.
 
         # Other options
@@ -230,6 +255,14 @@ def engine_iterate(self, num=1):
             else:
                 new_pr_grad.fill(0.)
 
+            # Wavefield preconditioner
+            if self.p.wavefield_precond:
+                self.ob_fln += self.p.wavefield_delta_object
+                self.pr_fln += self.p.wavefield_delta_probe
+                for name, s in new_ob_grad.storages.items():
+                    new_ob_grad.storages[name].data /= np.sqrt(self.ob_fln.storages[name].data)
+                    new_pr_grad.storages[name].data /= np.sqrt(self.pr_fln.storages[name].data)
+
             # Smoothing preconditioner
             if self.smooth_gradient:
                 self.smooth_gradient.sigma *= (1. - self.p.smooth_gradient_decay)
@@ -276,19 +309,30 @@ def engine_iterate(self, num=1):
             self.pr_grad << new_pr_grad
 
             dt = self.ptycho.FType
+
             # 3. Next conjugate
             self.ob_h *= bt / self.tmin
-
-            # Smoothing preconditioner
-            if self.smooth_gradient:
+            # Smoothing and wavefield preconditioners for the object
+            if self.smooth_gradient and self.p.wavefield_precond:
+                for name, s in self.ob_h.storages.items():
+                    s.data[:] -= self.smooth_gradient(self.ob_grad.storages[name].data / np.sqrt(self.ob_fln.storages[name].data))
+            elif self.p.wavefield_precond:
+                for name, s in self.ob_h.storages.items():
+                    s.data[:] -= self.ob_grad.storages[name].data / np.sqrt(self.ob_fln.storages[name].data)
+            elif self.smooth_gradient:
                 for name, s in self.ob_h.storages.items():
                     s.data[:] -= self.smooth_gradient(self.ob_grad.storages[name].data)
             else:
                 self.ob_h -= self.ob_grad
 
             self.pr_h *= bt / self.tmin
             self.pr_grad *= self.scale_p_o
-            self.pr_h -= self.pr_grad
+            # Wavefield preconditioner for the probe
+            if self.p.wavefield_precond:
+                for name, s in self.pr_h.storages.items():
+                    s.data[:] -= self.pr_grad.storages[name].data / np.sqrt(self.pr_fln.storages[name].data)
+            else:
+                self.pr_h -= self.pr_grad
 
             # In principle, the way things are now programmed this part
             # could be iterated over in a real Newton-Raphson style.
@@ -309,9 +353,9 @@ def engine_iterate(self, num=1):
                 self.tmin = dt(-0.5 * B[1] / B[2])
             else:
                 raise NotImplementedError("poly_line_coeffs should be 'quadratic' or 'all'")
-                
+
             tc += time.time() - t2
-                
+
             self.ob_h *= self.tmin
             self.pr_h *= self.tmin
             self.ob += self.ob_h
@@ -353,6 +397,11 @@ def engine_finalize(self):
         del self.pr_grad_new
         del self.ptycho.containers[self.pr_h.ID]
         del self.pr_h
+        if self.p.wavefield_precond:
+            del self.ptycho.containers[self.ob_fln.ID]
+            del self.ob_fln
+            del self.ptycho.containers[self.pr_fln.ID]
+            del self.pr_fln
 
         # Save floating intensities into runtime
         self.ptycho.runtime["float_intens"] = parallel.gather_dict(self.ML_model.float_intens_coeff)
@@ -377,6 +426,9 @@ def __init__(self, MLengine):
         self.ob = self.engine.ob
         self.ob_grad = self.engine.ob_grad_new
         self.pr_grad = self.engine.pr_grad_new
+        if self.p.wavefield_precond:
+            self.ob_fln = self.engine.ob_fln
+            self.pr_fln = self.engine.pr_fln
         self.pr = self.engine.pr
         self.float_intens_coeff = {}
 
@@ -490,6 +542,9 @@ def new_grad(self):
         """
         self.ob_grad.fill(0.)
         self.pr_grad.fill(0.)
+        if self.p.wavefield_precond:
+            self.ob_fln.fill(0.)
+            self.pr_fln.fill(0.)
 
         # We need an array for MPI
         LL = np.array([0.])
@@ -531,13 +586,21 @@ def new_grad(self):
                 self.ob_grad[pod.ob_view] += 2. * xi * pod.probe.conj()
                 self.pr_grad[pod.pr_view] += 2. * xi * pod.object.conj()
 
+                # Compute fluence maps for object and probe
+                if self.p.wavefield_precond:
+                    self.ob_fln[pod.ob_view] += u.abs2(pod.probe)
+                    self.pr_fln[pod.pr_view] += u.abs2(pod.object)
+
             diff_view.error = LLL
             error_dct[dname] = np.array([0, LLL / np.prod(DI.shape), 0])
             LL += LLL
 
         # MPI reduction of gradients
         self.ob_grad.allreduce()
         self.pr_grad.allreduce()
+        if self.p.wavefield_precond:
+            self.ob_fln.allreduce()
+            self.pr_fln.allreduce()
         parallel.allreduce(LL)
 
         # Object regularizer
@@ -733,6 +796,9 @@ def new_grad(self):
         """
         self.ob_grad.fill(0.)
         self.pr_grad.fill(0.)
+        if self.p.wavefield_precond:
+            self.ob_fln.fill(0.)
+            self.pr_fln.fill(0.)
 
         # We need an array for MPI
         LL = np.array([0.])
@@ -774,13 +840,21 @@ def new_grad(self):
                 self.ob_grad[pod.ob_view] += 2 * xi * pod.probe.conj()
                 self.pr_grad[pod.pr_view] += 2 * xi * pod.object.conj()
 
+                # Compute fluence maps for object and probe
+                if self.p.wavefield_precond:
+                    self.ob_fln[pod.ob_view] += u.abs2(pod.probe)
+                    self.pr_fln[pod.pr_view] += u.abs2(pod.object)
+
             diff_view.error = LLL
             error_dct[dname] = np.array([0, LLL / np.prod(DI.shape), 0])
             LL += LLL
 
         # MPI reduction of gradients
         self.ob_grad.allreduce()
         self.pr_grad.allreduce()
+        if self.p.wavefield_precond:
+            self.ob_fln.allreduce()
+            self.pr_fln.allreduce()
         parallel.allreduce(LL)
 
         # Object regularizer
@@ -989,6 +1063,9 @@ def new_grad(self):
         """
         self.ob_grad.fill(0.)
         self.pr_grad.fill(0.)
+        if self.p.wavefield_precond:
+            self.ob_fln.fill(0.)
+            self.pr_fln.fill(0.)
 
         # We need an array for MPI
         LL = np.array([0.])
@@ -1030,13 +1107,21 @@ def new_grad(self):
                 self.ob_grad[pod.ob_view] += 2. * xi * pod.probe.conj()
                 self.pr_grad[pod.pr_view] += 2. * xi * pod.object.conj()
 
+                # Compute fluence maps for object and probe
+                if self.p.wavefield_precond:
+                    self.ob_fln[pod.ob_view] += u.abs2(pod.probe)
+                    self.pr_fln[pod.pr_view] += u.abs2(pod.object)
+
             diff_view.error = LLL
             error_dct[dname] = np.array([0, LLL / np.prod(DA.shape), 0])
             LL += LLL
 
         # MPI reduction of gradients
         self.ob_grad.allreduce()
         self.pr_grad.allreduce()
+        if self.p.wavefield_precond:
+            self.ob_fln.allreduce()
+            self.pr_fln.allreduce()
         parallel.allreduce(LL)
 
         # Object regularizer

templates/notebooks/moonflower_ml_wavefield.ipynb

@@ -0,0 +1,147 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## PtyPy moonflower example\n",
+    "#### scan model: BlockGradFull\n",
+    "#### engine: Maximum Likelihood (ML)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from ptypy.core import Ptycho\n",
+    "from ptypy import utils as u"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# create parameter tree\n",
+    "p = u.Param()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# set verbose level to interactive\n",
+    "p.verbose_level = \"interactive\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# set home path and io settings (no files saved)\n",
+    "p.io = u.Param()\n",
+    "p.io.rfile = None\n",
+    "p.io.autosave = u.Param(active=False)\n",
+    "p.io.autoplot = u.Param(active=False)\n",
+    "p.io.interaction = u.Param(active=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# max 200 frames (128x128px) of diffraction data\n",
+    "p.scans = u.Param()\n",
+    "p.scans.MF = u.Param()\n",
+    "p.scans.MF.name = 'BlockGradFull'\n",
+    "p.scans.MF.data= u.Param()\n",
+    "p.scans.MF.data.name = 'MoonFlowerScan'\n",
+    "p.scans.MF.data.shape = 128\n",
+    "p.scans.MF.data.num_frames = 200\n",
+    "p.scans.MF.data.save = None\n",
+    "p.scans.MF.data.density = 0.2\n",
+    "p.scans.MF.data.photons = 1e8\n",
+    "p.scans.MF.data.psf = 0."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# maximum likelihood reconstrucion engine\n",
+    "p.engines = u.Param()\n",
+    "p.engines.engine00 = u.Param()\n",
+    "p.engines.engine00.name = 'ML'\n",
+    "p.engines.engine00.ML_type = 'Gaussian'\n",
+    "p.engines.engine00.wavefield_precond = True\n",
+    "#p.engines.engine00.reg_del2 = True\n",
+    "#p.engines.engine00.reg_del2_amplitude = 1.\n",
+    "#p.engines.engine00.scale_precond = True\n",
+    "#p.engines.engine00.smooth_gradient = 20.\n",
+    "#p.engines.engine00.smooth_gradient_decay = 1/50.\n",
+    "#p.engines.engine00.floating_intensities = False\n",
+    "p.engines.engine00.numiter = 200"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# prepare and run\n",
+    "P = Ptycho(p,level=5)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Plotting the results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import ptypy.utils.plot_client as pc\n",
+    "fig = pc.figure_from_ptycho(P)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.6"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}