From fbbecbd04c6becec68ab54720ffca9ae842c4f8e Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 23 Apr 2020 14:32:19 -0400 Subject: [PATCH] First work on the off-axis propagator --- CDTools/models/bragg_2d_ptycho.py | 380 +++++++++++++++++++++++++++++- CDTools/tools/propagators.py | 60 ++++- 2 files changed, 438 insertions(+), 2 deletions(-) diff --git a/CDTools/models/bragg_2d_ptycho.py b/CDTools/models/bragg_2d_ptycho.py index daf7fd2..8b1eb04 100644 --- a/CDTools/models/bragg_2d_ptycho.py +++ b/CDTools/models/bragg_2d_ptycho.py @@ -23,6 +23,23 @@ from copy import copy # 4) Include a correction for the thickness of the sample # +# +# How to do this properly? +# First thing to note is that the two corrections (probe propagation before +# interaction and high-NA correction for the final diffraction measurement) +# should be able to be turned on separately, since they show up in different +# situations. In fact, I would like to focus on the first aspect initially +# since I think that's the dominant issue we will contend with at CSX. +# + +# +# It also should be possible to choose an "auto" setting for the two +# corrections, since the geometry information given should be enough to +# decide if the correction is needed. +# Probably the automatic check will have to be very conservative for the +# probe propagation side since the model has no information about the +# expected numerical aperture of the probe. +# class Bragg2DPtycho(CDIModel): @@ -31,6 +48,367 @@ class Bragg2DPtycho(CDIModel): 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): + probe_support = None, obj_support=None, oversampling=1, + propagate_probe=None, correct_tilt=None): + + + super(FancyPtycho,self).__init__() + self.wavelength = t.Tensor([wavelength]) + self.detector_geometry = copy(detector_geometry) + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = t.Tensor(det_geo['distance']) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = t.Tensor(det_geo['basis']) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = t.Tensor(det_geo['corner']) + + self.min_translation = t.Tensor(min_translation) + + self.probe_basis = t.Tensor(probe_basis) + self.detector_slice = detector_slice + self.surface_normal = t.Tensor(surface_normal) + + self.saturation = saturation + + if mask is None: + self.mask = mask + else: + self.mask = t.BoolTensor(mask) + + # We rescale the probe here so it learns at the same rate as the + # object + if probe_guess.dim() > 3: + 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))) + + self.probe = t.nn.Parameter(probe_guess.to(t.float32) + / self.probe_norm) + + self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + + if background is None: + if detector_slice is not None: + background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1]) + else: + background = 1e-6 * t.ones(self.probe[0].shape[:-1]) + + + self.background = t.nn.Parameter(t.Tensor(background).to(t.float32)) + + if weights is None: + self.weights = None + else: + self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32)) + + if translation_offsets is None: + self.translation_offsets = None + else: + self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale) + + self.translation_scale = translation_scale + + if probe_support is not None: + self.probe_support = probe_support + else: + self.probe_support = t.ones_like(self.probe[0]) + + if obj_support is not None: + self.obj_support = obj_support + self.obj.data = self.obj * obj_support + else: + self.obj_support = t.ones_like(self.obj) + + 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 + + if not(correct_tilt is True or correct_tilt is False): + pass + else: + self.correct_tilt = correct_tilt + + + @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): + + wavelength = dataset.wavelength + det_basis = dataset.detector_geometry['basis'] + det_shape = dataset[0][1].shape + distance = dataset.detector_geometry['distance'] + + # always do this on the cpu + get_as_args = dataset.get_as_args + dataset.get_as(device='cpu') + (indices, translations), patterns = dataset[:] + dataset.get_as(*get_as_args[0],**get_as_args[1]) + + # Set to none to avoid issues with things outside the detector + if auto_center: + center = tools.image_processing.centroid(t.sum(patterns,dim=0)) + else: + center = None + + # Then, generate the probe geometry from the dataset + ewg = tools.initializers.exit_wave_geometry + probe_basis, probe_shape, det_slice = ewg(det_basis, + det_shape, + wavelength, + distance, + center=center, + padding=padding, + opt_for_fft=False, + oversampling=oversampling) + + + 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'}: + surface_normal = np.array([0.,0.,1.]) + elif scattering_mode in {'r', 'reflection'}: + outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) + outgoing_dir /= np.linalg.norm(outgoing_dir) + surface_normal = outgoing_dir + np.array([0.,0.,1.]) + surface_normal /= np.linalg.norm(outgoing_dir) + + + # 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) + + obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200) + + if hasattr(dataset, 'background') and dataset.background is not None: + background = t.sqrt(dataset.background) + else: + background = None + + # Finally, initialize the probe and object using this information + if probe_size is None: + probe = tools.initializers.SHARP_style_probe(dataset, probe_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) + + + # Now we initialize all the subdominant probe modes + probe_max = t.max(cmath.cabs(probe)) + probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)] + probe = t.stack([probe,] + probe_stack) + + obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) + + det_geo = dataset.detector_geometry + + translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) + + weights = t.ones(len(dataset)) + + if hasattr(dataset, 'mask') and dataset.mask is not None: + mask = dataset.mask.to(t.bool) + else: + mask = None + + if probe_support_radius is not None: + probe_support = t.zeros_like(probe[0].to(dtype=t.float32)) + p_cent = np.array(probe.shape[1:3]).astype(int) // 2 + psr = int(probe_support_radius) + probe_support[p_cent[0]-psr:p_cent[0]+psr, + p_cent[1]-psr:p_cent[1]+psr] = 1 + else: + probe_support = None; + + if restrict_obj != -1: + ro = restrict_obj + os = np.array(obj_size) + ps = np.array(probe_shape) + obj_support = t.zeros_like(obj.to(dtype=t.float32)) + obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2, + ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1 + else: + obj_support = None + + return cls(wavelength, det_geo, probe_basis, probe, obj, + detector_slice=det_slice, + surface_normal=surface_normal, + min_translation=min_translation, + 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) + + + def interaction(self, index, translations): + pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, + translations, + surface_normal=self.surface_normal) + pix_trans -= self.min_translation + + if self.translation_offsets is not None: + pix_trans += self.translation_scale * self.translation_offsets[index] + + 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[...,:,:] + + + if exit_waves.dim() == 4: + exit_waves = self.weights[index][:,None,None,None] * exit_waves + else: + exit_waves = self.weights[index] * exit_waves + + all_exit_waves.append(exit_waves) + + + return t.stack(all_exit_waves) + + + def forward_propagator(self, wavefields): + return tools.propagators.far_field(wavefields) + + + def backward_propagator(self, wavefields): + return tools.propagators.inverse_far_field(wavefields) + + + def measurement(self, wavefields): + return tools.measurements.quadratic_background(wavefields, + self.background, + detector_slice=self.detector_slice, + measurement=tools.measurements.incoherent_sum, + saturation=self.saturation, + oversampling=self.oversampling) + + + def loss(self, sim_data, real_data, mask=None): + return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) + + + def to(self, *args, **kwargs): + super(FancyPtycho, self).to(*args, **kwargs) + self.wavelength = self.wavelength.to(*args,**kwargs) + # move the detector geometry too + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = det_geo['distance'].to(*args,**kwargs) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = det_geo['basis'].to(*args,**kwargs) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = det_geo['corner'].to(*args,**kwargs) + + if self.mask is not None: + self.mask = self.mask.to(*args, **kwargs) + + + self.min_translation = self.min_translation.to(*args,**kwargs) + self.probe_basis = self.probe_basis.to(*args,**kwargs) + self.probe_norm = self.probe_norm.to(*args,**kwargs) + self.probe_support = self.probe_support.to(*args,**kwargs) + 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 + # over what metadata is saved (names, etc.) + + # First, I need to gather all the relevant data + # that needs to be added to the dataset + entry_info = {'program_name': 'CDTools', + 'instrument_n': 'Simulated Data', + 'start_time': datetime.now()} + + surface_normal = self.surface_normal.detach().cpu().numpy() + xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal) + xsurfacevec /= np.linalg.norm(xsurfacevec) + ysurfacevec = np.cross(surface_normal, xsurfacevec) + ysurfacevec /= np.linalg.norm(ysurfacevec) + orientation = np.array([xsurfacevec, ysurfacevec, surface_normal]) + + sample_info = {'description': 'A simulated sample', + 'orientation': orientation} + + + detector_geometry = self.detector_geometry + mask = self.mask + wavelength = self.wavelength + indices, translations = args_list + + # Then we simulate the results + data = self.forward(indices, translations) + + # And finally, we make the dataset + return Ptycho2DDataset(translations, data, + entry_info = entry_info, + sample_info = sample_info, + wavelength=wavelength, + detector_geometry=detector_geometry, + mask=mask) + + + def corrected_translations(self,dataset): + translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device) + t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal) + return translations + t_offset + + + # Needs to be updated to allow for plotting to an existing figure + plot_list = [ + ('Dominant Probe Amplitude', + lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)), + ('Dominant Probe Phase', + lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)), + ('Subdominant Probe Amplitude', + lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis), + lambda self: len(self.probe) >=2), + ('Subdominant Probe Phase', + lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis), + lambda self: len(self.probe) >=2), + ('Object Amplitude', + lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)), + ('Object Phase', + lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)), + ('Corrected Translations', + lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)), + ('Background', + lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)) + ] + + + def save_results(self, dataset): + basis = self.probe_basis.detach().cpu().numpy() + translations = self.corrected_translations(dataset).detach().cpu().numpy() + probe = cmath.torch_to_complex(self.probe.detach().cpu()) + probe = probe * self.probe_norm.detach().cpu().numpy() + obj = cmath.torch_to_complex(self.obj.detach().cpu()) + background = self.background.detach().cpu().numpy()**2 + weights = self.weights.detach().cpu().numpy() + + return {'basis':basis, 'translation':translations, + 'probe':probe,'obj':obj, + 'background':background, + 'weights':weights} diff --git a/CDTools/tools/propagators.py b/CDTools/tools/propagators.py index d66ea4d..98950ff 100644 --- a/CDTools/tools/propagators.py +++ b/CDTools/tools/propagators.py @@ -117,10 +117,68 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, * # the previous expression to ensure that complex frequencies # get mapped to values <1 instead of >1 propagator = complex_to_torch(np.conj(propagator)) - + return propagator.to(*args, **kwargs) +def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, propagation_vector, *args, **kwargs): + """Generates an angular-spectrum based near-field propagator from experimental quantities + + This function generates an angular-spectrum based near field + propagator that will work on torch Tensors. The function is structured + 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. + + Formally, this propagator is the complex conjugate of the fourier + transform of the convolution kernel for light propagation in free + space + + This function is written to work on any wavefield defined on any + array of parallelograms. In addition, there is an assumed phase ramp + applied to the wavefield before propagation, defined such that a feature + with uniform phase will propagate along the direction of the + defined propagation vector. This helps simplify + + + Parameters + ---------- + shape : array + The shape of the arrays to be propagated + spacing : array + 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 + The displacement to propagate the wavefield along. + tilt : float + The tilt, in radians, of the plane that the wavefield is defined on + + Returns + ------- + propagator : torch.Tensor + A phase mask which accounts for the phase change that each plane wave will undergo. + """ + + ki = 2 * np.pi * fftpack.fftfreq(shape[0],spacing[0]) + kj = 2 * np.pi * fftpack.fftfreq(shape[1],spacing[1]) + Kj, Ki = np.meshgrid(kj,ki) + + # Define this as complex so the square root properly gives + # k>k0 components imaginary frequencies + k0 = np.complex128((2*np.pi/wavelength)) + + propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z) + + # Take the conjugate explicitly here instead of negating + # the previous expression to ensure that complex frequencies + # get mapped to values <1 instead of >1 + propagator = complex_to_torch(np.conj(propagator)) + + return propagator.to(**kwargs) + + def near_field(wavefront, angular_spectrum_propagator): """ Propagates a wavefront via the angular spectrum method