tools.py

def add_credible_levels(
    filename: str,
    df: pd.DataFrame,
    pipe: bool = True
) -> None:
    """
    Calculate the minimum credible region containing each source
    and add to the dataframe in-place.

    Args:
        filename: Path to the healpix skymap file.
        df: Dataframe of sources.
        pipe: Whether the dataframe is from the pipeline. Defaults to True.

    Returns:
        None

    Raises:
        Exception: Path does not exist.
    """

    skymap = Path(filename)
    if not skymap.is_file():
        raise Exception("{} does not exist".format(filename))

    if pipe:
        ra_col = 'wavg_ra'
        dec_col = 'wavg_dec'
    else:
        ra_col = 'ra'
        dec_col = 'dec'

    hpx = hp.read_map(filename)
    nside = hp.get_nside(hpx)

    i = np.flipud(np.argsort(hpx))
    sorted_credible_levels = np.cumsum(hpx[i])
    credible_levels = np.empty_like(sorted_credible_levels)
    credible_levels[i] = sorted_credible_levels
    theta = 0.5 * np.pi - np.deg2rad(df[dec_col].values)
    phi = np.deg2rad(df[ra_col].values)

    ipix = hp.ang2pix(nside, theta, phi)

    df.loc[:, 'credible_level'] = credible_levels[ipix]

def add_obs_date(epoch: str,
                 image_type: str,
                 image_dir: str,
                 epoch_path: str = None
                 ) -> None:
    """
    Add datetime information to all fits files in a single epoch.

    Args:
        epoch: The epoch of interest.
        image_type: `COMBINED` or `TILES`.
        image_dir: The name of the folder containing the images to be updated.
            E.g. `STOKESI_IMAGES`.
        epoch_path: Full path to the folder containing the epoch.
            Defaults to None, which will set the value based on the
            `VAST_DATA_DIR` environment variable and `epoch`.

    Returns:
        None

    Raises:
        ValueError: When image_type is not 'TILES' or 'COMBINED'.
    """
    epoch_info = load_fields_file(epoch)

    if epoch_path is None:
        epoch_path = _set_epoch_path(epoch)

    raw_images = _get_epoch_images(epoch_path, image_type, image_dir)

    for filename in raw_images:
        split_name = filename.split("/")[-1].split(".")
        if image_type == 'TILES':
            field = split_name[4]
        elif image_type == 'COMBINED':
            field = split_name[0]
        else:
            raise ValueError(
                "Image type not recognised, "
                "must be either 'TILES' or 'COMBINED'."
            )

        field_info = epoch_info[epoch_info.FIELD_NAME == field].iloc[0]
        field_start = Time(field_info.DATEOBS)
        field_end = Time(field_info.DATEEND)
        duration = field_end - field_start

        hdu = fits.open(filename, mode="update")
        hdu_index = 0
        if filename.endswith('.fits.fz'):
            hdu_index = 1
        hdu[hdu_index].header["DATE-OBS"] = field_start.fits
        hdu[hdu_index].header["MJD-OBS"] = field_start.mjd
        hdu[hdu_index].header["DATE-BEG"] = field_start.fits
        hdu[hdu_index].header["DATE-END"] = field_end.fits
        hdu[hdu_index].header["MJD-BEG"] = field_start.mjd
        hdu[hdu_index].header["MJD-END"] = field_end.mjd
        hdu[hdu_index].header["TELAPSE"] = duration.sec
        hdu[hdu_index].header["TIMEUNIT"] = "s"
        hdu.close()

def create_fields_metadata(epoch_num: str,
                           db_path: str,
                           outdir: Union[str, Path] = '.'
                           ) -> None:
    """
    Create and write the fields csv and skycoord pickle for a single epoch.

    Args:
        epoch_num: Epoch number of interest.
        db_path: Path to the askap_surveys database.
        outdir: Path to the output directory.
            Defaults to the current directory.

    Returns:
        None

    Raises:
        Exception: Path does not exist.
    """

    if isinstance(outdir, str):
        outdir = Path(outdir)

    if not outdir.exists():
        raise Exception("{} does not exist!".format(outdir))

    fields_df = _create_fields_df(epoch_num, db_path)
    fields_sc = _create_fields_sc(fields_df)

    if len(epoch_num.rstrip('x')) == 1:
        epoch_num = f'0{epoch_num}'
    fields_outfile = f'vast_epoch{epoch_num}_info.csv'
    sc_outfile = f'vast_epoch{epoch_num}_fields_sc.pickle'

    fields_df.to_csv(outdir / fields_outfile, index=False)

    with open(outdir / sc_outfile, 'wb') as picklefile:
        pickle.dump(fields_sc, picklefile)

def find_in_moc(
    moc: MOC,
    df: pd.DataFrame,
    pipe: bool = True
) -> np.ndarray:
    """
    Find the sources that are contained within a MOC

    Args:
        moc: The MOC of interest.
        df: Dataframe of sources.
        pipe: Whether the dataframe is from the pipeline. Defaults to True.

    Returns:
        Indices of all sources contained within the MOC.
    """

    if pipe:
        ra_col = 'wavg_ra'
        dec_col = 'wavg_dec'
    else:
        ra_col = 'ra'
        dec_col = 'dec'

    ra = Angle(df[ra_col], unit='deg')
    dec = Angle(df[dec_col], unit='deg')

    return np.where(moc.contains(ra, dec))[0]

def gen_mocs_epoch(epoch: str,
                   image_type: str,
                   image_dir: str,
                   epoch_path: str = None,
                   outdir: Union[str, Path] = '.',
                   base_stmoc: Union[str, Path] = None
                   ) -> None:
    """
    Generate MOCs and STMOCs for all images in a single epoch, and create a new
    full pilot STMOC.

    Args:
        epoch: The epoch of interest.
        image_type: `COMBINED` or `TILES`.
        image_dir: The name of the folder containing the images to be updated.
            E.g. `STOKESI_IMAGES`.
        epoch_path: Full path to the folder containing the epoch.
            Defaults to None, which will set the value based on the
            `VAST_DATA_DIR` environment variable and `epoch`.
        outdir: Path to the output directory.
            Defaults to the current directory.
        base_stmoc: Path to the STMOC to use as the base. Defaults to `None`,
            in which case the VAST STMOC installed with vast-tools will
            be used.

    Returns:
        None

    Raises:
        Exception: Path does not exist.
    """

    if isinstance(outdir, str):
        outdir = Path(outdir)

    if not outdir.exists():
        raise Exception("{} does not exist".format(outdir))

    if base_stmoc is None:
        vtm = VASTMOCS()
        full_STMOC = vtm.load_pilot_stmoc()
    else:
        if isinstance(base_stmoc, str):
            base_stmoc = Path(base_stmoc)

        if not base_stmoc.exists():
            raise Exception("{} does not exist".format(base_stmoc))

        full_STMOC = STMOC.from_fits(base_stmoc)

    if epoch_path is None:
        epoch_path = _set_epoch_path(epoch)

    raw_images = _get_epoch_images(epoch_path, image_type, image_dir)

    for i, f in enumerate(raw_images):
        themoc, thestmoc = gen_mocs_image(f)

        if i == 0:
            mastermoc = themoc
            masterstemoc = thestmoc
        else:
            mastermoc = mastermoc.union(themoc)
            masterstemoc = masterstemoc.union(thestmoc)

    master_name = "VAST_PILOT_EPOCH{}.moc.fits".format(epoch)
    master_stmoc_name = master_name.replace("moc", "stmoc")

    mastermoc.write(outdir / master_name, overwrite=True)
    masterstemoc.write(outdir / master_stmoc_name, overwrite=True)

    full_STMOC = full_STMOC.union(masterstemoc)
    full_STMOC.write(outdir / 'VAST_PILOT.stmoc.fits', overwrite=True)

def gen_mocs_image(fits_file: str,
                   outdir: Union[str, Path] = '.',
                   write: bool = False
                   ) -> Union[MOC, STMOC]:
    """
    Generate a MOC and STMOC for a single fits file.

    Args:
        fits_file: path to the fits file.
        outdir: Path to the output directory.
            Defaults to the current directory.
        write: Write the MOC/STMOC to file.

    Returns:
        The MOC and STMOC.

    Raises:
        Exception: Path does not exist.
    """

    if isinstance(outdir, str):
        outdir = Path(outdir)

    if not outdir.exists():
        raise Exception("{} does not exist".format(outdir))

    if not Path(fits_file).exists():
        raise Exception("{} does not exist".format(fits_file))

    moc = create_moc_from_fits(fits_file)

    header = fits.getheader(fits_file, 0)
    start = Time([header['DATE-BEG']])
    end = Time([header['DATE-END']])
    stmoc = STMOC.from_spatial_coverages(
        start, end, [moc]
    )

    if write:
        filename = os.path.split(fits_file)[1]
        moc_name = filename.replace(".fits", ".moc.fits")
        stmoc_name = filename.replace(".fits", ".stmoc.fits")

        moc.write(outdir / moc_name, overwrite=True)
        stmoc.write(outdir / stmoc_name, overwrite=True)

    return moc, stmoc

def offset_postagestamp_axes(ax: plt.Axes,
                             centre_sc: SkyCoord,
                             ra_units: u.core.Unit = u.arcsec,
                             dec_units: u.core.Unit = u.arcsec,
                             ra_label: str = 'R.A. Offset',
                             dec_label: str = 'Dec. Offset',
                             major_tick_length: Union[int, float] = 6,
                             minor_tick_length: Union[int, float] = 3,
                             ) -> None:
    """
    Display axis ticks and labels as offsets from a given coordinate,
    rather than in absolute coordinates.

    Args:
        ax: The axis of interest
        centre_sc: SkyCoord to calculate offsets from
        ra_units: Right Ascension axis ticklabel units
        dec_units: Declination axis ticklabel units
        ra_label: Right Ascension axis label
        dec_label: Declination axis label
        major_tick_length: Major tick length in points
        minor_tick_length: Minor tick length in points

    Returns:
        None

    Raises:
        Exception: R.A. and Dec. units must be angles.
    """

    if ra_units.physical_type != 'angle' or dec_units.physical_type != 'angle':
        raise Exception("R.A. and Dec. units must be angles.")

    ra_offs, dec_offs = ax.get_coords_overlay(centre_sc.skyoffset_frame())
    plt.minorticks_on()
    ra_offs.set_coord_type('longitude', 180)
    ra_offs.set_format_unit(ra_units, decimal=True)
    dec_offs.set_format_unit(dec_units, decimal=True)
    ra_offs.tick_params(direction='in', color='black')
    ra_offs.tick_params(which='major', length=major_tick_length)
    ra_offs.tick_params(which='minor', length=minor_tick_length)

    dec_offs.tick_params(direction='in', color='black')
    dec_offs.tick_params(which='major', length=major_tick_length)
    dec_offs.tick_params(which='minor', length=minor_tick_length)

    ra, dec = ax.coords

    ra.set_ticks_visible(False)
    ra.set_ticklabel_visible(False)
    dec.set_ticks_visible(False)
    dec.set_ticklabel_visible(False)
    ra_offs.display_minor_ticks(True)
    dec_offs.display_minor_ticks(True)

    dec_offs.set_ticks_position('lr')
    dec_offs.set_ticklabel_position('l')
    dec_offs.set_axislabel_position('l')
    dec_offs.set_axislabel(dec_label)

    ra_offs.set_ticks_position('tb')
    ra_offs.set_ticklabel_position('b')
    ra_offs.set_axislabel_position('b')
    ra_offs.set_axislabel(ra_label)

    return

def skymap2moc(filename: str, cutoff: float) -> MOC:
    """
    Creates a MOC of the specified credible region of a given skymap.

    Args:
        filename: Path to the healpix skymap file.
        cutoff: Credible level cutoff.

    Returns:
        A MOC containing the credible region.

    Raises:
        ValueError: Credible level cutoff must be between 0 and 1.
        Exception: Path does not exist.
    """
    skymap = Path(filename)

    if not 0.0 <= cutoff <= 1.0:
        raise Exception("Credible level cutoff must be between 0 and 1")

    if not skymap.is_file():
        raise Exception("{} does not exist".format(skymap))

    hpx = hp.read_map(filename, nest=True)
    nside = hp.get_nside(hpx)
    level = np.log2(nside)

    i = np.flipud(np.argsort(hpx))
    sorted_credible_levels = np.cumsum(hpx[i])
    credible_levels = np.empty_like(sorted_credible_levels)
    credible_levels[i] = sorted_credible_levels

    idx = np.where(credible_levels < cutoff)[0]
    levels = np.ones(len(idx)) * level

    moc = MOC.from_healpix_cells(idx, depth=levels)

    return moc

def wise_color_color_plot(
    patch_style_overrides: Optional[Dict[WiseClass, WisePatchConfig]] = None,
    annotation_text_size: Union[float, str] = "x-small",
) -> matplotlib.figure.Figure:
    """Make an empty WISE color-color plot with common object classes drawn as
    patches. The patches have default styles that may be overridden. To
    override a patch style, pass in a dict containing the desired style and
    annotation settings. The overrides must be complete, i.e. a complete
    `WisePatchConfig` object must be provided for each `WiseClass` you wish to
    modify. Partial updates to the style or annotation of individual patches is
    not supported.

    For example, to change the color of the stars patch to blue:
    ```python
    fig = wise_color_color_plot({
        WiseClass.STARS: WisePatchConfig(
            style=dict(fc="blue", ec="none"),
            annotation_text="Stars",
            annotation_position=(0.5, 0.4),
        )
    })
    ```

    Args:
        patch_style_overrides (Optional[Dict[WiseClass, WisePatchConfig]],
            optional): Override the default patch styles for the given WISE
            object class. If None, use defaults for each patch. Defaults to
            None.
        annotation_text_size (Union[float, str]): Font size for the patch
            annotations. Accepts a font size (float) or a matplotlib font scale
            string (e.g. "xx-small", "medium", "xx-large"). Defaults to
            "x-small".
    Returns:
        `matplotlib.figure.Figure`: the WISE color-color figure. Access the
            axes with the `.axes` attribute.
    """
    # set the WISE object classification patch styles
    if patch_style_overrides is not None:
        patch_styles = WISE_DEFAULT_PATCH_CONFIGS.copy()
        patch_styles.update(patch_style_overrides)
    else:
        patch_styles = WISE_DEFAULT_PATCH_CONFIGS

    # parse the WISE color-color SVG that contains the object class patches
    with importlib.resources.path(
        "vasttools.data", "WISE-color-color.svg"
    ) as svg_path:
        doc = minidom.parse(str(svg_path))
    # define the transform from the SVG frame to the plotting frame
    transform = (
        matplotlib.transforms.Affine2D().scale(sx=1, sy=-1).translate(-1, 4)
    )

    fig, ax = plt.subplots()
    # add WISE object classification patches
    for svg_path in doc.getElementsByTagName("path"):
        name = svg_path.getAttribute("id")
        patch_style = patch_styles[getattr(WiseClass, name)]
        path_mpl = parse_path(svg_path.getAttribute("d")).transformed(
            transform
        )
        patch = matplotlib.patches.PathPatch(path_mpl, **patch_style.style)
        ax.add_patch(patch)
        ax.annotate(
            patch_style.annotation_text,
            patch_style.annotation_position,
            ha="center",
            fontsize=annotation_text_size,
        )
    ax.set_xlim(-1, 6)
    ax.set_ylim(-0.5, 4)
    ax.set_aspect(1)
    ax.set_xlabel("[4.6] - [12] (mag)")
    ax.set_ylabel("[3.4] - [4.6] (mag)")
    ax.xaxis.set_major_locator(matplotlib.ticker.MultipleLocator(1))
    ax.yaxis.set_major_locator(matplotlib.ticker.MultipleLocator(1))
    return fig