Add docstrings to the fits module

pyvisgen/fits/data.py

@@ -1,61 +1,129 @@
 from pathlib import Path
 
+import natsort
 import numpy as np
 from astropy.io import fits
+from numpy.typing import ArrayLike
 
 
 class fits_data:
-    def __init__(self, data_path):
-        """Handle h5 files simulated with radiosim package.
+    """Class that handles h5 files simulated
+    with the radiosim package.
+
+    Parameters
+    ----------
+    data_path: str Path
+        path to fits data directory
+    """
+
+    def __init__(self, data_path: str | Path) -> None:
+        """Handle h5 files simulated with the
+        radiosim package.
 
         Parameters
         ----------
-        data_path: str
+        data_path: str Path
             path to fits data directory
-
         """
         self.path = data_path
         self.files = self.get_files(Path(data_path))
         self.num_files = len(self.files)
 
-    def __call__(self):
+    def __call__(self) -> None:
         return print("This is the pyvisgen fits files data set class.")
 
-    def __len__(self):
+    def __len__(self) -> int:
         """
         Returns the total number of fits files in this dataset.
+
+        Returns
+        -------
+        num_files : int
+            Number of files in data set.
         """
         return self.num_files
 
     def __getitem__(self, i):
+        """Returns the file at index ``i`` of :py:attr:`~self.files`."""
         return self.open_file(i)
 
-    def get_files(self, path):
-        import natsort
+    def get_files(self, path: Path) -> list[str]:
+        """Returns a list of files in natural ordering.
+
+        Parameters
+        ----------
+        path : Path
+            Path to the data directory.
 
+        Returns
+        -------
+        fits_files_sorted : list
+            List of files in natural ordering.
+        """
         fits_files = [str(x) for x in path.iterdir()]
         fits_files_sorted = natsort.natsorted(fits_files)
+
         return fits_files_sorted
 
-    def get_uv_data(self, i):
+    def get_uv_data(self, i: int) -> ArrayLike:
+        """Loads uv data from the file at index ``i`` of
+        :py:attr:`~self.files`.
+
+        Parameters
+        ----------
+        i : int
+            Index of the file in :py:attr:`self.files`.
+
+        Returns
+        -------
+        uv_data : array_like
+            Array of uv data.
+        """
         with fits.open(self.files[i]) as hdul:
             uv_data = hdul[0].data
-            hdul.close()
-            return uv_data
 
-    def get_freq_data(self, i):
+        return uv_data
+
+    def get_freq_data(self, i: int) -> ArrayLike:
+        """Loads frequency data from the file at index ``i`` of
+        :py:attr:`~self.files`.
+
+        Parameters
+        ----------
+        i : int
+            Index of the file in :py:attr:`~self.files`.
+
+        Returns
+        -------
+        freq_data : array_like
+            Array of frequency data.
+        base_freq : float
+            Base frequency of the observing antenna array.
+        """
         with fits.open(self.files[i]) as hdul:
             base_freq = hdul[0].header["CRVAL4"]
             freq_data = hdul[2].data
-            hdul.close()
-            return freq_data, base_freq
 
-    def open_file(self, i):
+        return freq_data, base_freq
+
+    def open_file(self, i: int) -> ArrayLike:
+        """Opens the file and returns the file information.
+
+        Parameters
+        ----------
+        i : int
+            Index of the file in :py:attr:`~self.files`.
+
+        Returns
+        -------
+        :py:func:`~astropy.io.fits.info`
+            Summary information on the FITS file.
+        """
         if isinstance(i, np.ndarray):
             print("Only one file can be open at the same time!")
             return
 
         with fits.open(self.files[i]) as hdul:
             fits_file = hdul
-            hdul.close()
-            return fits_file.info()
+
+        return fits_file.info()

pyvisgen/fits/writer.py

@@ -6,11 +6,34 @@
 from astropy import wcs
 from astropy.io import fits
 from astropy.time import Time
+from astropy.utils import iers
 
 import pyvisgen.layouts.layouts as layouts
 
 
-def create_vis_hdu(data, obs, source_name="sim-source-0"):
+def create_vis_hdu(data, obs, source_name="sim-source-0") -> fits.GroupsHDU:
+    """Creates the visibility HDU for the given visibility
+    data and observation.
+
+    Parameters
+    ----------
+    data : :class:`~pyvisgen.simulation.Visibilities`
+        :class:`~pyvisgen.simulation.Visibilities` object
+        containing visibility data.
+    obs : :class:`~pyvisgen.simulation.Observation`
+        :class:`~pyvisgen.simulation.Observation` class object
+        containing information on the observation, such as
+        baselines.
+    source_name : str, optional
+        Source name saved to the ``OBJECT`` key inside the
+        FITS file. Default: ``'sim-source-0'``
+
+    Returns
+    -------
+    hdu_vis : :class:`~astropy.io.fits.GroupsHDU`
+        :class:`~astropy.io.fits.GroupsHDU` containing visibility
+        data for the FITS file.
+    """
     u = data.u
 
     v = data.v
@@ -129,7 +152,21 @@ def create_vis_hdu(data, obs, source_name="sim-source-0"):
     return hdu_vis
 
 
-def create_time_hdu(data):
+def create_time_hdu(data) -> fits.BinTableHDU:
+    """Creates the time HDU for the FITS file.
+
+    Parameters
+    ----------
+    data : :class:`~pyvisgen.simulation.Visibilities`
+        :class:`~pyvisgen.simulation.Visibilities` object
+        containing visibility and time data.
+
+    Returns
+    -------
+    hdu_vis : :class:`~astropy.io.fits.BinTableHDU`
+        :class:`~astropy.io.fits.BinTableHDU` containing time
+        data for the FITS file.
+    """
     TIME = np.array(
         [data.date.mean() - int(data.date.min())],
         dtype=">f4",
@@ -184,7 +221,22 @@ def create_time_hdu(data):
     return hdu_time
 
 
-def create_frequency_hdu(obs):
+def create_frequency_hdu(obs) -> fits.BinTableHDU:
+    """Creates the frequency HDU of the FITS file.
+
+    Parameters
+    ----------
+    obs : :class:`~pyvisgen.simulation.Observation`
+        :class:`~pyvisgen.simulation.Observation` class object
+        containing information on the observation, including
+        frequency data.
+
+    Returns
+    -------
+    hdu_freq : :class:`~astropy.io.fits.BinTableHDU`
+        :class:`~astropy.io.fits.BinTableHDU` containing
+        frequency data for the FITS file.
+    """
     FRQSEL = np.array([1], dtype=">i4")
     col1 = fits.Column(name="FRQSEL", format="1J", array=FRQSEL)
 
@@ -227,7 +279,22 @@ def create_frequency_hdu(obs):
     return hdu_freq
 
 
-def create_antenna_hdu(obs):
+def create_antenna_hdu(obs) -> fits.BinTableHDU:
+    """Creates the antenna HDU for the FITS file.
+
+    Parameters
+    ----------
+    obs : :class:`~pyvisgen.simulation.Observation`
+        :class:`~pyvisgen.simulation.Observation` class object
+        containing information on the observation, including
+        antenna data.
+
+    Returns
+    -------
+    hdu_ant : :class:`~astropy.io.fits.BinTableHDU`
+        :class:`~astropy.io.fits.BinTableHDU` containing
+        antenna data for the FITS file.
+    """
     array = layouts.get_array_layout(obs.layout, writer=True)
 
     ANNAME = np.chararray(len(array), itemsize=8, unicode=True)
@@ -296,8 +363,6 @@ def create_antenna_hdu(obs):
         obs.start.isot.split("T")[0] + "T0:00:00.000", format="isot", scale="utc"
     )
 
-    from astropy.utils import iers
-
     iers_b = iers.IERS_B.open()
 
     # add additional keywords
@@ -364,7 +429,26 @@ def create_antenna_hdu(obs):
     return hdu_ant
 
 
-def create_hdu_list(data, obs):
+def create_hdu_list(data, obs) -> fits.HDUList:
+    """Creates a :class:`~astropy.io.fits.HDUList` as the
+    top-level object for the FITS file.
+
+    Parameters
+    ----------
+    data : :class:`~pyvisgen.simulation.Visibilities`
+        :class:`~pyvisgen.simulation.Visibilities` object containing
+        data on visibilities and observation time.
+    obs : :class:`~pyvisgen.simulation.Observation`
+        :class:`~pyvisgen.simulation.Observation` object containing
+        data on the source position, baselines, antenna configuration,
+        and frequencies.
+
+    Returns
+    -------
+    hdu_list : :class:`~astropy.io.fits.HDUList`
+        :class:`~astropy.io.fits.HDUList` that comprises of
+        HDU objects.
+    """
     warnings.filterwarnings("ignore", module="astropy.io.fits")
     vis_hdu = create_vis_hdu(data, obs)
     time_hdu = create_time_hdu(data)