Implementation of Huber regularizer in maximum likelihood (ML) engine

ptypy/engines/ML.py

@@ -70,6 +70,23 @@ class ML(PositionCorrectionEngine):
     lowlim = 0.0
     help = Amplitude of the Gaussian prior if used
 
+    [reg_Huber]
+    default = False
+    type = bool
+    help = Whether to use a Huber regularizer
+
+    [reg_Huber_amplitude]
+    default = .01
+    type = float
+    lowlim = 0.0
+    help = Amplitude of the Huber regularizer if used
+
+    [reg_Huber_scale]
+    default = .01
+    type = float
+    lowlim = 0.0
+    help = Scale of the Huber regularizer if used
+
     [smooth_gradient]
     default = 0.0
     type = float
@@ -371,6 +388,17 @@ def prepare_regularizer(self):
             reg_del2_amplitude = self.p.reg_del2_amplitude * reg_rescale
             self.regularizer = Regul_del2(amplitude=reg_del2_amplitude)
 
+        elif self.p.reg_Huber:
+            obj_Npix = self.ob.size
+            expected_obj_var = obj_Npix / self.tot_power  # Poisson
+            reg_rescale = self.tot_measpts / (8. * obj_Npix * expected_obj_var)
+            logger.debug(
+                'Rescaling regularization amplitude using '
+                'the Poisson distribution assumption.')
+            logger.debug('Factor: %8.5g' % reg_rescale)
+            reg_Huber_amplitude = self.p.reg_Huber_amplitude * reg_rescale
+            self.regularizer = Regul_Huber(amplitude=reg_Huber_amplitude, scale=self.p.reg_Huber_scale)
+
     def __del__(self):
         """
         Clean up routine
@@ -792,6 +820,89 @@ def poly_line_coeffs(self, h, x=None):
         return self.coeff
 
 
+class Regul_Huber(object):
+    """\
+    Huber regularizer.
+
+    See https://en.wikipedia.org/wiki/Huber_loss#Pseudo-Huber_loss_function
+    This class applies to any numpy array.
+    """
+    def __init__(self, amplitude, scale, axes=[-2, -1]):
+        # Regul.__init__(self, axes)
+        self.axes = axes
+        self.amplitude = amplitude
+        self.scale = scale
+        self.delxy = None
+        self.fpxy = None
+        self.g = None
+        self.LL = None
+
+    def grad(self, x):
+        """
+        Compute and return the regularizer gradient given the array x.
+        """
+        ax0, ax1 = self.axes
+        del_xf = u.delxf(x, axis=ax0)
+        del_yf = u.delxf(x, axis=ax1)
+        del_xb = u.delxb(x, axis=ax0)
+        del_yb = u.delxb(x, axis=ax1)
+
+        fp_xb = np.sqrt(self.scale + u.abs2(del_xb))
+        fp_yb = np.sqrt(self.scale + u.abs2(del_yb))
+        fp_xf = np.sqrt(self.scale + u.abs2(del_xf))
+        fp_yf = np.sqrt(self.scale + u.abs2(del_yf))
+
+        self.delxy = [del_xf, del_yf, del_xb, del_yb]
+        self.fpxy = [fp_xf, fp_yf, fp_xb, fp_yb]
+
+        self.g = self.amplitude * (del_xb / fp_xb + del_yb / fp_yb - del_xf / fp_xf - del_yf / fp_yf)
+
+        self.LL = self.amplitude * (fp_xb.sum() + fp_yb.sum() + fp_xf.sum() + fp_yf.sum())
+
+        return self.g
+
+    def poly_line_coeffs(self, h, x=None):
+        ax0, ax1 = self.axes
+        if x is None:
+            del_xf, del_yf, del_xb, del_yb = self.delxy
+            fp_xf, fp_yf, fp_xb, fp_yb = self.fpxy
+            c0 = self.LL
+        else:
+            del_xf = u.delxf(x, axis=ax0)
+            del_yf = u.delxf(x, axis=ax1)
+            del_xb = u.delxb(x, axis=ax0)
+            del_yb = u.delxb(x, axis=ax1)
+            fp_xb = np.sqrt(self.scale + u.abs2(del_xb))
+            fp_yb = np.sqrt(self.scale + u.abs2(del_yb))
+            fp_xf = np.sqrt(self.scale + u.abs2(del_xf))
+            fp_yf = np.sqrt(self.scale + u.abs2(del_yf))
+            c0 = self.amplitude * (fp_xb.sum() + fp_yb.sum() + fp_xf.sum() + fp_yf.sum())
+
+        hdel_xf = u.delxf(h, axis=ax0)
+        hdel_yf = u.delxf(h, axis=ax1)
+        hdel_xb = u.delxb(h, axis=ax0)
+        hdel_yb = u.delxb(h, axis=ax1)
+
+        c1 = self.amplitude * np.real(np.vdot(del_xf / fp_xf, hdel_xf)
+                                          + np.vdot(del_yf / fp_yf, hdel_yf)
+                                          + np.vdot(del_xb / fp_xb, hdel_xb)
+                                          + np.vdot(del_yb / fp_yb, hdel_yb))
+
+        c2 = .5 * self.amplitude * np.real(np.vdot(hdel_xf / fp_xf, hdel_xf)
+                                           + np.vdot(hdel_yf / fp_yf, hdel_yf)
+                                           + np.vdot(hdel_xb / fp_xb, hdel_xb)
+                                           + np.vdot(hdel_yb / fp_yb, hdel_yb))
+
+        c2 -= .5 * self.amplitude * ((np.real(del_xf * hdel_xf.conj() / fp_xf) ** 2 / fp_xf).sum()
+                                     + (np.real(del_yf * hdel_yf.conj() / fp_yf) ** 2 / fp_yf).sum()
+                                     + (np.real(del_xb * hdel_xb.conj() / fp_xb) ** 2 / fp_xb).sum()
+                                     + (np.real(del_yb * hdel_yb.conj() / fp_yb) ** 2 / fp_yb).sum())
+
+        self.coeff = np.array([c0, c1, c2])
+
+        return self.coeff
+
+
 def prepare_smoothing_preconditioner(amplitude):
     """
     Factory for smoothing preconditioner.

templates/minimal_prep_and_run_ML_Gaussian_Huber.py

@@ -0,0 +1,56 @@
+"""
+This script is a test for ptychographic reconstruction in the absence
+of actual data. It uses the test Scan class
+`ptypy.core.data.MoonFlowerScan` to provide "data".
+"""
+#import ptypy
+from ptypy.core import Ptycho
+from ptypy import utils as u
+p = u.Param()
+
+# for verbose output
+p.verbose_level = 4
+
+# set home path
+p.io = u.Param()
+p.io.home = "/tmp/ptypy/"
+p.io.autosave = u.Param()
+p.io.autosave.active = False
+p.io.autoplot = u.Param()
+p.io.autoplot.active = True
+p.io.autoplot.dump = False
+
+# max 100 frames (128x128px) of diffraction data
+p.scans = u.Param()
+p.scans.MF = u.Param()
+p.scans.MF.name = 'Full'
+p.scans.MF.data= u.Param()
+p.scans.MF.data.name = 'MoonFlowerScan'
+p.scans.MF.data.shape = 128
+p.scans.MF.data.num_frames = 100
+p.scans.MF.data.save = None
+
+# position distance in fraction of illumination frame
+p.scans.MF.data.density = 0.2
+# total number of photon in empty beam
+p.scans.MF.data.photons = 1e8
+# Gaussian FWHM of possible detector blurring
+p.scans.MF.data.psf = 0.
+
+# attach a reconstrucion engine
+p.engines = u.Param()
+p.engines.engine00 = u.Param()
+p.engines.engine00.name = 'ML'
+p.engines.engine00.ML_type = 'Gaussian'
+p.engines.engine00.reg_Huber = True                      # Whether to use a Gaussian prior (smoothing) regularizer
+p.engines.engine00.reg_Huber_amplitude = 1.             # Amplitude of the Gaussian prior if used
+p.engines.engine00.reg_Huber_scale = 0.01             # Amplitude of the Gaussian prior if used
+p.engines.engine00.scale_precond = True
+#p.engines.engine00.scale_probe_object = 1.
+p.engines.engine00.smooth_gradient = 20.
+p.engines.engine00.smooth_gradient_decay = 1/50.
+p.engines.engine00.floating_intensities = False
+p.engines.engine00.numiter = 300
+
+# prepare and run
+P = Ptycho(p,level=5)