Add a tool for loading dark images from cxi files

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