Merge pull request #33 from AlecThomson/docs

Clean up the docs

Author: AlecThomson

tests/test_2d.py

@@ -3,6 +3,7 @@
 
 import subprocess as sp
 import unittest
+from typing import List, Tuple, Union
 
 import astropy.units as u
 import numpy as np
@@ -13,7 +14,15 @@
 from racs_tools import beamcon_2D
 
 
-def make_2d_image(beam):
+def make_2d_image(beam: Beam) -> str:
+    """Make a fake 2D image from with a Gaussian beam.
+
+    Args:
+        beam (Beam): Gaussian beam.
+
+    Returns:
+        str: Name of the output FITS file.
+    """
     pix_scale = 2.5 * u.arcsec
 
     data = beam.as_kernel(pixscale=pix_scale, x_size=100, y_size=100).array
@@ -44,7 +53,16 @@ def make_2d_image(beam):
     return outf
 
 
-def mirsmooth(outf, target_beam):
+def mirsmooth(outf: str, target_beam: Beam) -> Tuple[str, str, str]:
+    """Smooth a FITS image to a target beam using MIRIAD.
+
+    Args:
+        outf (str): FITS image to smooth.
+        target_beam (Beam): Target beam.
+
+    Returns:
+        Tuple[str, str, str]: Names of the output images.
+    """
     outim = outf.replace(".fits", ".im")
     cmd = f"fits op=xyin in={outf} out={outim}"
     sp.run(cmd.split())
@@ -60,20 +78,37 @@ def mirsmooth(outf, target_beam):
     return outim, smoothim, smoothfits
 
 
-def check_images(fname_1, fname_2):
+def check_images(fname_1: str, fname_2: str) -> bool:
+    """Compare two FITS images.
+
+    Args:
+        fname_1 (str): Image 1.
+        fname_2 (str): Image 2.
+
+    Returns:
+        bool: True if the images are the same.
+    """
     data_1 = fits.getdata(fname_1)
     data_2 = fits.getdata(fname_2)
 
     return np.allclose(data_1, data_2, atol=1e-5)
 
 
-def cleanup(files):
+def cleanup(files: List[str]):
+    """Remove files.
+
+    Args:
+        files (List[str]): List of files to remove.
+    """
     for f in files:
         sp.run(f"rm -rfv {f}".split())
 
 
 class test_Beamcon2D(unittest.TestCase):
+    """Test the 2D beam convolution."""
+
     def setUp(self) -> None:
+        """Set up the test."""
         self.orginal_beam = Beam(20 * u.arcsec, 10 * u.arcsec, 10 * u.deg)
         test_image = make_2d_image(self.orginal_beam)
 
@@ -90,6 +125,7 @@ def setUp(self) -> None:
         ]
 
     def test_robust(self):
+        """Test the robust convolution."""
         with schwimmbad.SerialPool() as pool:
             beamcon_2D.main(
                 pool=pool,
@@ -103,9 +139,12 @@ def test_robust(self):
 
         fname_beamcon = self.test_image.replace(".fits", ".robust.fits")
         self.files.append(fname_beamcon)
-        assert check_images(self.test_mir, fname_beamcon), "Beamcon does not match miriad"
+        assert check_images(
+            self.test_mir, fname_beamcon
+        ), "Beamcon does not match miriad"
 
     def test_astropy(self):
+        """Test the astropy convolution."""
         print(f"{self.test_image=}")
         with schwimmbad.SerialPool() as pool:
             beamcon_2D.main(
@@ -120,9 +159,12 @@ def test_astropy(self):
 
         fname_beamcon = self.test_image.replace(".fits", ".astropy.fits")
         self.files.append(fname_beamcon)
-        assert check_images(self.test_mir, fname_beamcon), "Beamcon does not match miriad"
+        assert check_images(
+            self.test_mir, fname_beamcon
+        ), "Beamcon does not match miriad"
 
     def test_scipy(self):
+        """Test the scipy convolution."""
         with schwimmbad.SerialPool() as pool:
             beamcon_2D.main(
                 pool=pool,
@@ -136,9 +178,12 @@ def test_scipy(self):
 
         fname_beamcon = self.test_image.replace(".fits", ".scipy.fits")
         self.files.append(fname_beamcon)
-        assert check_images(self.test_mir, fname_beamcon), "Beamcon does not match miriad"
+        assert check_images(
+            self.test_mir, fname_beamcon
+        ), "Beamcon does not match miriad"
 
     def tearDown(self) -> None:
+        """Clean up."""
         cleanup(self.files)
 
 

tests/test_3d.py

@@ -1,23 +1,32 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
-import subprocess as sp
-import unittest
 import os
 import shutil
+import subprocess as sp
+import unittest
 
 import astropy.units as u
 import numpy as np
 import schwimmbad
 from astropy.io import fits
 from astropy.table import Table
 from radio_beam import Beam, Beams
+from test_2d import check_images, cleanup, mirsmooth
 
 from racs_tools import beamcon_3D
-from test_2d import mirsmooth, check_images, cleanup
 
-def smoothcube(outf: str, target_beam: Beam):
 
+def smoothcube(outf: str, target_beam: Beam) -> str:
+    """Smooth a FITS cube to a target beam.
+
+    Args:
+        outf (str): FITS cube to smooth.
+        target_beam (Beam): Target beam.
+
+    Returns:
+        str: Output FITS cube.
+    """
     cube = np.squeeze(fits.getdata(outf))
     header = fits.getheader(outf)
     with fits.open(outf) as hdulist:
@@ -36,7 +45,6 @@ def smoothcube(outf: str, target_beam: Beam):
         shutil.rmtree(outim)
         shutil.rmtree(smoothim)
 
-
     smoothcube = np.array(smoothcube)
     cube_hdu = fits.PrimaryHDU(data=smoothcube, header=header)
     cube_hdu.header = target_beam.attach_to_header(cube_hdu.header)
@@ -47,7 +55,15 @@ def smoothcube(outf: str, target_beam: Beam):
     return smooth_outf
 
 
-def make_3d_image(beams: Beams):
+def make_3d_image(beams: Beams) -> str:
+    """Make a fake 3D image from with a Gaussian beam.
+
+    Args:
+        beams (Beams): Gaussian beams.
+
+    Returns:
+        str: FITS cube filename.
+    """
     pix_scale = 2.5 * u.arcsec
 
     freqs = np.linspace(1, 2, len(beams)) * u.GHz
@@ -60,7 +76,6 @@ def make_3d_image(beams: Beams):
 
     cube = np.array(cube)[np.newaxis]
 
-
     hdu = fits.PrimaryHDU(data=cube)
     hdu.header["BUNIT"] = "Jy/beam"
     hdu.header["CDELT1"] = -pix_scale.to(u.deg).value
@@ -117,8 +132,12 @@ def make_3d_image(beams: Beams):
 
     return outf
 
-class test_Beamcon2D(unittest.TestCase):
+
+class test_Beamcon3D(unittest.TestCase):
+    """Test the beamcon_3D script."""
+
     def setUp(self) -> None:
+        """Set up the test."""
         self.orginal_beams = Beams(
             major=np.linspace(50, 10, 10) * u.arcsec,
             minor=np.linspace(10, 10, 10) * u.arcsec,
@@ -137,6 +156,7 @@ def setUp(self) -> None:
         ]
 
     def test_robust(self):
+        """Test the robust mode."""
         beamcon_3D.main(
             infile=[self.test_image],
             suffix="robust",
@@ -152,6 +172,7 @@ def test_robust(self):
         check_images(self.test_cube, fname_beamcon), "Beamcon does not match Miriad"
 
     def test_astropy(self):
+        """Test the astropy mode."""
         beamcon_3D.main(
             infile=[self.test_image],
             suffix="astropy",
@@ -167,6 +188,7 @@ def test_astropy(self):
         check_images(self.test_cube, fname_beamcon), "Beamcon does not match Miriad"
 
     def test_scipy(self):
+        """Test the scipy mode."""
         beamcon_3D.main(
             infile=[self.test_image],
             suffix="scipy",
@@ -181,4 +203,11 @@ def test_scipy(self):
         check_images(self.test_cube, fname_beamcon), "Beamcon does not match Miriad"
 
     def tearDown(self) -> None:
-        cleanup(self.files)
+        """Tear down the test."""
+        cleanup(self.files)
+
+
+if __name__ == "__main__":
+    unittest.TestLoader.sortTestMethodsUsing = None
+    suite = unittest.TestLoader().loadTestsFromTestCase(test_Beamcon3D)
+    unittest.TextTestRunner(verbosity=1).run(suite)