From fe1f4db0c6a6d780094df41ca1a90137082648db Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Mon, 1 Jul 2019 15:02:14 -0400 Subject: [PATCH] Consistently get Bragg/specular geometry implemented across both models/datasets and dataloaders --- CDTools/models/fancy_ptycho.py | 33 +++-- CDTools/models/incoherent_ptycho.py | 188 ---------------------------- CDTools/models/simple_ptycho.py | 34 +++-- CDTools/tools/data.py | 12 +- CDTools/tools/interactions.py | 2 +- examples/specular_pinhole_ptycho.py | 10 +- 6 files changed, 69 insertions(+), 210 deletions(-) delete mode 100644 CDTools/models/incoherent_ptycho.py diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 37683bc..5fe6326 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -90,7 +90,7 @@ class FancyPtycho(CDIModel): @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): + 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): wavelength = dataset.wavelength det_basis = dataset.detector_geometry['basis'] @@ -123,7 +123,19 @@ class FancyPtycho(CDIModel): 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) @@ -202,9 +214,6 @@ class FancyPtycho(CDIModel): 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_round(self.probe[i], - # self.obj, - # pix_trans) exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr, self.obj_support * self.obj, pix_trans, @@ -263,8 +272,9 @@ class FancyPtycho(CDIModel): 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.) @@ -275,7 +285,16 @@ class FancyPtycho(CDIModel): 'instrument_n': 'Simulated Data', 'start_time': datetime.now()} - sample_info = {'description': 'A simulated sample'} + 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 diff --git a/CDTools/models/incoherent_ptycho.py b/CDTools/models/incoherent_ptycho.py deleted file mode 100644 index 2d405cc..0000000 --- a/CDTools/models/incoherent_ptycho.py +++ /dev/null @@ -1,188 +0,0 @@ -from __future__ import division, print_function, absolute_import - -import torch as t -from CDTools.models import CDIModel -from CDTools import tools -from copy import copy -import numpy as np - - -class IncoherentPtycho(CDIModel): - - def __init__(self, wavelength, detector_geometry, - probe_basis, detector_slice, - probe_guess, obj_guess, min_translation = t.Tensor([0,0]), - translation_offsets=None, - background = None, mask=None, weights = None): - - super(IncoherentPtycho,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 - - - if mask is None: - self.mask = mask - else: - self.mask = t.ByteTensor(mask) - - # We rescale the probe here so it learns at the same rate as the - # object - probe_norm = t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) - - self.probe = t.nn.Parameter(probe_guess.to(t.float32)/probe_norm) - - self.probe_norm = float(probe_norm.numpy()) - - self.obj = t.nn.Parameter(obj_guess.to(t.float32)) - - if background is None: - background = 1e-6 * t.ones(self.probe[(np.s_[0],)+self.detector_slice].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)) - - - - @classmethod - def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0): - 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 - center = tools.image_processing.centroid(t.sum(patterns,dim=0)) - - # 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) - - # Next generate the object geometry from the probe geometry and - # the translations - pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations) - - obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=20) - - # 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) - else: - probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size) - - translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5) - - # For incoherent probe mixing - probe = t.stack((probe,0.05*t.rand(probe.shape).to(probe.dtype))) - - obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5)) - - det_geo = dataset.detector_geometry - - weights = t.ones(len(dataset)) - - if hasattr(dataset, 'mask') and dataset.mask is not None: - mask = dataset.mask.to(t.uint8) - else: - mask = None - - return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, translation_offsets=translation_offsets, weights=weights, mask=mask) - - - def interaction(self, index, translations): - pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, - translations) - # The 10x term is to condition the translation offsets - pix_trans -= self.min_translation - pix_trans = pix_trans + self.translation_offsets[index] - - all_exit_waves = [] - for i in range(self.probe.shape[0]): - exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(self.probe[i], - self.obj, - pix_trans, - shift_probe=True) - - 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) - - - 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(IncoherentPtycho, 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) - - - def sim_to_dataset(self, args_list): - pass - - diff --git a/CDTools/models/simple_ptycho.py b/CDTools/models/simple_ptycho.py index 341a7c4..09ee69c 100644 --- a/CDTools/models/simple_ptycho.py +++ b/CDTools/models/simple_ptycho.py @@ -16,7 +16,7 @@ class SimplePtycho(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, detector_slice, probe_guess, obj_guess, min_translation = [0,0], - mask=None): + surface_normal=np.array([0.,0.,1.]), mask=None): super(SimplePtycho,self).__init__() self.wavelength = t.Tensor([wavelength]) @@ -33,6 +33,9 @@ class SimplePtycho(CDIModel): self.probe_basis = t.Tensor(probe_basis) self.detector_slice = detector_slice + + self.surface_normal = t.Tensor(surface_normal) + if mask is None: self.mask = None else: @@ -71,9 +74,16 @@ class SimplePtycho(CDIModel): distance, center=center) + 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.]) + # Next generate the object geometry from the probe geometry and # the translations - pix_translations = tools.interactions.translations_to_pixel(probe_basis, 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) # Finally, initialize the probe and object using this information @@ -82,17 +92,19 @@ class SimplePtycho(CDIModel): obj = t.ones(obj_size+(2,)) det_geo = dataset.detector_geometry + if hasattr(dataset, 'mask') and dataset.mask is not None: mask = dataset.mask.to(t.uint8) else: mask = None - return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, mask=mask) + return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, mask=mask, surface_normal=surface_normal) def interaction(self, index, translations): pix_trans = tools.interactions.translations_to_pixel(self.probe_basis, - translations) + translations, + surface_normal=self.surface_normal) pix_trans -= self.min_translation return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe, self.obj, @@ -130,11 +142,11 @@ class SimplePtycho(CDIModel): 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.surface_normal = self.surface_normal.to(*args, **kwargs) def sim_to_dataset(self, args_list): # In the future, potentially add more control @@ -146,7 +158,15 @@ class SimplePtycho(CDIModel): 'instrument_n': 'Simulated Data', 'start_time': datetime.now()} - sample_info = {'description': 'A simulated sample'} + 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 diff --git a/CDTools/tools/data.py b/CDTools/tools/data.py index aa01f55..5d70b43 100644 --- a/CDTools/tools/data.py +++ b/CDTools/tools/data.py @@ -154,8 +154,16 @@ def get_sample_info(cxi_file): orient = np.array(s1['geometry_1/orientation']).astype(np.float32) xvec = orient[:3] / np.linalg.norm(orient[:3]) yvec = orient[3:] / np.linalg.norm(orient[3:]) - metadata['orientation'] = np.array([xvec,yvec, - np.cross(xvec,yvec)]) + metadata['orientation'] = np.array([xvec,yvec, + np.cross(xvec,yvec)]) + + if 'geometry_1/surface_normal' in s1: + snorm = np.array(s1['geometry_1/surface_normal']).astype(np.float32) + xvec = np.cross(np.array([0.,1.,0.]), snorm) + xvec /= np.linalg.norm(xvec) + yvec = np.cross(snorm, xvec) + yvec /= np.linalg.norm(yvec) + metadata['orientation'] = np.array([xvec, yvec, snorm]) # Check if the metadata is empty if metadata == {}: diff --git a/CDTools/tools/interactions.py b/CDTools/tools/interactions.py index 808efd9..5de75ce 100644 --- a/CDTools/tools/interactions.py +++ b/CDTools/tools/interactions.py @@ -15,7 +15,7 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations', # -def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0,0,1])): +def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])): """Takes real space translations and outputs them in pixel space This works for any 2D ptychography geometry. It takes in diff --git a/examples/specular_pinhole_ptycho.py b/examples/specular_pinhole_ptycho.py index 6920753..7d058cb 100644 --- a/examples/specular_pinhole_ptycho.py +++ b/examples/specular_pinhole_ptycho.py @@ -22,11 +22,11 @@ with h5py.File(filename,'r') as f: dataset = CDTools.datasets.Ptycho_2D_Dataset.from_cxi(f) -# A hack for the specular geometry -dataset.detector_geometry['basis'] = np.array([[0,-30e-6],[30e-6,0],[0,0]]) - -# Figure out why the padding doesn't work here -model = CDTools.models.FancyPtycho.from_dataset(dataset,randomize_ang = np.pi/4, padding=0, translation_scale=10)#, n_modes=2, propagation_distance=-5e-4) +model = CDTools.models.FancyPtycho.from_dataset(dataset, + randomize_ang = np.pi/4, + padding=0, + translation_scale=10, + scattering_mode='reflection') # Uncomment these to use on the CPU