Optimise the method for distortion correction

httomolibgpu/prep/alignment.py

@@ -254,9 +254,9 @@ def distortion_correction_proj_discorpy(
     cp.ndarray
         3D array. Distortion-corrected image(s).
     """
-    # Check if it's a stack of 2D images, or only a single 2D image
-    if len(data.shape) == 2:
-        data = cp.expand_dims(data, axis=0)
+    # Check if it's a stack of 2D images only
+    if data.ndim != 3:
+        raise ValueError("Input data must be a stack of 2D images only")
 
     # Get info from metadata txt file
     xcenter, ycenter, list_fact = _load_metadata_txt(metadata_path)

tests/test_prep/test_alignment.py

@@ -1,26 +1,28 @@
 import os
+import time
 
 import cupy as cp
 import numpy as np
 import pytest
+from cupy.cuda import nvtx
+from imageio.v2 import imread
+from numpy.testing import assert_allclose
+
+from httomolibgpu import method_registry
 from httomolibgpu.prep.alignment import (
     distortion_correction_proj,
     distortion_correction_proj_discorpy,
 )
-from httomolibgpu import method_registry
-from imageio.v2 import imread
-from numpy.testing import assert_allclose
-
 from tests import MaxMemoryHook
 
 
 @cp.testing.gpu
 @pytest.mark.parametrize(
-    "image, max_value, mean_value",
+    "image, mean_value, mean_sum, sum",
     [
-        ("dot_pattern_03.tif", 255, 200.16733869461675),
-        ("peppers.tif", 228, 95.51871109008789),
-        ("cameraman.tif", 254, 122.2400016784668),
+        ("dot_pattern_03.tif", 200.16733869461675, 20016.733869461674, 214366729300),
+        ("peppers.tif", 95.51871109008789, 9551.871109008789, 2503965700),
+        ("cameraman.tif", 122.2400016784668, 12224.00016784668, 3204448300),
     ],
     ids=["dot_pattern_03", "peppers", "cameraman"],
 )
@@ -34,24 +36,42 @@ def test_correct_distortion(
     ensure_clean_memory,
     image,
     implementation,
-    max_value,
     mean_value,
+    mean_sum,
+    sum,
 ):
     distortion_coeffs_path = os.path.join(
         distortion_correction_path, "distortion-coeffs.txt"
     )
 
     path = os.path.join(distortion_correction_path, image)
-    im_host = imread(path)
-    im = cp.asarray(im_host)
+    im_host = np.asarray(imread(path))
+    im_host = np.expand_dims(im_host, axis=0)
+    # replicate into a stack of images
+    im_stack = cp.asarray(cp.tile(im_host, (100, 1, 1)))
+
+    preview = {
+        "starts": [0, 0],
+        "stops": [im_stack.shape[0], im_stack.shape[1]],
+        "steps": [1, 1],
+    }
+    corrected_data = implementation(im_stack, distortion_coeffs_path, preview).get()
 
-    preview = {"starts": [0, 0], "stops": [im.shape[0], im.shape[1]], "steps": [1, 1]}
-    corrected_data = implementation(im, distortion_coeffs_path, preview).get()
-
     assert_allclose(np.mean(corrected_data), mean_value)
-    assert np.max(corrected_data) == max_value
+    assert_allclose(np.mean(corrected_data, axis=(1, 2)).sum(), mean_sum, rtol=1e-6)
+    assert np.sum(corrected_data) == sum
 
     assert corrected_data.dtype == np.uint8
+    assert corrected_data.ndim == 3
+
+
+def test_correct_distortion_1d_raises(distortion_correction_path, ensure_clean_memory):
+    with pytest.raises(ValueError):
+        preview = {"starts": [0, 0], "stops": [1, 1], "steps": [1, 1]}
+        distortion_coeffs_path = os.path.join(
+            distortion_correction_path, "distortion-coeffs.txt"
+        )
+        distortion_correction_proj_discorpy(cp.ones(1), distortion_coeffs_path, preview)
 
 
 @pytest.mark.parametrize("stack_size", [20, 50, 100])
@@ -60,7 +80,9 @@ def test_correct_distortion(
     [distortion_correction_proj, distortion_correction_proj_discorpy],
     ids=["cupy", "tomopy"],
 )
-def test_distortion_correction_meta(distortion_correction_path, stack_size, implementation):
+def test_distortion_correction_meta(
+    distortion_correction_path, stack_size, implementation
+):
     distortion_coeffs_path = os.path.join(
         distortion_correction_path, "distortion-coeffs.txt"
     )
@@ -70,20 +92,68 @@ def test_distortion_correction_meta(distortion_correction_path, stack_size, impl
     im_host = np.expand_dims(im_host, axis=0)
     # replicate into a stack of images
     im_stack = cp.asarray(np.tile(im_host, (stack_size, 1, 1)))
+
     hook = MaxMemoryHook(im_stack.size * im_stack.itemsize)
-    preview = {"starts": [0, 0], "stops": [im_stack.shape[1], im_stack.shape[2]], "steps": [1, 1]}
-
+    preview = {
+        "starts": [0, 0],
+        "stops": [im_stack.shape[1], im_stack.shape[2]],
+        "steps": [1, 1],
+    }
+
     with hook:
         implementation(im_stack, distortion_coeffs_path, preview)
-    
+
     # make sure estimator function is within range (80% min, 100% max)
     max_mem = hook.max_mem
     actual_slices = im_stack.shape[0]
     estimated_slices, dtype_out, output_dims = implementation.meta.calc_max_slices(
-        0, (im_stack.shape[1], im_stack.shape[2]),
-        im_stack.dtype, max_mem, 
-        metadata_path=distortion_coeffs_path, preview=preview)
+        0,
+        (im_stack.shape[1], im_stack.shape[2]),
+        im_stack.dtype,
+        max_mem,
+        metadata_path=distortion_coeffs_path,
+        preview=preview,
+    )
     assert estimated_slices <= actual_slices
-    assert estimated_slices / actual_slices >= 0.8 
-
-    assert implementation.__name__ in method_registry['httomolibgpu']['prep']['alignment']
+    assert estimated_slices / actual_slices >= 0.8
+
+    assert (
+        implementation.__name__ in method_registry["httomolibgpu"]["prep"]["alignment"]
+    )
+
+
+@cp.testing.gpu
+@pytest.mark.perf
+def test_distortion_correction_performance(
+    distortion_correction_path, ensure_clean_memory
+):
+    dev = cp.cuda.Device()
+    distortion_coeffs_path = os.path.join(
+        distortion_correction_path, "distortion-coeffs.txt"
+    )
+
+    image_path = os.path.join(distortion_correction_path, "dot_pattern_03.tif")
+    im_host = np.asarray(imread(image_path))
+    im_host = np.expand_dims(im_host, axis=0)
+    # replicate into a stack of images
+    im_stack = cp.asarray(np.tile(im_host, (150, 3, 2)))
+
+    preview = {
+        "starts": [0, 0],
+        "stops": [im_stack.shape[0], im_stack.shape[1]],
+        "steps": [1, 1],
+    }
+
+    # cold run first
+    distortion_correction_proj_discorpy(im_stack, distortion_coeffs_path, preview)
+    dev.synchronize()
+
+    start = time.perf_counter_ns()
+    nvtx.RangePush("Core")
+    for _ in range(10):
+        distortion_correction_proj_discorpy(im_stack, distortion_coeffs_path, preview)
+    nvtx.RangePop()
+    dev.synchronize()
+    duration_ms = float(time.perf_counter_ns() - start) * 1e-6 / 10
+
+    assert "performance in ms" == duration_ms