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 way to save and load the quantum efficiency masks, and add test coverage

Browse Source
This commit is contained in:
levita_a
2025-05-13 16:05:16 +02:00
parent c86e505ae5
commit 2c2f2d2d93
6 changed files with 195 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/cdtools/datasets/base.py
+14 -5
View File
@@ -110,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)
@@ -205,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):
@@ -248,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/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
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
@@ -21,6 +21,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,',
@@ -28,7 +29,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_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)