From 2c2f2d2d935708a5b203ad4ffb434d159df68089 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 13 May 2025 16:05:16 +0200 Subject: [PATCH] Add a way to save and load the quantum efficiency masks, and add test coverage --- src/cdtools/datasets/base.py | 19 ++++++-- src/cdtools/tools/data/data.py | 75 ++++++++++++++++++++++++++++ tests/conftest.py | 14 ++++++ tests/models/test_fancy_ptycho.py | 3 +- tests/test_datasets.py | 81 ++++++++++++++++++++++++++----- tests/tools/test_data.py | 22 ++++++++- 6 files changed, 195 insertions(+), 19 deletions(-) diff --git a/src/cdtools/datasets/base.py b/src/cdtools/datasets/base.py index 85ec854..3f8ec8c 100644 --- a/src/cdtools/datasets/base.py +++ b/src/cdtools/datasets/base.py @@ -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) diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index b69d0c0..89879dd 100644 --- a/src/cdtools/tools/data/data.py +++ b/src/cdtools/tools/data/data.py @@ -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 diff --git a/tests/conftest.py b/tests/conftest.py index 51a58dc..781dd9e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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') diff --git a/tests/models/test_fancy_ptycho.py b/tests/models/test_fancy_ptycho.py index c9eaa1b..1591bf9 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -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) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index 6897991..3bb91be 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -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: diff --git a/tests/tools/test_data.py b/tests/tools/test_data.py index e637509..39b64dd 100644 --- a/tests/tools/test_data.py +++ b/tests/tools/test_data.py @@ -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)