Merge pull request #93 from cdtools-developers/bugfix/dataset_inspect_gpu

Add test coverage for dataset.inspect on GPU and fix the relevant bug

src/cdtools/datasets/ptycho_2d_dataset.py

@@ -250,7 +250,8 @@ class Ptycho2DDataset(CDataset):
         # once, but it avoids creating another self.patterns-sized array
         # as an intermediate step. This can be super important because
         # self.patterns can be more than half the available memory
-        nanomap_values = np.ones(self.translations.shape[0])
+        nanomap_values = t.ones(self.translations.shape[0],
+                                device=self.patterns.device)
 
         chunk_size = 10
         for i in range(0, self.translations.shape[0], chunk_size):

tests/test_datasets.py

@@ -7,6 +7,7 @@ import h5py
 import numpy as np
 import pytest
 import torch as t
+import matplotlib.pyplot as plt
 
 from cdtools.datasets import CDataset, Ptycho2DDataset
 from cdtools.tools import data as cdtdata
@@ -510,3 +511,32 @@ def test_Ptycho2DDataset_crop_translations(ptycho_cxi_1):
         assert t.allclose(copied_dataset.patterns, dataset.patterns[10:-10, :])
 
         assert t.allclose(copied_dataset.translations, dataset.translations[10:-10, :])
+
+# This isn't actually slow, but it will fail by default if there is no
+# gpu on the machine
+@pytest.mark.slow
+def test_Ptycho2DDataset_inspect(ptycho_cxi_1, reconstruction_device, show_plot):
+    cxi, expected = ptycho_cxi_1
+    dataset = Ptycho2DDataset.from_cxi(cxi)
+
+    # Test for failure in several cases
+
+    # First, no additional plots
+    dataset.inspect(plot_mean_pattern=False, plot_mask=False)
+    plt.close('all')
+    
+    # Then, with additional plots
+    dataset.inspect(plot_mean_pattern=True, plot_mask=True)
+    plt.close('all')
+
+    # Finally, if data is on a special device like the GPU
+    dataset.to(device=reconstruction_device)
+    dataset.get_as(device=reconstruction_device)    
+    dataset.inspect(plot_mean_pattern=True, plot_mask=True)
+
+    if show_plot:
+        plt.show()
+    plt.close('all')
+
+
+