Modules

A collection of utility functions for data manipulation.

This module contains a collection of utility functions for astronomical data manipulation.

catalog_quest(name, service='Vizier')

Fetch a catalog from a given astronomical service e.g. VizieR, Simbad.

:param name: The name of the catalog to be fetched. :type name: str

:param service: The name of the astronomical service to be used. :type service: str

:return: A pandas DataFrame containing the fetched catalog. :rtype: pd.DataFrame

:raises _UnsupportedServiceError: If an unsupported service is provided.

Source code in rgc/utils/data.py

def catalog_quest(name: str, service: str = "Vizier") -> pd.DataFrame:
    """
    Fetch a catalog from a given astronomical service e.g. VizieR, Simbad.

    :param name: The name of the catalog to be fetched.
    :type name: str

    :param service: The name of the astronomical service to be used.
    :type service: str

    :return: A pandas DataFrame containing the fetched catalog.
    :rtype: pd.DataFrame

    :raises _UnsupportedServiceError: If an unsupported service is provided.
    """
    if service == "Vizier":
        Vizier.ROW_LIMIT = -1
        catalog = Vizier.get_catalogs(name)
        return cast(pd.DataFrame, catalog[0].to_pandas())
    else:
        raise _UnsupportedServiceError()

celestial_capture(survey, ra, dec, filename)

Capture a celestial image using the SkyView service.

:param survey: The name of the survey to be used e.g. 'VLA FIRST (1.4 GHz)'. :type survey: str

:param ra: The right ascension of the celestial object. :type ra: Skycoord

:param dec: The declination of the celestial object. :type dec: Skycoord

:param filename: The name of the file to save the image. :type filename: str

Source code in rgc/utils/data.py

def celestial_capture(survey: str, ra: float, dec: float, filename: str) -> None:
    """
    Capture a celestial image using the SkyView service.

    :param survey: The name of the survey to be used e.g. 'VLA FIRST (1.4 GHz)'.
    :type survey: str

    :param ra: The right ascension of the celestial object.
    :type ra: Skycoord

    :param dec: The declination of the celestial object.
    :type dec: Skycoord

    :param filename: The name of the file to save the image.
    :type filename: str
    """
    image = SkyView.get_images(position=f"{ra}, {dec}", survey=survey, coordinates="J2000", pixels=(150, 150))[0]

    comment = str(image[0].header["COMMENT"])
    comment = comment.replace("\n", " ")
    comment = comment.replace("\t", " ")

    image[0].header.remove("comment", comment, True)
    image[0].header.add_comment(comment)

    folder_path = Path(filename).parent
    Path(folder_path).mkdir(parents=True, exist_ok=True)
    image.writeto(filename, overwrite=True)

celestial_capture_bulk(catalog, survey, img_dir, classes=None, cls_col=None)

Capture celestial images for a catalog of celestial objects.

:param catalog: A pandas DataFrame containing the catalog of celestial objects. :type catalog: pd.DataFrame

:param survey: The name of the survey to be used e.g. 'VLA FIRST (1.4 GHz)'. :type survey: str

:param img_dir: The path to the directory to save the images. :type img_dir: str

:param classes: A dictionary containing the classes of the celestial objects. :type classes: dict

:param cls_col: The name of the column containing the class labels.

:raises _InvalidCoordinatesError: If coordinates are invalid.

Source code in rgc/utils/data.py

def celestial_capture_bulk(
    catalog: pd.DataFrame, survey: str, img_dir: str, classes: Optional[dict] = None, cls_col: Optional[str] = None
) -> None:
    """
    Capture celestial images for a catalog of celestial objects.

    :param catalog: A pandas DataFrame containing the catalog of celestial objects.
    :type catalog: pd.DataFrame

    :param survey: The name of the survey to be used e.g. 'VLA FIRST (1.4 GHz)'.
    :type survey: str

    :param img_dir: The path to the directory to save the images.
    :type img_dir: str

    :param classes: A dictionary containing the classes of the celestial objects.
    :type classes: dict

    :param cls_col: The name of the column containing the class labels.

    :raises _InvalidCoordinatesError: If coordinates are invalid.
    """
    failed = pd.DataFrame(columns=catalog.columns)
    for _, entry in catalog.iterrows():
        try:
            tag = celestial_tag(entry)
            coordinate = SkyCoord(tag, unit=(u.hourangle, u.deg))

            right_ascension = coordinate.ra.deg
            declination = coordinate.dec.deg

            label = _get_class_labels(entry, classes, cls_col) if classes is not None and cls_col is not None else ""

            if "filename" in catalog.columns:
                filename = f'{img_dir}/{label}_{entry["filename"]}.fits'
            else:
                filename = f"{img_dir}/{label}_{tag}.fits"

            celestial_capture(survey, right_ascension, declination, filename)
        except Exception as err:
            series = entry.to_frame().T
            failed = pd.concat([failed, series], ignore_index=True)
            print(f"Failed to capture image. {err}")

celestial_tag(entry)

Generate a name tag for a celestial object based on its coordinates.

:param entry: A pandas Series entry of the catalog. :type entry: pd.Series

:return: A string containing the name tag. :rtype: str

:raises _NoValidCelestialCoordinatesError: If no valid celestial coordinates are found in the entry.

Source code in rgc/utils/data.py

def celestial_tag(entry: pd.Series) -> str:
    """
    Generate a name tag for a celestial object based on its coordinates.

    :param entry: A pandas Series entry of the catalog.
    :type entry: pd.Series

    :return: A string containing the name tag.
    :rtype: str

    :raises _NoValidCelestialCoordinatesError: If no valid celestial coordinates are found in the entry.
    """

    def format_dec(dec: str) -> str:
        sign = "+" if float(dec.replace(" ", "")) > 0 else ""
        return f"{sign}{dec}"

    if {"RAJ2000", "DEJ2000"}.issubset(entry.index):
        ra, dec = entry["RAJ2000"], entry["DEJ2000"]
    elif {"RA", "DEC"}.issubset(entry.index):
        ra, dec = entry["RA"], entry["DEC"]
    elif "filename" in entry.index:
        return f"{entry['filename']}"
    elif "FCG" in entry.index:
        return f"{entry['FCG']}"
    else:
        raise _NoValidCelestialCoordinatesError()

    return f"{ra}{format_dec(dec)}"

compute_mean_std(dataloader)

Compute the mean and standard deviation of the dataset.

:param dataloader: The dataloader for the dataset. :type dataloader: torch.utils.data.DataLoader

:return: The mean and standard deviation of the dataset. :rtype: tuple[torch.Tensor, torch.Tensor]

Source code in rgc/utils/data.py

def compute_mean_std(dataloader: torch.utils.data.DataLoader) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute the mean and standard deviation of the dataset.

    :param dataloader: The dataloader for the dataset.
    :type dataloader: torch.utils.data.DataLoader

    :return: The mean and standard deviation of the dataset.
    :rtype: tuple[torch.Tensor, torch.Tensor]
    """
    data = torch.tensor([])
    for batch in dataloader:
        data = torch.cat((data, batch[0]), 0)

    mean = torch.mean(data, dim=(0, 2, 3))
    std = torch.std(data, dim=(0, 2, 3))

    return mean, std

dataframe_to_html(catalog, save_dir)

Save the catalog as an HTML file.

:param catalog: Catalog of the astronomical objects :type catalog: pd.DataFrame :param save_dir: Path to the directory to save the HTML file :type save_dir: str

Source code in rgc/utils/data.py

def dataframe_to_html(catalog: pd.DataFrame, save_dir: str) -> None:
    """
    Save the catalog as an HTML file.

    :param catalog: Catalog of the astronomical objects
    :type catalog: pd.DataFrame
    :param save_dir: Path to the directory to save the HTML file
    :type save_dir: str
    """
    Path(save_dir).mkdir(parents=True, exist_ok=True)
    catalog.to_html(os.path.join(save_dir, "catalog.html"))

fits_to_png(fits_file, img_size=None)

Convert a FITS file to a PNG image.

:param fits_file: The path to the FITS file. :type fits_file: str

:param img_size: The size of the output image. :type img_size: Optional[tuple[int, int]]

:return: A PIL Image object containing the PNG image. :rtype: Image.Image

:raises _FileNotFoundError: If the FITS file is not found.

Source code in rgc/utils/data.py

def fits_to_png(fits_file: str, img_size: Optional[tuple[int, int]] = None) -> Image.Image:
    """
    Convert a FITS file to a PNG image.

    :param fits_file: The path to the FITS file.
    :type fits_file: str

    :param img_size: The size of the output image.
    :type img_size: Optional[tuple[int, int]]

    :return: A PIL Image object containing the PNG image.
    :rtype: Image.Image

    :raises _FileNotFoundError: If the FITS file is not found.
    """
    try:
        image = fits.getdata(fits_file)
        header = fits.getheader(fits_file)
    except FileNotFoundError as err:
        raise _FileNotFoundError(fits_file) from err

    if img_size is not None:
        width, height = img_size
    else:
        width, height = header["NAXIS1"], header["NAXIS2"]

    image = np.reshape(image, (height, width))
    image[np.isnan(image)] = np.nanmin(image)

    image = (image - np.nanmin(image)) / (np.nanmax(image) - np.nanmin(image)) * 255
    image = image.astype(np.uint8)
    image = Image.fromarray(image, mode="L")

    return cast(Image.Image, image)

fits_to_png_bulk(fits_dir, png_dir, img_size=None)

Convert a directory of FITS files to PNG images.

:param fits_dir: The path to the directory containing the FITS files. :type fits_dir: str

:param png_dir: The path to the directory to save the PNG images. :type png_dir: str

:param img_size: The size of the output image. :type img_size: Optional[tuple[int, int]]

Source code in rgc/utils/data.py

def fits_to_png_bulk(fits_dir: str, png_dir: str, img_size: Optional[tuple[int, int]] = None) -> None:
    """
    Convert a directory of FITS files to PNG images.

    :param fits_dir: The path to the directory containing the FITS files.
    :type fits_dir: str

    :param png_dir: The path to the directory to save the PNG images.
    :type png_dir: str

    :param img_size: The size of the output image.
    :type img_size: Optional[tuple[int, int]]
    """
    fits_files = Path(fits_dir).rglob("*.fits")
    for fits_file in fits_files:
        image = fits_to_png(str(fits_file), img_size)

        png_file = os.path.join(png_dir, fits_file.stem)
        Path(png_file).parent.mkdir(parents=True, exist_ok=True)

        if image is not None:
            image.save(png_file)

generate_mask(image_path, mask_dir, freq, beam, dilation, threshold_pixel=5.0, threshold_island=3.0)

Detect sources in the image and generate a mask.

:param image_path: Path to the image file :type image_path: str

:param mask_dir: Path to the directory to save the mask :type mask_dir: str

:param freq: Frequency of the image in MHz :type freq: float

:param beam: Beam size of the image in arcsec :type beam: tuple

:param dilation: Dilation factor for the mask :type dilation: int

:param threshold_pixel: Threshold for island peak in number of sigma above the mean :type threshold_pixel: float

:param threshold_island: Threshold for island detection in number of sigma above the mean :type threshold_island: float

Source code in rgc/utils/data.py

def generate_mask(
    image_path: str,
    mask_dir: str,
    freq: float,
    beam: tuple[float, float, float],
    dilation: int,
    threshold_pixel: float = 5.0,
    threshold_island: float = 3.0,
) -> None:
    """
    Detect sources in the image and generate a mask.

    :param image_path: Path to the image file
    :type image_path: str

    :param mask_dir: Path to the directory to save the mask
    :type mask_dir: str

    :param freq: Frequency of the image in MHz
    :type freq: float

    :param beam: Beam size of the image in arcsec
    :type beam: tuple

    :param dilation: Dilation factor for the mask
    :type dilation: int

    :param threshold_pixel: Threshold for island peak in number of sigma above the mean
    :type threshold_pixel: float

    :param threshold_island: Threshold for island detection in number of sigma above the mean
    :type threshold_island: float
    """
    try:
        image = bdsf.process_image(
            image_path,
            beam=beam,
            thresh_isl=threshold_island,
            thresh_pix=threshold_pixel,
            frequency=freq,
        )

        mask_file = Path(mask_dir) / Path(image_path).name
        Path(mask_file).parent.mkdir(parents=True, exist_ok=True)

        image.export_image(
            img_type="island_mask",
            outfile=mask_file,
            clobber=True,
            mask_dilation=dilation,
        )

    except Exception:
        print("Failed to generate mask.")
        return None

generate_mask_bulk(catalog, img_dir, mask_dir, freq, beam)

Generate masks for a catalog of celestial objects.

:param catalog: A pandas DataFrame containing the catalog of celestial objects. :type catalog: pd.DataFrame

:param img_dir: The path to the directory containing the images. :type img_dir: str

:param mask_dir: The path to the directory to save the masks. :type mask_dir: str

:param freq: Frequency of the image in MHz :type freq: float

:param beam: Beam size of the image in arcsec :type beam: tuple

Source code in rgc/utils/data.py

def generate_mask_bulk(
    catalog: pd.DataFrame, img_dir: str, mask_dir: str, freq: float, beam: tuple[float, float, float]
) -> None:
    """
    Generate masks for a catalog of celestial objects.

    :param catalog: A pandas DataFrame containing the catalog of celestial objects.
    :type catalog: pd.DataFrame

    :param img_dir: The path to the directory containing the images.
    :type img_dir: str

    :param mask_dir: The path to the directory to save the masks.
    :type mask_dir: str

    :param freq: Frequency of the image in MHz
    :type freq: float

    :param beam: Beam size of the image in arcsec
    :type beam: tuple
    """
    for _, entry in catalog.iterrows():
        try:
            filename = entry["filename"]
            image_path = os.path.join(img_dir, f"{filename}.fits")
            dilation = entry["dilation"]
            threshold_pixel = entry["background sigma"]
            threshold_island = entry["foreground sigma"]

            generate_mask(
                image_path,
                mask_dir,
                freq,
                beam,
                dilation,
                threshold_pixel,
                threshold_island,
            )

        except Exception as err:
            print(f"Failed to generate mask. {err}")
            return None

mask_image(image, mask)

Mask an image with a given mask image.

:param image: The image to be masked. :type image: Image.Image

:param mask: The mask image. :type mask: Image.Image

:return: A PIL Image object containing the masked image. :rtype: Image.Image

:raises _ImageMaskDimensionError: If the dimensions of the image and mask do not match.

Source code in rgc/utils/data.py

def mask_image(image: Image.Image, mask: Image.Image) -> Image.Image:
    """
    Mask an image with a given mask image.

    :param image: The image to be masked.
    :type image: Image.Image

    :param mask: The mask image.
    :type mask: Image.Image

    :return: A PIL Image object containing the masked image.
    :rtype: Image.Image

    :raises _ImageMaskDimensionError: If the dimensions of the image and mask do not match.
    """
    image_array = np.array(image)
    mask_array = np.array(mask)

    if image_array.shape != mask_array.shape:
        raise _ImageMaskDimensionError()

    masked_array = np.where(mask_array == 0, 0, image_array)
    masked_image = Image.fromarray(masked_array, mode="L")

    return cast(Image.Image, masked_image)

mask_image_bulk(image_dir, mask_dir, masked_dir)

Mask a directory of images with a directory of mask images.

:param image_dir: The path to the directory containing the images. :type image_dir: str

:param mask_dir: The path to the directory containing the mask images. :type mask_dir: str

:param masked_dir: The path to the directory to save the masked images. :type masked_dir: str

:raises _FileNotFoundError: If no images or masks are found in the directories. :raises _ImageMaskCountMismatchError: If the number of images and masks do not match.

Source code in rgc/utils/data.py

def mask_image_bulk(image_dir: str, mask_dir: str, masked_dir: str) -> None:
    """
    Mask a directory of images with a directory of mask images.

    :param image_dir: The path to the directory containing the images.
    :type image_dir: str

    :param mask_dir: The path to the directory containing the mask images.
    :type mask_dir: str

    :param masked_dir: The path to the directory to save the masked images.
    :type masked_dir: str

    :raises _FileNotFoundError: If no images or masks are found in the directories.
    :raises _ImageMaskCountMismatchError: If the number of images and masks do not match.
    """
    image_paths = sorted(Path(image_dir).glob("*.png"))
    mask_paths = sorted(Path(mask_dir).glob("*.png"))

    if len(image_paths) == 0 or len(mask_paths) == 0:
        raise _FileNotFoundError()

    if len(image_paths) != len(mask_paths):
        raise _ImageMaskCountMismatchError() from None

    os.makedirs(masked_dir, exist_ok=True)

    for image_path in image_paths:
        mask_path = Path(mask_dir) / image_path.name

        if not mask_path.exists():
            print(f"Skipping {image_path.name} due to missing mask.")
            continue

        image = Image.open(image_path)
        mask = Image.open(mask_path)

        if image.size != mask.size:
            print(f"Skipping {image_path.name} due to mismatched dimensions.")
            continue
        else:
            masked_image = mask_image(image, mask)

        masked_image.save(Path(masked_dir) / image_path.name)

remove_artifacts(folder, extension)

Remove files with the given extensions from a folder.

:param folder: Path to the folder to clear :type folder: str :param extension: List of file with the given extensions to keep :type extension: list

Source code in rgc/utils/data.py

def remove_artifacts(folder: str, extension: list[str]) -> None:
    """
    Remove files with the given extensions from a folder.

    :param folder: Path to the folder to clear
    :type folder: str
    :param extension: List of file with the given extensions to keep
    :type extension: list
    """
    for file in os.listdir(folder):
        if not file.endswith(tuple(extension)):
            os.remove(os.path.join(folder, file))

    print(f"Artifacts removed from {folder} with extensions {', '.join(extension)}")