Add Fourier and Interp shift functions

ptypy/core/classes.py

@@ -369,8 +369,10 @@ def __init__(self, container, ID=None, data=None, shape=DEFAULT_SHAPE,
         padding: int
             Number of pixels (voxels) to add as padding around the area defined
             by the views.
+
         """
         super(Storage, self).__init__(container, ID)
+        self.shift_type = None
 
         #: Default fill value
         self.fill_value = fill if fill is not None else 0.
@@ -381,7 +383,6 @@ def __init__(self, container, ID=None, data=None, shape=DEFAULT_SHAPE,
 
         # Additional padding around tight field of view
         self.padding = padding
-
         # dimensionality suggestion from container
         ndim = container.ndim if container.ndim is not None else 2
 
@@ -399,10 +400,12 @@ def __init__(self, container, ID=None, data=None, shape=DEFAULT_SHAPE,
             logger.warning('Storage view access dimension %d is not in regular '
                         'scope (2,3). Behavior is untested.' % len(shape[1:]))
 
+
         self.shape = shape
         self.data = np.empty(self.shape, self.dtype)
         self.data.fill(self.fill_value)
 
+
         """
         # Set data buffer
         if data is None:
@@ -425,6 +428,8 @@ def __init__(self, container, ID=None, data=None, shape=DEFAULT_SHAPE,
             self.shape = self.data.shape
         """
 
+
+
         if layermap is None:
             layermap = list(range(len(self.data)))
         self.layermap = layermap
@@ -1030,14 +1035,32 @@ def formatted_report(self, table_format=None, offset=8, align='right',
 
         return ''.join(fstring), dct
 
+    def subpixel_shift(self, data, sp):
+        """
+        HOOK - to be implemented and replaced.
+        Apply subpixel shift sp to data.
+
+
+        Parameters
+        ----------
+        data: numpy array
+        sp: tuple representing subpixel shift
+
+        Returns
+        -------
+        shifted array the same shape and type as data
+        """
+        return data
+
     def __getitem__(self, v):
         """
-        Storage[v]
+        Storage[v], Storage[v, sp=True]
 
         Returns
         -------
         ndarray
             The view to internal data buffer corresponding to View or layer `v`
+            Apply subpixel shift if sp=True.
         """
 
         # Here things could get complicated.
@@ -1048,30 +1071,44 @@ def __getitem__(self, v):
         # return shift(self.data[v.slayer, v.roi[0, 0]:v.roi[1, 0],
         #             v.roi[0, 1]:v.roi[1, 1]], v.sp)
         if isinstance(v, View):
+            # Default: return slice from 'data' array (without subpixel shift)
             if self.ndim == 2:
-                return shift(self.data[
-                             v.dlayer, v.dlow[0]:v.dhigh[0], v.dlow[1]:v.dhigh[1]],
-                             v.sp)
+                return self.data[v.dlayer, v.dlow[0]:v.dhigh[0], v.dlow[1]:v.dhigh[1]]
             elif self.ndim == 3:
-                return shift(self.data[
-                             v.dlayer, v.dlow[0]:v.dhigh[0], v.dlow[1]:v.dhigh[1],
-                             v.dlow[2]:v.dhigh[2]], v.sp)
+                return self.data[
+                       v.dlayer, v.dlow[0]:v.dhigh[0], v.dlow[1]:v.dhigh[1],
+                       v.dlow[2]:v.dhigh[2]]
         elif v in self.layermap:
+            # Return layer
             return self.data[self.layermap.index(v)]
+        elif isinstance(v, tuple):
+            # Call included second argument (subpixel switch)
+            v, sp = v
+            if self.ndim == 2 and sp and np.any(v.sp != 0.0):
+                return self.subpixel_shift(self.data[
+                                           v.dlayer, v.dlow[0]:v.dhigh[0], v.dlow[1]:v.dhigh[1]],
+                                           v.sp)
+            else:
+                return self[v]
         else:
             raise ValueError("View or layer '%s' is not present in storage %s"
                              % (v, self.ID))
 
+
     def __setitem__(self, v, newdata):
         """
         Storage[v] = newdata
+        or
+        Storage[v, sp] = newdata
 
         Set internal data buffer to `newdata` for the region of view `v`.
 
         Parameters
         ----------
         v : View
             A View for this storage
+        sp (optional): bool
+            Apply subpixel shift
 
         newdata : ndarray
             Two-dimensional array that fits the view's shape
@@ -1088,28 +1125,36 @@ def __setitem__(self, v, newdata):
             # right, but returns copies and not views.
             if self.ndim == 2:
                 self.data[v.dlayer,
-                          v.dlow[0]:v.dhigh[0],
-                          v.dlow[1]:v.dhigh[1]] = (shift(newdata, -v.sp))
+                      v.dlow[0]:v.dhigh[0],
+                      v.dlow[1]:v.dhigh[1]] = newdata
             elif self.ndim == 3:
                 self.data[v.dlayer,
                           v.dlow[0]:v.dhigh[0],
                           v.dlow[1]:v.dhigh[1],
-                          v.dlow[2]:v.dhigh[2]] = (shift(newdata, -v.sp))
+                          v.dlow[2]:v.dhigh[2]] = newdata
             elif self.ndim == 4:
                 self.data[v.dlayer,
                           v.dlow[0]:v.dhigh[0],
                           v.dlow[1]:v.dhigh[1],
                           v.dlow[2]:v.dhigh[2],
-                          v.dlow[3]:v.dhigh[3]] = (shift(newdata, -v.sp))
+                          v.dlow[3]:v.dhigh[3]] = newdata
             elif self.ndim == 5:
                 self.data[v.dlayer,
                           v.dlow[0]:v.dhigh[0],
                           v.dlow[1]:v.dhigh[1],
                           v.dlow[2]:v.dhigh[2],
                           v.dlow[3]:v.dhigh[3],
-                          v.dlow[4]:v.dhigh[4]] = (shift(newdata, -v.sp))
+                          v.dlow[4]:v.dhigh[4]] = newdata
         elif v in self.layermap:
             self.data[self.layermap.index(v)] = newdata
+        elif isinstance(v, tuple):
+            v, sp = v
+            if (self.ndim == 2) and sp and np.any(v.sp != 0.0):
+                self.data[v.dlayer,
+                      v.dlow[0]:v.dhigh[0],
+                      v.dlow[1]:v.dhigh[1]] = self.subpixel_shift(newdata, -v.sp)
+            else:
+                self[v] = newdata
         else:
             raise ValueError("View or layer '%s' is not present in storage %s"
                              % (v, self.ID))
@@ -1123,13 +1168,6 @@ def __str__(self):
         return info + ' psize=%(_psize)s center=%(_center)s' % self.__dict__
 
 
-def shift(v, sp):
-    """
-    Placeholder for future subpixel shifting method.
-    """
-    return v
-
-
 class View(Base):
     """
     A 'window' on a Container.
@@ -1385,6 +1423,20 @@ def data(self, v):
         """
         self.storage[self] = v
 
+    @property
+    def data_sp(self):
+        """
+        The view content in data buffer of associated storage.
+        """
+        return self.storage[self, True]
+
+    @data_sp.setter
+    def data_sp(self, v):
+        """
+        Set the view content in data buffer of associated storage.
+        """
+        self.storage[self, True] = v
+
     @property
     def shape(self):
         """
@@ -2192,6 +2244,23 @@ def object(self, v):
         if not self.is_empty:
             self.ob_view.data = v
 
+    @property
+    def object_sp(self):
+        """
+        Convenience property that links to slice of object :any:`Storage`,
+        including eventual pixel shift.
+        """
+        if not self.is_empty:
+            return self.ob_view.data_sp
+        else:
+            # Empty probe means no object (perfect transmission)
+            return np.ones(self.geometry.shape, dtype=self.owner.CType)
+
+    @object_sp.setter
+    def object_sp(self, v):
+        if not self.is_empty:
+            self.ob_view.data_sp = v
+
     @property
     def probe(self):
         """
@@ -2204,6 +2273,18 @@ def probe(self):
     def probe(self, v):
         self.pr_view.data = v
 
+    @property
+    def probe_sp(self):
+        """
+        Convenience property that links to slice of probe :any:`Storage`,
+        including subpixel.
+        """
+        return self.pr_view.data_sp
+
+    @probe_sp.setter
+    def probe_sp(self, v):
+        self.pr_view.data_sp = v
+
     @property
     def exit(self):
         """

ptypy/core/data.py

@@ -46,7 +46,7 @@
          EOS: 'End of scan reached'}
 
 __all__ = ['PtyScan', 'PTYD', 'PtydScan',
-           'MoonFlowerScan']
+           'MoonFlowerScan', 'NullScan']
 
 
 @defaults_tree.parse_doc('scandata.PtyScan')
@@ -1496,6 +1496,11 @@ class MoonFlowerScan(PtyScan):
     type = float
     help = Point spread function of the detector
 
+    [use_subpix]
+    default = False
+    type = bool
+    help = allows the simulation to use subpixel offsets
+
     """
 
     def __init__(self, pars=None, **kwargs):
@@ -1537,6 +1542,13 @@ def __init__(self, pars=None, **kwargs):
         self.pixel = np.round(pixel).astype(int) + 10
         frame = self.pixel.max(0) + 10 + geo.shape
         self.geo = geo
+        self.subpix = np.round(pixel) - pixel
+        print(self.subpix)
+        if p.use_subpix:
+            from ..utils import shift_fourier
+            self.probe_shift = shift_fourier
+        else:
+            self.probe_shift = lambda x, y: x
 
         try:
             self.obj = resources.flower_obj(frame)
@@ -1575,7 +1587,7 @@ def load(self, indices):
 
         for k in indices:
             intensity_j = u.abs2(self.geo.propagator.fw(
-                self.pr * self.obj[p[k][0]:p[k][0] + s[0],
+                self.probe_shift(self.pr, self.subpix[k]) * self.obj[p[k][0]:p[k][0] + s[0],
                                    p[k][1]:p[k][1] + s[1]]))
 
             if self.p.psf > 0.:
@@ -1589,6 +1601,107 @@ def load(self, indices):
         return raw, {}, {}
 
 
+@defaults_tree.parse_doc('scandata.NullScan')
+class NullScan(PtyScan):
+    """
+    Test PtyScan class producing zero frames. Meant for testing.
+
+    Override parent class default:
+
+    Defaults:
+
+    [name]
+    default = NullScan
+    type = str
+    help =
+    doc =
+
+    [num_frames]
+    default = 100
+    type = int
+    help = Number of frames
+    doc =
+
+    [shape]
+    type = int, tuple
+    default = 128
+    help = Shape of the region of interest cropped from the raw data.
+    doc = Cropping dimension of the diffraction frame
+      Can be None, (dimx, dimy), or dim. In the latter case shape will be (dim, dim).
+    userlevel = 1
+
+    [density]
+    default = 0.2
+    type = float
+    help = Position distance in fraction of illumination frame
+
+    [model]
+    default = 'round'
+    type = str
+    help = The scan pattern
+    """
+
+    def __init__(self, pars=None, **kwargs):
+        """
+        Parent pars are for the
+        """
+
+        p = self.DEFAULT.copy(depth=99)
+        p.update(pars)
+
+        # Initialize parent class
+        super(NullScan, self).__init__(p, **kwargs)
+
+        # Derive geometry from input
+        geo = geometry.Geo(pars=self.meta)
+
+        # Derive scan pattern
+        if p.model is 'raster':
+            pos = u.Param()
+            pos.spacing = geo.resolution * geo.shape * p.density
+            pos.steps = np.int(np.round(np.sqrt(self.num_frames))) + 1
+            pos.extent = pos.steps * pos.spacing
+            pos.model = p.model
+            self.num_frames = pos.steps ** 2
+            pos.count = self.num_frames
+            self.pos = xy.raster_scan(pos.spacing[0], pos.spacing[1], pos.steps, pos.steps)
+        else:
+            pos = u.Param()
+            pos.spacing = geo.resolution * geo.shape * p.density
+            pos.steps = np.int(np.round(np.sqrt(self.num_frames))) + 1
+            pos.extent = pos.steps * pos.spacing
+            pos.model = p.model
+            pos.count = self.num_frames
+            self.pos = xy.from_pars(pos)
+
+        # Calculate pixel positions
+        pixel = self.pos / geo.resolution
+        pixel -= pixel.min(0)
+        self.pixel = np.round(pixel).astype(int) + 10
+        frame = self.pixel.max(0) + 10 + geo.shape
+        self.geo = geo
+
+        self.obj = np.ones_like(frame, dtype=complex)
+
+        self.pr = np.zeros(self.geo.shape, dtype=complex)
+
+        self.load_common_in_parallel = True
+
+        self.p = p
+
+    def load_positions(self):
+        return self.pos
+
+    def load_weight(self):
+        return np.ones(self.pr.shape)
+
+    def load(self, indices):
+        p = self.pixel
+        s = self.geo.shape
+        raw = {k: np.zeros(self.geo.shape) for k in indices}
+        return raw, {}, {}
+
+
 @defaults_tree.parse_doc('scandata.QuickScan')
 class QuickScan(PtyScan):
     """