From 26de79ca923f575e605171f6cc46421897fdf861 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Thu, 25 Apr 2019 16:06:59 -0400 Subject: [PATCH] Include ability to read background from a .cxi file and include it in SHARP-style initialization --- CDTools/datasets.py | 17 ++++++++++++++--- CDTools/models/fancy_ptycho.py | 11 +++++++---- CDTools/tools/initializers.py | 4 ++++ examples/gold_ball_ptycho.py | 2 +- tests/test_datasets.py | 19 ++++++++++++++++++- 5 files changed, 44 insertions(+), 9 deletions(-) diff --git a/CDTools/datasets.py b/CDTools/datasets.py index 5eece18..2a5edf3 100644 --- a/CDTools/datasets.py +++ b/CDTools/datasets.py @@ -20,7 +20,8 @@ class CDataset(torchdata.Dataset): def __init__(self, entry_info=None, sample_info=None, wavelength=None, - detector_geometry=None, mask=None): + detector_geometry=None, mask=None, + background=None): # Force pass-by-value-like behavior to stop strangeness self.entry_info = copy(entry_info) @@ -31,6 +32,10 @@ class CDataset(torchdata.Dataset): self.mask = t.tensor(mask) else: self.mask = None + if background is not None: + self.background = t.Tensor(background) + else: + self.background = None self.get_as(device='cpu') @@ -45,6 +50,8 @@ class CDataset(torchdata.Dataset): if self.mask is not None: self.mask = self.mask.to(*args,**mask_kwargs) + if self.background is not None: + self.background = self.background.to(*args,**kwargs) def get_as(self, *args, **kwargs): @@ -83,11 +90,12 @@ class CDataset(torchdata.Dataset): 'basis' : basis, 'corner' : corner} mask = cdtdata.get_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) + mask=mask, background=dark) def to_cxi(self, cxi_file): @@ -108,6 +116,8 @@ class CDataset(torchdata.Dataset): corner = corner) if self.mask is not None: cdtdata.add_mask(cxi_file, self.mask) + if self.background is not None: + cdtdata.add_dark(cxi_file, self.background) @@ -152,6 +162,7 @@ class Ptycho_2D_Dataset(CDataset): 'basis' : basis, 'corner' : corner} mask = cdtdata.get_mask(cxi_file) + dark = cdtdata.get_dark(cxi_file) patterns, axes = cdtdata.get_data(cxi_file) translations = cdtdata.get_ptycho_translations(cxi_file) @@ -160,7 +171,7 @@ class Ptycho_2D_Dataset(CDataset): sample_info = sample_info, wavelength=wavelength, detector_geometry=detector_geometry, - mask=mask) + mask=mask, background=dark) def to_cxi(self, cxi_file): diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 1bc3c39..4247759 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -110,6 +110,11 @@ class FancyPtycho(CDIModel): obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=50) + if hasattr(dataset, 'background'): + background = t.sqrt(dataset.background) + else: + background = None + # Finally, initialize the probe and object using this information if probe_size is None: probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice) @@ -145,9 +150,7 @@ class FancyPtycho(CDIModel): probe_support = t.ones_like(probe[0].to(dtype=t.float32)) - - #return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, translation_offsets = translation_offsets) - return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, translation_offsets = translation_offsets, weights=weights, mask=mask, translation_scale=translation_scale, saturation=saturation, probe_support=probe_support) + return cls(wavelength, det_geo, probe_basis, det_slice, probe, obj, min_translation=min_translation, translation_offsets = translation_offsets, weights=weights, mask=mask, background=background, translation_scale=translation_scale, saturation=saturation, probe_support=probe_support) def interaction(self, index, translations): @@ -264,7 +267,7 @@ class FancyPtycho(CDIModel): probe = cmath.torch_to_complex(self.probe.detach().cpu()) probe = probe * self.probe_norm.detach().cpu().numpy() obj = cmath.torch_to_complex(self.obj.detach().cpu()) - background = self.background.detach().cpu().numpy() + background = self.background.detach().cpu().numpy()**2 weights = self.weights.detach().cpu().numpy() return {'basis':basis, 'translation':translations, diff --git a/CDTools/tools/initializers.py b/CDTools/tools/initializers.py index 1810f4e..0acaed9 100644 --- a/CDTools/tools/initializers.py +++ b/CDTools/tools/initializers.py @@ -241,6 +241,10 @@ def SHARP_style_probe(dataset, shape, det_slice): intensities[det_slice] += im.cpu().numpy() intensities /= len(dataset) + # Subtract off a known background if it's stored + if hasattr(dataset, 'background'): + intensities[det_slice] = np.clip(intensities[det_slice] - dataset.background.cpu().numpy(), a_min=0,a_max=None) + probe_fft = cmath.complex_to_torch(np.sqrt(intensities)) probe_guess = cmath.torch_to_complex(inverse_far_field(probe_fft)) diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index 1e3cee4..79c3783 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -36,7 +36,7 @@ for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)): for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100, lr=0.001)): print(i,loss) -for i, loss in enumerate(model.Adam_optimize(50, dataset, batch_size=100, lr=0.0001)): +for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100, lr=0.0001)): print(i,loss) diff --git a/tests/test_datasets.py b/tests/test_datasets.py index aa0d6d8..91d38c1 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -63,6 +63,11 @@ 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['dark'] is not None: + assert t.all(t.eq(t.Tensor(expected['dark']),dataset.background)) + + + def test_CDataset_to_cxi(test_ptycho_cxis, tmp_path): for cxi, expected in test_ptycho_cxis: @@ -95,7 +100,10 @@ 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.background is not None: + assert t.all(t.eq(dataset.background, read_dataset.background)) + def test_CDataset_to(ptycho_cxi_1): @@ -107,6 +115,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.background.device == t.device('cuda:0') @@ -176,6 +185,10 @@ def test_Ptycho_2D_Dataset_from_cxi(test_ptycho_cxis): if expected['mask'] is not None: assert t.all(t.eq(t.tensor(expected['mask']),dataset.mask)) + if expected['dark'] is not None: + assert t.all(t.eq(t.Tensor(expected['dark']),dataset.background)) + + assert t.allclose(t.tensor(expected['data']),dataset.patterns) assert t.allclose(t.tensor(expected['translations']),dataset.translations) @@ -214,6 +227,9 @@ def test_Ptycho_2D_Dataset_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.background is not None: + assert t.all(t.eq(dataset.background, read_dataset.background)) + assert t.allclose(dataset.patterns, read_dataset.patterns) assert t.allclose(dataset.translations, read_dataset.translations) @@ -229,6 +245,7 @@ def test_Ptycho_2D_Dataset_to(ptycho_cxi_1): if t.cuda.is_available(): dataset.to(device='cuda:0') assert dataset.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')