Source code for pyxel.util.timing

#  Copyright (c) European Space Agency, 2020.
#
#  This file is subject to the terms and conditions defined in file 'LICENCE.txt', which
#  is part of this Pyxel package. No part of the package, including
#  this file, may be copied, modified, propagated, or distributed except according to
#  the terms contained in the file ‘LICENCE.txt’.


"""Timing functions."""

import timeit
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    import pandas as pd

    from pyxel.detectors import Detector
    from pyxel.pipelines import DetectionPipeline, ModelFunction, ModelGroup


def _list_of_times_to_dataframe(times: list, model_names: list) -> "pd.DataFrame":
    """Convert a list of model times to a dataframe.

    Parameters
    ----------
    times: list
    model_names: list

    Returns
    -------
    df: pd.DataFrame
    """
    # Late import to speedup start-up time
    import pandas as pd

    times_array = np.array(times) * 1000  # to milliseconds
    times_sum = np.sum(times_array)
    percentages = times_array * 100 / times_sum

    final_times = np.round(np.append(times_array, times_sum), 2)
    percentages = np.round(np.append(percentages, 100.0), 4)

    model_names.append("TOTAL:")

    df = pd.DataFrame({"time [ms]": final_times, "%time": percentages})
    df.index = model_names

    return df




def time_pipeline(
    detector: "Detector", pipeline: "DetectionPipeline"
) -> "pd.DataFrame":
    """Time a single pipeline.

    Parameters
    ----------
    detector : Detector
    pipeline : DetectionPipeline

    Returns
    -------
    df: DataFrame
    """

    times = []
    model_names = []

    for group_name in pipeline.model_group_names:
        models_grp: ModelGroup | None = getattr(pipeline, group_name)

        if models_grp:
            model: ModelFunction
            for model in models_grp.models:
                if model.enabled:
                    model_start = timeit.default_timer()
                    model(detector)
                    model_end = timeit.default_timer()
                    times.append(model_end - model_start)
                    model_names.append(model.name)

    df = _list_of_times_to_dataframe(times, model_names)

    return df