Include ability to read background from a .cxi file and include it in SHARP-style initialization

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):

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,

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))

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)
 
 

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')