Load probes

ptypy/core/illumination.py

@@ -11,9 +11,11 @@
 
 """
 import numpy as np
+from scipy.ndimage.interpolation import map_coordinates
 
 from .. import utils as u
 from ..core import geometry
+from ..core import Storage
 from ..utils.verbose import logger
 from .. import resources
 from ..utils.descriptor import EvalDescriptor
@@ -120,7 +122,7 @@
 
     [model] 
     default = None
-    type = str, ndarray
+    type = str, ndarray, Storage
     help = Type of illumination model
     doc = One of:
     	 - ``None`` : model initialitziation defaults to flat array filled with the specified number of photons
@@ -280,7 +282,7 @@ def aperture(A, grids=None, pars=None, **kwargs):
     return np.resize(ap, sh)
 
 
-def init_storage(storage, pars, energy=None, **kwargs):
+def init_storage(storage, pars, energy=None):
     """
     Initializes :any:`Storage` `storage` with parameters from `pars`
 
@@ -366,18 +368,37 @@ def init_storage(storage, pars, energy=None, **kwargs):
     elif type(p.model) is np.ndarray:
         model = p.model
     elif p.model in resources.probes:
-        model = resources.probes[p.model](s.shape)
+        model = resources.probes[p.model](s.shape[1:])
+    elif type(p.model) is Storage:
+        # resample the incoming probe storage based on the current pixel
+        #size and shape
+        layer = p.recon.get('layer')
+        layer = 0 if layer is None else layer
+        if (np.allclose(p.model.psize, s.psize) and
+                np.allclose(p.model.shape, s.shape)):
+            model = p.model.data[layer]
+        else:
+            logger.info(
+                prefix +
+                'Attempt to resample layer %s from Storage %s'
+                % (str(layer), str(p.model.ID)))
+            y, x = np.array(s.grids())[:, layer]
+            ii = (y - p.model.origin[0]) / p.model.psize[0]
+            jj = (x - p.model.origin[1]) / p.model.psize[1]
+            model = np.empty(shape=ii.shape, dtype=np.complex128)
+            model[:] = map_coordinates(np.abs(p.model.data[layer]), np.array((ii, jj)), mode='constant', cval=0)
+            model[:] = model * np.exp(1j * map_coordinates(np.angle(p.model.data[layer]), (ii, jj), mode='constant', cval=0))
     elif str(p.model) == 'recon':
         # Loading from a reconstruction file
-        layer = p.recon.get('layer')
         ID = p.recon.get('ID')
         logger.info(
             prefix +
-            'Attempt to load layer `%s` of probe storage with ID `%s` from `%s`'
-            % (str(layer), str(ID), p.recon.rfile))
-        model = u.load_from_ptyr(p.recon.rfile, 'probe', ID, layer)
-        # This could be more sophisticated,
-        # i.e. matching the real space grids etc.
+            'Attempt to load probe storage with ID `%s` from `%s`'
+            % (str(ID), p.recon.rfile))
+        model = u.load_from_ptyr(p.recon.rfile, 'probe', how='storage', ID=ID)
+        p.model = model
+        init_storage(s, p, energy)
+        return
     elif str(p.model) == 'stxm':
         logger.info(
             prefix + 'Probe initialization using averaged diffraction data.')

ptypy/utils/scripts.py

@@ -659,7 +659,7 @@ def stxm_init(storage, probe=None):
     trans, dpc_row, dpc_col = stxm_analysis(storage,probe)
     storage.data = trans * np.exp(-1j*phase_from_dpc(dpc_row, dpc_col))
 
-def load_from_ptyr(filename, what='probe', ID=None, layer=None):
+def load_from_ptyr(filename, what='probe', how='array', ID=None, layer=None):
     """
     Convenience script to extract data from ``*.ptyr``-file.
 
@@ -671,19 +671,20 @@ def load_from_ptyr(filename, what='probe', ID=None, layer=None):
     what : str
         Type of container to retrieve. Only `'probe'` and `'obj'` makes
         sense. Default is `'probe'`
+    how : str
+        Type of object to return. Only ``array`` and ``storage`` are
+        accepted. Default is ``array``.
     ID : str
         ID of storage in chosen container. If ``None`` the first stored
         storage is chosen.
     layer : int, optional
         If an integer, the data buffer of chosen storage gets sliced
-        with `layer` for its first index.
+        with `layer` for its first index. Has no effect if how is
+        ``storage``.
 
     Returns
     -------
-    data : ndarray
-        If `layer` is provided, that layer ``storage,data[layer]``
-        will be sliced from the 3d data buffer, else the whole buffer
-        ``storage.data`` will be returned.
+    data : ndarray or Storage
     """
     from .. import io
     header = io.h5read(filename, 'header')['header']
@@ -698,10 +699,16 @@ def load_from_ptyr(filename, what='probe', ID=None, layer=None):
             address = 'content/' + str(what)
             conti = io.h5read(filename, address)[address]
             storage = list(conti.values())[0]
-        if layer is None:
-            return storage['data']
-        else:
-            return storage['data'][layer]
+        if how == 'array':
+            if layer is None:
+                return storage['data']
+            else:
+                return storage['data'][layer]
+        elif how == 'storage':
+            from ..core import Storage, Container
+            C = Container(data_dims=2, data_type='complex')
+            storage['owner'] = C
+            return Storage._from_dict(storage)
 
 def phase_from_dpc(dpc_row, dpc_col):
     """

ptypy/utils/verbose.py

@@ -238,6 +238,8 @@ def _format_None(key,level, obj):
     def _format(key,level, obj):
         if hasattr(obj,'items'):
             stringout = _format_dict(key,level, obj)
+        elif hasattr(obj, 'views'):
+            stringout = _format(key, level, str(obj))
         elif type(obj) is np.ndarray:
             stringout = _format_numpy(key,level, obj)
         elif str(obj)==obj: