expose rfft rescale utilities (#5)

teamtomo · Jul 11, 2024 · 264bf64 · 264bf64
1 parent ab9cc41
commit 264bf64
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/src/torch_fourier_rescale/__init__.py b/src/torch_fourier_rescale/__init__.py
@@ -1,2 +1,2 @@
-from .fourier_rescale_2d import fourier_rescale_2d
-from .fourier_rescale_3d import fourier_rescale_3d
+from .fourier_rescale_2d import fourier_rescale_2d, fourier_rescale_rfft_2d
+from .fourier_rescale_3d import fourier_rescale_3d, fourier_rescale_rfft_3d
diff --git a/src/torch_fourier_rescale/fourier_rescale_2d.py b/src/torch_fourier_rescale/fourier_rescale_2d.py
@@ -45,7 +45,7 @@ def fourier_rescale_2d(
     dft = torch.fft.fftshift(dft, dim=(-2,))
 
     # Fourier pad/crop
-    dft, new_nyquist = _rescale_rfft_2d(
+    dft, new_nyquist = fourier_rescale_rfft_2d(
         dft=dft,
         image_shape=image.shape[-2:],
         source_spacing=source_spacing,
@@ -65,7 +65,7 @@ def fourier_rescale_2d(
     return rescaled_image, tuple(new_spacing)
 
 
-def _rescale_rfft_2d(
+def fourier_rescale_rfft_2d(
     dft: torch.Tensor,
     image_shape: tuple[int, int],
     source_spacing: tuple[float, float],

diff --git a/src/torch_fourier_rescale/fourier_rescale_3d.py b/src/torch_fourier_rescale/fourier_rescale_3d.py
@@ -43,7 +43,7 @@ def fourier_rescale_3d(
     dft = torch.fft.fftshift(dft, dim=(-3, -2))
 
     # Fourier pad/crop
-    dft, new_nyquist = _rescale_rfft_3d(
+    dft, new_nyquist = fourier_rescale_rfft_3d(
         dft=dft,
         image_shape=image.shape[-3:],
         source_spacing=source_spacing,
@@ -63,7 +63,7 @@ def fourier_rescale_3d(
     return rescaled_image, tuple(new_spacing)
 
 
-def _rescale_rfft_3d(
+def fourier_rescale_rfft_3d(
     dft: torch.Tensor,
     image_shape: tuple[int, int, int],
     source_spacing: tuple[float, float, float],