From a09cd4b39ee1d45a87c3cff027a77610d66e494e Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Wed, 22 May 2019 20:20:43 -0400 Subject: [PATCH] Add ability to load/use Bragg geometry ptycho data, also add a 1D convolve to try with the terrible data from CSX --- CDTools/models/fancy_ptycho.py | 36 ++++++++++++++---- CDTools/tools/data.py | 23 +++++++++--- CDTools/tools/image_processing.py | 55 +++++++++++++++++++++++++++- CDTools/tools/initializers.py | 5 ++- tests/conftest.py | 5 ++- tests/tools/test_image_processing.py | 31 +++++++++++++++- 6 files changed, 138 insertions(+), 17 deletions(-) diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 9d83724..0b398c4 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -14,7 +14,9 @@ class FancyPtycho(CDIModel): def __init__(self, wavelength, detector_geometry, probe_basis, detector_slice, - probe_guess, obj_guess, min_translation = t.Tensor([0,0]), + probe_guess, obj_guess, + surface_normal=np.array([0.,0.,1.]), + min_translation = t.Tensor([0,0]), background = None, translation_offsets=None, mask=None, weights = None, translation_scale = 1, saturation=None, probe_support = None, obj_support=None): @@ -34,7 +36,8 @@ class FancyPtycho(CDIModel): 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: @@ -45,9 +48,9 @@ class FancyPtycho(CDIModel): # We rescale the probe here so it learns at the same rate as the # object if probe_guess.dim() > 3: - self.probe_norm = t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) + self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32))) else: - self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32))) + 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) @@ -110,10 +113,18 @@ class FancyPtycho(CDIModel): center=center, padding=padding, opt_for_fft=False) + + + 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, padding=50) @@ -166,12 +177,21 @@ class FancyPtycho(CDIModel): else: obj_support = None - return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, 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) + return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, + 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) 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 if self.translation_offsets is not None: @@ -249,7 +269,7 @@ class FancyPtycho(CDIModel): 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) + 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 diff --git a/CDTools/tools/data.py b/CDTools/tools/data.py index 0a2f133..aa01f55 100644 --- a/CDTools/tools/data.py +++ b/CDTools/tools/data.py @@ -135,12 +135,13 @@ def get_sample_info(cxi_file): metadata[attr] = str(s1[attr][()].decode()) except AttributeError as e: metadata[attr] = str(np.array(s1[attr][:])[0].decode()) - + float_attrs = ['concentration', 'mass', 'temperature', 'thickness', 'unit_cell_volume'] + for attr in float_attrs: if attr in s1: metadata[attr] = np.float32(s1[attr][()]) @@ -148,11 +149,14 @@ def get_sample_info(cxi_file): if 'unit_cell' in s1: metadata['unit_cell'] = np.array(s1['unit_cell']).astype(np.float32) - # TODO: Add my nonstandard "surface normal" attribute here - - # TODO: I should add the sample geometry as a valid metadata that can - # be copied over + if 'geometry_1/orientation' in s1: + 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)]) + # Check if the metadata is empty if metadata == {}: metadata = None @@ -444,7 +448,16 @@ def add_sample_info(cxi_file, metadata): cxi_file['entry_1'].create_group('sample_1') s1 = cxi_file['entry_1/sample_1'] + if 'orientation' in metadata: + if 'geometry_1' not in s1: + s1.create_group('geometry_1') + # Only store the part of this matrix as defined in the CXI file spec + s1['geometry_1'].create_dataset('orientation', + data=metadata['orientation'].ravel()[:6]) + for key, value in metadata.items(): + if key == 'orientation': + continue # this is a special case if isinstance(value,(str,bytes)): s1[key] = np.string_(value) elif isinstance(value, datetime.datetime): diff --git a/CDTools/tools/image_processing.py b/CDTools/tools/image_processing.py index 5eea286..03a2c57 100644 --- a/CDTools/tools/image_processing.py +++ b/CDTools/tools/image_processing.py @@ -4,7 +4,8 @@ import torch as t from CDTools.tools import cmath __all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift', - 'find_subpixel_shift', 'find_pixel_shift', 'find_shift'] + 'find_subpixel_shift', 'find_pixel_shift', 'find_shift', + 'convolve_1d'] def centroid(im, dims=2): @@ -206,3 +207,55 @@ def find_shift(im1, im2, resolution=10): resolution=resolution) return subpixel_shift + + +def convolve_1d(image, kernel, dim=0, fftshift_kernel=True): + """Convolves an image with a 1d kernel along a specified dimension + + The convolution is a circular convolution calculated using a Fourier + transform. The calculation is done so the input remains differentiable + with respect to the output. + + If the image has a final dimension of 2, it is assumed to be complex. + Otherwise, the image is assumed to be real. The image and kernel + must either both be real or both be complex. + + Args: + image (torch.Tensor) : The image to convolve + kernel (torch.Tensor) : The 1d kernel to convolve with + dim (int) : Default 0, the dimension to convolve along + fftshift_kernel (bool) : Default True, whether to fftshift the kernel first. + + Returns: + (torch.Tensor) : The convolved image + """ + + complex_things = 2 + if image.shape[-1] != 2: + image = t.stack((image,t.zeros_like(image)),dim=-1) + complex_things -= 1 + + if kernel.shape[-1] != 2: + kernel = t.stack((kernel,t.zeros_like(kernel)),dim=-1) + complex_things -= 1 + + # Take a correlation + if fftshift_kernel: + kernel = cmath.ifftshift(kernel) + + + # We have to transpose the relevant dimension to -2 before using the fft, + # which expects to operate on the final non-complex dimension + trans_im = t.transpose(image, dim, -2) + + fft_im = t.fft(trans_im, 1) + fft_kernel = t.fft(kernel, 1) + trans_conv = t.ifft(cmath.cmult(fft_im,fft_kernel), 1) + + conv_im = t.transpose(trans_conv, dim, -2) + + # If nothing was input as complex, the result should be returned as real + if complex_things == 0: + return conv_im[...,0] + else: + return conv_im diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index 5ff497e..33e2112 100644 --- a/CDTools/tools/initializers.py +++ b/CDTools/tools/initializers.py @@ -240,7 +240,10 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None): # to use the mask or not? intensities = np.zeros(shape) for params, im in dataset: - intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy() + if hasattr(dataset,'mask') and dataset.mask is not None: + intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy() + else: + intensities[det_slice] += im.cpu().numpy() intensities /= len(dataset) # Subtract off a known background if it's stored diff --git a/tests/conftest.py b/tests/conftest.py index a1b6a0c..f1abff3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -123,8 +123,11 @@ def ptycho_cxi_1(): expected['axes'] = ['translation','y','x'] g1f = s1f.create_group('geometry_1') + orientation = np.array([1.,0,0,0,1,0]) + g1f.create_dataset('orientation', data=orientation) + s1e['orientation'] = np.array([[1.,0,0],[0,1,0],[0,0,1]]) translations = np.arange(300).reshape((100,3)).astype(np.float32) - g1f.create_dataset('translation',data=translations) + g1f.create_dataset('translation',data=translations) data1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') d1f['translation'] = h5py.SoftLink('/entry_1/sample_1/geometry_1/translation') expected['translations'] = -translations diff --git a/tests/tools/test_image_processing.py b/tests/tools/test_image_processing.py index 7fcd736..e93869b 100644 --- a/tests/tools/test_image_processing.py +++ b/tests/tools/test_image_processing.py @@ -7,7 +7,7 @@ import torch as t from CDTools.tools import image_processing, cmath, initializers, interactions from scipy import ndimage - +from scipy.signal import fftconvolve def test_centroid(): # Test single im @@ -108,4 +108,33 @@ def test_find_shift(): retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50) # tolerance of 0.03 on this measurement assert t.all(t.abs(shift + t.Tensor((40,6)) - retrieved_shift) < 0.03) + +def test_convolve_1d(): + from matplotlib import pyplot as plt + test_image = np.random.rand(400,300) + #test_image = np.hstack((np.ones((400,150)),np.zeros((400,150)))) + xs = np.linspace(-100,100,300) + kernel = 1/(1+xs**2) + + # First, we test with everything real, dim=1 + convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(kernel),dim=1) + + np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1)) + assert np.allclose(convolved.numpy(),np_result) + + + xs = np.linspace(-100,100,400) + kernel = 1/(1+xs**2) + + # Then with dim=0, and a non-fftshifted kernel + convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(np.fft.ifftshift(kernel)), fftshift_kernel=False) + + np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)) + assert np.allclose(convolved.numpy(),np_result) + + # And finally with complex input + convolved = cmath.torch_to_complex(image_processing.convolve_1d(cmath.complex_to_torch(test_image),cmath.complex_to_torch(kernel))) + + np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0) + assert np.allclose(convolved,np_result)