Further develop the 2D Bragg ptychography code and edit some tests to work with a newer version of h5py

CDTools/models/__init__.py

@@ -52,7 +52,7 @@ from matplotlib.widgets import Slider
 from matplotlib import ticker
 import numpy as np
 
-__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho']
+__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho']
 
 
 class CDIModel(t.nn.Module):
@@ -470,3 +470,4 @@ class CDIModel(t.nn.Module):
 from CDTools.models.simple_ptycho import SimplePtycho
 from CDTools.models.fancy_ptycho import FancyPtycho
 from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
+from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho

CDTools/models/bragg_2d_ptycho.py

@@ -2,10 +2,11 @@ from __future__ import division, print_function, absolute_import
 
 import torch as t
 from CDTools.models import CDIModel
-from CDTools.datasets import Ptycho_2D_Dataset
+from CDTools.datasets import Ptycho2DDataset
 from CDTools import tools
 from CDTools.tools import cmath
 from CDTools.tools import plotting as p
+from CDTools.tools.propagators import generate_generalized_angular_spectrum_propagator as ggasp
 from matplotlib import pyplot as plt
 from datetime import datetime
 import numpy as np
@@ -46,16 +47,24 @@ class Bragg2DPtycho(CDIModel):
     def __init__(self, wavelength, detector_geometry,
                  probe_basis, probe_guess, obj_guess,
                  detector_slice=None,
-                 surface_normal=np.array([0.,0.,1.]),
                  min_translation = t.Tensor([0,0]),
                  median_propagation = t.Tensor(data=[0]),
                  background = None, translation_offsets=None, mask=None,
                  weights = None, translation_scale = 1, saturation=None,
                  probe_support = None, obj_support=None, oversampling=1,
-                 propagate_probe=None, correct_tilt=None):
+                 propagate_probe=True, correct_tilt=True):
 
         
-        super(FancyPtycho,self).__init__()
+        # We need the detector geometry
+        # We need the probe basis (but in this case, we don't need the surface
+        # normal because it comes implied by the probe basis
+        # we do need the detector slice I suppose
+        # The min translation is also needed
+        # The median propagation should be needed as well
+        # translation_offsets can stay 2D for now
+        # propagate_probe and correct_tilt are important!
+        
+        super(Bragg2DPtycho,self).__init__()
         self.wavelength = t.Tensor([wavelength])
         self.detector_geometry = copy(detector_geometry)
         det_geo = self.detector_geometry
@@ -67,9 +76,15 @@ class Bragg2DPtycho(CDIModel):
             det_geo['corner'] = t.Tensor(det_geo['corner'])
 
         self.min_translation = t.Tensor(min_translation)
-
+        self.median_propagation = median_propagation
+        
         self.probe_basis = t.Tensor(probe_basis)
         self.detector_slice = detector_slice
+
+        # calculate the surface normal from the probe basis
+        surface_normal =  np.cross(np.array(probe_basis)[:,1],
+                           np.array(probe_basis)[:,0])
+        surface_normal /= np.linalg.norm(surface_normal)
         self.surface_normal = t.Tensor(surface_normal)
         
         self.saturation = saturation
@@ -85,6 +100,16 @@ class Bragg2DPtycho(CDIModel):
             self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
         else:
             self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))         
+
+        # Not strictly necessary but otherwise it will return
+        # a probe with the stuff outside of the support unchanged after
+        # optimization
+        if probe_support is not None:
+            probe_guess = probe_guess * probe_support
+            # This seems dumb, but otherwise it winds up with a mixture
+            # of negative and positive zeros and it's super annoying when
+            # you look at the phase map
+            probe_guess[probe_guess == 0] = 0
             
         self.probe = t.nn.Parameter(probe_guess.to(t.float32)
                                     / self.probe_norm)
@@ -125,21 +150,27 @@ class Bragg2DPtycho(CDIModel):
 
         self.oversampling = oversampling
 
-        # Here we need to implement a simple condition to choose whether
-        # to propagate the probe or not
-        if not( propagate_probe is True or propagate_probe is False):
-            pass
-        else:
-            self.propagate_probe = propagate_probe
+        self.propagate_probe = propagate_probe
+        self.correct_tilt = correct_tilt
+        if correct_tilt:
+            # recall that here we always want the shape of the detector
+            # before it's cut down by the detector slice to match the
+            # physical detector region
+            det_shape = self.probe[0].shape[:-1]
 
-        if not(correct_tilt is True or correct_tilt is False):
-            pass
+            self.k_map, self.intensity_map = \
+                tools.propagators.generate_high_NA_k_intensity_map(
+                    self.probe_basis, self.detector_geometry['basis'],
+                    det_shape,
+                    self.detector_geometry['distance'],
+                    self.wavelength,dtype=t.float32)
         else:
-            self.correct_tilt = correct_tilt
+            self.k_map = None
+            self.intensity_map = None
         
         
     @classmethod
-    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True):
+    def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True):
         
         wavelength = dataset.wavelength
         det_basis = dataset.detector_geometry['basis']
@@ -158,9 +189,9 @@ class Bragg2DPtycho(CDIModel):
         else:
             center = None
             
-        # Then, generate the probe geometry from the dataset
+        # Then, generate the exit wave geometry from the dataset
         ewg = tools.initializers.exit_wave_geometry
-        probe_basis, probe_shape, det_slice =  ewg(det_basis,
+        ew_basis, ew_shape, det_slice =  ewg(det_basis,
                                                    det_shape,
                                                    wavelength,
                                                    distance,
@@ -169,15 +200,14 @@ class Bragg2DPtycho(CDIModel):
                                                    opt_for_fft=False,
                                                    oversampling=oversampling)
 
-
+        # now we grab the sample surface normal
         if hasattr(dataset, 'sample_info') and \
            dataset.sample_info is not None and \
            'orientation' in dataset.sample_info:
             surface_normal = dataset.sample_info['orientation'][2]
         else:
             surface_normal = np.array([0.,0.,1.])
-
-
+            
         # If this information is supplied when the function is called,
         # then we override the information in the .cxi file
         if scattering_mode in {'t', 'transmission'}:
@@ -188,12 +218,42 @@ class Bragg2DPtycho(CDIModel):
             surface_normal = outgoing_dir + np.array([0.,0.,1.])
             surface_normal /= np.linalg.norm(outgoing_dir)
 
+        # and we use that to generate the probe basis
 
+        ew_normal = np.cross(np.array(ew_basis)[:,1],
+                           np.array(ew_basis)[:,0])
+        ew_normal /= np.linalg.norm(ew_normal)
+
+        # This is a bit of an odd way to do a projection but I think it's
+        # the most compact and reliable. We set up two matrix-vector equations
+        # to enforce the two conditions (on the plane normal to the surface
+        # normal and in the ew_normal direction from the original point
+
+        mat = np.vstack([np.eye(3) - np.outer(ew_normal,ew_normal),
+                         surface_normal])
+
+        # We know that this matrix multiplied by the final result must
+        # equal the input vector with a trailing 0, so we can do the
+        # projection with a pseudoinverse and removing the last column
+
+        projector = np.linalg.pinv(mat)[:,:3]
+        
+        
+        probe_basis = t.Tensor(np.dot(projector, ew_basis))
+        
+
+        # Now we need a much better way to handle the translations here
+        # than translations_to_pixel
+        
         # Next generate the object geometry from the probe geometry and
         # the translations
-        pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
+        p2s = tools.interactions.project_translations_to_sample
+        pix_translations, propagations = p2s(probe_basis, translations)
 
-        obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
+
+        obj_size, min_translation = tools.initializers.calc_object_setup(ew_shape, pix_translations, padding=200)
+
+        median_propagation = t.median(propagations)
 
         if hasattr(dataset, 'background') and dataset.background is not None:
             background = t.sqrt(dataset.background)
@@ -201,10 +261,14 @@ class Bragg2DPtycho(CDIModel):
             background = None
 
         # Finally, initialize the probe and  object using this information
+        # Because the grid we defined on the sample is projected from the
+        # detector conjugate space, we can pretend that the grid is just in
+        # that space and use the standard initializations anyway
+        
         if probe_size is None:
-            probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
+            probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
         else:
-            probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
+            probe = tools.initializers.gaussian_probe(dataset, ew_basis, ew_shape, probe_size, propagation_distance=propagation_distance)
 
 
         # Now we initialize all the subdominant probe modes
@@ -244,42 +308,72 @@ class Bragg2DPtycho(CDIModel):
         else:
             obj_support = None
 
+
+        # Here we need to implement a simple condition to choose whether
+        # to propagate the probe or not
+        if not( propagate_probe is True or propagate_probe is False):
+            raise NotImplementedError('No auto option implemented yet')
+
+        if not(correct_tilt is True or correct_tilt is False):
+            raise NotImplementedError('No auto option implemented yet')
+
+            
         return cls(wavelength, det_geo, probe_basis, probe, obj,
                    detector_slice=det_slice,
-                   surface_normal=surface_normal,
                    min_translation=min_translation,
+                   median_propagation =median_propagation,
                    translation_offsets = translation_offsets,
                    weights=weights, mask=mask, background=background,
                    translation_scale=translation_scale,
                    saturation=saturation,
                    probe_support=probe_support,
                    obj_support=obj_support,
-                   oversampling=oversampling)
+                   oversampling=oversampling,
+                   propagate_probe=propagate_probe,
+                   correct_tilt=correct_tilt)
                    
     
     def interaction(self, index, translations):
-        pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
-                                                             translations,
-                                                             surface_normal=self.surface_normal)
+        pix_trans, props = tools.interactions.project_translations_to_sample(
+            self.probe_basis, translations)
+
         pix_trans -= self.min_translation
+        props -= self.median_propagation
 
         if self.translation_offsets is not None:
             pix_trans += self.translation_scale * self.translation_offsets[index]
+
+        single_translation = False
+        if translations.dim() == 1:
+            translations = translations[None,:]
+            single_translation = True
             
         all_exit_waves = []
         for i in range(self.probe.shape[0]):
             pr = self.probe[i] * self.probe_support
-            exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
-                                                           self.obj_support * self.obj,
-                                                           pix_trans,
-                                                           shift_probe=True)
-            exit_waves = exit_waves * self.probe_support[...,:,:]
+
+            exit_waves = []
+            for j in range(translations.size()[0]):
+                if self.propagate_probe:
+                    propagator = ggasp(pr.shape, self.probe_basis, self.wavelength,
+                                       t.Tensor([0,0,5*props[j]]),
+                                       dtype=pr.dtype,device=pr.device, propagate_along_offset=True)
+                    prop_pr = tools.propagators.near_field(pr, propagator)
+
+                else:
+                    prop_pr = pr
+                exit_waves.append(self.probe_norm *
+                                  tools.interactions.ptycho_2D_sinc(prop_pr,
+                                            self.obj_support * self.obj,
+                                            pix_trans[j],
+                                            shift_probe=True))
+            exit_waves = t.stack(exit_waves)
 
 
-            if exit_waves.dim() == 4:
+            if not single_translation:
                 exit_waves =  self.weights[index][:,None,None,None] * exit_waves
             else:
-                exit_waves =  self.weights[index] * exit_waves
+                exit_waves =  self.weights[index] * exit_waves[0,:,:,:]
                 
             all_exit_waves.append(exit_waves)
 
@@ -288,11 +382,18 @@ class Bragg2DPtycho(CDIModel):
     
         
     def forward_propagator(self, wavefields):
-        return tools.propagators.far_field(wavefields)
+        if self.correct_tilt:
+            return tools.propagators.high_NA_far_field(
+                wavefields,self.k_map,intensity_map=self.intensity_map)
+        else:
+            return tools.propagators.far_field(wavefields)
 
 
     def backward_propagator(self, wavefields):
-        return tools.propagators.inverse_far_field(wavefields)
+        if self.correct_tilt:
+            assert NotImplementedError('Backward propagator not defined with tilt correction')
+        else:
+            return tools.propagators.inverse_far_field(wavefields)
 
     
     def measurement(self, wavefields):
@@ -309,7 +410,7 @@ class Bragg2DPtycho(CDIModel):
 
     
     def to(self, *args, **kwargs):
-        super(FancyPtycho, self).to(*args, **kwargs)
+        super(Bragg2DPtycho, self).to(*args, **kwargs)
         self.wavelength = self.wavelength.to(*args,**kwargs)
         # move the detector geometry too
         det_geo = self.detector_geometry
@@ -323,6 +424,10 @@ class Bragg2DPtycho(CDIModel):
         if self.mask is not None:
             self.mask = self.mask.to(*args, **kwargs)
 
+        if self.k_map is not None:
+            self.k_map = self.k_map.to(*args,**kwargs)
+        if self.intensity_map is not None:
+            self.intensity_map = self.intensity_map.to(*args,**kwargs)
             
         self.min_translation = self.min_translation.to(*args,**kwargs)
         self.probe_basis = self.probe_basis.to(*args,**kwargs)
@@ -331,6 +436,7 @@ class Bragg2DPtycho(CDIModel):
         self.obj_support = self.obj_support.to(*args,**kwargs)
         self.surface_normal = self.surface_normal.to(*args, **kwargs)
 
+
         
     def sim_to_dataset(self, args_list):
         # In the future, potentially add more control

CDTools/tools/interactions.py

@@ -13,7 +13,9 @@ import numpy as np
 
 
 __all__ = ['translations_to_pixel', 'pixel_to_translations',
-           'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc']
+           'project_translations_to_sample',
+           'ptycho_2D_round','ptycho_2D_linear','ptycho_2D_sinc',
+           'ptycho_2D_propagate']
 
 
 
@@ -127,7 +129,92 @@ def pixel_to_translations(basis, pixel_translations, surface_normal=t.Tensor([0,
     else:
         return translations
 
+
+
+def project_translations_to_sample(sample_basis, translations):
+    """Takes real space translations and outputs them in pixels in a sample basis
     
+    This projection function is designed for the Bragg2DPtycho class. More
+    broadly, it works to take a set of translations in the lab frame and
+    convert each one into two values. First, an (i,j) value in pixels
+    describing the location of the probe's intersection with the sample
+    plane, assuming the basis found in sample_basis is used. Second, an
+    amount that the probe must be propagated along the z-axis to reach
+    that location on the sample plane.
+    
+    The assumed geometry is incoming radiation with a wavevector parallel
+    to the +z axis, [0,0,1].
+
+    Parameters
+    ----------
+    sample_basis : torch.Tensor
+        The real space basis the wavefields are defined in
+    translations : torch.Tensor 
+        A Jx3 stack of real-space translations, or a single translation
+
+    Returns
+    -------
+    pixel_translations : torch.Tensor
+        A Jx2 stack of translations in internal (i,j) pixel-space, or a single translation
+    propagations : torch.Tensor
+        A length-J vector of propagation distances, in meters
+    """
+
+    # Must do this all in pytorch unfortunately
+    # First we calculate the surface normal for the projection onto the sample
+    surface_normal =  t.cross(sample_basis[:,1],sample_basis[:,0])
+    surface_normal /= t.norm(surface_normal)
+
+
+    # Then we calculate a matrix which can do the projection
+    
+    propagation_dir = t.Tensor(np.array([0,0,1])).to(
+        device=surface_normal.device,
+        dtype=surface_normal.dtype)
+
+    I = t.eye(3).to(
+        device=surface_normal.device,
+        dtype=surface_normal.dtype)
+    
+    # Here we're setting up a matrix-vector equation mat*answer=input
+    mat = t.cat((I - t.ger(propagation_dir,propagation_dir),
+                 surface_normal.unsqueeze(0)))
+    
+    # And we invert the matrix to do the projection
+    projector = t.pinverse(mat)[:,:3].to(device=translations.device,
+                                         dtype=translations.dtype)
+
+    # Finally, we need to get the result in the right basis, so this will
+    # do that conversion for us
+    basis_vectors_inv = t.pinverse(sample_basis).to(device=translations.device,
+                                                    dtype=translations.dtype)
+
+    # Same for the propagation direction to extract how far it must propagate
+    propagation_dir_inv = propagation_dir.unsqueeze(0).to(
+        device=translations.device,
+        dtype=translations.dtype)
+
+    
+    sample_projection = t.mm(basis_vectors_inv, projector).t()
+    prop_projection = t.mm(propagation_dir_inv, projector).t()
+
+
+    single_translation = False
+    if len(translations.shape) == 1:
+        translations = translations[None,:]
+        single_translation = True
+
+    pixel_translations = t.mm(translations, sample_projection)
+    propagations = t.mm(translations, prop_projection)
+    
+    if single_translation:
+        return pixel_translations[0], propagations[0]
+    else:
+        return pixel_translations, propagations
+    
+
+    
+
 def ptycho_2D_round(probe, obj, translations):
     """Returns a stack of exit waves without accounting for subpixel shifts
 
@@ -334,3 +421,6 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
         return exit_waves[0]
     else:
         return t.stack(exit_waves)
+
+
+

CDTools/tools/propagators.py

@@ -7,8 +7,10 @@ from __future__ import division, print_function, absolute_import
 
 from CDTools.tools.cmath import *
 import torch as t
+from torch.nn.functional import grid_sample
 from scipy import fftpack
 import numpy as np
+from matplotlib import pyplot as plt
 
 __all__ = ['far_field', 'near_field',
            'generate_angular_spectrum_propagator',
@@ -72,6 +74,198 @@ def inverse_far_field(wavefront):
     return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
 
 
+def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, **kwargs):
+    """Generates k-space and intensity maps to allow for high-NA far-field propagation of light
+
+    At high numerical apertures or for very tilted samples, the simple
+    linear map between location on the detector and location in k-space
+    starts to break down. In addition, at angles above roughly 15 degrees,
+    a correction is needed to account for the decreasing solid angle
+    intersected by each pixel on the detector.
+    
+    This function generates a map which can be used to apply both corrections
+    via the high_NA_far_field propagator. The k-map which is output is
+    defined as a map between pixel location on the detector and locations
+    in the output of the standard, low-NA far-field propagated wavefield.
+    The output coordinate system is defined to run from -1 to 1 in both
+    directions - this allows for compatibility with pytorch's grid_sample
+    function. Some detector pixels may be mapped to values outside the
+    rangel [-1,1], depending on the respective sample and detector geometries.
+    This is most likely to be the case if the major correction is due to a
+    tilted sample.
+
+    The intensity map is simply an object, the shape of the detector, which
+    encodes intensity corrections between 0 and 1 per pixel.
+
+    Parameters
+    ----------
+    sample_basis: array
+        The 3x2 sample basis, in real space
+    det_basis: array
+        The 3x2 detector basis, in real space
+    det_shape: array
+        The length-2 shape of the detector array, (N,M)
+    distance: float
+        The sample-to-detector distance
+    wavelength: float
+        The wavelength of light being propagated
+
+
+    Returns
+    -------
+    k_map : torch.Tensor
+        An NxMx2 tensor mapping detector pixels to locations in the low-NA propagated wavefield
+    intensity_map : torch.Tensor
+        An NxM tensor encoding the high-NA intensity correction
+    """
+
+    # First we convert things to numpy
+    try:
+        sample_basis = sample_basis.detach().cpu().numpy()
+    except:
+        pass
+
+    try:
+        det_basis = det_basis.detach().cpu().numpy()
+    except:
+        pass
+
+    det_shape = np.array(tuple(det_shape))
+
+    try:
+        distance = distance.detach().cpu().numpy()[0]
+    except:
+        pass
+
+    try:
+        wavelength = wavelength.detach().cpu().numpy()[0]
+    except:
+        pass
+
+    # The next order of business is to calculate the k values associated with
+    # each pixel.
+
+    i_arr = np.arange(det_shape[0])
+    i_arr = i_arr - np.mean(i_arr)
+    j_arr = np.arange(det_shape[1])
+    j_arr = j_arr - np.mean(j_arr)
+    Is, Js = np.meshgrid(i_arr,j_arr,indexing='ij')
+    samp_det_vec = np.cross(det_basis[:,0],det_basis[:,1])
+    samp_det_vec *= distance / np.linalg.norm(samp_det_vec)
+
+    Rs = np.tensordot(det_basis,np.stack([Is,Js]),axes=1) \
+        + samp_det_vec[:,None,None]
+
+    k0 = 2*np.pi/wavelength
+    
+    Ks = k0 * Rs / np.linalg.norm(Rs, axis=0)
+
+    # This is the cosine of the angle with the detector normal
+    intensity_map = np.tensordot(samp_det_vec/(k0*distance),Ks,axes=1)
+    intensity_map = t.Tensor(intensity_map).to(*args, **kwargs)
+
+
+    # This accounts for the implied phase ramp along the exit wave direction
+    # In other words, it prevents the diffraction pattern from sliding off the
+    # detector when the sample is tilted but represented by an object with
+    # uniform phase.
+    Ks -= k0 * samp_det_vec[:,None,None] / distance
+
+    # Now we move on to finding the conversion into k-space
+    # for the sample grid. It turns out we can do this by multiplying
+    # them with the real space basis (dual of the reciprocal space
+    # basis is the real space basis). In fact, because we want to return
+    # values scaled to the overall size of the k-space window, we don't
+    # even need the shape of the sample array
+
+    k_map = np.tensordot(2*sample_basis.transpose()[::-1,:] / (2*np.pi),Ks,axes=1)
+    k_map = t.Tensor(np.moveaxis(k_map,0,2)).to(*args, **kwargs)
+    
+    # Potentially we need a correction to account for the discrete nature
+    # of the FFT
+    
+    # And finally, we need to convert the results to pytorch
+    
+    return k_map, intensity_map
+
+
+
+
+def high_NA_far_field(wavefront, k_map, intensity_map=None):
+    """Performs a far-field propagation step including a correction for high-NA scenarios
+
+    Two major corrections need to be performed when propagating light fields
+    into the far field at high numerical aperture or when the sample is
+    tilted as compared to the detector. The first correction is a deviation
+    from the linear relationship between detector position and spatial
+    frequency in the near field. This is accounted for with the k_map
+    argument, as generated by the generate_high_NA_k_intensity_map 
+    function.
+
+    The second correction is the change in the solid angle which each pixel
+    subtends at high NA. This is accounted for with an optional intensity
+    map. This is kept optional because some detectors - specifically, those
+    for penetrating radiation - may either not need a correction or need
+    a different correction due to the volumetric nature of the pixels.
+    
+    If the k-map map any pixels on the detector to pixels outside of the
+    k-space range of the wavefront, these will be set to zero. This is in
+    keeping with the typical assumption that the sample is band-limited to
+    the Nyquist frequency for the array on which it is sampled.
+
+    Parameters
+    ----------
+    wavefront : torch.Tensor
+        The JxNxMx2 stack of complex wavefronts propagated to the far-field
+    k_map : torch.Tensor
+        The NxMx2 map accounting for high NA distortion, as generated by generate_high_NA_k_intensity_map
+    intensity_map : torch.Tensor
+        The optional NxM tensor accounting for the intensity variation across the detector
+
+    
+    Returns
+    -------
+    propagated : torch.Tensor
+        The JxNxMx2 exit wavefield
+
+
+    """
+    low_NA_wavefield = far_field(wavefront)
+    # I'm going to need to separately interpolate the real and complex parts
+    # This can be done
+
+    k_map = k_map[None,:,:,:]
+    # Will only work for a 4D wavefile stack.
+    def process_wavefield_stack(low_NA_wavefield):
+        real_output = grid_sample(low_NA_wavefield[None,:,:,:,0],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
+        imag_output = grid_sample(low_NA_wavefield[None,:,:,:,1],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
+
+        result = t.stack((real_output[0,:,:,:],imag_output[0,:,:,:]),dim=3)
+    
+        if intensity_map is not None:
+            result = result * intensity_map[None,:,:,None]
+
+        return result
+
+    original_dim = wavefront.dim()
+    if original_dim == 3:
+        result = process_wavefield_stack(low_NA_wavefield[None,:,:,:])
+        return result[0,:,:,:]
+    if original_dim == 4:
+        result = process_wavefield_stack(low_NA_wavefield)
+        return result
+    if original_dim == 5:
+        result = []
+        for i in range(low_NA_wavefield.size()[0]):
+            result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:,:]))
+        return t.stack(result)
+    else:
+        raise IndexError('Wavefield had incorrect number of dimensions')
+
+
+
+
+
 def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, **kwargs):
     """Generates an angular-spectrum based near-field propagator from experimental quantities
 
@@ -80,7 +274,7 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
     this way - to generate the propagator first - because the
     generation of the propagation mask is a bit expensive and if this
     propagator is used in a reconstruction program, then it will be best
-    to calculate this mask once and close over it.
+    to calculate this mask once and reuse it.
 
     Formally, this propagator is the complex conjugate of the fourier
     transform of the convolution kernel for light propagation in free
@@ -163,7 +357,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
         The (2x3) set of basis vectors describing the array to be propagated
     wavelength : float
         The wavelength of light to simulate propagation of
-    propagation_vector : array
+    offset_vector : array
         The displacement to propagate the wavefield along.
     propagate_along_offset : bool
         Optional, whether to include an implied phase ramp to propagate uniform phase features along the offset direction
@@ -174,7 +368,25 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
         A phase mask which accounts for the phase change that each plane wave will undergo.
     """
 
-    
+    # These check for any pytorch inputs and convert them
+    try:
+        basis = basis.detach().cpu().numpy()
+    except:
+        pass
+
+    shape = np.array(tuple(shape))
+
+    try:
+        offset_vector = offset_vector.detach().cpu().numpy()
+    except:
+        pass
+
+    try:
+        wavelength = wavelength.detach().cpu().numpy()[0]
+    except:
+        pass
+
+
     # First we calculate a dual basis for the real space grid
     inv_basis =  np.linalg.pinv(basis).transpose()
 
@@ -207,15 +419,23 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
     perpendicular_dir /= np.linalg.norm(perpendicular_dir)
     offset_perpendicular = np.dot(perpendicular_dir, offset_vector)
     offset_parallel = offset_vector - perpendicular_dir * offset_perpendicular
-    print(offset_perpendicular)
-    print(offset_parallel)
-    k0 = 2*np.pi/wavelength
-    k_offset = offset_parallel * k0 / np.sqrt(offset_perpendicular**2 +
-                                    np.linalg.norm(offset_parallel)**2)
 
+    
+    k0 = 2*np.pi/wavelength
+
+    if np.linalg.norm(offset_vector) == 0:
+        # for numerical stability if this is exactly zero we need a special case
+        k_offset = np.array([0,0,0]) 
+    else:  
+        k_offset = offset_parallel * k0 / np.linalg.norm(offset_vector)
+
+
+    # There apparently is a sign correction that I need to apply
+    sign_correction = np.sign(np.dot(perpendicular_dir,offset_vector))
+    
     # Only implement the shift if the flag is set to True
     if propagate_along_offset:
-        K_xyz = K_xyz + k_offset[:,None,None]
+        K_xyz = K_xyz + k_offset[:,None,None] * sign_correction
 
         # we also need to remove the z-dependence on the phase
         # This time, though, the z-dependence actually has to do with

conda_requirements.txt

@@ -2,7 +2,7 @@ numpy>=1.0
 scipy>=1.0
 matplotlib>=2.0
 python-dateutil
-pytorch>=1.2.0
+pytorch>=1.3.0
 h5py>=2.1
 pytest
 sphinx

examples/lab_bragg_2d_ptycho.py

@@ -0,0 +1,36 @@
+from __future__ import division, print_function, absolute_import
+
+import CDTools
+from matplotlib import pyplot as plt
+import pickle
+
+
+#
+# This is the testing file for the Bragg ptychography code that includes
+# corrections for probe propagation and for off-axis far field diffraction
+# I'm starting to build it up using the lab ptychography forward data so I
+# can initially test that it works on "softball" data before moving to the
+# actual Bragg geometry data
+#
+
+#filename = 'example_data/lab_ptycho_data.cxi'
+filename = '/media/Data Bank/Lab Ptycho/Zone Plate Bragg 633.cxi'
+dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
+
+#dataset.inspect()
+
+model = CDTools.models.Bragg2DPtycho.from_dataset(dataset,probe_support_radius=60)#propagate_probe=False)#, n_modes=8)
+model.to(device='cuda')
+dataset.get_as(device='cuda')
+
+
+model.translation_offsets.requires_grad = False
+for i, loss in enumerate(model.Adam_optimize(100, dataset)):
+    model.inspect(dataset)
+    print(i,loss)
+    
+#with open('example_reconstructions/lab_bragg_2d_ptycho.pickle', 'wb') as f:
+#    pickle.dump(model.save_results(dataset),f)
+
+#model.compare(dataset)
+plt.show()

setup.py

@@ -17,7 +17,7 @@ setuptools.setup(
         "scipy>=1.0",
         "matplotlib>=2.0",
         "python-dateutil",
-        "torch>=1.2.0", #1.2.0 introduced boolean tensors in a breaking way, we use the boolean tensors here for masking
+        "torch>=1.3.0", #1.3.0 adds the align_corners option for grid_sample which is used for high NA far-field propagation. 1.2.0 introduced boolean tensors in a breaking way, we use the boolean tensors here for masking.
         "h5py>=2.1",
         "pathlib2 ; python_version<'3.4'"],
     extras_require={

tests/conftest.py

@@ -34,7 +34,7 @@ def ptycho_cxi_1():
     """
 
     expected = {}
-    f = h5py.File('ptycho_cxi_1',driver='core',backing_store=False)
+    f = h5py.File('ptycho_cxi_1','w',driver='core',backing_store=False)
 
     # Start by defining the basic structure
     f.create_dataset('cxi_version', data=150)
@@ -153,7 +153,7 @@ def ptycho_cxi_2():
     """
 
     expected = {}
-    f = h5py.File('ptycho_cxi_2',driver='core',backing_store=False)
+    f = h5py.File('ptycho_cxi_2','w',driver='core',backing_store=False)
 
     # Start by defining the basic structure
     f.create_dataset('cxi_version', data=150)
@@ -234,7 +234,7 @@ def ptycho_cxi_3():
     """
 
     expected = {}
-    f = h5py.File('ptycho_cxi_3',driver='core',backing_store=False)
+    f = h5py.File('ptycho_cxi_3','w',driver='core',backing_store=False)
 
     # Start by defining the basic structure
     f.create_dataset('cxi_version', data=150)

tests/tools/test_data.py

@@ -100,7 +100,7 @@ def test_get_ptycho_translations(test_ptycho_cxis):
 
 def test_create_cxi(tmp_path):
     data.create_cxi(tmp_path / 'test_create.cxi')
-    with h5py.File(tmp_path / 'test_create.cxi') as f:
+    with h5py.File(tmp_path / 'test_create.cxi','r') as f:
         assert f['cxi_version'][()] == 160
         assert 'entry_1' in f
 
@@ -114,7 +114,7 @@ def test_add_entry_info(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_entry_info.cxi') as f:
         data.add_entry_info(f, entry_info)
 
-    with h5py.File(tmp_path / 'test_add_entry_info.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_entry_info.cxi','r') as f:
         read_entry_info = data.get_entry_info(f)
 
     for key in entry_info:
@@ -135,7 +135,7 @@ def test_add_sample_info(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_sample_info.cxi') as f:
         data.add_sample_info(f, sample_info)
 
-    with h5py.File(tmp_path / 'test_add_sample_info.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_sample_info.cxi','r') as f:
         read_sample_info = data.get_sample_info(f)    
 
     for key in sample_info:
@@ -154,7 +154,7 @@ def test_add_source(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_source.cxi') as f:
         data.add_source(f, wavelength)
 
-    with h5py.File(tmp_path / 'test_add_source.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_source.cxi','r') as f:
         # Check this directly since we want to make sure it saved
         # the wavelength and energy
         read_wavelength = f['entry_1/instrument_1/source_1/wavelength'][()]
@@ -173,7 +173,7 @@ def test_add_detector(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_detector.cxi') as f:
         data.add_detector(f, distance, basis, corner=corner)
 
-    with h5py.File(tmp_path / 'test_add_detector.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_detector.cxi','r') as f:
         # Check this directly since we want to make sure it saved
         # the pixel sizes
         d1 = f['entry_1/instrument_1/detector_1']
@@ -196,7 +196,7 @@ def test_add_mask(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_mask.cxi') as f:
         data.add_mask(f, mask)
 
-    with h5py.File(tmp_path / 'test_add_mask.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_mask.cxi','r') as f:
         read_mask = data.get_mask(f)
 
     assert np.all(mask == read_mask)
@@ -208,7 +208,7 @@ def test_add_dark(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_dark.cxi') as f:
         data.add_dark(f, dark)
 
-    with h5py.File(tmp_path / 'test_add_dark.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_dark.cxi','r') as f:
         read_dark = data.get_dark(f)
 
     print(dark.shape)
@@ -224,14 +224,13 @@ def test_add_data(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_data.cxi') as f:
         data.add_data(f, fake_data, axes)
 
-    with h5py.File(tmp_path / 'test_add_data.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_data.cxi','r') as f:
         # Check this directly since we want to make sure it saved
         # it in all the places it should have
         read_data_1 = np.array(f['entry_1/data_1/data'])
         read_data_2 = np.array(f['entry_1/instrument_1/detector_1/data'])
         read_axes = str(f['entry_1/instrument_1/detector_1/data'].attrs['axes'].decode())
     
-
     assert np.allclose(fake_data, read_data_1)
     assert np.allclose(fake_data, read_data_2)
     assert 'translation:y:x' == read_axes
@@ -242,7 +241,7 @@ def test_add_data(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_data_torch.cxi') as f:
         data.add_data(f, fake_data)
 
-    with h5py.File(tmp_path / 'test_add_data_torch.cxi') as f:    
+    with h5py.File(tmp_path / 'test_add_data_torch.cxi','r') as f:    
         read_data, axes = data.get_data(f)
 
     assert np.allclose(fake_data.numpy(),read_data)
@@ -255,7 +254,7 @@ def test_add_ptycho_translations(tmp_path):
     with data.create_cxi(tmp_path / 'test_add_ptycho_translations.cxi') as f:
         data.add_ptycho_translations(f, translations)
     
-    with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi') as f:
+    with h5py.File(tmp_path / 'test_add_ptycho_translations.cxi','r') as f:
         # Check this directly since we want to make sure it saved
         # it in all the places it should have
         read_translations_1 = np.array(f['entry_1/data_1/translation'])

tests/tools/test_interactions.py

@@ -95,7 +95,11 @@ def test_pixel_to_translations():
                                                 surface_normal=surface_normal)
     assert t.allclose(translations, roundtrip)
 
-    
+
+
+def test_project_translations_to_sample():
+    # Needs to be tested
+    assert 0
 
 def test_ptycho_2D_round(random_probe, random_obj):
     # Test a stack of images

tests/tools/test_propagators.py

@@ -42,7 +42,90 @@ def test_inverse_far_field(exit_waves_1):
     assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
 
 
+def test_generate_high_NA_k_intensity_map():
 
+    # We need to generate a plausible scenario. I will start
+    # by using the initializer to generate a reasonable exit wave geometry
+    # and detector pair
+    basis = t.Tensor([[0,-30e-6,0],
+                      [-20e-6,0,0]]).transpose(0,1)
+    shape = t.Size([478,573])
+    #shape = t.Size([3,5])
+    wavelength = 1e-9
+    distance = 1#6e-3
+    rs_basis, full_shape, det_slice = \
+        initializers.exit_wave_geometry(basis, shape, wavelength,
+                                        distance, opt_for_fft=False)
+
+    k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
+        rs_basis, basis, shape, distance, wavelength,
+        dtype=t.float32)
+
+    # generate a good test exit wave
+    i = (np.arange(478) - 240)
+    j = (np.arange(573) - 270)
+    Is,Js = np.meshgrid(i,j,indexing='ij')
+    wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128)
+    t_wavefield = cmath.complex_to_torch(wavefield).to(dtype=t.float32)
+
+    high_NA_propagated = propagators.high_NA_far_field(
+        t_wavefield, k_map, intensity_map=intensity_map)
+    low_NA_propagated = propagators.far_field(t_wavefield)
+
+    low_NA = cmath.torch_to_complex(low_NA_propagated)
+    high_NA = cmath.torch_to_complex(high_NA_propagated)
+
+    # Checking first that for a low-NA propagation they give the same result
+    # 1e-4 tolerance seems to be reasonable in this comparison given my
+    # exploration with the code
+    #assert np.max(np.abs(high_NA-low_NA))/np.max(np.abs(low_NA)) < 1e-4
+
+    # Now I will explore some results with a tilted sample
+    #print(rs_basis)
+    #print(rs_basis_tilted)
+    distance = 0.01#6e-3
+    rs_basis, full_shape, det_slice = \
+        initializers.exit_wave_geometry(basis, shape, wavelength,
+                                        distance, opt_for_fft=False)
+    rs_basis_tilted = rs_basis.clone()
+    rs_basis_tilted[2,1] = rs_basis_tilted[0,1]
+
+    
+    k_map, intensity_map = propagators.generate_high_NA_k_intensity_map(
+        rs_basis_tilted, basis, shape, distance, wavelength,
+        dtype=t.float32)
+    
+    high_NA_propagated = propagators.high_NA_far_field(
+        t_wavefield, k_map, intensity_map=intensity_map)
+    low_NA_propagated = propagators.far_field(t_wavefield)
+
+    low_NA = cmath.torch_to_complex(low_NA_propagated)
+    high_NA = cmath.torch_to_complex(high_NA_propagated)
+    
+    #plt.close('all')
+    #plt.imshow(np.abs(low_NA)) 
+    #plt.figure()
+    #plt.imshow(np.abs(high_NA))
+    #plt.colorbar()
+    #plt.imshow(np.abs(wavefield))
+    #plt.show()
+
+
+    # Now I want to test that it doesn't crash for wavefields of various shapes
+    propagators.high_NA_far_field(t_wavefield.unsqueeze(0),
+                                  k_map, intensity_map=intensity_map)
+    propagators.high_NA_far_field(t_wavefield.unsqueeze(0).unsqueeze(0),
+                                  k_map, intensity_map=intensity_map)
+
+    # I believe this works, but I still would like to get a second method for
+    # simulating at least one diffraction pattern as an independent check
+    
+    assert 0
+    
+    
+
+    
+    
 def test_near_field():
 
     # The strategy is to compare the propagation of a gaussian beam to