cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-09 21:12:42 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add an option to simulate finite pixel sizes

Browse Source
This commit is contained in:
Abe Levitan
2022-09-29 14:14:43 -07:00
parent 036fa8ea2c
commit 978c9d3206
3 changed files with 80 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/models/fancy_ptycho.py
+11 -4
View File
@@ -33,7 +33,8 @@ class FancyPtycho(CDIModel):
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False
simulate_probe_translation=False,
simulate_finite_pixels=False,
):
super(FancyPtycho, self).__init__()
@@ -124,6 +125,8 @@ class FancyPtycho(CDIModel):
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
self.I_phase = 2 * np.pi* Is
self.J_phase = 2 * np.pi* Js
self.simulate_finite_pixels = simulate_finite_pixels
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
@@ -155,7 +158,8 @@ class FancyPtycho(CDIModel):
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False
simulate_probe_translation=False,
simulate_finite_pixels=False,
):
wavelength = dataset.wavelength
@@ -209,6 +213,7 @@ class FancyPtycho(CDIModel):
# 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)
@@ -293,7 +298,8 @@ class FancyPtycho(CDIModel):
fourier_probe=fourier_probe,
oversampling=oversampling,
loss=loss, units=units,
simulate_probe_translation=simulate_probe_translation)
simulate_probe_translation=simulate_probe_translation,
simulate_finite_pixels=simulate_finite_pixels)
def interaction(self, index, translations, *args):
@@ -376,7 +382,8 @@ class FancyPtycho(CDIModel):
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
oversampling=self.oversampling,
simulate_finite_pixels=self.simulate_finite_pixels)
# Note: No "loss" function is defined here, because it is added
src/cdtools/models/rpi.py
+17 -6
View File
@@ -62,7 +62,8 @@ class RPI(CDIModel):
def __init__(self, wavelength, detector_geometry, probe_basis,
probe, obj_guess, detector_slice=None,
background=None, mask=None, saturation=None,
obj_support=None, oversampling=1, weight_matrix=False):
obj_support=None, oversampling=1, weight_matrix=False,
simulate_finite_pixels=False):
super(RPI, self).__init__()
@@ -134,9 +135,11 @@ class RPI(CDIModel):
self.oversampling = oversampling
self.simulate_finite_pixels=simulate_finite_pixels
@classmethod
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0):
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0, simulate_finite_pixels=False):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -213,7 +216,13 @@ class RPI(CDIModel):
# on a coarser grid needs to be accounted for here that is not
# accounted for yet
scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+list(obj_size)))
elif initialization.lower().strip() == 'uniform':
# I think something to do with the fact that the object is defined
# on a coarser grid needs to be accounted for here that is not
# accounted for yet
scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
obj_guess = scale * t.ones([n_modes,]+list(obj_size), dtype=probe.dtype)
elif initialization.lower().strip() == 'spectral':
if background is not None:
obj_guess = initializers.RPI_spectral_init(
@@ -226,7 +235,7 @@ class RPI(CDIModel):
else:
raise KeyError('Initialization "' + str(initialization) + \
'" invalid - use "spectral" or "random"')
'" invalid - use "spectral", "uniform", or "random"')
probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
probe_fft = tools.propagators.far_field(probe_intensity)
@@ -244,7 +253,8 @@ class RPI(CDIModel):
probe, obj_guess, detector_slice=det_slice,
background=background, mask=mask, saturation=saturation,
obj_support=obj_support, oversampling=oversampling,
weight_matrix=weight_matrix)
weight_matrix=weight_matrix,
simulate_finite_pixels=simulate_finite_pixels)
@classmethod
def from_calibration(cls, calibration, obj_size=None, n_modes=1, saturation=None):
@@ -353,7 +363,8 @@ class RPI(CDIModel):
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
oversampling=self.oversampling,
simulate_finite_pixels=self.simulate_finite_pixels)
return m
def loss(self, sim_data, real_data, mask=None):
src/cdtools/tools/measurements/measurements.py
+52 -8
View File
@@ -17,7 +17,7 @@ from torch.nn.functional import avg_pool2d
__all__ = ['intensity', 'incoherent_sum', 'quadratic_background']
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False):
"""Returns the intensity of a wavefield
The intensity is defined as the magnitude squared of the
@@ -40,8 +40,30 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
sim_patterns : torch.Tensor
A real MxN array storing the wavefield's intensities
"""
output = t.abs(wavefield)**2
if simulate_finite_pixels:
inverse_fft = t.fft.fftshift(t.fft.ifft2(wavefield), dim=(-2,-1))
pad1l = wavefield.shape[-2]//2
pad1r = wavefield.shape[-2] - pad1l
pad2l = wavefield.shape[-1]//2
pad2r = wavefield.shape[-1] - pad2l
padded = t.nn.functional.pad(inverse_fft, (pad1l, pad1r, pad2l, pad2r))
upsampled_field = t.fft.fft2(t.fft.ifftshift(padded, dim=(-2,-2)))
upsampled_intensity = t.abs(upsampled_field)**2
ifft_intensity = t.fft.fftshift(t.fft.ifft2(upsampled_intensity), dim=(-2,-1))
# Now we take a sinc function
xs = t.arange(ifft_intensity.shape[-2])
xs = (xs / t.max(xs)) * 2 - 1
ys = t.arange(ifft_intensity.shape[-1])
ys = (ys / t.max(ys)) * 2 - 1
Xs, Ys = t.meshgrid(xs, ys, indexing='ij')
mask = (t.special.sinc(Xs) * t.special.sinc(Ys)).to(device=ifft_intensity.device)
blurred_intensity = t.fft.fft2(t.fft.ifftshift(mask * ifft_intensity, dim=(-2,-2)))
output = blurred_intensity[...,::2,::2]
else:
output = t.abs(wavefield)**2
# Now we apply oversampling
if oversampling != 1:
if wavefield.dim() == 2:
@@ -63,7 +85,7 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
return t.clamp(output + epsilon,0,saturation)
def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False):
"""Returns the incoherent sum of the intensities of the wavefields
The intensity is defined as the sum of the magnitudes squared of
@@ -91,8 +113,27 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
sim_patterns : torch.Tensor
A real LXMxN array storing the incoherently summed intensities
"""
output = t.sum(t.abs(wavefields)**2,dim=-3)
if simulate_finite_pixels:
inverse_fft = t.fft.fftshift(t.fft.ifft2(wavefields), dim=(-2,-1))
pad1l = wavefields.shape[-2]//2
pad1r = wavefields.shape[-2] - pad1l
pad2l = wavefields.shape[-1]//2
pad2r = wavefields.shape[-1] - pad2l
padded = t.nn.functional.pad(inverse_fft, (pad1l, pad1r, pad2l, pad2r))
upsampled_field = t.fft.fft2(t.fft.ifftshift(padded, dim=(-2,-2)))
upsampled_intensity = t.sum(t.abs(upsampled_field)**2, dim=-3)
ifft_intensity = t.fft.fftshift(t.fft.ifft2(upsampled_intensity), dim=(-2,-1))
# Now we take a sinc function
xs = t.arange(ifft_intensity.shape[-2])
xs = (xs / t.max(xs)) * 2 - 1
ys = t.arange(ifft_intensity.shape[-1])
ys = (ys / t.max(ys)) * 2 - 1
Xs, Ys = t.meshgrid(xs, ys, indexing='ij')
mask = (t.special.sinc(Xs) * t.special.sinc(Ys)).to(device=ifft_intensity.device)
blurred_intensity = t.fft.fft2(t.fft.ifftshift(mask * ifft_intensity, dim=(-2,-2)))
output = t.abs(blurred_intensity[...,::2,::2])
else:
output = t.sum(t.abs(wavefields)**2,dim=-3)
# Now we apply oversampling
if oversampling != 1:
@@ -114,7 +155,7 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
return t.clamp(output + epsilon,0,saturation)
def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1):
def quadratic_background(wavefield, background, *args, detector_slice=None, measurement=intensity, epsilon=1e-7, saturation=None, oversampling=1, simulate_finite_pixels=False):
"""Returns the intensity of a wavefield plus a background
The intensity is calculated via the given measurment function
@@ -145,10 +186,13 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
if detector_slice is None:
output = measurement(wavefield, *args, epsilon=epsilon,
oversampling=oversampling) + background**2
oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels) \
+ background**2
else:
output = measurement(wavefield, *args, detector_slice=detector_slice,
epsilon=epsilon, oversampling=oversampling) \
epsilon=epsilon, oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels) \
+ background**2
# This has to be done after the background is added, hence we replicate