adding docstring and linting

Author: dkazanc

httomolibgpu/prep/phase.py

@@ -74,6 +74,8 @@ def paganin_filter(
         Beam energy in keV.
     ratio_delta_beta : float
         The ratio of delta/beta, where delta is the phase shift and real part of the complex material refractive index and beta is the absorption.
+    calc_peak_gpu_mem: bool
+        Parameter to support memory estimation in HTTomo. Irrelevant to the method itself and can be ignored by user.
 
     Returns
     -------
@@ -125,24 +127,33 @@ def paganin_filter(
     indy = _reciprocal_coord(pixel_size, dx)
 
     if mem_stack:
-        mem_stack.malloc(indx.size * indx.dtype.itemsize) # cp.asarray(indx)
-        mem_stack.malloc(indx.size * indx.dtype.itemsize) # cp.square
-        mem_stack.free(indx.size * indx.dtype.itemsize) # cp.asarray(indx)
+        mem_stack.malloc(indx.size * indx.dtype.itemsize)  # cp.asarray(indx)
+        mem_stack.malloc(indx.size * indx.dtype.itemsize)  # cp.square
+        mem_stack.free(indx.size * indx.dtype.itemsize)  # cp.asarray(indx)
         mem_stack.malloc(indy.size * indy.dtype.itemsize)  # cp.asarray(indy)
-        mem_stack.malloc(indy.size * indy.dtype.itemsize) # cp.square
+        mem_stack.malloc(indy.size * indy.dtype.itemsize)  # cp.square
         mem_stack.free(indy.size * indy.dtype.itemsize)  # cp.asarray(indy)
 
-        mem_stack.malloc(indx.size * indy.size * indx.dtype.itemsize) # cp.add.outer
-        mem_stack.free(indx.size * indx.dtype.itemsize) # cp.square
-        mem_stack.free(indy.size * indy.dtype.itemsize) # cp.square
-        mem_stack.malloc(indx.size * indy.size * indx.dtype.itemsize) # phase_filter
-        mem_stack.free(indx.size * indy.size * indx.dtype.itemsize) # cp.add.outer
-        mem_stack.free(indx.size * indy.size * indx.dtype.itemsize) # phase_filter
+        mem_stack.malloc(indx.size * indy.size * indx.dtype.itemsize)  # cp.add.outer
+        mem_stack.free(indx.size * indx.dtype.itemsize)  # cp.square
+        mem_stack.free(indy.size * indy.dtype.itemsize)  # cp.square
+        mem_stack.malloc(indx.size * indy.size * indx.dtype.itemsize)  # phase_filter
+        mem_stack.free(indx.size * indy.size * indx.dtype.itemsize)  # cp.add.outer
+        mem_stack.free(indx.size * indy.size * indx.dtype.itemsize)  # phase_filter
 
     else:
         # Build Lorentzian-type filter
         phase_filter = fftshift(
-            1.0 / (1.0 + alpha * (cp.add.outer(cp.square(cp.asarray(indx)), cp.square(cp.asarray(indy)))))
+            1.0
+            / (
+                1.0
+                + alpha
+                * (
+                    cp.add.outer(
+                        cp.square(cp.asarray(indx)), cp.square(cp.asarray(indy))
+                    )
+                )
+            )
         )
 
         phase_filter = phase_filter / phase_filter.max()  # normalisation
@@ -152,15 +163,17 @@ def paganin_filter(
         del phase_filter
 
     # Apply filter and take inverse FFT
-    ifft_input = fft_tomo if not mem_stack else cp.empty(padded_tomo, dtype=cp.complex64)
+    ifft_input = (
+        fft_tomo if not mem_stack else cp.empty(padded_tomo, dtype=cp.complex64)
+    )
     ifft_plan = get_fft_plan(ifft_input, axes=(-2, -1))
     if mem_stack:
         mem_stack.malloc(ifft_plan.work_area.mem.size)
         mem_stack.free(ifft_plan.work_area.mem.size)
     else:
         with ifft_plan:
             ifft_filtered_tomo = ifft2(fft_tomo, axes=(-2, -1), overwrite_x=True).real
-        del fft_tomo    
+        del fft_tomo
     del ifft_plan
     del ifft_input
 
@@ -172,9 +185,13 @@ def paganin_filter(
     )
 
     if mem_stack:
-        mem_stack.malloc(np.prod(tomo) * np.float32().itemsize) # astype(cp.float32)
-        mem_stack.free(np.prod(padded_tomo) * np.complex64().itemsize) # ifft_filtered_tomo
-        mem_stack.malloc(np.prod(tomo) * np.float32().itemsize) # return _log_kernel(tomo)
+        mem_stack.malloc(np.prod(tomo) * np.float32().itemsize)  # astype(cp.float32)
+        mem_stack.free(
+            np.prod(padded_tomo) * np.complex64().itemsize
+        )  # ifft_filtered_tomo
+        mem_stack.malloc(
+            np.prod(tomo) * np.float32().itemsize
+        )  # return _log_kernel(tomo)
         return mem_stack.highwater
 
     # crop the padded filtered data:
@@ -232,8 +249,7 @@ def _calculate_pad_size(datashape: tuple) -> list:
 
 
 def _pad_projections_to_second_power(
-    tomo: cp.ndarray,
-    mem_stack: Optional[_DeviceMemStack]
+    tomo: cp.ndarray, mem_stack: Optional[_DeviceMemStack]
 ) -> Tuple[cp.ndarray, Tuple[int, int]]:
     """
     Performs padding of each projection to the next power of 2.
@@ -255,7 +271,9 @@ def _pad_projections_to_second_power(
     pad_list = _calculate_pad_size(full_shape_tomo)
 
     if mem_stack:
-        padded_tomo = [sh + pad[0] + pad[1] for sh, pad in zip(full_shape_tomo, pad_list)]
+        padded_tomo = [
+            sh + pad[0] + pad[1] for sh, pad in zip(full_shape_tomo, pad_list)
+        ]
         mem_stack.malloc(np.prod(padded_tomo) * np.float32().itemsize)
     else:
         padded_tomo = cp.pad(tomo, tuple(pad_list), "edge")
@@ -317,11 +335,20 @@ def paganin_filter_savu_legacy(
         Beam energy in keV.
     ratio_delta_beta : float
         The ratio of delta/beta, where delta is the phase shift and real part of the complex material refractive index and beta is the absorption.
+    calc_peak_gpu_mem: bool
+        Parameter to support memory estimation in HTTomo. Irrelevant to the method itself and can be ignored by user.
 
     Returns
     -------
     cp.ndarray
         The 3D array of Paganin phase-filtered projection images.
     """
 
-    return paganin_filter(tomo, pixel_size, distance, energy, ratio_delta_beta / 4, calc_peak_gpu_mem=calc_peak_gpu_mem)
+    return paganin_filter(
+        tomo,
+        pixel_size,
+        distance,
+        energy,
+        ratio_delta_beta / 4,
+        calc_peak_gpu_mem=calc_peak_gpu_mem,
+    )

httomolibgpu/recon/_phase_cross_correlation.py

@@ -36,9 +36,8 @@
 import cupyx.scipy.ndimage as ndi
 import numpy as np
 
-def _upsampled_dft(
-    data, upsampled_region_size, upsample_factor=1, axis_offsets=None
-):
+
+def _upsampled_dft(data, upsampled_region_size, upsample_factor=1, axis_offsets=None):
     """
     Upsampled DFT by matrix multiplication.
 
@@ -148,9 +147,7 @@ def _compute_error(cross_correlation_max, src_amp, target_amp):
         )
 
     with np.errstate(invalid="ignore"):
-        error = 1.0 - cross_correlation_max * cross_correlation_max.conj() / (
-            amp
-        )
+        error = 1.0 - cross_correlation_max * cross_correlation_max.conj() / (amp)
 
     return cp.sqrt(cp.abs(error))
 
@@ -192,9 +189,7 @@ def _disambiguate_shift(reference_image, moving_image, shift):
     negative_shift = [shift_i - s for shift_i, s in zip(positive_shift, shape)]
     subpixel = any(s % 1 != 0 for s in shift)
     interp_order = 3 if subpixel else 0
-    shifted = ndi.shift(
-        moving_image, shift, mode="grid-wrap", order=interp_order
-    )
+    shifted = ndi.shift(moving_image, shift, mode="grid-wrap", order=interp_order)
     indices = tuple(round(s) for s in positive_shift)
     splits_per_dim = [(slice(0, i), slice(i, None)) for i in indices]
     max_corr = -1.0
@@ -217,9 +212,7 @@ def _disambiguate_shift(reference_image, moving_image, shift):
         )
         return shift
     real_shift_acc = []
-    for sl, pos_shift, neg_shift in zip(
-        max_slice, positive_shift, negative_shift
-    ):
+    for sl, pos_shift, neg_shift in zip(max_slice, positive_shift, negative_shift):
         real_shift_acc.append(pos_shift if sl.stop is None else neg_shift)
     if not subpixel:
         real_shift = tuple(map(int, real_shift_acc))
@@ -359,16 +352,12 @@ def phase_cross_correlation(
         # Initial shift estimate in upsampled grid
         # shift = cp.around(shift * upsample_factor) / upsample_factor
         upsample_factor = float(upsample_factor)
-        shift = tuple(
-            round(s * upsample_factor) / upsample_factor for s in shift
-        )
+        shift = tuple(round(s * upsample_factor) / upsample_factor for s in shift)
         upsampled_region_size = math.ceil(upsample_factor * 1.5)
         # Center of output array at dftshift + 1
         dftshift = float(upsampled_region_size // 2)
         # Matrix multiply DFT around the current shift estimate
-        sample_region_offset = tuple(
-            dftshift - s * upsample_factor for s in shift
-        )
+        sample_region_offset = tuple(dftshift - s * upsample_factor for s in shift)
         cross_correlation = _upsampled_dft(
             image_product.conj(),
             upsampled_region_size,
@@ -394,9 +383,7 @@ def phase_cross_correlation(
 
     # If its only one row or column the shift along that dimension has no
     # effect. We set to zero.
-    shift = tuple(
-        s if axis_size != 1 else 0 for s, axis_size in zip(shift, shape)
-    )
+    shift = tuple(s if axis_size != 1 else 0 for s, axis_size in zip(shift, shape))
 
     if disambiguate:
         if space.lower() != "real":
@@ -406,10 +393,7 @@ def phase_cross_correlation(
 
     # Redirect user to masked_phase_cross_correlation if NaNs are observed
     if cp.isnan(CCmax) or cp.isnan(src_amp) or cp.isnan(target_amp):
-        raise ValueError(
-            "NaN values found, please remove NaNs from your "
-            "input data"
-        )
+        raise ValueError("NaN values found, please remove NaNs from your " "input data")
 
     return (
         shift,

tests/test_prep/test_phase.py

@@ -84,6 +84,7 @@ def test_paganin_filter_performance(ensure_clean_memory):
 
     assert "performance in ms" == duration_ms
 
+
 @pytest.mark.parametrize("slices", [3, 7, 32, 61, 109, 120, 150])
 @pytest.mark.parametrize("dim_x", [128, 140])
 def test_paganin_filter_calc_mem(slices, dim_x, ensure_clean_memory):
@@ -95,27 +96,21 @@ def test_paganin_filter_calc_mem(slices, dim_x, ensure_clean_memory):
     actual_mem_peak = hook.max_mem
 
     try:
-        estimated_mem_peak = paganin_filter(
-            data.shape, calc_peak_gpu_mem=True
-        )
+        estimated_mem_peak = paganin_filter(data.shape, calc_peak_gpu_mem=True)
     except cp.cuda.memory.OutOfMemoryError:
         pytest.skip("Not enough GPU memory to estimate memory peak")
 
     assert actual_mem_peak * 0.99 <= estimated_mem_peak
     assert estimated_mem_peak <= actual_mem_peak * 1.01
 
 
-@pytest.mark.parametrize(
-    "slices", [38, 177, 268, 320, 490, 607, 803, 859, 902, 951]
-)
+@pytest.mark.parametrize("slices", [38, 177, 268, 320, 490, 607, 803, 859, 902, 951])
 @pytest.mark.parametrize("dims", [(900, 1280), (1801, 1540), (1801, 2560)])
 def test_paganin_filter_calc_mem_big(slices, dims, ensure_clean_memory):
     dim_y, dim_x = dims
     data_shape = (slices, dim_x, dim_y)
     try:
-        estimated_mem_peak = paganin_filter(
-            data_shape, calc_peak_gpu_mem=True
-        )
+        estimated_mem_peak = paganin_filter(data_shape, calc_peak_gpu_mem=True)
     except cp.cuda.memory.OutOfMemoryError:
         pytest.skip("Not enough GPU memory to estimate memory peak")
     av_mem = cp.cuda.Device().mem_info[0]