savu paganin work

Author: dkazanc

httomolibgpu/prep/phase.py

@@ -145,6 +145,16 @@ def _calculate_alpha(energy, distance_micron, ratio_delta_beta):
     ) * ratio_delta_beta
 
 
+# the scaling is different here and doesn't follow the original formula
+def _calculate_alpha_savu(energy, distance_micron, ratio_delta_beta):
+    return (
+        ((1240.0 / energy) * 10.0 ** (-9))
+        * distance_micron
+        * ratio_delta_beta
+        * math.pi
+    )
+
+
 def _shift_bit_length(x: int) -> int:
     return 1 << (x - 1).bit_length()
 
@@ -264,10 +274,64 @@ def paganin_filter_savu_legacy(
     cp.ndarray
         The 3D array of Paganin phase-filtered projection images.
     """
+    # Check the input data is valid
+    if tomo.ndim != 3:
+        raise ValueError(
+            f"Invalid number of dimensions in data: {tomo.ndim},"
+            " please provide a stack of 2D projections."
+        )
+
+    dz_orig, dy_orig, dx_orig = tomo.shape
 
-    # the scaling is different here and doesn't follow the original formula
-    # ((1240.0 / energy) * 10.0 ** (-9)) * distance_micron * ratio_delta_beta * math.pi
+    # Perform padding to the power of 2 as FFT is O(n*log(n)) complexity
+    # TODO: adding other options of padding?
+    padded_tomo, pad_tup = _pad_projections_to_second_power(tomo)
 
-    return paganin_filter(
-        tomo, pixel_size, distance * 1e-03, energy, ratio_delta_beta / (4 * 2 * np.pi)
+    dz, dy, dx = padded_tomo.shape
+
+    # 3D FFT of tomo data
+    padded_tomo = cp.asarray(padded_tomo, dtype=cp.complex64)
+    fft_tomo = fft2(padded_tomo, axes=(-2, -1), overwrite_x=True)
+
+    # calculate alpha constant as Savu does
+    micron = 10 ** (-6)
+    KeV = 1000.0
+    alpha = _calculate_alpha_savu(energy * KeV, distance * micron, ratio_delta_beta)
+
+    # Compute the reciprocal grid
+    # NOTE: pixel_size is ALREADY in microns, why it rescaled to micron again in Savu?
+    indx = _reciprocal_coord(pixel_size * micron, dy)
+    indy = _reciprocal_coord(pixel_size * micron, dx)
+
+    # Build Lorentzian-type filter
+    phase_filter = fftshift(
+        1.0 / (1.0 + alpha * (cp.add.outer(cp.square(indx), cp.square(indy))))
     )
+
+    phase_filter = phase_filter / phase_filter.max()  # normalisation
+
+    # Filter projections
+    fft_tomo *= phase_filter
+
+    # Apply filter and take inverse FFT
+    ifft_filtered_tomo = ifft2(fft_tomo, axes=(-2, -1), overwrite_x=True).real
+
+    # slicing indices for cropping
+    slc_indices = (
+        slice(pad_tup[0][0], pad_tup[0][0] + dz_orig, 1),
+        slice(pad_tup[1][0], pad_tup[1][0] + dy_orig, 1),
+        slice(pad_tup[2][0], pad_tup[2][0] + dx_orig, 1),
+    )
+
+    # crop the padded filtered data:
+    tomo = ifft_filtered_tomo[slc_indices].astype(cp.float32)
+
+    # taking the negative log
+    _log_kernel = cp.ElementwiseKernel(
+        "C tomo",
+        "C out",
+        "out = -log(tomo)",
+        name="log_kernel",
+    )
+
+    return _log_kernel(tomo)