Merge branch 'master' of github.mit.edu:Scattering/CDTools

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
     

CDTools/tools/measurements.py

@@ -47,22 +47,22 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
     the wavefields. If a detector slice is given, the returned array
     will only include that slice from the simulated wavefronts.
 
-    The first index is the index of the diffraction pattern to measure,
-    the second index is the set of incoherently adding patterns, and
-    the next two indices index the wavefield. The final index is the complex
+    The first index is the set of incoherently adding patterns, and
+    the second index is the index of the diffraction pattern to measure.
+    The next two indices index the wavefield. The final index is the complex
     index.
     
     Args:
-        wavefields (torch.Tensor) : A JxLxMxNx2 stack of complex wavefields
+        wavefields (torch.Tensor) : An LxJxMxNx2 stack of complex wavefields
         detector_slice (slice) : Optional, a slice or tuple of slices defining a section of the simulation to return
         saturation (float) : Optional, a maximum saturation value to clamp the resulting intensities to
 
     Returns:
-        torch.Tensor : A real MxN array storing the incoherently summed intensities
+        torch.Tensor : A real JXMxN array storing the incoherently summed intensities
     """
     # This syntax just adds an axis to the slice to preserve the J direction
     if detector_slice is None:
-        output = t.sum(cmath.cabssq(wavefields),dim=-3) + epsilon
+        output = t.sum(cmath.cabssq(wavefields),dim=0) + epsilon
     else:
         if wavefields.dim() == 4:
             output =  t.sum(cmath.cabssq(wavefields[(np.s_[:],)+detector_slice]),dim=0) + epsilon

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

tests/tools/test_measurements.py

@@ -31,11 +31,11 @@ def test_intensity():
 def test_incoherent_sum():
     wavefields = t.rand((5,4,10,10,2))
     epsilon=1e-6
-    np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=1) + epsilon
+    np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) + epsilon
     assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon),
                       t.tensor(np_result))
     # Test single field case
-    assert t.allclose(measurements.incoherent_sum(wavefields[0],epsilon=epsilon),
+    assert t.allclose(measurements.incoherent_sum(wavefields[:,0],epsilon=epsilon),
                       t.tensor(np_result[0]))
 
     
@@ -43,7 +43,7 @@ def test_incoherent_sum():
     assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon),
                       t.tensor(np_result[(np.s_[:],)+det_slice]))
     # Test single field case
-    assert t.allclose(measurements.incoherent_sum(wavefields[0],det_slice,epsilon=epsilon),
+    assert t.allclose(measurements.incoherent_sum(wavefields[:,0],det_slice,epsilon=epsilon),
                       t.tensor(np_result[0][det_slice]))