diff --git a/pyvisgen/fits/data.py b/pyvisgen/fits/data.py index 1fdfd090..d23bdcfc 100644 --- a/pyvisgen/fits/data.py +++ b/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() diff --git a/pyvisgen/fits/writer.py b/pyvisgen/fits/writer.py index ab00fc93..8bc75bb3 100644 --- a/pyvisgen/fits/writer.py +++ b/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)