From d740e0484e0d5cd5a6062948811de4466a9aa3b5 Mon Sep 17 00:00:00 2001
From: Abe Levitan <abraham.levitan@psi.ch>
Date: Tue, 4 Aug 2026 02:24:11 +0200
Subject: [PATCH] Add test coverage for dataset.inspect on GPU and fix the
 relevant bug

---
 src/cdtools/datasets/ptycho_2d_dataset.py |  3 ++-
 tests/test_datasets.py                    | 30 +++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/cdtools/datasets/ptycho_2d_dataset.py b/src/cdtools/datasets/ptycho_2d_dataset.py
index 5a84ca14..9c347e74 100644
--- a/src/cdtools/datasets/ptycho_2d_dataset.py
+++ b/src/cdtools/datasets/ptycho_2d_dataset.py
@@ -250,7 +250,8 @@ def get_images(idx):
         # 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):
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index 4d17f0b3..d57eeaac 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -7,6 +7,7 @@
 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')
+
+
+