utils.py

def _get_skyregion_relations(
    row: pd.Series,
    coords: SkyCoord,
    ids: pd.core.indexes.numeric.Int64Index
) -> List[int]:
    '''
    For each sky region row a list is returned that
    contains the ids of other sky regions that overlap
    with the row sky region (including itself).

    Args:
        row:
            A row from the dataframe containing all the sky regions of the run.
            Contains the 'id', 'centre_ra', 'centre_dec' and 'xtr_radius'
            columns.
        coords: A SkyCoord holding the coordinates of all sky regions.
        ids: The sky regions ids that match the coords.

    Returns:
        A list of other sky regions (including self) that are within the
        'xtr_radius' of the sky region in the row.
    '''
    target = SkyCoord(
        row['centre_ra'],
        row['centre_dec'],
        unit=(u.deg, u.deg)
    )

    seps = target.separation(coords)

    # place a slight buffer on the radius to make sure
    # any neighbouring fields are caught
    mask = seps <= row['xtr_radius'] * 1.1 * u.deg

    related_ids = ids[mask].to_list()

    return related_ids

def _load_measurements(
    image: Image,
    cols: List[str],
    start_id: int = 0,
    ini_df: bool = False
) -> pd.DataFrame:
    """
    Load the measurements for an image from the parquet file.

    Args:
        image:
            The object representing the image for which to load the
            measurements.
        cols:
            The columns to load.
        start_id:
            The number to start from when setting the source ids (when
            'ini_df' is 'True'). Defaults to 0.
        ini_df:
            Boolean to indicate whether these sources are part of the initial
            source list creation for association. If 'True' the source ids are
            reset ready for the first iteration. Defaults to 'False'.

    Returns:
        The measurements of the image with some extra values set ready for
        association .
    """
    image_centre = SkyCoord(
        image.ra,
        image.dec,
        unit=(u.deg, u.deg)
    )

    df = pd.read_parquet(image.measurements_path, columns=cols)
    df['image'] = image.name
    df['datetime'] = image.datetime
    # these are the first 'sources' if ini_df is True.
    df['source'] = df.index + start_id + 1 if ini_df else -1
    df['ra_source'] = df['ra']
    df['dec_source'] = df['dec']
    df['d2d'] = 0.
    df['dr'] = 0.
    df['related'] = None

    sources_sc = SkyCoord(df['ra'], df['dec'], unit=(u.deg, u.deg))

    seps = sources_sc.separation(image_centre).degree
    df['dist_from_centre'] = seps

    del sources_sc
    del seps

    return df

def add_new_many_to_one_relations(row: pd.Series) -> List[int]:
    """
    This handles the relation information being created from the
    many_to_one function in advanced association.
    It is a lot simpler than the one_to_many case as it purely just adds
    the new relations to the relation column, taking into account if it is
    already a list of relations or not (i.e. no previous relations).

    Args:
        row:
            The relation information Series from the association dataframe.
            Only the columns ['related_skyc1', 'new_relations'] are required.

    Returns:
        The new related field for the source in question, containing the
        appended ids.
    """
    out = row['new_relations'].copy()

    if isinstance(row['related_skyc1'], list):
        out += row['related_skyc1'].copy()

    return out

def add_new_one_to_many_relations(
    row: pd.Series, advanced: bool = False,
    source_ids: Optional[pd.DataFrame] = None
) -> List[int]:
    """
    This handles the relation information being created from the
    one_to_many functions in association.

    Args:
        row:
            The relation information Series from the association dataframe.
            Only the columns ['related_skyc1', 'source_skyc1'] are required
            for advanced, these are instead called ['related', 'source']
            for basic.
        advanced:
            Whether advanced association is being used which changes the names
            of the columns involved.
        source_ids:
            A dataframe that contains the other ids to append to related for
            each original source.
            +----------------+--------+
            |   source_skyc1 | 0      |
            |----------------+--------|
            |            122 | [5542] |
            |            254 | [5543] |
            |            262 | [5544] |
            |            405 | [5545] |
            |            656 | [5546] |
            +----------------+--------+

    Returns:
        The new related field for the source in question, containing the
        appended ids.
    """
    if source_ids is None:
        source_ids = pd.DataFrame()

    related_col = 'related_skyc1' if advanced else 'related'
    source_col = 'source_skyc1' if advanced else 'source'

    # this is the not_original case where the original source id is appended.
    if source_ids.empty:
        if isinstance(row[related_col], list):
            out = row[related_col]
            out.append(row[source_col])
        else:
            out = [row[source_col],]

    else:  # the original case to append all the new ids.
        source_ids = source_ids.loc[row[source_col]].iloc[0]
        if isinstance(row[related_col], list):
            out = row[related_col] + source_ids
        else:
            out = source_ids

    return out

def backup_parquets(p_run_path: str) -> None:
    """
    Backups up all the existing parquet files in a pipeline run directory.
    Backups are named with a '.bak' suffix in the pipeline run directory.

    Args:
        p_run_path:
            The path of the pipeline run where the parquets are stored.

    Returns:
        None
    """
    parquets = (
        glob.glob(os.path.join(p_run_path, "*.parquet"))
        # TODO Remove arrow when arrow files are no longer required.
        + glob.glob(os.path.join(p_run_path, "*.arrow")))

    for i in parquets:
        backup_name = i + '.bak'
        if os.path.isfile(backup_name):
            logger.debug(f'Removing old backup file: {backup_name}.')
            os.remove(backup_name)
        shutil.copyfile(i, backup_name)

def calc_ave_coord(grp: pd.DataFrame) -> pd.Series:
    """
    Calculates the average coordinate of the grouped by sources dataframe for
    each unique group, along with defining the image and epoch list for each
    unique source (group).

    Args:
        grp: The current group dataframe (unique source) of the grouped by
            dataframe being acted upon.

    Returns:
        A pandas series containing the average coordinate along with the
        image and epoch lists.
    """
    d = {}
    grp = grp.sort_values(by='datetime')
    d['img_list'] = grp['image'].values.tolist()
    d['epoch_list'] = grp['epoch'].values.tolist()
    d['wavg_ra'] = grp['interim_ew'].sum() / grp['weight_ew'].sum()
    d['wavg_dec'] = grp['interim_ns'].sum() / grp['weight_ns'].sum()

    return pd.Series(d)

def calculate_m_metric(flux_a: float, flux_b: float) -> float:
    """Calculate the m variability metric which is the modulation index between two fluxes.
    This is proportional to the fractional variability.
    See Section 5 of Mooley et al. (2016) for details, DOI: 10.3847/0004-637X/818/2/105.

    Args:
        flux_a (float): flux value "A".
        flux_b (float): flux value "B".

    Returns:
        float: the m metric for flux values "A" and "B".
    """
    return 2 * ((flux_a - flux_b) / (flux_a + flux_b))

def calculate_measurement_pair_metrics(df: pd.DataFrame) -> pd.DataFrame:
    """Generate a DataFrame of measurement pairs and their 2-epoch variability metrics
    from a DataFrame of measurements. For more information on the variability metrics, see
    Section 5 of Mooley et al. (2016), DOI: 10.3847/0004-637X/818/2/105.

    Args:
        df (pd.DataFrame): Input measurements. Must contain columns: id, source, flux_int,
            flux_int_err, flux_peak, flux_peak_err, has_siblings.

    Returns:
        pd.DataFrame: Measurement pairs and 2-epoch metrics. Will contain columns:
            source - the source ID
            id_a, id_b - the measurement IDs
            flux_int_a, flux_int_b - measurement integrated fluxes in mJy
            flux_int_err_a, flux_int_err_b - measurement integrated flux errors in mJy
            flux_peak_a, flux_peak_b - measurement peak fluxes in mJy/beam
            flux_peak_err_a, flux_peak_err_b - measurement peak flux errors in mJy/beam
            vs_peak, vs_int - variability t-statistic
            m_peak, m_int - variability modulation index
    """
    n_cpu = cpu_count() - 1

    """Create a DataFrame containing all measurement ID combinations per source.
    Resultant DataFrame will have a MultiIndex(["source", RangeIndex]) where "source" is
    the source ID and RangeIndex is an unnamed temporary ID for each measurement pair,
    unique only together with source.
    DataFrame will have columns [0, 1], each containing a measurement ID. e.g.
                       0      1
        source
        1       0      1   9284
                1      1  17597
                2      1  26984
                3   9284  17597
                4   9284  26984
        ...          ...    ...
        11105   2  11845  19961
        11124   0   3573  12929
                1   3573  21994
                2  12929  21994
        11128   0   6216  23534
    """
    measurement_combinations = (
        dd.from_pandas(df, n_cpu)
        .groupby("source")["id"]
        .apply(
            lambda x: pd.DataFrame(list(combinations(x, 2))
        ), meta={0: "i", 1: "i"},)
        .compute(num_workers=n_cpu, scheduler="processes")
    )

    """Drop the RangeIndex from the MultiIndex as it isn't required and rename the columns.
    Example resultant DataFrame:
               source   id_a   id_b
        0           1      1   9284
        1           1      1  17597
        2           1      1  26984
        3           1   9284  17597
        4           1   9284  26984
        ...       ...    ...    ...
        33640   11105  11845  19961
        33641   11124   3573  12929
        33642   11124   3573  21994
        33643   11124  12929  21994
        33644   11128   6216  23534
    Where source is the source ID, id_a and id_b are measurement IDs.
    """
    measurement_combinations = measurement_combinations.reset_index(
        level=1, drop=True
    ).rename(columns={0: "id_a", 1: "id_b"}).astype(int).reset_index()

    # Dask has a tendency to swap which order the measurement pairs are
    # defined in, even if the dataframe is pre-sorted. We want the pairs to be
    # in date order (a < b) so the code below corrects any that are not.
    measurement_combinations = measurement_combinations.join(
        df[['source', 'id', 'datetime']].set_index(['source', 'id']),
        on=['source', 'id_a'],
    )

    measurement_combinations = measurement_combinations.join(
        df[['source', 'id', 'datetime']].set_index(['source', 'id']),
        on=['source', 'id_b'], lsuffix='_a', rsuffix='_b'
    )

    to_correct_mask = (
        measurement_combinations['datetime_a']
        > measurement_combinations['datetime_b']
    )

    if np.any(to_correct_mask):
        logger.debug('Correcting measurement pairs order')
        (
            measurement_combinations.loc[to_correct_mask, 'id_a'],
            measurement_combinations.loc[to_correct_mask, 'id_b']
        ) = np.array([
            measurement_combinations.loc[to_correct_mask, 'id_b'].values,
            measurement_combinations.loc[to_correct_mask, 'id_a'].values
        ])

    measurement_combinations = measurement_combinations.drop(
        ['datetime_a', 'datetime_b'], axis=1
    )

    # add the measurement fluxes and errors
    association_fluxes = df.set_index(["source", "id"])[
        ["flux_int", "flux_int_err", "flux_peak", "flux_peak_err", "image"]
    ].rename(columns={"image": "image_name"})
    measurement_combinations = measurement_combinations.join(
        association_fluxes,
        on=["source", "id_a"],
    ).join(
        association_fluxes,
        on=["source", "id_b"],
        lsuffix="_a",
        rsuffix="_b",
    )

    # calculate 2-epoch metrics
    measurement_combinations["vs_peak"] = calculate_vs_metric(
        measurement_combinations.flux_peak_a,
        measurement_combinations.flux_peak_b,
        measurement_combinations.flux_peak_err_a,
        measurement_combinations.flux_peak_err_b,
    )
    measurement_combinations["vs_int"] = calculate_vs_metric(
        measurement_combinations.flux_int_a,
        measurement_combinations.flux_int_b,
        measurement_combinations.flux_int_err_a,
        measurement_combinations.flux_int_err_b,
    )
    measurement_combinations["m_peak"] = calculate_m_metric(
        measurement_combinations.flux_peak_a,
        measurement_combinations.flux_peak_b,
    )
    measurement_combinations["m_int"] = calculate_m_metric(
        measurement_combinations.flux_int_a,
        measurement_combinations.flux_int_b,
    )

    return measurement_combinations

def calculate_vs_metric(
    flux_a: float, flux_b: float, flux_err_a: float, flux_err_b: float
) -> float:
    """Calculate the Vs variability metric which is the t-statistic that the provided
    fluxes are variable. See Section 5 of Mooley et al. (2016) for details,
    DOI: 10.3847/0004-637X/818/2/105.

    Args:
        flux_a (float): flux value "A".
        flux_b (float): flux value "B".
        flux_err_a (float): error of `flux_a`.
        flux_err_b (float): error of `flux_b`.

    Returns:
        float: the Vs metric for flux values "A" and "B".
    """
    return (flux_a - flux_b) / np.hypot(flux_err_a, flux_err_b)

def check_primary_image(row: pd.Series) -> bool:
    """
    Checks whether the primary image of the ideal source
    dataframe is in the image list for the source.

    Args:
        row:
            Input dataframe row, with columns ['primary'] and ['img_list'].

    Returns:
        True if primary in image list else False.
    """
    return row['primary'] in row['img_list']

def create_measurement_pairs_arrow_file(p_run: Run) -> None:
    """
    Creates a measurement_pairs.arrow file using the parquet outputs
    of a pipeline run.

    Args:
        p_run:
            Pipeline model instance.

    Returns:
        None
    """
    logger.info('Creating measurement_pairs.arrow for run %s.', p_run.name)

    measurement_pairs_df = pd.read_parquet(
        os.path.join(
            p_run.path,
            'measurement_pairs.parquet'
        )
    )

    logger.debug('Optimising dataframe.')
    measurement_pairs_df = optimize_ints(optimize_floats(measurement_pairs_df))

    logger.debug("Loading to pyarrow table.")
    measurement_pairs_df = pa.Table.from_pandas(measurement_pairs_df)

    logger.debug("Exporting to arrow file.")
    outname = os.path.join(p_run.path, 'measurement_pairs.arrow')

    local = pa.fs.LocalFileSystem()

    with local.open_output_stream(outname) as file:
       with pa.RecordBatchFileWriter(file, measurement_pairs_df.schema) as writer:
          writer.write_table(measurement_pairs_df)

def create_measurements_arrow_file(p_run: Run) -> None:
    """
    Creates a measurements.arrow file using the parquet outputs
    of a pipeline run.

    Args:
        p_run:
            Pipeline model instance.

    Returns:
        None
    """
    logger.info('Creating measurements.arrow for run %s.', p_run.name)

    associations = pd.read_parquet(
        os.path.join(
            p_run.path,
            'associations.parquet'
        )
    )
    images = pd.read_parquet(
        os.path.join(
            p_run.path,
            'images.parquet'
        )
    )

    m_files = images['measurements_path'].tolist()

    m_files += glob.glob(os.path.join(
        p_run.path,
        'forced*.parquet'
    ))

    logger.debug('Loading %i files...', len(m_files))
    measurements = dd.read_parquet(m_files, engine='pyarrow').compute()

    measurements = measurements.loc[
        measurements['id'].isin(associations['meas_id'].values)
    ]

    measurements = (
        associations.loc[:, ['meas_id', 'source_id']]
        .set_index('meas_id')
        .merge(
            measurements,
            left_index=True,
            right_on='id'
        )
        .rename(columns={'source_id': 'source'})
    )

    logger.debug('Optimising dataframes.')
    measurements = optimize_ints(optimize_floats(measurements))

    logger.debug("Loading to pyarrow table.")
    measurements = pa.Table.from_pandas(measurements)

    logger.debug("Exporting to arrow file.")
    outname = os.path.join(p_run.path, 'measurements.arrow')

    local = pa.fs.LocalFileSystem()

    with local.open_output_stream(outname) as file:
       with pa.RecordBatchFileWriter(file, measurements.schema) as writer:
          writer.write_table(measurements)

def cross_join(left: pd.DataFrame, right: pd.DataFrame) -> pd.DataFrame:
    """
    A convenience function to merge two dataframes.

    Args:
        left: The base pandas DataFrame to merge.
        right: The pandas DataFrame to merge to the left.

    Returns:
        The resultant merged DataFrame.

    """
    return (
        left.assign(key=1)
        .merge(right.assign(key=1), on='key')
        .drop('key', axis=1)
    )

def get_create_img(
    p_run: Run, band_id: int, image: Image
) -> Tuple[Image, SkyRegion, bool]:
    """
    Function to fetch or create the Image and Sky Region objects for the images
    in the pipeline run.

    Args:
        p_run: The pipeline run Django ORM object.
        band_id: The integer database id value of the frequency band of the
            image.
        image: The image Django ORM object.

    Returns:
        The resulting image django ORM object, the sky region Django ORM
        object and a bool value denoting if the image already existed in the
        database.
    """
    img = Image.objects.filter(name__exact=image.name)
    if img.exists():
        img = img.get()
        # Add background path if not originally provided
        if image.background_path and not img.background_path:
            img.background_path = image.background_path
            img.save()
        skyreg = get_create_skyreg(p_run, img)
        # check and add the many to many if not existent
        if not Image.objects.filter(
            id=img.id, run__id=p_run.id
        ).exists():
            img.run.add(p_run)

        return (img, skyreg, True)

    # at this stage, measurement parquet file is not created but
    # assume location
    img_folder_name = image.name.replace('.', '_')
    measurements_path = os.path.join(
        settings.PIPELINE_WORKING_DIR,
        'images',
        img_folder_name,
        'measurements.parquet'
        )
    img = Image(
        band_id=band_id,
        measurements_path=measurements_path
    )
    # set the attributes and save the image,
    # by selecting only valid (not hidden) attributes
    # FYI attributs and/or method starting with _ are hidden
    # and with __ can't be modified/called
    for fld in img._meta.get_fields():
        if getattr(fld, 'attname', None) and (
            getattr(image, fld.attname, None) is not None
        ):
            setattr(img, fld.attname, getattr(image, fld.attname))

    # get create the sky region and associate with image
    skyreg = get_create_skyreg(p_run, img)
    img.skyreg = skyreg

    img.rms_median, img.rms_min, img.rms_max = get_rms_noise_image_values(
        img.noise_path
    )

    img.save()
    img.run.add(p_run)

    return (img, skyreg, False)

def get_create_img_band(image: Image) -> Band:
    '''
    Return the existing Band row for the given FitsImage.
    An image is considered to belong to a band if its frequency is within some
    tolerance of the band's frequency.
    Returns a Band row or None if no matching band.

    Args:
        image: The image Django ORM object.

    Returns:
        The band Django ORM object.
    '''
    # For now we match bands using the central frequency.
    # This assumes that every band has a unique frequency,
    # which is true for the data we've used so far.
    freq = int(image.freq_eff * 1.e-6)
    freq_band = int(image.freq_bw * 1.e-6)
    # TODO: refine the band query
    for band in Band.objects.all():
        diff = abs(freq - band.frequency) / float(band.frequency)
        if diff < 0.02:
            return band

    # no band has been found so create it
    band = Band(name=str(freq), frequency=freq, bandwidth=freq_band)
    logger.info('Adding new frequency band: %s', band)
    band.save()

    return band

def get_create_p_run(
    name: str, path: str, description: str=None, user: User=None
) -> Tuple[Run, bool]:
    '''
    Get or create a pipeline run in db, return the run django object and
    a flag True/False if has been created or already exists.

    Args:
        name: The name of the pipeline run.
        path: The system path to the pipeline run folder which contains the
            configuration file and where outputs will be saved.
        description: An optional description of the pipeline run.
        user: The Django user that launched the pipeline run.

    Returns:
        The pipeline run object and a boolean object representing whether
        the pipeline run already existed ('True') or not ('False').
    '''
    p_run = Run.objects.filter(name__exact=name)
    if p_run:
        return p_run.get(), True

    description = "" if description is None else description
    p_run = Run(name=name, description=description, path=path)
    if user:
        p_run.user = user
    p_run.save()

    return p_run, False

def get_create_skyreg(p_run: Run, image: Image) -> SkyRegion:
    '''
    This creates a Sky Region object in Django ORM given the related
    image object.

    Args:
        p_run: The pipeline run Django ORM object.
        image: The image Django ORM object.

    Returns:
        The sky region Django ORM object.
    '''
    # In the calculations below, it is assumed the image has square
    # pixels (this pipeline has been designed for ASKAP images, so it
    # should always be square). It will likely give wrong results if not
    skyr = SkyRegion.objects.filter(
        centre_ra=image.ra,
        centre_dec=image.dec,
        xtr_radius=image.fov_bmin
    )
    if skyr:
        skyr = skyr.get()
        logger.info('Found sky region %s', skyr)
        if p_run not in skyr.run.all():
            logger.info('Adding %s to sky region %s', p_run, skyr)
            skyr.run.add(p_run)
        return skyr

    x, y, z = eq_to_cart(image.ra, image.dec)
    skyr = SkyRegion(
        centre_ra=image.ra,
        centre_dec=image.dec,
        width_ra=image.physical_bmin,
        width_dec=image.physical_bmaj,
        xtr_radius=image.fov_bmin,
        x=x,
        y=y,
        z=z,
    )
    skyr.save()
    logger.info('Created sky region %s', skyr)
    skyr.run.add(p_run)
    logger.info('Adding %s to sky region %s', p_run, skyr)

    return skyr

def get_eta_metric(
    row: Dict[str, float], df: pd.DataFrame, peak: bool=False
) -> float:
    '''
    Calculates the eta variability metric of a source.
    Works on the grouped by dataframe using the fluxes
    of the associated measurements.

    Args:
        row: Dictionary containing statistics for the current source.
        df: The grouped by sources dataframe of the measurements containing all
            the flux and flux error information,
        peak: Whether to use peak_flux for the calculation. If False then the
            integrated flux is used.

    Returns:
        The calculated eta value.

    '''
    if row['n_meas'] == 1:
        return 0.

    suffix = 'peak' if peak else 'int'
    weights = 1. / df[f'flux_{suffix}_err'].values**2
    fluxes = df[f'flux_{suffix}'].values
    eta = (row['n_meas'] / (row['n_meas']-1)) * (
        (weights * fluxes**2).mean() - (
            (weights * fluxes).mean()**2 / weights.mean()
        )
    )
    return eta

def get_image_list_diff(row: pd.Series) -> List[str]:
    """
    Calculate the difference between the ideal coverage image list of a source
    and the actual observed image list. Also checks whether an epoch does in
    fact contain a detection but is not in the expected 'ideal' image for that
    epoch.

    Args:
        row: The row from the sources dataframe that is being iterated over.

    Returns:
        A list of the images missing from the observed image list. Will be
        returned as '-1' integer value if there are no missing images.
    """
    out = list(
        filter(lambda arg: arg not in row['img_list'], row['skyreg_img_list'])
    )

    # set empty list to -1
    if not out:
        return -1

    # Check that an epoch has not already been seen (just not in the 'ideal'
    # image)
    out_epochs = [
        row['skyreg_epoch'][pair[0]] for pair in enumerate(
            row['skyreg_img_list']
        ) if pair[1] in out
    ]

    out = [
        out[pair[0]] for pair in enumerate(
            out_epochs
        ) if pair[1] not in row['epoch_list']
    ]

    if not out:
        out = -1

    return out

def get_names_and_epochs(grp: pd.DataFrame) -> pd.Series:
    """
    Convenience function to group together the image names, epochs and
    datetimes into one list object which is then returned as a pandas series.
    This is necessary for easier processing in the ideal coverage analysis.

    Args:
        grp: A group from the grouped by sources DataFrame.

    Returns:
        Pandas series containing the list object that contains the lists of the
        image names, epochs and datetimes.
    """
    d = {}
    d['skyreg_img_epoch_list'] = [[[x,], y, z] for x,y,z in zip(
        grp['name'].values.tolist(),
        grp['epoch'].values.tolist(),
        grp['datetime'].values.tolist()
    )]

    return pd.Series(d)

def get_parallel_assoc_image_df(
    images: List[Image], skyregion_groups: pd.DataFrame
) -> pd.DataFrame:
    """
    Merge the sky region groups with the images and skyreg_ids.

    Args:
        images:
            A list of the Image objects.
        skyregion_groups:
            The sky region group of each skyregion id.
            +----+----------------+
            |    |   skyreg_group |
            |----+----------------|
            |  2 |              1 |
            |  3 |              1 |
            |  1 |              2 |
            +----+----------------+

    Returns:
        Dataframe containing the merged images and skyreg_id and skyreg_group.
        +----+-------------------------------+-------------+----------------+
        |    | image                         |   skyreg_id |   skyreg_group |
        |----+-------------------------------+-------------+----------------|
        |  0 | VAST_2118+00A.EPOCH01.I.fits  |           2 |              1 |
        |  1 | VAST_2118-06A.EPOCH01.I.fits  |           3 |              1 |
        |  2 | VAST_0127-73A.EPOCH01.I.fits  |           1 |              2 |
        |  3 | VAST_2118-06A.EPOCH03x.I.fits |           3 |              1 |
        |  4 | VAST_2118-06A.EPOCH02.I.fits  |           3 |              1 |
        |  5 | VAST_2118-06A.EPOCH05x.I.fits |           3 |              1 |
        |  6 | VAST_2118-06A.EPOCH06x.I.fits |           3 |              1 |
        |  7 | VAST_0127-73A.EPOCH08.I.fits  |           1 |              2 |
        +----+-------------------------------+-------------+----------------+
    """
    skyreg_ids = [i.skyreg_id for i in images]

    images_df = pd.DataFrame({
        'image_dj': images,
        'skyreg_id': skyreg_ids,
    })

    images_df = images_df.merge(
        skyregion_groups,
        how='left',
        left_on='skyreg_id',
        right_index=True
    )

    images_df['image_name'] = images_df['image_dj'].apply(
        lambda x: x.name
    )

    images_df['image_datetime'] = images_df['image_dj'].apply(
        lambda x: x.datetime
    )

    return images_df

def get_rms_noise_image_values(rms_path: str) -> Tuple[float, float, float]:
    '''
    Open the RMS noise FITS file and compute the median, max and min
    rms values to be added to the image model and then used in the
    calculations.

    Args:
        rms_path: The system path to the RMS FITS image.

    Returns:
        The median, minimum and maximum values of the RMS image.

    Raises:
        IOError: Raised when the RMS FITS file cannot be found.
    '''
    logger.debug('Extracting Image RMS values from Noise file...')
    med_val = min_val = max_val = 0.
    try:
        with fits.open(rms_path) as f:
            data = f[0].data
            data = data[np.logical_not(np.isnan(data))]
            data = data[data != 0]
            med_val = np.median(data) * 1e+3
            min_val = np.min(data) * 1e+3
            max_val = np.max(data) * 1e+3
            del data
    except Exception:
        raise IOError(f'Could not read this RMS FITS file: {rms_path}')

    return med_val, min_val, max_val

def get_src_skyregion_merged_df(
    sources_df: pd.DataFrame, images_df: pd.DataFrame, skyreg_df: pd.DataFrame
) -> pd.DataFrame:
    """
    Analyses the current sources_df to determine what the 'ideal coverage'
    for each source should be. In other words, what images is the source
    missing in when it should have been seen.

    Args:
        sources_df:
            The output of the association  step containing the
            measurements associated into sources.
        images_df:
            Contains the images of the pipeline run. I.e. all image
            objects for the run loaded into a dataframe.
        skyreg_df:
            Contains the sky regions of the pipeline run. I.e. all
            sky region objects for the run loaded into a dataframe.

    Returns:
        DataFrame containing missing image information. Output format:
        +----------+----------------------------------+-----------+------------+
        |   source | img_list                         |   wavg_ra |   wavg_dec |
        |----------+----------------------------------+-----------+------------+
        |      278 | ['VAST_0127-73A.EPOCH01.I.fits'] |  22.2929  |   -71.8717 |
        |      702 | ['VAST_0127-73A.EPOCH01.I.fits'] |  28.8125  |   -69.3547 |
        |      844 | ['VAST_0127-73A.EPOCH01.I.fits'] |  17.3152  |   -72.346  |
        |      934 | ['VAST_0127-73A.EPOCH01.I.fits'] |   9.75754 |   -72.9629 |
        |     1290 | ['VAST_0127-73A.EPOCH01.I.fits'] |  20.8455  |   -76.8269 |
        +----------+----------------------------------+-----------+------------+
        ------------------------------------------------------------------+
         skyreg_img_list                                                  |
        ------------------------------------------------------------------+
         ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
         ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
         ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
         ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
         ['VAST_0127-73A.EPOCH01.I.fits', 'VAST_0127-73A.EPOCH08.I.fits'] |
        ------------------------------------------------------------------+
        ----------------------------------+------------------------------+
         img_diff                         | primary                      |
        ----------------------------------+------------------------------+
         ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
         ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
         ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
         ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
         ['VAST_0127-73A.EPOCH08.I.fits'] | VAST_0127-73A.EPOCH01.I.fits |
        ----------------------------------+------------------------------+
        ------------------------------+--------------+
         detection                    | in_primary   |
        ------------------------------+--------------|
         VAST_0127-73A.EPOCH01.I.fits | True         |
         VAST_0127-73A.EPOCH01.I.fits | True         |
         VAST_0127-73A.EPOCH01.I.fits | True         |
         VAST_0127-73A.EPOCH01.I.fits | True         |
         VAST_0127-73A.EPOCH01.I.fits | True         |
        ------------------------------+--------------+
    """
    logger.info("Creating ideal source coverage df...")

    merged_timer = StopWatch()

    skyreg_df = skyreg_df.drop(
        ['x', 'y', 'z', 'width_ra', 'width_dec'], axis=1
    )

    images_df['name'] = images_df['image_dj'].apply(
        lambda x: x.name
    )
    images_df['datetime'] = images_df['image_dj'].apply(
        lambda x: x.datetime
    )

    skyreg_df = skyreg_df.join(
        pd.DataFrame(
            images_df.groupby('skyreg_id').apply(
                get_names_and_epochs
            )
        ),
        on='id'
    )

    sources_df = sources_df.sort_values(by='datetime')
    # calculate some metrics on sources
    # compute only some necessary metrics in the groupby
    timer = StopWatch()
    srcs_df = parallel_groupby_coord(sources_df)
    logger.debug('Groupby-apply time: %.2f seconds', timer.reset())

    del sources_df

    # create dataframe with all skyregions and sources combinations
    src_skyrg_df = cross_join(
        srcs_df.drop(['epoch_list', 'img_list'], axis=1).reset_index(),
        skyreg_df.drop('skyreg_img_epoch_list', axis=1)
    )

    skyreg_df = skyreg_df.drop(
        ['centre_ra', 'centre_dec', 'xtr_radius'],
        axis=1
    ).set_index('id')

    src_skyrg_df['sep'] = np.rad2deg(
        on_sky_sep(
            np.deg2rad(src_skyrg_df['wavg_ra'].values),
            np.deg2rad(src_skyrg_df['centre_ra'].values),
            np.deg2rad(src_skyrg_df['wavg_dec'].values),
            np.deg2rad(src_skyrg_df['centre_dec'].values),
        )
    )

    # select rows where separation is less than sky region radius
    # drop not more useful columns and groupby source id
    # compute list of images
    src_skyrg_df = (
        src_skyrg_df.loc[
            src_skyrg_df.sep < src_skyrg_df.xtr_radius,
            ['source', 'id', 'sep']
        ].merge(skyreg_df, left_on='id', right_index=True)
        .drop('id', axis=1)
        .explode('skyreg_img_epoch_list')
    )

    del skyreg_df

    src_skyrg_df[
        ['skyreg_img_list', 'skyreg_epoch', 'skyreg_datetime']
    ] = pd.DataFrame(
        src_skyrg_df['skyreg_img_epoch_list'].tolist(),
        index=src_skyrg_df.index
    )

    src_skyrg_df = src_skyrg_df.drop('skyreg_img_epoch_list', axis=1)

    src_skyrg_df = (
        src_skyrg_df.sort_values(
            ['source', 'sep']
        )
        .drop_duplicates(['source', 'skyreg_epoch'])
        .sort_values(by='skyreg_datetime')
        .drop(
            ['sep', 'skyreg_datetime'],
            axis=1
        )
    )
    # annoyingly epoch needs to be not a list to drop duplicates
    # but then we need to sum the epochs into a list
    src_skyrg_df['skyreg_epoch'] = src_skyrg_df['skyreg_epoch'].apply(
        lambda x: [x,]
    )

    src_skyrg_df = (
        src_skyrg_df.groupby('source')
        .agg('sum') # sum because we need to preserve order
    )

    # merge into main df and compare the images
    srcs_df = srcs_df.merge(
        src_skyrg_df, left_index=True, right_index=True
    )

    del src_skyrg_df

    srcs_df['img_diff'] = srcs_df[
        ['img_list', 'skyreg_img_list', 'epoch_list', 'skyreg_epoch']
    ].apply(
        get_image_list_diff, axis=1
    )

    srcs_df = srcs_df.loc[
        srcs_df['img_diff'] != -1
    ]

    srcs_df = srcs_df.drop(
        ['epoch_list', 'skyreg_epoch'],
        axis=1
    )

    srcs_df['primary'] = srcs_df[
        'skyreg_img_list'
    ].apply(lambda x: x[0])

    srcs_df['detection'] = srcs_df[
        'img_list'
    ].apply(lambda x: x[0])

    srcs_df['in_primary'] = srcs_df[
        ['primary', 'img_list']
    ].apply(
        check_primary_image,
        axis=1
    )

    srcs_df = srcs_df.drop(['img_list', 'skyreg_img_list', 'primary'], axis=1)

    logger.info(
        'Ideal source coverage time: %.2f seconds', merged_timer.reset()
    )

    return srcs_df

def group_skyregions(df: pd.DataFrame) -> pd.DataFrame:
    """
    Logic to group sky regions into overlapping groups.
    Returns a dataframe containing the sky region id as
    the index and a column containing a list of the
    sky region group number it belongs to.

    Args:
        df:
            A dataframe containing all the sky regions of the run. Only the
            'id', 'centre_ra', 'centre_dec' and 'xtr_radius' columns are
            required.
            +------+-------------+--------------+--------------+
            |   id |   centre_ra |   centre_dec |   xtr_radius |
            |------+-------------+--------------+--------------|
            |    2 |    319.652  |    0.0030765 |      6.72488 |
            |    3 |    319.652  |   -6.2989    |      6.7401  |
            |    1 |     21.8361 |  -73.121     |      7.24662 |
            +------+-------------+--------------+--------------+

    Returns:
        The sky region group of each skyregion id.
        +----+----------------+
        |    |   skyreg_group |
        |----+----------------|
        |  2 |              1 |
        |  3 |              1 |
        |  1 |              2 |
        +----+----------------+
    """
    sr_coords = SkyCoord(
        df['centre_ra'],
        df['centre_dec'],
        unit=(u.deg, u.deg)
    )

    df = df.set_index('id')

    results = df.apply(
        _get_skyregion_relations,
        args=(sr_coords, df.index),
        axis=1
    )

    skyreg_groups = {}

    master_done = []  # keep track of all checked ids in master done

    for skyreg_id, neighbours in results.iteritems():

        if skyreg_id not in master_done:
            local_done = []   # a local done list for the sky region group.
            # add the current skyreg_id to both master and local done.
            master_done.append(skyreg_id)
            local_done.append(skyreg_id)
            # Define the new group number based on the existing ones.
            skyreg_group = len(skyreg_groups) + 1
            # Add all the ones that we know are neighbours that were obtained
            # from _get_skyregion_relations.
            skyreg_groups[skyreg_group] = list(neighbours)

            # Now the sky region group is extended out to include all those sky
            # regions that overlap with the neighbours.
            # Each neighbour is checked and added to the local done list.
            # Checked means that for each neighbour, it's own neighbours are
            # added to the current group if not in already.
            # When the local done is equal to the skyreg group we know that
            # we have exhausted all possible neighbours and that results in a
            # sky region group.
            while sorted(local_done) != sorted(skyreg_groups[skyreg_group]):
                # Loop over each neighbour
                for other_skyreg_id in skyreg_groups[skyreg_group]:
                    # If we haven't checked this neighbour locally proceed.
                    if other_skyreg_id not in local_done:
                        # Add it to the local checked.
                        local_done.append(other_skyreg_id)
                        # Get the neighbours neighbour and add these.
                        new_vals = results.loc[other_skyreg_id]
                        for k in new_vals:
                            if k not in skyreg_groups[skyreg_group]:
                                skyreg_groups[skyreg_group].append(k)

            # Reached the end of the group so append all to the master
            # done list
            for j in skyreg_groups[skyreg_group]:
                master_done.append(j)
        else:
            # continue if already placed in group
            continue

    # flip the dictionary around
    skyreg_group_ids = {}
    for i in skyreg_groups:
        for j in skyreg_groups[i]:
            skyreg_group_ids[j]=i

    skyreg_group_ids = pd.DataFrame.from_dict(
        skyreg_group_ids, orient='index'
    ).rename(columns={0: 'skyreg_group'})

    return skyreg_group_ids

def groupby_funcs(df: pd.DataFrame) -> pd.Series:
    '''
    Performs calculations on the unique sources to get the
    lightcurve properties. Works on the grouped by source
    dataframe.

    Args:
        df: The current iteration dataframe of the grouped by sources
            dataframe.

    Returns:
        Pandas series containing the calculated metrics of the source.

    '''
    # calculated average ra, dec, fluxes and metrics
    d = {}
    d['img_list'] = df['image'].values.tolist()
    d['n_meas_forced'] = df['forced'].sum()
    d['n_meas'] = df['id'].count()
    d['n_meas_sel'] = d['n_meas'] - d['n_meas_forced']
    d['n_sibl'] = df['has_siblings'].sum()
    if d['n_meas_forced'] > 0:
        non_forced_sel = df['forced'] != True
        d['wavg_ra'] = (
            df.loc[non_forced_sel, 'interim_ew'].sum() /
            df.loc[non_forced_sel, 'weight_ew'].sum()
        )
        d['wavg_dec'] = (
            df.loc[non_forced_sel, 'interim_ns'].sum() /
            df.loc[non_forced_sel, 'weight_ns'].sum()
        )
        d['avg_compactness'] = df.loc[
            non_forced_sel, 'compactness'
        ].mean()
        d['min_snr'] = df.loc[
            non_forced_sel, 'snr'
        ].min()
        d['max_snr'] = df.loc[
            non_forced_sel, 'snr'
        ].max()

    else:
        d['wavg_ra'] = df['interim_ew'].sum() / df['weight_ew'].sum()
        d['wavg_dec'] = df['interim_ns'].sum() / df['weight_ns'].sum()
        d['avg_compactness'] = df['compactness'].mean()
        d['min_snr'] = df['snr'].min()
        d['max_snr'] = df['snr'].max()

    d['wavg_uncertainty_ew'] = 1. / np.sqrt(df['weight_ew'].sum())
    d['wavg_uncertainty_ns'] = 1. / np.sqrt(df['weight_ns'].sum())
    for col in ['avg_flux_int', 'avg_flux_peak']:
        d[col] = df[col.split('_', 1)[1]].mean()
    for col in ['max_flux_peak', 'max_flux_int']:
        d[col] = df[col.split('_', 1)[1]].max()
    for col in ['min_flux_peak', 'min_flux_int']:
        d[col] = df[col.split('_', 1)[1]].min()
    for col in ['min_flux_peak_isl_ratio', 'min_flux_int_isl_ratio']:
        d[col] = df[col.split('_', 1)[1]].min()

    for col in ['flux_int', 'flux_peak']:
        d[f'{col}_sq'] = (df[col]**2).mean()
    d['v_int'] = df['flux_int'].std() / df['flux_int'].mean()
    d['v_peak'] = df['flux_peak'].std() / df['flux_peak'].mean()
    d['eta_int'] = get_eta_metric(d, df)
    d['eta_peak'] = get_eta_metric(d, df, peak=True)
    # remove not used cols
    for col in ['flux_int_sq', 'flux_peak_sq']:
        d.pop(col)

    # get unique related sources
    list_uniq_related = list(set(
        chain.from_iterable(
            lst for lst in df['related'] if isinstance(lst, list)
        )
    ))
    d['related_list'] = list_uniq_related if list_uniq_related else -1

    return pd.Series(d).fillna(value={"v_int": 0.0, "v_peak": 0.0})