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

Merge branch 'cdtools-developers:master' into bugfix/rpi_compare

Browse Source
This commit is contained in:
Dayne Yoshiki Sasaki
2025-06-11 17:03:06 -07:00
committed by GitHub
parent 4c784de922 b3a70ec325
commit 31318819d0
11 changed files with 432 additions and 78 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
+49 -12
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]
@@ -269,9 +270,12 @@ class Ptycho2DDataset(CDataset):
"""Plots the mean diffraction pattern across the dataset
The output is normalized so that the summed intensity on the
detector is equal to the total intensity of light that passed
detector is roughly equal to the total intensity of light that passed
through the sample within each detector conjugate field of view.
If the scan points are colinear (which causes issues for this
estimation), the mean pattern is displayed unscaled.
The plot is plotted as log base 10 of the output plus log_offset.
By default, log_offset is set equal to 1, which is a good level for
shot-noise limited data captured in units of photons. More
@@ -372,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
----------
@@ -383,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
+85 -25
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,
@@ -597,15 +629,40 @@ class FancyPtycho(CDIModel):
def center_probes(self, iterations=4):
"""Centers the probes
"""Centers the probes in real space
Takes the current guess of the illumination function and centers it
using a shift with periodic boundary conditions. It uses
cdtools.tools.image_processing.center internally to do the centering.
Multiple iterations of an algorithm are run, which is helpful if the
illumination is reconstructed near the corners and "wraps around" the
probe field of view.
Note that the centering is always performed in real space, even if
the probe array is defined in Fourier space.
Note that this does not compensate for the centering by adjusting
Note also that this does not compensate for the centering by adjusting
the object, so it's a good idea to reset the object after centering
the probes
Parameters
----------
iterations : int
Default 4, how many iterations of the centering algorithm to run
"""
centered_probe = tools.image_processing.center(
self.probe.data.cpu(), iterations=iterations)
self.probe.data = centered_probe.to(device=self.probe.data.device)
if self.fourier_probe:
prs = tools.propagators.inverse_far_field(self.probe.detach()).cpu()
else:
prs = self.probe.detach().cpu()
centered_prs = tools.image_processing.center(prs, iterations=iterations)
if self.fourier_probe:
self.probe.data = tools.propagators.far_field(
centered_prs.to(device=self.probe.data.device))
else:
self.probe.data = centered_prs.to(device=self.probe.data.device)
def tidy_probes(self):
@@ -840,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/analysis/analysis.py
+19 -5
View File
@@ -14,6 +14,7 @@ from scipy import linalg as sla
from scipy import special
from scipy import optimize as opt
from scipy import spatial
import warnings
__all__ = [
'product_svd',
@@ -1443,6 +1444,13 @@ def calc_spectral_info(dataset, nbins=50):
the scan pattern whose area matches one detector conjugate field of
view.
This estimation will start to deviate from the truth if the scan area
is not significantly larger than the illumination function, because
the nonzero size of the illumination function is not taken into account.
Furthermore, in the edge case where all the scan points are colinear,
the estimate will fail, and the mean diffraction pattern will be returned
instead
Parameters
----------
dataset : Ptycho2DDataset
@@ -1461,10 +1469,12 @@ def calc_spectral_info(dataset, nbins=50):
"""
scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
scan_area = scan_hull.volume
try:
scan_hull = spatial.ConvexHull(dataset.translations[:,:2].cpu().numpy())
scan_area = scan_hull.volume
except spatial._qhull.QhullError as e:
scan_area = None
ewg = cdtools.tools.initializers.exit_wave_geometry
obj_basis = ewg(
dataset.detector_geometry['basis'],
@@ -1477,8 +1487,12 @@ def calc_spectral_info(dataset, nbins=50):
np.cross(obj_basis[:,0]*dataset.patterns.shape[-2],
obj_basis[:,1]*dataset.patterns.shape[-1])
)
scale_factor = det_conj_fov_area / scan_area
if scan_area is not None:
scale_factor = det_conj_fov_area / scan_area
else:
warnings.warn("The scan points in this dataset are all colinear. The mean pattern will be calculated rather than a scaled mean based on the scanned area.")
scale_factor = 1/len(dataset)
mask = dataset.mask.cpu().numpy().astype(int)
sum_pattern = dataset.mask * t.sum(dataset.patterns, dim=0) * scale_factor
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
+41 -1
View File
@@ -5,6 +5,45 @@ import torch as t
import cdtools
from matplotlib import pyplot as plt
def test_center_probe(lab_ptycho_cxi):
dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
fourier_probe=False
)
base_probe = model.probe.detach().clone()
model.center_probes()
centered_probe = model.probe.detach().clone()
fourier_model = cdtools.models.FancyPtycho.from_dataset(
dataset,
n_modes=3,
fourier_probe=True,
)
fourier_model.probe.data = cdtools.tools.propagators.far_field(
base_probe
)
fourier_base_probe = fourier_model.probe.detach().clone()
fourier_model.center_probes()
fourier_centered_probe = fourier_model.probe.detach().clone()
ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
fourier_centered_probe)
# So we know the code had to do something
assert not t.allclose(base_probe, centered_probe)
# And checking that they both do the same thing, whether or not
# fourier_probe was set to True
assert t.allclose(
centered_probe,
ifft_fourier_centered_probe,
atol=1e-4,
rtol=1e-3
)
@pytest.mark.slow
def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
@@ -21,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,',
@@ -28,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)