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

get it running in principle, but it is not good

Browse Source
This commit is contained in:
levita_a
2025-05-12 16:55:26 +02:00
parent d853a141bf
commit ac1499f511
2 changed files with 83 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/models/fancy_ptycho.py
+48 -20
View File
@@ -28,6 +28,7 @@ class FancyPtycho(CDIModel):
probe_fourier_shifts=None,
mask=None,
weights=None,
qe_mask=None,
translation_scale=1,
saturation=None,
probe_support=None,
@@ -87,7 +88,14 @@ class FancyPtycho(CDIModel):
else:
self.register_buffer('mask',
t.as_tensor(mask, dtype=t.bool))
if qe_mask is None:
self.qe_mask = None
else:
self.qe_mask = t.nn.Parameter(
t.as_tensor(qe_mask, dtype=dtype))
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
@@ -202,6 +210,7 @@ class FancyPtycho(CDIModel):
dm_rank=None,
translation_scale=1,
saturation=None,
use_qe_mask=False,
probe_support_radius=None,
probe_fourier_crop=None,
propagation_distance=None,
@@ -376,6 +385,11 @@ class FancyPtycho(CDIModel):
else:
mask = None
if use_qe_mask:
qe_mask = t.ones(mask.shape, dtype=t.float32)
else:
qe_mask = None
if probe_support_radius is not None:
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
@@ -389,24 +403,34 @@ class FancyPtycho(CDIModel):
else:
probe_support = None
return cls(wavelength, det_geo, obj_basis, probe, obj,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_basis=probe_basis,
probe_support=probe_support,
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,
exponentiate_obj=exponentiate_obj,
obj_view_crop=obj_view_crop)
return cls(
wavelength,
det_geo,
obj_basis,
probe,
obj,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets=translation_offsets,
weights=Ws,
mask=mask,
background=background,
qe_mask=qe_mask,
translation_scale=translation_scale,
saturation=saturation,
probe_basis=probe_basis,
probe_support=probe_support,
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,
exponentiate_obj=exponentiate_obj,
obj_view_crop=obj_view_crop
)
def interaction(self, index, translations, *args):
@@ -521,6 +545,7 @@ class FancyPtycho(CDIModel):
wavefields,
self.background,
measurement=tools.measurements.incoherent_sum,
qe_mask=self.qe_mask,
saturation=self.saturation,
oversampling=self.oversampling,
simulate_finite_pixels=self.simulate_finite_pixels,
@@ -840,7 +865,10 @@ class FancyPtycho(CDIModel):
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: p.plot_amplitude(self.background**2, fig=fig))
lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)),
('Quantum Efficiency Mask',
lambda self, fig: p.plot_amplitude(self.qe_mask, fig=fig),
lambda self: (hasattr(self, 'qe_mask') and self.qe_mask is not None))
]
src/cdtools/tools/measurements/measurements.py
+35 -12
View File
@@ -155,7 +155,18 @@ 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, simulate_finite_pixels=False):
def quadratic_background(
wavefield,
background,
*args,
detector_slice=None,
measurement=intensity,
epsilon=1e-7,
qe_mask=None,
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
@@ -173,6 +184,8 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
Optional, a slice or tuple of slices defining a section of the simulation to return
measurement : function
Default is measurements.intensity, the measurement function to use.
qe_mask : torch.Tensor
A tensor storing the per-pixel quantum efficiency (up to an unknown global scaling factor)
saturation : float
Optional, a maximum saturation value to clamp the resulting intensities to
oversampling : int
@@ -183,18 +196,28 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
sim_patterns : torch.Tensor
A real MxN array storing the wavefield's intensities
"""
if detector_slice is None:
output = measurement(wavefield, *args, epsilon=epsilon,
oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels) \
+ background**2
else:
output = measurement(wavefield, *args, detector_slice=detector_slice,
epsilon=epsilon, oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels) \
+ background**2
if detector_slice is None:
raw_intensity = measurement(
wavefield,
*args,
epsilon=epsilon,
oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels)
else:
raw_intensity = measurement(
wavefield,
*args,
detector_slice=detector_slice,
epsilon=epsilon,
oversampling=oversampling,
simulate_finite_pixels=simulate_finite_pixels)
if qe_mask is None:
output = raw_intensity + background**2
else:
output = (qe_mask * raw_intensity) + background**2
# This has to be done after the background is added, hence we replicate
# it here
if saturation is None: