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 a fourier translation to fancy ptycho

Browse Source
This commit is contained in:
levita_a
2024-09-19 18:09:37 +02:00
parent ad9420c3f2
commit 6b2358df36
2 changed files with 107 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/models/fancy_ptycho.py
+26 -4
View File
@@ -25,6 +25,7 @@ class FancyPtycho(CDIModel):
background=None,
probe_basis=None,
translation_offsets=None,
probe_fourier_shifts=None,
mask=None,
weights=None,
translation_scale=1,
@@ -134,6 +135,13 @@ class FancyPtycho(CDIModel):
t_o = t.as_tensor(translation_offsets, dtype=t.float32)
t_o = t_o / translation_scale
self.translation_offsets = t.nn.Parameter(t_o)
if probe_fourier_shifts is None:
self.probe_fourier_shifts = None
else:
self.probe_fourier_shifts = t.nn.Parameter(
t.as_tensor(translation_offsets, dtype=t.float32)
)
self.register_buffer('translation_scale',
t.as_tensor(translation_scale, dtype=dtype))
@@ -152,7 +160,7 @@ class FancyPtycho(CDIModel):
t.as_tensor(simulate_probe_translation, dtype=bool)
)
if simulate_probe_translation:
if simulate_probe_translation or (self.probe_fourier_shifts is not None):
Is = t.arange(self.probe.shape[-2], dtype=dtype)
Js = t.arange(self.probe.shape[-1], dtype=dtype)
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
@@ -195,6 +203,7 @@ class FancyPtycho(CDIModel):
fourier_probe=False,
loss='amplitude mse',
units='um',
allow_probe_fourier_shifts=False,
simulate_probe_translation=False,
simulate_finite_pixels=False,
exponentiate_obj=False,
@@ -327,6 +336,11 @@ class FancyPtycho(CDIModel):
translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
if allow_probe_fourier_shifts:
probe_fourier_shifts = t.zeros((len(dataset), 2), dtype=t.float32)
else:
probe_fourier_shifts = None
if dm_rank is not None and dm_rank != 0:
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
@@ -380,6 +394,7 @@ class FancyPtycho(CDIModel):
fourier_probe=fourier_probe,
oversampling=oversampling,
loss=loss, units=units,
probe_fourier_shifts=probe_fourier_shifts,
simulate_probe_translation=simulate_probe_translation,
simulate_finite_pixels=simulate_finite_pixels,
phase_only=phase_only,
@@ -431,12 +446,19 @@ class FancyPtycho(CDIModel):
# Maybe this can be done with a matmul now?
prs = t.sum(Ws[..., None, None] * basis_prs, axis=-3)
if self.simulate_probe_translation:
det_pix_trans = tools.interactions.translations_to_pixel(
if self.simulate_probe_translation or (self.probe_fourier_shifts is not None):
if self.probe_fourier_shifts is not None:
det_pix_trans = self.probe_fourier_shifts[index]
else:
det_pix_trans = t.zeros_like(translations)
if self.simulate_probe_translation:
det_pix_trans = det_pix_trans + tools.interactions.translations_to_pixel(
self.det_basis,
translations,
surface_normal=self.surface_normal)
probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
self.I_phase[None,...] +
det_pix_trans[:,1,None,None] *
src/cdtools/tools/analysis/analysis.py
+81
View File
@@ -13,6 +13,7 @@ import cdtools
from scipy import linalg as sla
from scipy import special
from scipy import optimize as opt
from scipy import spatial
__all__ = [
'product_svd',
@@ -31,6 +32,7 @@ __all__ = [
'remove_amplitude_exponent',
'standardize_reconstruction_set',
'standardize_reconstruction_pair',
'calc_spectral_info',
]
@@ -1432,3 +1434,82 @@ def standardize_reconstruction_pair(
return results
def calc_spectral_info(dataset, nbins=50):
"""Makes a properly normalized sum diffraction pattern
This returns a scaled version of sum of all the diffraction patterns
within the dataset. The scaling is defined so that the total intensity
in the final image is equal to the intensity arising from a region of
the scan pattern whose area matches one detector conjugate field of
view.
Parameters
----------
dataset : Ptycho2DDataset
A ptychography dataset to use
nbins : int
The number of bins to use for the SNR curve
Returns
-------
spectrum : t.tensor
An image of the spectral signal rate
freqs : t.tensor
The frequencies at which the SSNR is estimated
SSNR : t.tensor
The estimated SSNR
"""
scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
scan_area = scan_hull.volume
ewg = cdtools.tools.initializers.exit_wave_geometry
obj_basis = ewg(
dataset.detector_geometry['basis'],
dataset[0][1].shape,
dataset.wavelength,
dataset.detector_geometry['distance'],
)
det_conj_fov_area = np.linalg.norm(
np.cross(obj_basis[:,0]*dataset.patterns.shape[-2],
obj_basis[:,1]*dataset.patterns.shape[-1])
)
scale_factor = det_conj_fov_area / scan_area
mask = dataset.mask.cpu().numpy().astype(int)
sum_pattern = dataset.mask * t.sum(dataset.patterns, dim=0) * scale_factor
sum_pattern = sum_pattern.cpu().numpy()
# TODO this assumes orthogonal axes
pix_sizes = np.linalg.norm(obj_basis, axis=0)
i_freqs = np.fft.fftshift(np.fft.fftfreq(
sum_pattern.shape[0],d=pix_sizes[0]))
j_freqs = np.fft.fftshift(np.fft.fftfreq(
sum_pattern.shape[1],d=pix_sizes[1]))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
max_i = np.max(i_freqs)
max_j = np.max(j_freqs)
frc_range = [0, max(max_i,max_j)]
sum_spectrum, frc_bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=sum_pattern)
sum_spectrum_sq, frc_bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=sum_pattern**2)
n_pix, frc_bins = np.histogram(Rs, bins=nbins, range=frc_range,
weights=mask)
mean_spectrum = sum_spectrum / n_pix
pattern_snr = sum_spectrum_sq / sum_spectrum
return sum_pattern, frc_bins[:-1], mean_spectrum