import torch as t from CDTools.models import CDIModel from CDTools.datasets import Ptycho2DDataset from CDTools import tools 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 from copy import copy __all__ = ['Bragg2DPtycho'] # # Key ideas: # 1) To a first approximation, do the reconstruction on a parallelogram # shaped grid on the sample which is conjugate to the detector coordinates # 2) Simulate a probe in those same coordinates, but propagate it back and # forth (along with the translations), using the angular spectrum method # 3) Apply a correction to the simulated data to account for the tilt of the # sample with respect to the detector # 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. # # # I'm worried that the propagation doesn't happen along the correct # direction, if a phase ramp is expected to be baked in to the # retrieved focal spot. Unclear if this is the case though. # # Retrieved focal spot should have the implicit phase ramp subtracted # (that is, it should be the focal spot along the sample plane, but with # the e^ikz dependence removed). Therefore, the original e^ikz dependence # should be easy to re-add jusy by using the propagate_along feature # in ggasp. So I believe this should not be a problem # class Bragg2DPtycho(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, probe_guess, obj_guess, detector_slice=None, min_translation=t.tensor([0, 0], dtype=t.float32), median_propagation=t.tensor(0, dtype=t.float32), background=None, translation_offsets=None, mask=None, weights=None, translation_scale=1, saturation=None, probe_support=None, oversampling=1, propagate_probe=True, correct_tilt=True, lens=False): # 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 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.median_propagation = t.tensor(median_propagation) self.probe_basis = t.tensor(probe_basis) self.detector_slice = copy(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 if mask is None: self.mask = mask else: self.mask = t.tensor(mask, dtype=t.bool) probe_guess = t.tensor(probe_guess, dtype=t.complex64) obj_guess = t.tensor(obj_guess, dtype=t.complex64) # We rescale the probe here so it learns at the same rate as the # object if probe_guess.dim() > 2: self.probe_norm = 1 * t.max(t.abs(probe_guess[0])) else: self.probe_norm = 1 * t.max(t.abs(probe_guess)) # 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: self.probe_support = t.tensor(probe_support, dtype=t.bool) 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 else: self.probe_support = t.ones(self.probe[0].shape, dtype=t.bool) self.probe = t.nn.Parameter(probe_guess / self.probe_norm) self.obj = t.nn.Parameter(obj_guess) if background is None: if detector_slice is not None: background = 1e-6 * t.ones( self.probe[0][self.detector_slice].shape, dtype=t.float32) else: background = 1e-6 * t.ones(self.probe[0].shape, dtype=t.float32) self.background = t.nn.Parameter(background) if weights is None: self.weights = None else: # No incoherent + unstable here yet self.weights = t.nn.Parameter(t.tensor(weights, dtype=t.float32)) if translation_offsets is None: self.translation_offsets = None else: t_o = t.tensor(translation_offsets, dtype=t.float32) t_o = t_o / translation_scale self.translation_offsets = t.nn.Parameter(t_o) self.translation_scale = translation_scale self.oversampling = oversampling 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 probe_shape = self.probe[0] self.k_map, self.intensity_map = \ tools.propagators.generate_high_NA_k_intensity_map( self.probe_basis, self.detector_geometry['basis'] / oversampling, probe_shape, self.detector_geometry['distance'], self.wavelength,dtype=t.float32, lens=lens) else: self.k_map = None self.intensity_map = None self.prop_dir = t.tensor([0, 0, 1], dtype=t.float32) # This propagator should be able to be multiplied by the propagation # distance each time to get a propagator self.universal_propagator = t.angle(ggasp( self.probe.shape[1:], self.probe_basis, self.wavelength, t.tensor([0, 0, self.wavelength/(2*np.pi)], dtype=t.float32), propagation_vector=self.prop_dir, dtype=t.complex64, propagate_along_offset=True)) @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, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True, lens=False, opt_for_fft=False): 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 exit wave geometry from the dataset ewg = tools.initializers.exit_wave_geometry ew_basis, ew_shape, det_slice = ewg(det_basis, det_shape, wavelength, distance, center=center, padding=padding, opt_for_fft=opt_for_fft, 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'}: 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) # 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 p2s = tools.interactions.project_translations_to_sample pix_translations, propagations = p2s(probe_basis, translations) 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) else: 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, ew_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling) else: probe = tools.initializers.gaussian_probe(dataset, ew_basis, ew_shape, probe_size, propagation_distance=propagation_distance) # Now we initialize all the subdominant probe modes probe_max = t.max(t.abs(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 = t.exp(1j*(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(probe[0].shape, dtype=t.bool) p_cent = np.array(probe[0].shape).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; # 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, 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, oversampling=oversampling, propagate_probe=propagate_probe, correct_tilt=correct_tilt, lens=lens) def interaction(self, index, translations): 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] Ws = self.weights[index] prs = Ws[...,None,None,None] * self.probe # Now we need to propagate each of the probes for j in range(prs.shape[0]): # I believe this -1 sign is in error, but I need a dataset with # well understood geometry to figure it out propagator = t.exp( 1j*(props[j]*(2*np.pi)/self.wavelength) * self.universal_propagator) prs[j] = tools.propagators.near_field(prs[j], propagator) exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc( prs, self.obj,pix_trans, shift_probe=True, multiple_modes=True) return exit_waves def forward_propagator(self, 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): 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): 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(Bragg2DPtycho, 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) 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) self.probe_norm = self.probe_norm.to(*args,**kwargs) self.probe_support = self.probe_support.to(*args,**kwargs) self.surface_normal = self.surface_normal.to(*args, **kwargs) self.prop_dir = self.prop_dir.to(*args, **kwargs) self.universal_propagator = self.universal_propagator.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.real.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 plot_list = [ ('Dominant Probe Amplitude', lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig)), ('Dominant Probe Phase', lambda self, fig: p.plot_phase(self.probe[0], fig=fig)), ('Subdominant Probe Amplitude', lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig), lambda self: len(self.probe) >=2), ('Subdominant Probe Phase', lambda self, fig: p.plot_phase(self.probe[1], fig=fig), lambda self: len(self.probe) >=2), ('Object Amplitude', lambda self, fig: p.plot_amplitude(self.obj, fig=fig)), ('Object Phase', lambda self, fig: p.plot_phase(self.obj, fig=fig)), ('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 = self.probe.detach().cpu().numpy() probe = probe * self.probe_norm.detach().cpu().numpy() obj = self.obj.detach().cpu().numpy() 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}