diff --git a/ptypy/core/illumination.py b/ptypy/core/illumination.py index ce4430400..0967d7374 100644 --- a/ptypy/core/illumination.py +++ b/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.') diff --git a/ptypy/utils/scripts.py b/ptypy/utils/scripts.py index afbf74c53..dadfda346 100644 --- a/ptypy/utils/scripts.py +++ b/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): """ diff --git a/ptypy/utils/verbose.py b/ptypy/utils/verbose.py index 6e6ae300c..d35c1239b 100644 --- a/ptypy/utils/verbose.py +++ b/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: