diff --git a/src/cdtools/datasets/base.py b/src/cdtools/datasets/base.py index 9dc4801..3f8ec8c 100644 --- a/src/cdtools/datasets/base.py +++ b/src/cdtools/datasets/base.py @@ -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) diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py index 971baf1..6631cd6 100644 --- a/src/cdtools/datasets/ptycho_2d_dataset.py +++ b/src/cdtools/datasets/ptycho_2d_dataset.py @@ -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 x 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, diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 8854d41..6220299 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -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)) ] 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/src/cdtools/tools/measurements/measurements.py b/src/cdtools/tools/measurements/measurements.py index e726149..c926568 100644 --- a/src/cdtools/tools/measurements/measurements.py +++ b/src/cdtools/tools/measurements/measurements.py @@ -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: 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 4bcbac1..7ef429f 100644 --- a/tests/models/test_fancy_ptycho.py +++ b/tests/models/test_fancy_ptycho.py @@ -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) 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_analysis.py b/tests/tools/test_analysis.py index 27c6a65..f0a1cf1 100644 --- a/tests/tools/test_analysis.py +++ b/tests/tools/test_analysis.py @@ -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, 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)