refactor: add docstring and make utils functions more shape-independent

Author: aknierim

src/pyvisgen/dataset/utils.py

@@ -1,24 +1,40 @@
 import numpy as np
 
 
-def calc_truth_fft(sky_dist):
-    truth_fft = np.fft.fftshift(
-        np.fft.fft2(np.fft.fftshift(sky_dist, axes=(2, 3)), axes=(2, 3)),
-        axes=(2, 3),
+def calc_truth_fft(image):
+    """Calculates the Fourier transform (image space -> uv-space)
+    for a single image or a batch of images.
+
+    This is shape independent as long as the last two axes are
+    height and width, i.e. ``(..., H, W)``.
+
+    Parameters
+    ----------
+    image : array_like, shape (..., H, W)
+        (True) sky distribution.
+
+    Returns
+    -------
+    fft_image : array_like, shape (..., H, W)
+        Complex type array of the fft of the input image.
+    """
+    fft_image = np.fft.fftshift(
+        np.fft.fft2(np.fft.fftshift(image, axes=(-2, -1)), axes=(-2, -1)),
+        axes=(-2, -1),
     )
-    return truth_fft
+    return fft_image
 
 
 def convert_amp_phase(data, sky_sim=False):
     if sky_sim:
         amp = np.abs(data)
         phase = np.angle(data)
-        data = np.concatenate((amp, phase), axis=1)
+        data = np.concatenate((amp, phase), axis=-3)
     else:
         test = data[:, 0] + 1j * data[:, 1]
         amp = np.abs(test)
         phase = np.angle(test)
-        data = np.stack((amp, phase), axis=1)
+        data = np.stack((amp, phase), axis=-3)
 
     return data
 
@@ -28,11 +44,11 @@ def convert_real_imag(data, sky_sim=False):
         real = data.real
         imag = data.imag
 
-        data = np.concatenate((real, imag), axis=1)
+        data = np.concatenate((real, imag), axis=-3)
     else:
         real = data[:, 0]
         imag = data[:, 1]
 
-        data = np.stack((real, imag), axis=1)
+        data = np.stack((real, imag), axis=-3)
 
     return data