From 501eb4d00ec3ff656266baa1e70d102d444d1c99 Mon Sep 17 00:00:00 2001 From: Jari Date: Fri, 6 Dec 2024 13:42:22 +0000 Subject: [PATCH] Make requested changes to ML.py --- ptypy/engines/ML.py | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/ptypy/engines/ML.py b/ptypy/engines/ML.py index 9d43efbad..f369e3b44 100644 --- a/ptypy/engines/ML.py +++ b/ptypy/engines/ML.py @@ -154,9 +154,8 @@ def __init__(self, ptycho_parent, pars=None): self.pr_grad_new = None # Object and probe fluence maps - if self.p.wavefield_precond: - self.ob_fln = None - self.pr_fln = None + self.ob_fln = None + self.pr_fln = None # Other self.tmin = None @@ -257,11 +256,9 @@ def engine_iterate(self, num=1): # 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) + new_ob_grad.storages[name].data /= np.sqrt(self.ob_fln.storages[name].data + self.p.wavefield_delta_object) + new_pr_grad.storages[name].data /= np.sqrt(self.pr_fln.storages[name].data + self.p.wavefield_delta_probe) # Smoothing preconditioner if self.smooth_gradient: @@ -312,14 +309,17 @@ def engine_iterate(self, num=1): # 3. Next conjugate self.ob_h *= bt / self.tmin - # 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: + # Wavefield preconditioner for the object (with and without smoothing preconditioner) + if 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: + if self.smooth_gradient: + s.data[:] -= self.smooth_gradient(self.ob_grad.storages[name].data + / np.sqrt(self.ob_fln.storages[name].data + self.p.wavefield_delta_object)) + else: + s.data[:] -= (self.ob_grad.storages[name].data + / np.sqrt(self.ob_fln.storages[name].data + self.p.wavefield_delta_object)) + # Smoothing preconditioner for the object + if self.smooth_gradient: for name, s in self.ob_h.storages.items(): s.data[:] -= self.smooth_gradient(self.ob_grad.storages[name].data) else: @@ -330,7 +330,8 @@ 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 / np.sqrt(self.pr_fln.storages[name].data) + s.data[:] -= (self.pr_grad.storages[name].data + / np.sqrt(self.pr_fln.storages[name].data + self.p.wavefield_delta_probe)) else: self.pr_h -= self.pr_grad @@ -426,9 +427,8 @@ 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.ob_fln = self.engine.ob_fln + self.pr_fln = self.engine.pr_fln self.pr = self.engine.pr self.float_intens_coeff = {}