From 1884eaa5a0c92b1c8b51b1f9a2e7828076d4d054 Mon Sep 17 00:00:00 2001 From: Jaroslav Fowkes Date: Thu, 12 Sep 2024 11:10:15 +0200 Subject: [PATCH 1/4] Add Lipschitz preconditioner --- ptypy/engines/ML.py | 73 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/ptypy/engines/ML.py b/ptypy/engines/ML.py index 9d5bb7d1..aff77fa3 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' + [lipschitz_precond] + default = False + type = bool + help = Whether to use the Lipschitz preconditioner + doc = This parameter can give faster convergence. + + [lipschitz_delta_object] + default = 0.1 + type = float + help = Lipschitz preconditioner damping constant for the object. + + [lipschitz_delta_probe] + default = 0.1 + type = float + help = Lipschitz 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 normalisation + if self.p.lipschitz_precond: + self.ob_nrm = None + self.pr_nrm = 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 normalisation + if self.p.lipschitz_precond: + self.ob_nrm = self.ob.copy(self.ob.ID + '_nrm', fill=0., dtype='real') + self.pr_nrm = self.pr.copy(self.pr.ID + '_nrm', fill=0., dtype='real') + self.tmin = 1. # Other options @@ -290,6 +315,13 @@ def engine_iterate(self, num=1): self.pr_grad *= self.scale_p_o self.pr_h -= self.pr_grad + # Lipschitz preconditioner + if self.p.lipschitz_precond: + self.ob_nrm += self.p.lipschitz_delta_object + self.ob_h /= self.ob_nrm + self.pr_nrm += self.p.lipschitz_delta_probe + self.pr_h /= self.pr_nrm + # In principle, the way things are now programmed this part # could be iterated over in a real Newton-Raphson style. t2 = time.time() @@ -353,6 +385,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.lipschitz_precond: + del self.ptycho.containers[self.ob_nrm.ID] + del self.ob_nrm + del self.ptycho.containers[self.pr_nrm.ID] + del self.pr_nrm # Save floating intensities into runtime self.ptycho.runtime["float_intens"] = parallel.gather_dict(self.ML_model.float_intens_coeff) @@ -377,6 +414,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.lipschitz_precond: + self.ob_nrm = self.engine.ob_nrm + self.pr_nrm = self.engine.pr_nrm self.pr = self.engine.pr self.float_intens_coeff = {} @@ -490,6 +530,9 @@ def new_grad(self): """ self.ob_grad.fill(0.) self.pr_grad.fill(0.) + if self.p.lipschitz_precond: + self.ob_nrm.fill(0.) + self.pr_nrm.fill(0.) # We need an array for MPI LL = np.array([0.]) @@ -531,6 +574,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 normalisations for object and probe + if self.p.lipschitz_precond: + self.ob_nrm[pod.ob_view] += u.abs2(pod.probe) + self.pr_nrm[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 +586,9 @@ def new_grad(self): # MPI reduction of gradients self.ob_grad.allreduce() self.pr_grad.allreduce() + if self.p.lipschitz_precond: + self.ob_nrm.allreduce() + self.pr_nrm.allreduce() parallel.allreduce(LL) # Object regularizer @@ -733,6 +784,9 @@ def new_grad(self): """ self.ob_grad.fill(0.) self.pr_grad.fill(0.) + if self.p.lipschitz_precond: + self.ob_nrm.fill(0.) + self.pr_nrm.fill(0.) # We need an array for MPI LL = np.array([0.]) @@ -774,6 +828,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 normalisations for object and probe + if self.p.lipschitz_precond: + self.ob_nrm[pod.ob_view] += u.abs2(pod.probe) + self.pr_nrm[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 +840,9 @@ def new_grad(self): # MPI reduction of gradients self.ob_grad.allreduce() self.pr_grad.allreduce() + if self.p.lipschitz_precond: + self.ob_nrm.allreduce() + self.pr_nrm.allreduce() parallel.allreduce(LL) # Object regularizer @@ -989,6 +1051,9 @@ def new_grad(self): """ self.ob_grad.fill(0.) self.pr_grad.fill(0.) + if self.p.lipschitz_precond: + self.ob_nrm.fill(0.) + self.pr_nrm.fill(0.) # We need an array for MPI LL = np.array([0.]) @@ -1030,6 +1095,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 normalisations for object and probe + if self.p.lipschitz_precond: + self.ob_nrm[pod.ob_view] += u.abs2(pod.probe) + self.pr_nrm[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 +1107,9 @@ def new_grad(self): # MPI reduction of gradients self.ob_grad.allreduce() self.pr_grad.allreduce() + if self.p.lipschitz_precond: + self.ob_nrm.allreduce() + self.pr_nrm.allreduce() parallel.allreduce(LL) # Object regularizer From 4b3c21fc53b84ed88c021dbeea8f0e3e7d176447 Mon Sep 17 00:00:00 2001 From: Jaroslav Fowkes Date: Thu, 12 Sep 2024 19:14:28 +0200 Subject: [PATCH 2/4] Move Lipschitz preconditioner before smoothing preconditioner --- ptypy/engines/ML.py | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/ptypy/engines/ML.py b/ptypy/engines/ML.py index aff77fa3..db5107f5 100644 --- a/ptypy/engines/ML.py +++ b/ptypy/engines/ML.py @@ -255,6 +255,13 @@ def engine_iterate(self, num=1): else: new_pr_grad.fill(0.) + # Lipschitz preconditioner + if self.p.lipschitz_precond: + self.ob_nrm += self.p.lipschitz_delta_object + self.pr_nrm += self.p.lipschitz_delta_probe + new_ob_grad /= self.ob_nrm + new_pr_grad /= self.pr_nrm + # Smoothing preconditioner if self.smooth_gradient: self.smooth_gradient.sigma *= (1. - self.p.smooth_gradient_decay) @@ -301,10 +308,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 + # Smoothing and Lipschitz preconditioners for the object + # if self.smooth_gradient and self.p.lipschitz_precond: + # for name, s in self.ob_h.storages.items(): + # s.data[:] -= self.smooth_gradient(self.ob_grad.storages[name].data / self.ob_nrm.storages[name].data) + # elif self.p.lipschitz_precond: + # for name, s in self.ob_h.storages.items(): + # s.data[:] -= self.ob_grad.storages[name].data / self.ob_nrm.storages[name].data + # elif self.smooth_gradient: if self.smooth_gradient: for name, s in self.ob_h.storages.items(): s.data[:] -= self.smooth_gradient(self.ob_grad.storages[name].data) @@ -313,15 +327,13 @@ def engine_iterate(self, num=1): self.pr_h *= bt / self.tmin self.pr_grad *= self.scale_p_o + # # Lipschitz preconditioner for the probe + # if self.p.lipschitz_precond: + # for name, s in self.pr_h.storages.items(): + # s.data[:] -= self.pr_grad.storages[name].data / self.pr_nrm.storages[name].data + # else: self.pr_h -= self.pr_grad - # Lipschitz preconditioner - if self.p.lipschitz_precond: - self.ob_nrm += self.p.lipschitz_delta_object - self.ob_h /= self.ob_nrm - self.pr_nrm += self.p.lipschitz_delta_probe - self.pr_h /= self.pr_nrm - # In principle, the way things are now programmed this part # could be iterated over in a real Newton-Raphson style. t2 = time.time() From 8701960bf331ae7a1f83494243722e2177b1dff6 Mon Sep 17 00:00:00 2001 From: Jaroslav Fowkes Date: Wed, 18 Sep 2024 07:38:52 +0100 Subject: [PATCH 3/4] Use sqrt of fluence and apply twice --- ptypy/engines/ML.py | 149 +++++++++--------- .../notebooks/moonflower_ml_wavefield.ipynb | 147 +++++++++++++++++ 2 files changed, 221 insertions(+), 75 deletions(-) create mode 100644 templates/notebooks/moonflower_ml_wavefield.ipynb diff --git a/ptypy/engines/ML.py b/ptypy/engines/ML.py index db5107f5..827fe119 100644 --- a/ptypy/engines/ML.py +++ b/ptypy/engines/ML.py @@ -106,21 +106,21 @@ class ML(PositionCorrectionEngine): help = How many coefficients to be used in the the linesearch doc = choose between the 'quadratic' approximation (default) or 'all' - [lipschitz_precond] + [wavefield_precond] default = False type = bool - help = Whether to use the Lipschitz preconditioner + help = Whether to use the wavefield preconditioner doc = This parameter can give faster convergence. - [lipschitz_delta_object] + [wavefield_delta_object] default = 0.1 type = float - help = Lipschitz preconditioner damping constant for the object. + help = Wavefield preconditioner damping constant for the object. - [lipschitz_delta_probe] + [wavefield_delta_probe] default = 0.1 type = float - help = Lipschitz preconditioner damping constant for the probe. + help = Wavefield preconditioner damping constant for the probe. """ @@ -153,10 +153,10 @@ def __init__(self, ptycho_parent, pars=None): # Probe gradient self.pr_grad_new = None - # Object and probe normalisation - if self.p.lipschitz_precond: - self.ob_nrm = None - self.pr_nrm = None + # Object and probe fluence maps + if self.p.wavefield_precond: + self.ob_fln = None + self.pr_fln = None # Other self.tmin = None @@ -192,10 +192,10 @@ 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 normalisation - if self.p.lipschitz_precond: - self.ob_nrm = self.ob.copy(self.ob.ID + '_nrm', fill=0., dtype='real') - self.pr_nrm = self.pr.copy(self.pr.ID + '_nrm', fill=0., dtype='real') + # 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. @@ -255,12 +255,12 @@ def engine_iterate(self, num=1): else: new_pr_grad.fill(0.) - # Lipschitz preconditioner - if self.p.lipschitz_precond: - self.ob_nrm += self.p.lipschitz_delta_object - self.pr_nrm += self.p.lipschitz_delta_probe - new_ob_grad /= self.ob_nrm - new_pr_grad /= self.pr_nrm + # Wavefield preconditioner + if self.p.wavefield_precond: + self.ob_fln += self.p.wavefield_delta_object + self.pr_fln += self.p.wavefield_delta_probe + new_ob_grad /= self.ob_fln + new_pr_grad /= self.pr_fln # Smoothing preconditioner if self.smooth_gradient: @@ -311,15 +311,14 @@ def engine_iterate(self, num=1): # 3. Next conjugate self.ob_h *= bt / self.tmin - # Smoothing and Lipschitz preconditioners for the object - # if self.smooth_gradient and self.p.lipschitz_precond: - # for name, s in self.ob_h.storages.items(): - # s.data[:] -= self.smooth_gradient(self.ob_grad.storages[name].data / self.ob_nrm.storages[name].data) - # elif self.p.lipschitz_precond: - # for name, s in self.ob_h.storages.items(): - # s.data[:] -= self.ob_grad.storages[name].data / self.ob_nrm.storages[name].data - # elif self.smooth_gradient: - 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 / 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 / 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: @@ -327,12 +326,12 @@ def engine_iterate(self, num=1): self.pr_h *= bt / self.tmin self.pr_grad *= self.scale_p_o - # # Lipschitz preconditioner for the probe - # if self.p.lipschitz_precond: - # for name, s in self.pr_h.storages.items(): - # s.data[:] -= self.pr_grad.storages[name].data / self.pr_nrm.storages[name].data - # else: - 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 / 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. @@ -353,9 +352,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 @@ -397,11 +396,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.lipschitz_precond: - del self.ptycho.containers[self.ob_nrm.ID] - del self.ob_nrm - del self.ptycho.containers[self.pr_nrm.ID] - del self.pr_nrm + 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) @@ -426,9 +425,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.lipschitz_precond: - self.ob_nrm = self.engine.ob_nrm - self.pr_nrm = self.engine.pr_nrm + 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 = {} @@ -542,9 +541,9 @@ def new_grad(self): """ self.ob_grad.fill(0.) self.pr_grad.fill(0.) - if self.p.lipschitz_precond: - self.ob_nrm.fill(0.) - self.pr_nrm.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.]) @@ -586,10 +585,10 @@ 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 normalisations for object and probe - if self.p.lipschitz_precond: - self.ob_nrm[pod.ob_view] += u.abs2(pod.probe) - self.pr_nrm[pod.pr_view] += u.abs2(pod.object) + # Compute fluence maps for object and probe + if self.p.wavefield_precond: + self.ob_fln[pod.ob_view] += np.sqrt(u.abs2(pod.probe)) + self.pr_fln[pod.pr_view] += np.sqrt(u.abs2(pod.object)) diff_view.error = LLL error_dct[dname] = np.array([0, LLL / np.prod(DI.shape), 0]) @@ -598,9 +597,9 @@ def new_grad(self): # MPI reduction of gradients self.ob_grad.allreduce() self.pr_grad.allreduce() - if self.p.lipschitz_precond: - self.ob_nrm.allreduce() - self.pr_nrm.allreduce() + if self.p.wavefield_precond: + self.ob_fln.allreduce() + self.pr_fln.allreduce() parallel.allreduce(LL) # Object regularizer @@ -796,9 +795,9 @@ def new_grad(self): """ self.ob_grad.fill(0.) self.pr_grad.fill(0.) - if self.p.lipschitz_precond: - self.ob_nrm.fill(0.) - self.pr_nrm.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.]) @@ -840,10 +839,10 @@ 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 normalisations for object and probe - if self.p.lipschitz_precond: - self.ob_nrm[pod.ob_view] += u.abs2(pod.probe) - self.pr_nrm[pod.pr_view] += u.abs2(pod.object) + # Compute fluence maps for object and probe + if self.p.wavefield_precond: + self.ob_fln[pod.ob_view] += np.sqrt(u.abs2(pod.probe)) + self.pr_fln[pod.pr_view] += np.sqrt(u.abs2(pod.object)) diff_view.error = LLL error_dct[dname] = np.array([0, LLL / np.prod(DI.shape), 0]) @@ -852,9 +851,9 @@ def new_grad(self): # MPI reduction of gradients self.ob_grad.allreduce() self.pr_grad.allreduce() - if self.p.lipschitz_precond: - self.ob_nrm.allreduce() - self.pr_nrm.allreduce() + if self.p.wavefield_precond: + self.ob_fln.allreduce() + self.pr_fln.allreduce() parallel.allreduce(LL) # Object regularizer @@ -1063,9 +1062,9 @@ def new_grad(self): """ self.ob_grad.fill(0.) self.pr_grad.fill(0.) - if self.p.lipschitz_precond: - self.ob_nrm.fill(0.) - self.pr_nrm.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.]) @@ -1107,10 +1106,10 @@ 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 normalisations for object and probe - if self.p.lipschitz_precond: - self.ob_nrm[pod.ob_view] += u.abs2(pod.probe) - self.pr_nrm[pod.pr_view] += u.abs2(pod.object) + # Compute fluence maps for object and probe + if self.p.wavefield_precond: + self.ob_fln[pod.ob_view] += np.sqrt(u.abs2(pod.probe)) + self.pr_fln[pod.pr_view] += np.sqrt(u.abs2(pod.object)) diff_view.error = LLL error_dct[dname] = np.array([0, LLL / np.prod(DA.shape), 0]) @@ -1119,9 +1118,9 @@ def new_grad(self): # MPI reduction of gradients self.ob_grad.allreduce() self.pr_grad.allreduce() - if self.p.lipschitz_precond: - self.ob_nrm.allreduce() - self.pr_nrm.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 +} From f6b3437f85de091ebd69a726ca95fe3198145051 Mon Sep 17 00:00:00 2001 From: Jaroslav Fowkes Date: Wed, 18 Sep 2024 08:58:23 +0100 Subject: [PATCH 4/4] Take sqrt of the sum --- ptypy/engines/ML.py | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/ptypy/engines/ML.py b/ptypy/engines/ML.py index 827fe119..9d43efba 100644 --- a/ptypy/engines/ML.py +++ b/ptypy/engines/ML.py @@ -259,8 +259,9 @@ def engine_iterate(self, num=1): if self.p.wavefield_precond: self.ob_fln += self.p.wavefield_delta_object self.pr_fln += self.p.wavefield_delta_probe - new_ob_grad /= self.ob_fln - new_pr_grad /= self.pr_fln + 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: @@ -314,10 +315,10 @@ def engine_iterate(self, num=1): # 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 / self.ob_fln.storages[name].data) + 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 / self.ob_fln.storages[name].data + 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) @@ -329,7 +330,7 @@ def engine_iterate(self, num=1): # 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 / self.pr_fln.storages[name].data + s.data[:] -= self.pr_grad.storages[name].data / np.sqrt(self.pr_fln.storages[name].data) else: self.pr_h -= self.pr_grad @@ -587,8 +588,8 @@ def new_grad(self): # Compute fluence maps for object and probe if self.p.wavefield_precond: - self.ob_fln[pod.ob_view] += np.sqrt(u.abs2(pod.probe)) - self.pr_fln[pod.pr_view] += np.sqrt(u.abs2(pod.object)) + 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]) @@ -841,8 +842,8 @@ def new_grad(self): # Compute fluence maps for object and probe if self.p.wavefield_precond: - self.ob_fln[pod.ob_view] += np.sqrt(u.abs2(pod.probe)) - self.pr_fln[pod.pr_view] += np.sqrt(u.abs2(pod.object)) + 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]) @@ -1108,8 +1109,8 @@ def new_grad(self): # Compute fluence maps for object and probe if self.p.wavefield_precond: - self.ob_fln[pod.ob_view] += np.sqrt(u.abs2(pod.probe)) - self.pr_fln[pod.pr_view] += np.sqrt(u.abs2(pod.object)) + 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])