mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 13:32:40 +02:00
Add ability to load/use Bragg geometry ptycho data, also add a 1D convolve to try with the terrible data from CSX
This commit is contained in:
Abe Levitan
@@ -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
|
||||
|
||||
|
||||
|
||||
+18
-5
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+4
-1
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user