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

Merge pull request #30 from cdtools-developers/variable_qe

Browse Source 
Add the ability for fancy_ptycho to model detectors with spatially varying quantum efficiency
This commit is contained in:
Dayne Yoshiki Sasaki
2025-06-10 08:44:22 -07:00
committed by GitHub
parent da695c7c06 4ecc782c28
commit b3a70ec325
10 changed files with 340 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/datasets/base.py
+30 -9
View File
@@ -34,10 +34,16 @@ class CDataset(torchdata.Dataset):
needed to allow for easy mixing of data on the CPU and GPU.
"""
def __init__(self, entry_info=None, sample_info=None,
wavelength=None,
detector_geometry=None, mask=None,
background=None):
def __init__(
self,
entry_info=None,
sample_info=None,
wavelength=None,
detector_geometry=None,
mask=None,
qe_mask=None,
background=None,
):
"""The __init__ function allows construction from python objects.
@@ -73,6 +79,12 @@ class CDataset(torchdata.Dataset):
self.mask = t.tensor(mask, dtype=t.bool)
else:
self.mask = None
if qe_mask is not None:
self.qe_mask = t.as_tensor(qe_mask, dtype=t.float32)
else:
self.qe_mask = None
if background is not None:
self.background = t.tensor(background, dtype=t.float32)
else:
@@ -98,6 +110,8 @@ class CDataset(torchdata.Dataset):
if self.mask is not None:
self.mask = self.mask.to(*args,**mask_kwargs)
if self.qe_mask is not None:
self.qe_mask = self.qe_mask.to(*args,**kwargs)
if self.background is not None:
self.background = self.background.to(*args,**kwargs)
@@ -193,12 +207,17 @@ class CDataset(torchdata.Dataset):
'basis' : basis,
'corner' : corner}
mask = cdtdata.get_mask(cxi_file)
qe_mask = cdtdata.get_qe_mask(cxi_file)
dark = cdtdata.get_dark(cxi_file)
return cls(entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask, background=dark)
return cls(
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask,
background=dark,
)
def to_cxi(self, cxi_file):
@@ -236,6 +255,8 @@ class CDataset(torchdata.Dataset):
corner = corner)
if self.mask is not None:
cdtdata.add_mask(cxi_file, self.mask)
if self.qe_mask is not None:
cdtdata.add_qe_mask(cxi_file, self.qe_mask)
if self.background is not None:
cdtdata.add_dark(cxi_file, self.background)
src/cdtools/datasets/ptycho_2d_dataset.py
+45 -11
View File
@@ -77,6 +77,7 @@ class Ptycho2DDataset(CDataset):
self.intensities = t.as_tensor(intensities, dtype=t.float32)
else:
self.intensities = None
def __len__(self):
return self.patterns.shape[0]
@@ -375,10 +376,19 @@ class Ptycho2DDataset(CDataset):
equal to the sum of a <factor> x <factor> region of pixels in the
input pattern. This summation is done by pytorch.functional.avg_pool2d.
Any mask and background data which is stored with the dataset is
downsampled with the data. The background is downsampled using the same
method as the data. The mask is expanded so that any output pixel
containing a masked pixel will be masked.
Any mask, quantum efficiency, and background data which is stored with
the dataset is downsampled with the data. The background is downsampled
using the same method as the data.
If there is no quantum efficiency mask, then the mask is downsampled so
that any output pixel containing a masked pixel will be masked. If there
is a quantum efficiency mask, then the quantum efficiency mask is
downsampled using the same method as the data, and the mask is
downsampled to include any pixels for which there is at least one valid
pixel.
To avoid leakage of data from masked pixels, the data is first
multiplied by the mask before downsampling.
Parameters
----------
@@ -386,17 +396,41 @@ class Ptycho2DDataset(CDataset):
Default 2, the factor to downsample by
"""
self.patterns = t.nn.functional.avg_pool2d(
self.patterns.unsqueeze(0), factor, divisor_override=1)[0]
self.mask = t.logical_not(t.nn.functional.max_pool2d(
(1-self.mask.to(dtype=t.uint8)).unsqueeze(0).unsqueeze(0),
factor
)[0,0].to(dtype=t.bool))
if hasattr(self, 'mask') and self.mask is not None:
self.patterns = t.nn.functional.avg_pool2d(
(self.mask * self.patterns).unsqueeze(0),
factor, divisor_override=1)[0]
else:
self.patterns = t.nn.functional.avg_pool2d(
self.patterns.unsqueeze(0),
factor, divisor_override=1)[0]
# If we have a QE mask, we want to include all pixels for which at
# least one of the input pixels was unmasked, because we can account
# for the masked pixels through quantum efficiency
if hasattr(self, 'qe_mask') and self.qe_mask is not None:
self.qe_mask = t.nn.functional.avg_pool2d(
(self.mask * self.qe_mask).unsqueeze(0).unsqueeze(0),
factor)[0,0]
self.mask = t.nn.functional.max_pool2d(
self.mask.to(dtype=t.uint8).unsqueeze(0).unsqueeze(0),
factor)[0,0].to(dtype=t.bool)
# But if there is no QE mask, we need to only preserve pixels for
# which all input pixels were unmasked
elif hasattr(self, 'mask') and self.mask is not None:
self.mask = t.logical_not(t.nn.functional.max_pool2d(
(1-self.mask.to(dtype=t.uint8)).unsqueeze(0).unsqueeze(0),
factor
)[0,0].to(dtype=t.bool))
self.detector_geometry['basis'] = \
self.detector_geometry['basis'] * factor
if self.background is not None:
if hasattr(self, 'background') and self.background is not None:
self.background = t.nn.functional.avg_pool2d(
self.background.unsqueeze(0).unsqueeze(0),
factor,
src/cdtools/models/fancy_ptycho.py
+55 -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,18 @@ 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))
# I want the ability to optimize over this, but experience shows
# that it is wildly unstable, so I think it's best to keep
# gradients turned off by default
self.qe_mask.requires_grad=False
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
obj_guess = t.as_tensor(obj_guess, dtype=t.complex64)
@@ -202,6 +214,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 +389,14 @@ class FancyPtycho(CDIModel):
else:
mask = None
if use_qe_mask:
if hasattr(dataset, 'qe_mask') and dataset.qe_mask is not None:
qe_mask = t.as_tensor(dataset.qe_mask, dtype=t.float32)
else:
qe_mask = t.ones(dataset.patterns.shape[-2:], 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 +410,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 +552,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,
@@ -865,7 +897,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/data/data.py
+75
View File
@@ -22,6 +22,7 @@ __all__ = ['get_entry_info',
'get_wavelength',
'get_detector_geometry',
'get_mask',
'get_qe_mask',
'get_dark',
'get_data',
'get_shot_to_shot_info',
@@ -32,6 +33,7 @@ __all__ = ['get_entry_info',
'add_source',
'add_detector',
'add_mask',
'add_qe_mask',
'add_dark',
'add_data',
'add_shot_to_shot_info',
@@ -300,6 +302,42 @@ def get_mask(cxi_file):
return None
def get_qe_mask(cxi_file):
"""Returns the quantum efficiency mask defined in the cxi file object
There is no way to store a quantum efficiency mask (a.k.a. a flat-field
image) in the .cxi file specification, but experience has indicated that
this is often a valuable thing to store, because just correcting for a
flatfield with e.g. a division will mess up the photon counting statistics.
Because there is no specification, I have simply chosen to store the
quantum efficiency mask as a float32 array in the same location as the
mask is, i.e. `entry_1/instrument_1/detector_1/qe_mask`.
The stored quantum efficiency mask should be defined as the mask that
a simulated intensity pattern needs to be multiplied by to realize the
measured image. In other words, it should be a flat-field image, not the
inverse of a flat-field image.
Parameters
----------
cxi_file : h5py.File
A file object to be read
Returns
-------
qe_mask : np.array
A float32 array storing the quantum efficiency mask from the cxi file
"""
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1/qe_mask' in i1:
qe_mask = i1['detector_1/qe_mask'][()].astype(np.float32)
return qe_mask
else:
return None
def get_dark(cxi_file):
"""Returns an array with a dark image to use for initialization of a background model
@@ -635,6 +673,43 @@ def add_mask(cxi_file, mask):
d1.create_dataset('mask',data=mask_to_save)
def add_qe_mask(cxi_file, qe_mask):
"""Adds the specified quantum efficiency mask to the cxi file
There is no way to store a quantum efficiency mask (a.k.a. a flat-field
image) in the .cxi file specification, but experience has indicated that
this is often a valuable thing to store, because just correcting for a
flatfield with e.g. a division will mess up the photon counting statistics.
Because there is no specification, I have simply chosen to store the
quantum efficiency mask as an array in the same location as the
mask is, i.e. `entry_1/instrument_1/detector_1/qe_mask`.
The stored quantum efficiency mask should be defined as the mask that
a simulated intensity pattern needs to be multiplied by to realize the
measured image. In other words, it should be a flat-field image, not the
inverse of a flat-field image.
Parameters
----------
cxi_file : h5py.File
The file to add the mask to
qe_mask : array
The quantum efficiency mask to save out to the file
"""
if 'entry_1/instrument_1' not in cxi_file:
cxi_file['entry_1'].create_group('instrument_1')
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1' not in i1:
i1.create_group('detector_1')
d1 = i1['detector_1']
if isinstance(qe_mask, t.Tensor):
qe_mask = qe_mask.detach().cpu().numpy()
d1.create_dataset('qe_mask',data=qe_mask)
def add_dark(cxi_file, dark):
"""Adds the specified dark image to a cxi file
src/cdtools/tools/measurements/measurements.py
+27 -11
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
@@ -184,17 +197,20 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
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
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:
tests/conftest.py
+14
View File
@@ -137,9 +137,16 @@ def ptycho_cxi_1():
# Remember the format for the CXI file differs from the format used
# internally
mask = np.zeros((256,256)).astype(np.int32)
mask[5,8] = 1
expected['mask'] = np.ones((256,256)).astype(bool)
expected['mask'][5,8] = 0
d1f.create_dataset('mask',data=mask)
# There is no specification for this in the CXI file format :(
qe_mask = np.ones((256,256)).astype(np.float32)
expected['qe_mask'] = qe_mask
d1f.create_dataset('qe_mask',data=qe_mask)
# Create an initial background
dark = np.ones((256,256)) * 0.01
expected['dark'] = dark
@@ -228,6 +235,8 @@ def ptycho_cxi_2():
# internally
expected['mask'] = None
expected['qe_mask'] = None
# Test with a set of dark images
dark = np.ones((10,256,256)) * 0.01
expected['dark'] = np.nanmean(dark,axis=0)
@@ -305,8 +314,13 @@ def ptycho_cxi_3():
# Remember the format for the CXI file differs from the format used
# internally
mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
mask[15,47] = 38
expected['mask'] = np.ones((256,256)).astype(bool)
expected['mask'][15,47] = 0
d1f.create_dataset('mask',data=mask)
expected['qe_mask'] = None
expected['dark'] = None
data1f = e1f.create_group('data_1')
tests/models/test_fancy_ptycho.py
+2 -1
View File
@@ -60,6 +60,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
propagation_distance=5e-3,
units='mm',
obj_view_crop=-50,
use_qe_mask=True, # test this in the case where no qe mask is defined
)
print('Running reconstruction on provided reconstruction_device,',
@@ -67,7 +68,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
model.to(device=reconstruction_device)
dataset.get_as(device=reconstruction_device)
for loss in model.Adam_optimize(50, dataset, lr=0.02, batch_size=10):
for loss in model.Adam_optimize(70, dataset, lr=0.02, batch_size=10):
print(model.report())
if show_plot and model.epoch % 10 == 0:
model.inspect(dataset)
tests/test_datasets.py
+69 -12
View File
@@ -62,6 +62,9 @@ def test_CDataset_from_cxi(test_ptycho_cxis):
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
@@ -101,6 +104,9 @@ def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path):
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
@@ -115,6 +121,7 @@ def test_CDataset_to(ptycho_cxi_1):
if t.cuda.is_available():
dataset.to(device='cuda:0')
assert dataset.mask.device == t.device('cuda:0')
assert dataset.qe_mask.device == t.device('cuda:0')
assert dataset.background.device == t.device('cuda:0')
@@ -135,6 +142,7 @@ def test_Ptycho2DDataset_init():
[-20e-6,0,0]]).transpose(),
'corner': np.array((2550e-6,3825e-6,0.3))}
mask = np.ones((256,256))
qe_mask = 1.2*np.ones((256,256), dtype=np.float32)
patterns = np.random.rand(20,256,256)
translations = np.random.rand(20,3)
@@ -153,6 +161,24 @@ def test_Ptycho2DDataset_init():
assert t.allclose(dataset.patterns, t.as_tensor(patterns))
assert t.allclose(dataset.translations, t.as_tensor(translations))
# Also test one with a qe_mask
dataset = Ptycho2DDataset(translations, patterns,
entry_info=entry_info,
sample_info=sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask,
qe_mask=qe_mask)
assert t.all(t.eq(dataset.mask,t.BoolTensor(mask)))
assert t.all(t.eq(dataset.qe_mask,t.as_tensor(qe_mask)))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
assert dataset.detector_geometry == detector_geometry
assert t.allclose(dataset.patterns, t.as_tensor(patterns))
assert t.allclose(dataset.translations, t.as_tensor(translations))
def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
@@ -182,6 +208,9 @@ def test_Ptycho2DDataset_from_cxi(test_ptycho_cxis):
if expected['mask'] is not None:
assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask))
if expected['qe_mask'] is not None:
assert t.all(t.eq(t.tensor(expected['qe_mask']),dataset.qe_mask))
if expected['dark'] is not None:
assert t.all(t.eq(t.as_tensor(expected['dark'], dtype=t.float32),
dataset.background))
@@ -221,11 +250,12 @@ def test_Ptycho2DDataset_to_cxi(test_ptycho_cxis, tmp_path):
if dataset.detector_geometry['corner'] is not None:
assert 'corner' in read_dataset.detector_geometry
if dataset.mask is not None:
assert t.all(t.eq(dataset.mask,read_dataset.mask))
if dataset.qe_mask is not None:
assert t.all(t.eq(dataset.qe_mask,read_dataset.qe_mask))
if dataset.background is not None:
assert t.all(t.eq(dataset.background, read_dataset.background))
@@ -238,12 +268,14 @@ def test_Ptycho2DDataset_to(ptycho_cxi_1):
dataset.to(dtype=t.float64)
assert dataset.mask.dtype == t.bool
assert dataset.qe_mask.dtype == t.float64
assert dataset.patterns.dtype == t.float64
assert dataset.translations.dtype == t.float64
# If cuda is available, check that moving the mask to CUDA works.
if t.cuda.is_available():
dataset.to(device='cuda:0')
assert dataset.mask.device == t.device('cuda:0')
assert dataset.qe_mask.device == t.device('cuda:0')
assert dataset.background.device == t.device('cuda:0')
assert dataset.patterns.device == t.device('cuda:0')
assert dataset.translations.device == t.device('cuda:0')
@@ -291,24 +323,49 @@ def test_Ptycho2DDataset_downsample(test_ptycho_cxis):
# May start failing if the test datasets are changed to include
# a dataset with any dimension not even. That's a problem with the
# test, not the code. Sorry! -Abe
masked_patterns = dataset.mask * dataset.patterns
assert t.allclose(
copied_dataset.patterns,
dataset.patterns[:,::2,::2] +
dataset.patterns[:,1::2,::2] +
dataset.patterns[:,::2,1::2] +
dataset.patterns[:,1::2,1::2]
masked_patterns[:,::2,::2] +
masked_patterns[:,1::2,::2] +
masked_patterns[:,::2,1::2] +
masked_patterns[:,1::2,1::2]
)
assert t.allclose(
copied_dataset.mask,
t.logical_and(
if dataset.qe_mask is None:
manually_downsampled_mask = t.logical_and(
t.logical_and(dataset.mask[::2,::2],
dataset.mask[1::2,::2]),
t.logical_and(dataset.mask[::2,1::2],
dataset.mask[1::2,1::2]),
dataset.mask[1::2,1::2])
)
assert t.allclose(
copied_dataset.mask,
manually_downsampled_mask,
)
else:
manually_downsampled_mask = t.logical_or(
t.logical_or(dataset.mask[::2,::2],
dataset.mask[1::2,::2]),
t.logical_or(dataset.mask[::2,1::2],
dataset.mask[1::2,1::2])
)
assert t.allclose(
copied_dataset.mask,
manually_downsampled_mask
)
)
masked_qe_mask = dataset.mask * dataset.qe_mask
manually_downsampled_qe_mask = (
masked_qe_mask[::2,::2] + masked_qe_mask[1::2,::2]
+ masked_qe_mask[::2,1::2] + masked_qe_mask[1::2,1::2]
) / 4
assert t.allclose(
copied_dataset.qe_mask,
manually_downsampled_qe_mask
)
if dataset.background is not None:
tests/tools/test_analysis.py
+2 -2
View File
@@ -160,7 +160,7 @@ def test_standardize():
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.angle(np.sum(probe)) < 1e-7
assert np.angle(np.sum(probe)) < 2e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
@@ -232,7 +232,7 @@ def test_synthesize_reconstructions():
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.abs(np.angle(np.sum(probe))) < 1e-7
assert np.abs(np.angle(np.sum(probe))) < 2e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
tests/tools/test_data.py
+21 -1
View File
@@ -60,7 +60,15 @@ def test_get_mask(test_ptycho_cxis):
mask = data.get_mask(cxi)
if expected['mask'] is None and mask is None:
continue
assert np.all(data.get_mask(cxi) == expected['mask'])
assert np.all(mask == expected['mask'])
def test_get_qe_mask(test_ptycho_cxis):
for cxi, expected in test_ptycho_cxis:
qe_mask = data.get_qe_mask(cxi)
if expected['qe_mask'] is None and qe_mask is None:
continue
assert np.allclose(qe_mask, expected['qe_mask'])
def test_get_dark(test_ptycho_cxis):
@@ -207,6 +215,18 @@ def test_add_mask(tmp_path):
assert np.all(mask == read_mask)
def test_add_qe_mask(tmp_path):
qe_mask = np.random.rand(350,199).astype(np.float32)
with data.create_cxi(tmp_path / 'test_add_qe_mask.cxi') as f:
data.add_qe_mask(f, qe_mask)
with h5py.File(tmp_path / 'test_add_qe_mask.cxi','r') as f:
read_qe_mask = data.get_qe_mask(f)
assert np.allclose(qe_mask, read_qe_mask)
def test_add_dark(tmp_path):
dark = np.random.rand(350,620)