Add a tool for loading dark images from cxi files

CDTools/tools/data.py

@@ -13,6 +13,7 @@ __all__ = ['get_entry_info',
            'get_wavelength',
            'get_detector_geometry',
            'get_mask',
+           'get_dark',
            'get_data',
            'get_ptycho_translations',
            'create_cxi',
@@ -21,6 +22,7 @@ __all__ = ['get_entry_info',
            'add_source',
            'add_detector',
            'add_mask',
+           'add_dark',
            'add_data',
            'add_ptycho_translations']
 
@@ -284,6 +286,35 @@ def get_mask(cxi_file):
         return None
 
 
+def get_dark(cxi_file):
+    """Returns an array with a dark image to use for initialization of a background model
+
+    This looks for a set of dark images at
+    entry_1/instrument_1/detector_1/data_dark. If the darks exist, it will
+    return the mean of the array along all axes but the last two. That is,
+    if the dark image is a single image, it will return that image. If it
+    is a stack of images, it will return the mean along the stack axis.
+
+    If the darks do not exist, it will return None
+
+    Args:
+        cxi_file (h5py.File) : a file object to be read
+    
+    Returns:
+        np.array : An array storing the dark image
+    """
+    i1 = cxi_file['entry_1/instrument_1']
+    if 'detector_1/data_dark' in i1:
+        darks = np.array(i1['detector_1/data_dark'])
+        dims = tuple(range(len(darks.shape) - 2))
+        darks = np.nanmean(darks,axis=dims)
+    else:
+        darks = None
+
+    return darks
+
+
+    
 def get_data(cxi_file, cut_zeroes = True):
     """Returns an array with the full stack of detector data defined in the cxi file object
 
@@ -508,6 +539,28 @@ def add_mask(cxi_file, mask):
     d1.create_dataset('mask',data=mask_to_save)
 
 
+def add_dark(cxi_file, dark):
+    """Adds the specified dark image to a cxi file
+
+    It places the dark image data into the data_dark dataset under 
+    entry_1/instrument_1/detector_1.
+    
+    Args:
+        cxi_file (h5py.File) : The file to add the mask to    
+        dark (array_like) : The dark image(s) 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(dark, t.Tensor):
+        dark = dark.detach().cpu().numpy()
+
+    d1.create_dataset('data_dark',data=dark)
+    
+    
 def add_data(cxi_file, data, axes=None):
     """Adds the specified data to the cxi file
     

tests/conftest.py

@@ -103,10 +103,15 @@ def ptycho_cxi_1():
 
     # Remember the format for the CXI file differs from the format used
     # internally
-    mask = np.zeros((100,256,256)).astype(np.uint32)
-    expected['mask'] = np.ones((100,256,256)).astype(np.uint8)
+    mask = np.zeros((256,256)).astype(np.uint32)
+    expected['mask'] = np.ones((256,256)).astype(np.uint8)
     d1f.create_dataset('mask',data=mask)
 
+    # Create an initial background
+    dark = np.ones((256,256)) * 0.01
+    expected['dark'] = dark
+    d1f.create_dataset('data_dark', data=dark)
+
     data1f = e1f.create_group('data_1')
 
     data = np.random.rand(100,256,256).astype(np.float32)
@@ -187,6 +192,12 @@ def ptycho_cxi_2():
     # internally
     expected['mask'] = None
 
+    # Test with a set of dark images
+    dark = np.ones((10,256,256)) * 0.01
+    expected['dark'] = np.nanmean(dark,axis=0)
+    d1f.create_dataset('data_dark', data=dark)
+
+
     data1f = e1f.create_group('data_1')
 
     data = np.random.rand(100,256,256).astype(np.float32)
@@ -257,10 +268,11 @@ def ptycho_cxi_3():
 
     # Remember the format for the CXI file differs from the format used
     # internally
-    mask = np.ones((100,256,256)).astype(np.uint32) * 0x00001000
-    expected['mask'] = np.ones((100,256,256)).astype(np.uint8)
+    mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
+    expected['mask'] = np.ones((256,256)).astype(np.uint8)
     d1f.create_dataset('mask',data=mask)
-
+    expected['dark'] = None
+    
     data1f = e1f.create_group('data_1')
 
     data = np.random.rand(100,256,256).astype(np.float32)

tests/tools/test_data.py

@@ -64,7 +64,16 @@ def test_get_mask(test_ptycho_cxis):
             continue
         assert np.all(data.get_mask(cxi) == expected['mask'])
 
-    
+
+def test_get_dark(test_ptycho_cxis):
+    for cxi, expected in test_ptycho_cxis:
+        dark = data.get_dark(cxi)
+        if dark is None:
+            assert expected['dark'] is None
+        else:
+            assert np.allclose(dark, expected['dark'])
+            
+        
 def test_get_data(test_ptycho_cxis):
     for cxi, expected in test_ptycho_cxis:
         patterns, axes = data.get_data(cxi)
@@ -188,7 +197,21 @@ def test_add_mask(tmp_path):
         read_mask = data.get_mask(f)
 
     assert np.all(mask == read_mask)
+
     
+def test_add_dark(tmp_path):
+    dark = np.random.rand(350,620)
+
+    with data.create_cxi(tmp_path / 'test_add_dark.cxi') as f:
+        data.add_dark(f, dark)
+
+    with h5py.File(tmp_path / 'test_add_dark.cxi') as f:
+        read_dark = data.get_dark(f)
+
+    print(dark.shape)
+    assert np.allclose(dark, read_dark)
+
+
 
 def test_add_data(tmp_path):
     # First test from numpy, with axes