From c5d224cba8f02d62e7449152ea29a247aff2304c Mon Sep 17 00:00:00 2001 From: Jaroslav Fowkes Date: Thu, 21 Nov 2024 14:53:54 +0000 Subject: [PATCH] Add Poisson and Euclid models to ML_separate_grads --- ptypy/custom/ML_separate_grads.py | 736 +++++++++++++++++++++++++++++- 1 file changed, 729 insertions(+), 7 deletions(-) diff --git a/ptypy/custom/ML_separate_grads.py b/ptypy/custom/ML_separate_grads.py index abb9e45f5..3c77e4b54 100644 --- a/ptypy/custom/ML_separate_grads.py +++ b/ptypy/custom/ML_separate_grads.py @@ -4,7 +4,7 @@ TODO. - * Implement other ML models (Poisson/Euclid) + * Implement other regularizers This file is part of the PTYPY package. @@ -181,10 +181,10 @@ def _initialize_model(self): # Create noise model if self.p.ML_type.lower() == "gaussian": self.ML_model = GaussianModel(self) - # elif self.p.ML_type.lower() == "poisson": - # self.ML_model = PoissonModel(self) - # elif self.p.ML_type.lower() == "euclid": - # self.ML_model = EuclidModel(self) + elif self.p.ML_type.lower() == "poisson": + self.ML_model = PoissonModel(self) + elif self.p.ML_type.lower() == "euclid": + self.ML_model = EuclidModel(self) else: raise RuntimeError("Unsupported ML_type: '%s'" % self.p.ML_type) @@ -707,7 +707,7 @@ def poly_line_coeffs_pr(self, pr_h): def poly_line_all_coeffs_ob(self, ob_h): """ Compute all the coefficients of the polynomial for line minimization - in direction h + in direction h for the object """ B = np.zeros((5,), dtype=np.longdouble) @@ -777,7 +777,7 @@ def poly_line_all_coeffs_ob(self, ob_h): def poly_line_all_coeffs_pr(self, pr_h): """ Compute all the coefficients of the polynomial for line minimization - in direction h + in direction h for the probe """ B = np.zeros((5,), dtype=np.longdouble) @@ -839,6 +839,728 @@ def poly_line_all_coeffs_pr(self, pr_h): return B +class PoissonModel(BaseModel): + """ + Poisson noise model. + """ + + def __init__(self, MLengine): + """ + Core functions for ML computation using a Poisson model. + """ + BaseModel.__init__(self, MLengine) + from scipy import special + self.LLbase = {} + for name, di_view in self.di.views.items(): + if not di_view.active: + continue + self.LLbase[name] = special.gammaln(di_view.data+1).sum() + + def new_grad(self): + """ + Compute a new gradient direction according to a Poisson noise model. + + Note: The negative log-likelihood and local errors are also computed + here. + """ + self.ob_grad.fill(0.) + self.pr_grad.fill(0.) + + # We need an array for MPI + LL = np.array([0.]) + error_dct = {} + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Mask and intensities for this view + I = diff_view.data + m = diff_view.pod.ma_view.data + + Imodel = np.zeros_like(I) + f = {} + + # First pod loop: compute total intensity + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f[name] = pod.fw(pod.probe * pod.object) + Imodel += u.abs2(f[name]) + + # Floating intensity option + if self.p.floating_intensities: + self.float_intens_coeff[dname] = I.sum() / Imodel.sum() + Imodel *= self.float_intens_coeff[dname] + + Imodel += 1e-6 + DI = m * (1. - I / Imodel) + + # Second pod loop: gradients computation + LLL = self.LLbase[dname] + (m * (Imodel - I * np.log(Imodel))).sum().astype(np.float64) + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + xi = pod.bw(DI * f[name]) + self.ob_grad[pod.ob_view] += 2 * xi * pod.probe.conj() + self.pr_grad[pod.pr_view] += 2 * xi * pod.object.conj() + + 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() + parallel.allreduce(LL) + + # Object regularizer + if self.regularizer: + for name, s in self.ob.storages.items(): + self.ob_grad.storages[name].data += self.regularizer.grad( + s.data) + LL += self.regularizer.LL + + self.LL = LL / self.tot_measpts + + return error_dct + + def poly_line_coeffs_ob(self, ob_h): + """ + Compute the coefficients of the polynomial for line minimization + in direction h for the object + """ + B = np.zeros((3,), dtype=np.longdouble) + Brenorm = 1/(self.tot_measpts * self.LL[0])**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and intensities for this view + I = diff_view.data + m = diff_view.pod.ma_view.data + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pod.probe * ob_h[pod.ob_view]) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-6 + DI = 1. - I/A0 + + B[0] += (self.LLbase[dname] + (m * (A0 - I * np.log(A0))).sum().astype(np.float64)) * Brenorm + B[1] += np.dot(m.flat, (A1*DI).flat) * Brenorm + B[2] += (np.dot(m.flat, (A2*DI).flat) + 0.5*np.dot(m.flat, (I*(A1/A0)**2).flat)) * Brenorm + + parallel.allreduce(B) + + # Object regularizer + if self.regularizer: + for name, s in self.ob.storages.items(): + B += Brenorm * self.regularizer.poly_line_coeffs( + ob_h.storages[name].data, s.data) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + + def poly_line_coeffs_pr(self, pr_h): + """ + Compute the coefficients of the polynomial for line minimization + in direction h for the probe + """ + B = np.zeros((3,), dtype=np.longdouble) + Brenorm = 1/(self.tot_measpts * self.LL[0])**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and intensities for this view + I = diff_view.data + m = diff_view.pod.ma_view.data + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pr_h[pod.pr_view] * pod.object) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-6 + DI = 1. - I/A0 + + B[0] += (self.LLbase[dname] + (m * (A0 - I * np.log(A0))).sum().astype(np.float64)) * Brenorm + B[1] += np.dot(m.flat, (A1*DI).flat) * Brenorm + B[2] += (np.dot(m.flat, (A2*DI).flat) + 0.5*np.dot(m.flat, (I*(A1/A0)**2).flat)) * Brenorm + + parallel.allreduce(B) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + + def poly_line_all_coeffs_ob(self, ob_h): + """ + Compute all the coefficients of the polynomial for line minimization + in direction h for the object + """ + B = np.zeros((5,), dtype=np.longdouble) + Brenorm = 1/(self.tot_measpts * self.LL[0])**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and intensities for this view + I = diff_view.data + m = diff_view.pod.ma_view.data + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pod.probe * ob_h[pod.ob_view]) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-6 + DI = 1. - I/A0 + + B[0] += (self.LLbase[dname] + (m * (A0 - I * np.log(A0))).sum().astype(np.float64)) * Brenorm + B[1] += np.dot(m.flat, (A1*DI).flat) * Brenorm + B[2] += (np.dot(m.flat, (A2*DI).flat) + 0.5*np.dot(m.flat, (I*(A1/A0)**2).flat)) * Brenorm + B[3] += (0.5*np.dot(m.flat, (I*((2*A1*A2)/A0**2)).flat)) * Brenorm + B[4] += (0.5*np.dot(m.flat, (I*((A2**2)/A0**2)).flat)) * Brenorm + + parallel.allreduce(B) + + # Object regularizer + if self.regularizer: + for name, s in self.ob.storages.items(): + B[:3] += Brenorm * self.regularizer.poly_line_coeffs( + ob_h.storages[name].data, s.data) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + +def poly_line_all_coeffs_pr(self, pr_h): + """ + Compute all the coefficients of the polynomial for line minimization + in direction h for the probe + """ + B = np.zeros((5,), dtype=np.longdouble) + Brenorm = 1/(self.tot_measpts * self.LL[0])**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and intensities for this view + I = diff_view.data + m = diff_view.pod.ma_view.data + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pr_h[pod.pr_view] * pod.object) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-6 + DI = 1. - I/A0 + + B[0] += (self.LLbase[dname] + (m * (A0 - I * np.log(A0))).sum().astype(np.float64)) * Brenorm + B[1] += np.dot(m.flat, (A1*DI).flat) * Brenorm + B[2] += (np.dot(m.flat, (A2*DI).flat) + 0.5*np.dot(m.flat, (I*(A1/A0)**2).flat)) * Brenorm + B[3] += (0.5*np.dot(m.flat, (I*((2*A1*A2)/A0**2)).flat)) * Brenorm + B[4] += (0.5*np.dot(m.flat, (I*((A2**2)/A0**2)).flat)) * Brenorm + + parallel.allreduce(B) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + + +class EuclidModel(BaseModel): + """ + Euclidean (Amplitude) noise model. + TODO: feed actual statistical weights instead of using a fixed variance. + """ + + def __init__(self, MLengine): + """ + Core functions for ML computation using a Euclidean model. + """ + BaseModel.__init__(self, MLengine) + + # Euclidean model requires weights + # TODO: update this part of the code once actual weights are passed in the PODs + self.weights = self.engine.di.copy(self.engine.di.ID + '_weights') + # FIXME: This part needs to be updated once statistical weights are properly + # supported in the data preparation. + for name, di_view in self.di.views.items(): + if not di_view.active: + continue + self.weights[di_view] = di_view.pod.ma_view.data # just the mask for now + #self.weights[di_view] = (di_view.pod.ma_view.data + # / (1. + stat_weights/di_view.data)) + + def __del__(self): + """ + Clean up routine + """ + BaseModel.__del__(self) + del self.engine.ptycho.containers[self.weights.ID] + del self.weights + + def new_grad(self): + """ + Compute a new gradient direction according to a Euclidean noise model. + + Note: The negative log-likelihood and local errors are also computed + here. + """ + self.ob_grad.fill(0.) + self.pr_grad.fill(0.) + + # We need an array for MPI + LL = np.array([0.]) + error_dct = {} + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and amplitudes for this view + w = self.weights[diff_view] + A = np.sqrt(diff_view.data) + + Amodel = np.zeros_like(A) + f = {} + + # First pod loop: compute total amplitude + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f[name] = pod.fw(pod.probe * pod.object) + Amodel += np.sqrt(u.abs2(f[name])) + + # Floating intensity option + if self.p.floating_intensities: + self.float_intens_coeff[dname] = A.sum() / Amodel.sum() + Amodel *= self.float_intens_coeff[dname] + + Amodel += 1e-6 # cf Poisson model + DA = (1. - A / Amodel) + + # Second pod loop: gradients computation + LLL = np.sum((w * (Amodel - A)**2).astype(np.float64)) + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + xi = pod.bw(w*DA * f[name]) + self.ob_grad[pod.ob_view] += 2. * xi * pod.probe.conj() + self.pr_grad[pod.pr_view] += 2. * xi * pod.object.conj() + + 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() + parallel.allreduce(LL) + + # Object regularizer + if self.regularizer: + for name, s in self.ob.storages.items(): + self.ob_grad.storages[name].data += self.regularizer.grad( + s.data) + LL += self.regularizer.LL + self.LL = LL / self.tot_measpts + + return error_dct + + def poly_line_coeffs_ob(self, ob_h): + """ + Compute the coefficients of the polynomial for line minimization + in direction h for the object + """ + + B = np.zeros((3,), dtype=np.longdouble) + Brenorm = 1. / self.LL[0]**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and amplitudes for this view + w = self.weights[diff_view] + A = np.sqrt(diff_view.data) + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pod.probe * ob_h[pod.ob_view]) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-12 # cf Poisson model sqrt(1e-12) = 1e-6 + DA = 1. - A/np.sqrt(A0) + + B[0] += np.dot(w.flat, ((np.sqrt(A0) - A)**2).flat) * Brenorm + B[1] += np.dot(w.flat, (A1*DA).flat) * Brenorm + B[2] += (np.dot(w.flat, (A2*DA).flat) + 0.25*np.dot(w.flat, (A1**2 * A/A0**(3/2)).flat)) * Brenorm + + parallel.allreduce(B) + + # Object regularizer + if self.regularizer: + for name, s in self.ob.storages.items(): + B += Brenorm * self.regularizer.poly_line_coeffs( + ob_h.storages[name].data, s.data) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + + def poly_line_coeffs_pr(self, pr_h): + """ + Compute the coefficients of the polynomial for line minimization + in direction h for the probe + """ + + B = np.zeros((3,), dtype=np.longdouble) + Brenorm = 1. / self.LL[0]**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and amplitudes for this view + w = self.weights[diff_view] + A = np.sqrt(diff_view.data) + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pr_h[pod.pr_view] * pod.object) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-12 # cf Poisson model sqrt(1e-12) = 1e-6 + DA = 1. - A/np.sqrt(A0) + + B[0] += np.dot(w.flat, ((np.sqrt(A0) - A)**2).flat) * Brenorm + B[1] += np.dot(w.flat, (A1*DA).flat) * Brenorm + B[2] += (np.dot(w.flat, (A2*DA).flat) + 0.25*np.dot(w.flat, (A1**2 * A/A0**(3/2)).flat)) * Brenorm + + parallel.allreduce(B) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + + def poly_line_all_coeffs_ob(self, ob_h): + """ + Compute all the coefficients of the polynomial for line minimization + in direction h for the object + """ + + B = np.zeros((9,), dtype=np.longdouble) + Brenorm = 1. / self.LL[0]**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and amplitudes for this view + w = self.weights[diff_view] + A = np.sqrt(diff_view.data) + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pod.probe * ob_h[pod.ob_view]) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-12 # cf Poisson model sqrt(1e-12) = 1e-6 + DA = 1. - A/np.sqrt(A0) + DA32 = A/A0**(3/2) + + B[0] += np.dot(w.flat, ((np.sqrt(A0) - A)**2).flat) * Brenorm + B[1] += np.dot(w.flat, (A1*DA).flat) * Brenorm + B[2] += (np.dot(w.flat, (A2*DA).flat) + 0.25*np.dot(w.flat, (A1**2 * DA32).flat)) * Brenorm + B[3] += (0.25*np.dot(w.flat, (2*A1*A2 * DA32).flat) - 0.125*np.dot(w.flat, (A1**3/A0**2).flat)) * Brenorm + B[4] += (0.25*np.dot(w.flat, ((A2**2) * DA32).flat) - 0.125*np.dot(w.flat, ((3*A1**2*A2)/A0**2).flat) + + 0.015625*np.dot(w.flat, (A1**4/A0**3).flat)) * Brenorm + B[5] += (- 0.125*np.dot(w.flat, ((3*A1*A2**2)/A0**2).flat) + + 0.015625*np.dot(w.flat, ((4*A1**3*A2)/A0**3).flat)) * Brenorm + B[6] += (- 0.125*np.dot(w.flat, ((A2**3)/A0**2).flat) + + 0.015625*np.dot(w.flat, ((6*A1**2*A2**2)/A0**3).flat)) * Brenorm + B[7] += (0.015625*np.dot(w.flat, ((4*A1*A2**3)/A0**3).flat)) * Brenorm + B[8] += (0.015625*np.dot(w.flat, ((A2**4)/A0**3).flat)) * Brenorm + + parallel.allreduce(B) + + # Object regularizer + if self.regularizer: + for name, s in self.ob.storages.items(): + B[:3] += Brenorm * self.regularizer.poly_line_coeffs( + ob_h.storages[name].data, s.data) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + +def poly_line_all_coeffs_pr(self, pr_h): + """ + Compute all the coefficients of the polynomial for line minimization + in direction h for the probe + """ + + B = np.zeros((9,), dtype=np.longdouble) + Brenorm = 1. / self.LL[0]**2 + + # Outer loop: through diffraction patterns + for dname, diff_view in self.di.views.items(): + if not diff_view.active: + continue + + # Weights and amplitudes for this view + w = self.weights[diff_view] + A = np.sqrt(diff_view.data) + + A0 = None + A1 = None + A2 = None + + for name, pod in diff_view.pods.items(): + if not pod.active: + continue + f = pod.fw(pod.probe * pod.object) + a = pod.fw(pr_h[pod.pr_view] * pod.object) + + if A0 is None: + A0 = u.abs2(f).astype(np.longdouble) + A1 = 2 * np.real(f * a.conj()).astype(np.longdouble) + A2 = u.abs2(a).astype(np.longdouble) + else: + A0 += u.abs2(f) + A1 += 2 * np.real(f * a.conj()) + A2 += u.abs2(a) + + if self.p.floating_intensities: + A0 *= self.float_intens_coeff[dname] + A1 *= self.float_intens_coeff[dname] + A2 *= self.float_intens_coeff[dname] + + A0 += 1e-12 # cf Poisson model sqrt(1e-12) = 1e-6 + DA = 1. - A/np.sqrt(A0) + DA32 = A/A0**(3/2) + + B[0] += np.dot(w.flat, ((np.sqrt(A0) - A)**2).flat) * Brenorm + B[1] += np.dot(w.flat, (A1*DA).flat) * Brenorm + B[2] += (np.dot(w.flat, (A2*DA).flat) + 0.25*np.dot(w.flat, (A1**2 * DA32).flat)) * Brenorm + B[3] += (0.25*np.dot(w.flat, (2*A1*A2 * DA32).flat) - 0.125*np.dot(w.flat, (A1**3/A0**2).flat)) * Brenorm + B[4] += (0.25*np.dot(w.flat, ((A2**2) * DA32).flat) - 0.125*np.dot(w.flat, ((3*A1**2*A2)/A0**2).flat) + + 0.015625*np.dot(w.flat, (A1**4/A0**3).flat)) * Brenorm + B[5] += (- 0.125*np.dot(w.flat, ((3*A1*A2**2)/A0**2).flat) + + 0.015625*np.dot(w.flat, ((4*A1**3*A2)/A0**3).flat)) * Brenorm + B[6] += (- 0.125*np.dot(w.flat, ((A2**3)/A0**2).flat) + + 0.015625*np.dot(w.flat, ((6*A1**2*A2**2)/A0**3).flat)) * Brenorm + B[7] += (0.015625*np.dot(w.flat, ((4*A1*A2**3)/A0**3).flat)) * Brenorm + B[8] += (0.015625*np.dot(w.flat, ((A2**4)/A0**3).flat)) * Brenorm + + parallel.allreduce(B) + + if np.isinf(B).any() or np.isnan(B).any(): + logger.warning( + 'Warning! inf or nan found! Trying to continue...') + B[np.isinf(B)] = 0. + B[np.isnan(B)] = 0. + + self.B = B + + return B + + class Regul_del2(object): """\ Squared gradient regularizer (Gaussian prior).