Adding rudimentary testing of kernels

src/tests/test_dssc_kernels.py

+from Correction import cuda_pipeline
+import numpy as np
+import pycuda.autoinit
+import pycuda.compiler
+import pycuda.driver
+import pycuda.gpuarray
+import pytest
+
+
+# TODO: test "splitter" part
+# TODO: test pulse filter
+
+pixels_x = 512
+pixels_y = 128
+memory_cells = 400
+image_data = np.random.randint(
+    low=0, high=2000, size=(pixels_x, pixels_y, memory_cells), dtype=np.uint16
+)
+offset_map = (
+    np.random.random(size=(pixels_x, pixels_y, memory_cells)).astype(np.float32) * 20
+)
+cell_table = np.arange(memory_cells, dtype=np.uint16)
+pulse_filter = np.arange(memory_cells, dtype=np.uint16)
+corrected = (image_data.astype(np.float32) - offset_map[..., cell_table]).astype(
+    np.float16
+)
+
+# TODO: test non-contiguous memory cells
+# TODO: test graceful handling of cells not covered by correction map
+
+kernel_runner = cuda_pipeline.PyCudaPipeline(
+    pixels_x,
+    pixels_y,
+    memory_cells,
+    pulse_filter,
+    output_buffer_name="dssc_correction_test",
+    output_buffer_size=1,
+    input_data_dtype=np.uint16,
+    output_data_dtype=np.float16,
+)
+# TODO: initialize with map (avoid reallocation of buffer, recompilation of kernel)
+
+
+def test_correction():
+    kernel_runner.load_constants(offset_map)
+    # TODO: test something about the shmem buffer
+    _, corrected_on_gpu = kernel_runner.correct(
+        pycuda.gpuarray.to_gpu(image_data), pycuda.gpuarray.to_gpu(cell_table)
+    )
+    assert np.allclose(corrected_on_gpu, corrected)
+
+
+def test_preview_max():
+    # can it find max intensity frame?
+    # note: in case correction failed, still test this separately
+    kernel_runner.gpu_result.set(corrected)
+    preview_raw, preview_corrected = kernel_runner.compute_preview(
+        raw_data=pycuda.gpuarray.to_gpu(image_data),
+        preview_index=-1,
+        reuse_corrected=True,
+        cell_table=None,
+    )
+    assert np.allclose(
+        preview_raw,
+        image_data[
+            ...,
+            np.argmax(
+                np.sum(image_data.astype(np.float16), axis=(0, 1), dtype=np.float32)
+            ),
+        ].astype(np.float32),
+    )
+    assert np.allclose(
+        preview_corrected,
+        corrected[
+            ..., np.argmax(np.sum(corrected, axis=(0, 1), dtype=np.float32))
+        ].astype(np.float32),
+    )
+
+
+def test_preview_mean():
+    kernel_runner.gpu_result.set(corrected)
+    preview_raw, preview_corrected = kernel_runner.compute_preview(
+        raw_data=pycuda.gpuarray.to_gpu(image_data),
+        preview_index=-2,
+        reuse_corrected=True,
+        cell_table=None,
+    )
+    assert np.allclose(preview_raw, np.mean(image_data, axis=2, dtype=np.float32))
+    assert np.allclose(preview_corrected, np.mean(corrected, axis=2, dtype=np.float32))
+
+
+def test_preview_sum():
+    preview_raw, preview_corrected = kernel_runner.compute_preview(
+        raw_data=pycuda.gpuarray.to_gpu(image_data),
+        preview_index=-3,
+        reuse_corrected=True,
+        cell_table=None,
+    )
+    assert np.allclose(preview_raw, np.sum(image_data, axis=2, dtype=np.float32))
+    assert np.allclose(preview_corrected, np.sum(corrected, axis=2, dtype=np.float32))
+
+
+def test_preview_std():
+    preview_raw, preview_corrected = kernel_runner.compute_preview(
+        raw_data=pycuda.gpuarray.to_gpu(image_data),
+        preview_index=-4,
+        reuse_corrected=True,
+        cell_table=None,
+    )
+    assert np.allclose(preview_raw, np.std(image_data, axis=2, dtype=np.float64))
+    assert np.allclose(preview_corrected, np.std(corrected, axis=2, dtype=np.float64))
+
+
+def test_preview_valid_index():
+    with pytest.raises(ValueError):
+        kernel_runner.compute_preview(
+            raw_data=pycuda.gpuarray.to_gpu(image_data),
+            preview_index=-5,
+            reuse_corrected=True,
+            cell_table=None,
+        )
+    with pytest.raises(ValueError):
+        kernel_runner.compute_preview(
+            raw_data=pycuda.gpuarray.to_gpu(image_data),
+            preview_index=-5,
+            reuse_corrected=True,
+            cell_table=None,
+        )