Introduced new parameter for changing method for smoothing kernel

ptypy/accelerate/base/engines/ML_serial.py

@@ -23,13 +23,23 @@
 from ptypy.engines import register
 from ptypy.accelerate.base.kernels import GradientDescentKernel, AuxiliaryWaveKernel, PoUpdateKernel, PositionCorrectionKernel
 from ptypy.accelerate.base import address_manglers
+from ptypy.accelerate.base.array_utils import complex_gaussian_filter, complex_gaussian_filter_fft
 
 
 __all__ = ['ML_serial']
 
 @register()
 class ML_serial(ML):
 
+    """
+    Defaults:
+
+    [smooth_gradient.method]
+    default = convolution
+    type = str
+    help = Method to be used for smoothing the gradient, choose between ```convolution``` or ```fft```.
+    """
+
     def __init__(self, ptycho_parent, pars=None):
         """
         Maximum likelihood reconstruction engine.
@@ -143,10 +153,12 @@ def engine_prepare(self):
         self.ML_model.prepare()
 
     def _get_smooth_gradient(self, data, sigma):
-        return self.smooth_gradient(data)
-
-    def _get_smooth_gradient_fft(self, data, sigma):
-        return self.smooth_gradient(data)
+        if self.p.smooth_gradient.method == "convolution":
+            return complex_gaussian_filter(data, sigma)
+        elif self.p.smooth_gradient_method == "fft":
+            return complex_gaussian_filter_fft(data, sigma)
+        else:
+            raise NotImplementedError("smooth_gradient.method can only be ```convolution``` or ```fft```.")
 
     def _replace_ob_grad(self):
         new_ob_grad = self.ob_grad_new

ptypy/accelerate/cuda_pycuda/engines/ML_pycuda.py

@@ -90,10 +90,12 @@ def engine_initialize(self):
         self.qu_htod = cuda.Stream()
         self.qu_dtoh = cuda.Stream()
 
-        self.GSK = GaussianSmoothingKernel(queue=self.queue)
-        self.GSK.tmp = None
+        if self.p.smooth_gradient.method == "convolution":
+            self.GSK = GaussianSmoothingKernel(queue=self.queue)
+            self.GSK.tmp = None
 
-        self.FGSK = FFTGaussianSmoothingKernel(queue=self.queue)
+        if self.p.smooth_gradient.method == "fft":
+            self.FGSK = FFTGaussianSmoothingKernel(queue=self.queue)
 
         # Real/Fourier Support Kernel
         self.RSK = {}
@@ -257,13 +259,14 @@ def _set_pr_ob_ref_for_data(self, dev='gpu', container=None, sync_copy=False):
                 self._set_pr_ob_ref_for_data(dev=dev, container=container, sync_copy=sync_copy)
 
     def _get_smooth_gradient(self, data, sigma):
-        if self.GSK.tmp is None:
-            self.GSK.tmp = gpuarray.empty(data.shape, dtype=np.complex64)
-        self.GSK.convolution(data, [sigma, sigma], tmp=self.GSK.tmp)
-        return data
-
-    def _get_smooth_gradient_fft(self, data, sigma):
-        self.FGSK.filter(data, sigma)
+        if self.p.smooth_gradient.method == "convolution":
+            if self.GSK.tmp is None:
+                self.GSK.tmp = gpuarray.empty(data.shape, dtype=np.complex64)
+            self.GSK.convolution(data, [sigma, sigma], tmp=self.GSK.tmp)
+        elif self.p.smooth_gradient.method == "fft":
+            self.FGSK.filter(data, sigma)
+        else:
+            raise NotImplementedError("smooth_gradient.method can only be ```convolution``` or ```fft```.")
         return data
 
     def _replace_ob_grad(self):