Add analysis submodule

docs/changes/8.feature.rst

@@ -0,0 +1,5 @@
+- Add analysis submodule
+  - Add ``radiotools.analysis.get_source_rms`` function
+  - Add ``radiotools.analysis.dynamic_range`` function that calculates the dynamic range as a function of the maximum pixel value of an image and the RMS of the source
+
+- Add ``rms`` function to ``radiotools.utils``

radiotools/analysis/__init__.py

@@ -0,0 +1,3 @@
+from .analysis import dynamic_range, get_source_rms
+
+__all__ = ["dynamic_range", "get_source_rms"]

radiotools/analysis/analysis.py

@@ -0,0 +1,74 @@
+import torch
+from astropy.io import fits
+from numpy.typing import ArrayLike
+
+from radiotools.utils import rms
+
+
+def get_source_rms(image: ArrayLike, center: tuple[int, int], *, offset=75):
+    """Get the root mean square (RMS) of a given source.
+
+    Parameters
+    ----------
+    image : array_like, shape (N, M)
+        Array containing image data of the source.
+    center : tuple[int, int]
+        Tuple of the center coordinates where the RMS
+        is calculated for a region spanned by ``offset``.
+    offset : int, optional
+        Offset that spans a region around ``center``:
+
+        .. code::
+            x0 = center_x - offset
+            x1 = center_x + offset
+
+            y0 = center_y - offset
+            y1 = center_y + offset
+
+        Default: 75
+
+    Returns
+    -------
+    rms : float
+        RMS for the given source.
+    """
+    center_x = int(center[0])
+    center_y = int(center[1])
+
+    x0 = center_x - offset
+    x1 = center_x + offset
+
+    y0 = center_y - offset
+    y1 = center_y + offset
+
+    _rms = torch.from_numpy(rms(image[x0:x1, y0:y1]))
+
+    return torch.mean(_rms)
+
+
+def dynamic_range(path: str, *, offset: int = 75):
+    """Computes the dynamic range of a source.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path to a FITS file containing image data of the source.
+    offset : int, optional
+        Offset for the region around the center pixel of the source.
+        Used to compute the RMS value of the source. Default: 75
+
+    Returns
+    -------
+    dr : float
+        Dynamic range of the source.
+    """
+    hdu = fits.open(path)
+
+    image = hdu[0].data[0, 0, ...]
+    center = hdu[0].header["CRPIX1"], hdu[0].header["CRPIX2"]
+
+    _rms = get_source_rms(image, center, offset=75)
+
+    dr = image.max() / _rms
+
+    return dr

radiotools/utils/__init__.py

@@ -1,3 +1,7 @@
-from .utils import get_array_names, img2jansky
+from .utils import get_array_names, img2jansky, rms
 
-__all__ = ["get_array_names", "img2jansky"]
+__all__ = [
+    "get_array_names",
+    "img2jansky",
+    "rms",
+]

radiotools/utils/utils.py

@@ -35,6 +35,27 @@ def get_array_names(url: str) -> list[str]:
     return layouts
 
 
+def rms(a: ArrayLike, *, axis: int | None = 0):
+    """Return an array of the root-mean-square (RMS) value of
+    the passed array.
+
+    Parameters
+    ----------
+    a : array_like
+        Array of which to find the RMS.
+    axis : int or None
+
+    Returns
+    -------
+    rms : np.ndarray
+        Array of rms values.
+    """
+    if np.ndim(a) == 0:
+        axis = None
+
+    return np.sqrt(np.mean(a**2, axis=axis))
+
+
 def img2jansky(image: ArrayLike, header: fits.Header):
     """Converts an image from Jy/beam to Jy/px.
 

tests/test_analysis.py

@@ -0,0 +1,54 @@
+import numpy as np
+from numpy.testing import assert_allclose
+
+from radiotools.utils import rms
+
+
+class TestRMS:
+    def test_1d_array(self):
+        """Test 1d array."""
+        a = np.array([10, 0, 150, 40, 30, 20, -50, -60, 5])
+        expected = np.sqrt(1 / 9 * 31625)
+
+        assert_allclose(rms(a), expected)
+
+    def test_2d_array_axis_none(self):
+        """Test 2d array with axis=None."""
+        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+        expected = 5.627314338711377
+
+        assert_allclose(rms(a, axis=None), expected)
+
+    def test_2d_array_axis_0(self):
+        """Test 2d array with axis=0."""
+        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+        expected = np.array(
+            [
+                4.69041575982343,
+                5.5677643628300215,
+                6.48074069840786,
+            ]
+        )
+
+        assert_allclose(rms(a, axis=0), expected)
+
+    def test_2d_array_axis_1(self):
+        """Test 2d array with axis=1."""
+        a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+        expected = np.array(
+            [
+                2.160246899469287,
+                5.066228051190222,
+                8.04155872120988,
+            ]
+        )
+
+        assert_allclose(rms(a, axis=1), expected)
+
+    def test_input_ndim_0(self):
+        """Test input with ndim=0"""
+        a = np.int16(10)
+
+        expected = 10.0
+
+        assert rms(a) == expected