diff --git a/ptypy/engines/ML.py b/ptypy/engines/ML.py index 9d5bb7d1..9d43efba 100644 --- a/ptypy/engines/ML.py +++ b/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,11 +309,17 @@ 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: @@ -288,7 +327,12 @@ def engine_iterate(self, num=1): 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,6 +586,11 @@ 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 @@ -538,6 +598,9 @@ def new_grad(self): # 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,6 +840,11 @@ 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 @@ -781,6 +852,9 @@ def new_grad(self): # 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,6 +1107,11 @@ 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 @@ -1037,6 +1119,9 @@ def new_grad(self): # 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 diff --git a/templates/notebooks/moonflower_ml_wavefield.ipynb b/templates/notebooks/moonflower_ml_wavefield.ipynb new file mode 100644 index 00000000..201d0c4a --- /dev/null +++ b/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 +}