Fix RMSF and add return vals

.github/workflows/python-package.yml

@@ -31,9 +31,7 @@ jobs:
       shell: bash -l {0}
       run: |
         conda install -c conda-forge pymultinest
-        pip install nestle
-        conda install pip pytest
-        pip install .[QUfitting,parallel] --upgrade
+        pip install -U .[QUfitting,parallel] nestle pytest numpy
     - name: Test with pytest
       shell: bash -l {0}
       run: |

RMtools_1D/calculate_RMSF.py

@@ -174,7 +174,7 @@ def determine_RMSF_parameters(
         -1 * phi_max / 2, phi_max / 2 + 1e-6, dphi
     )  # division by two accounts for how RMSF is always twice as wide as FDF.
 
-    RMSFcube, phi2Arr, fwhmRMSFArr, statArr = get_rmsf_planes(
+    rmsf_results = get_rmsf_planes(
         lambda2_array,
         phi_array,
         weightArr=weights_array,
@@ -190,7 +190,11 @@ def determine_RMSF_parameters(
             3.8 / (l2_max - l2_min)
         )
     )
-    print("Measured FWHM:                       {:.4g} rad m^-2".format(fwhmRMSFArr))
+    print(
+        "Measured FWHM:                       {:.4g} rad m^-2".format(
+            rmsf_results.fwhmRMSFArr
+        )
+    )
     print("Theoretical largest FD scale probed: {:.4g} rad m^-2".format(np.pi / l2_min))
     print(
         "Theoretical maximum FD*:             {:.4g} rad m^-2".format(
@@ -208,15 +212,21 @@ def determine_RMSF_parameters(
     # finds the highest amplitude one, and calls that the first sidelobe.
     try:
         x = np.diff(
-            np.sign(np.diff(np.abs(RMSFcube[RMSFcube.size // 2 :])))
+            np.sign(
+                np.diff(
+                    np.abs(rmsf_results.RMSFcube[rmsf_results.RMSFcube.size // 2 :])
+                )
+            )
         )  # -2=local max, +2=local min
         y = (
             1 + np.where(x == -2)[0]
         )  # indices of peaks, +1 is because of offset from double differencing
-        peaks = np.abs(RMSFcube[RMSFcube.size // 2 :])[y]
+        peaks = np.abs(rmsf_results.RMSFcube[rmsf_results.RMSFcube.size // 2 :])[y]
         print(
             "First sidelobe FD and amplitude:     {:.4g} rad m^-2".format(
-                phi2Arr[phi2Arr.size // 2 :][y[np.argmax(peaks)]]
+                rmsf_results.phi2Arr[rmsf_results.phi2Arr.size // 2 :][
+                    y[np.argmax(peaks)]
+                ]
             )
         )
         print(
@@ -235,9 +245,15 @@ def determine_RMSF_parameters(
         plt.title("Simulated RMSF")
     else:
         plt.title(plotname)
-    plt.plot(phi2Arr, np.real(RMSFcube), "b-", label="Stokes Q")
-    plt.plot(phi2Arr, np.imag(RMSFcube), "r--", label="Stokes U")
-    plt.plot(phi2Arr, np.abs(RMSFcube), "k-", label="Amplitude")
+    plt.plot(
+        rmsf_results.phi2Arr, np.real(rmsf_results.RMSFcube), "b-", label="Stokes Q"
+    )
+    plt.plot(
+        rmsf_results.phi2Arr, np.imag(rmsf_results.RMSFcube), "r--", label="Stokes U"
+    )
+    plt.plot(
+        rmsf_results.phi2Arr, np.abs(rmsf_results.RMSFcube), "k-", label="Amplitude"
+    )
     plt.legend()
     plt.xlabel(r"Faraday depth (rad m$^{-2}$)")
     plt.ylabel("RMSF (unitless)")
@@ -259,7 +275,7 @@ def determine_RMSF_parameters(
             + "# of channels:                                {:.4g}\n"
         ).format(
             3.8 / (l2_max - l2_min),
-            fwhmRMSFArr,
+            rmsf_results.fwhmRMSFArr,
             np.pi / l2_min,
             np.sqrt(3.0) / dl2,
             np.min(freq_array) / 1e9,
@@ -282,7 +298,10 @@ def determine_RMSF_parameters(
                 + "First sidelobe FD and amplitude:     {:.4g} rad m^-2\n"
                 + "                                     {:.4g} % of peak"
             ).format(
-                phi2Arr[phi2Arr.size // 2 :][y[np.argmax(peaks)]], np.max(peaks) * 100
+                rmsf_results.phi2Arr[rmsf_results.phi2Arr.size // 2 :][
+                    y[np.argmax(peaks)]
+                ],
+                np.max(peaks) * 100,
             ),
             family="monospace",
             horizontalalignment="left",
@@ -298,7 +317,7 @@ def determine_RMSF_parameters(
     #    ax.text(0.,0.22,'First sidelobe FD and amplitude:       {:.4g} rad m^-2'.format(phi2Arr[phi2Arr.size//2:][y[np.argmax(peaks)]]))
     #    ax.text(0.,0.1,'                                                           {:.4g} % of peak'.format(np.max(peaks)*100))
 
-    if plotfile != None:
+    if plotfile is not None:
         plt.savefig(plotfile, bbox_inches="tight")
     else:
         plt.show()

RMutils/util_RM.py

@@ -63,6 +63,7 @@
 import gc
 import math as m
 import sys
+from typing import NamedTuple
 
 import finufft
 import numpy as np
@@ -86,6 +87,15 @@
 
 
 # -----------------------------------------------------------------------------#
+class RMsynthResults(NamedTuple):
+    """Results of the RM-synthesis calculation"""
+
+    FDFcube: np.ndarray
+    """The Faraday dispersion function cube"""
+    lam0Sq_m2: float
+    """The reference lambda^2 value"""
+
+
 def do_rmsynth_planes(
     dataQ,
     dataU,
@@ -96,7 +106,7 @@ def do_rmsynth_planes(
     nBits=32,
     eps=1e-6,
     log=print,
-):
+) -> RMsynthResults:
     """Perform RM-synthesis on Stokes Q and U cubes (1,2 or 3D). This version
     of the routine loops through spectral planes and is faster than the pixel-
     by-pixel code. This version also correctly deals with isolated clumps of
@@ -214,10 +224,25 @@ def do_rmsynth_planes(
     # Remove redundant dimensions in the FDF array
     FDFcube = np.squeeze(FDFcube)
 
-    return FDFcube, lam0Sq_m2
+    return RMsynthResults(FDFcube=FDFcube, lam0Sq_m2=lam0Sq_m2)
 
 
 # -----------------------------------------------------------------------------#
+class RMSFResults(NamedTuple):
+    """Results of the RMSF calculation"""
+
+    RMSFcube: np.ndarray
+    """The RMSF cube"""
+    phi2Arr: np.ndarray
+    """The (double length) Faraday depth array"""
+    fwhmRMSFArr: np.ndarray
+    """The FWHM of the RMSF main lobe"""
+    statArr: np.ndarray
+    """The status of the RMSF fit"""
+    lam0Sq_m2: float
+    """The reference lambda^2 value"""
+
+
 def get_rmsf_planes(
     lambdaSqArr_m2,
     phiArr_radm2,
@@ -231,7 +256,7 @@ def get_rmsf_planes(
     eps=1e-6,
     verbose=False,
     log=print,
-):
+) -> RMSFResults:
     """Calculate the Rotation Measure Spread Function from inputs. This version
     returns a cube (1, 2 or 3D) of RMSF spectra based on the shape of a
     boolean mask array, where flagged data are True and unflagged data False.
@@ -423,10 +448,29 @@ def get_rmsf_planes(
         fwhmRMSFArr = np.reshape(fwhmRMSFArr, (old_data_shape[1], old_data_shape[2]))
         statArr = np.reshape(statArr, (old_data_shape[1], old_data_shape[2]))
 
-    return RMSFcube, phi2Arr, fwhmRMSFArr, statArr, lam0Sq_m2
+    return RMSFResults(
+        RMSFcube=RMSFcube,
+        phi2Arr=phi2Arr,
+        fwhmRMSFArr=fwhmRMSFArr,
+        statArr=statArr,
+        lam0Sq_m2=lam0Sq_m2,
+    )
 
 
 # -----------------------------------------------------------------------------#
+class RMCleanResults(NamedTuple):
+    """Results of the RM-CLEAN calculation"""
+
+    cleanFDF: np.ndarray
+    """The cleaned Faraday dispersion function cube"""
+    ccArr: np.ndarray
+    """The clean components cube"""
+    iterCountArr: np.ndarray
+    """The number of iterations for each pixel"""
+    residFDF: np.ndarray
+    """The residual Faraday dispersion function cube"""
+
+
 def do_rmclean_hogbom(
     dirtyFDF,
     phiArr_radm2,
@@ -444,7 +488,7 @@ def do_rmclean_hogbom(
     chunksize=None,
     log=print,
     window=0,
-):
+) -> RMCleanResults:
     """Perform Hogbom CLEAN on a cube of complex Faraday dispersion functions
     given a cube of rotation measure spread functions.
 
@@ -593,10 +637,19 @@ def do_rmclean_hogbom(
     iterCountArr = np.squeeze(iterCountArr)
     residFDF = np.squeeze(residFDF)
 
-    return cleanFDF, ccArr, iterCountArr, residFDF
+    return RMCleanResults(cleanFDF, ccArr, iterCountArr, residFDF)
 
 
 # -----------------------------------------------------------------------------#
+class CleanLoopResults(NamedTuple):
+    """Results of the RM-CLEAN loop"""
+
+    cleanFDF: np.ndarray
+    """The cleaned Faraday dispersion function cube"""
+    residFDF: np.ndarray
+    """The residual Faraday dispersion function cube"""
+    ccArr: np.ndarray
+    """The clean components cube"""
 
 
 class RMcleaner:
@@ -630,10 +683,10 @@ def __init__(
         self.nbits = nbits
         self.window = window
 
-    def cleanloop(self, args):
+    def cleanloop(self, args) -> CleanLoopResults:
         return self._cleanloop(*args)
 
-    def _cleanloop(self, yi, xi, dirtyFDF):
+    def _cleanloop(self, yi, xi, dirtyFDF) -> CleanLoopResults:
         dirtyFDF = dirtyFDF[:, yi, xi]
         # Initialise arrays to hold the residual FDF, clean components, clean FDF
         residFDF = dirtyFDF.copy()
@@ -724,7 +777,7 @@ def _cleanloop(self, yi, xi, dirtyFDF):
         residFDF = np.squeeze(residFDF)
         ccArr = np.squeeze(ccArr)
 
-        return cleanFDF, residFDF, ccArr
+        return CleanLoopResults(cleanFDF=cleanFDF, residFDF=residFDF, ccArr=ccArr)
 
 
 # -----------------------------------------------------------------------------#

setup.py

@@ -17,6 +17,7 @@
 
 REQUIRED = [
     "numpy<2",
+    "numpy>1.22;python_version=='3.8'",
     "scipy",
     "matplotlib>=3.4.0",
     "astropy",