diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 8dc91a7..c147817 100644 --- a/CDTools/models/__init__.py +++ b/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 diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index 8b1eb04..be7cee1 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/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 diff --git a/CDTools/tools/interactions.py b/CDTools/tools/interactions.py index e7365ea..be203da 100644 --- a/CDTools/tools/interactions.py +++ b/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) + + + diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py index fbc6ba0..bc827d4 100644 --- a/CDTools/tools/propagators.py +++ b/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 diff --git a/conda_requirements.txt b/conda_requirements.txt index 1a535df..da4dcb2 100644 --- a/conda_requirements.txt +++ b/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 diff --git a/examples/lab_bragg_2d_ptycho.py b/examples/lab_bragg_2d_ptycho.py new file mode 100644 index 0000000..968ce4b --- /dev/null +++ b/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() diff --git a/setup.py b/setup.py index 7e2494f..44ee5c6 100644 --- a/setup.py +++ b/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={ diff --git a/tests/conftest.py b/tests/conftest.py index 0894750..6b31585 100644 --- a/tests/conftest.py +++ b/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) diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index 0732bbd..5e0b7d7 100644 --- a/tests/tools/test_data.py +++ b/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']) diff --git a/tests/tools/test_interactions.py b/tests/tools/test_interactions.py index dcc9a5d..765dbfc 100644 --- a/tests/tools/test_interactions.py +++ b/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 diff --git a/tests/tools/test_propagators.py b/tests/tools/test_propagators.py index a6c3f7e..adb960f 100644 --- a/tests/tools/test_propagators.py +++ b/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