Metrics¶

def mosqito_metric_1ch(
    s: Signal,
    metric: Literal[
        "loudness_zwtv",
        "roughness_dw",
        "sharpness_din_from_loudness",
        "sharpness_din_perseg",
        "sharpness_din_tv",
    ],
    statistics: tuple[int | str, ...] = (
        5,
        10,
        50,
        90,
        95,
        "avg",
        "max",
        "min",
        "kurt",
        "skew",
    ),
    label: str | None = None,
    *,
    as_df: bool = False,
    return_time_series: bool = False,
    **kwargs: Unpack[_MosqitoMetricParams],
) -> dict | pd.DataFrame:
    """
    Calculate a MoSQITo psychoacoustic metric for a single channel signal.

    Parameters
    ----------
    s
        Single channel signal object to analyze.
    metric
        Name of the metric to calculate. Options are "loudness_zwtv",
        "roughness_dw", "sharpness_din_from_loudness", "sharpness_din_perseg",
        or "sharpness_din_tv".
    statistics
        Statistics to calculate on the metric results.
    label
        Label to use for the metric in the results. If None, uses a default label.
    as_df
        If True, return results as a pandas DataFrame. Otherwise, return a dictionary.
    return_time_series
        If True, include the full time series in the results.
    **kwargs
        Additional arguments to pass to the underlying MoSQITo function.

    Returns
    -------
    :
        Results of the metric calculation and statistics.

    Raises
    ------
    ValueError
        If the input signal is not single-channel
        or if an unrecognized metric is specified.

    Examples
    --------
    >>> # doctest: +SKIP
    >>> from soundscapy.audio import Binaural
    >>> signal = Binaural.from_wav("audio.wav", resample=480000)
    >>> results = mosqito_metric_1ch(signal[0], "loudness_zwtv", as_df=True)

    """
    logger.debug(f"Calculating MoSQITo metric: {metric}")

    # Checks and warnings
    if s.channels != 1:
        logger.error("Signal must be single channel")
        msg = "Signal must be single channel"
        raise ValueError(msg)
    try:
        label = label or DEFAULT_LABELS[metric]
    except KeyError as e:
        logger.error(f"Metric {metric} not recognized")
        msg = f"Metric {metric} not recognized."
        raise ValueError(msg) from e
    if as_df and return_time_series:
        logger.warning(
            "Cannot return both a dataframe and time series. Returning dataframe only."
        )
        return_time_series = False

    # Start the calc
    res = {}
    try:
        if metric == "loudness_zwtv":
            # Prepare args specifically for loudness_zwtv
            loudness_args = {}
            if "field_type" in kwargs:
                loudness_args["field_type"] = kwargs["field_type"]
            # Call with filtered args
            N, N_spec, _, time_axis = loudness_zwtv(s, s.fs, **loudness_args)  # noqa: N806
            # TODO(MitchellAcoustics): Add the bark_axis back in
            # when we implement time series calcs
            # https://github.com/MitchellAcoustics/Soundscapy/issues/113
            res = _stat_calcs(label, N, res, statistics)
            if return_time_series:
                res[f"{label}_ts"] = (time_axis, N)

        elif metric == "roughness_dw":
            # Prepare args specifically for roughness_dw
            roughness_args = {}
            if "overlap" in kwargs:
                roughness_args["overlap"] = kwargs["overlap"]
            # Call with filtered args
            R, _, _, time_axis = roughness_dw(s, s.fs, **roughness_args)  # noqa: N806
            # TODO(MitchellAcoustics): Add the R_spec and bark_axis back in
            # when we implement time series calcs
            # https://github.com/MitchellAcoustics/Soundscapy/issues/113
            if isinstance(R, float | int):
                res[label] = R
            elif isinstance(R, np.ndarray) and len(R) == 1:
                res[label] = R[0]
            else:
                res = _stat_calcs(label, R, res, statistics)
            if return_time_series:
                res[f"{label}_ts"] = (time_axis, R)

        elif metric == "sharpness_din_from_loudness":
            # Prepare args for loudness_zwtv (needed first)
            loudness_args = {}
            if "field_type" in kwargs:
                loudness_args["field_type"] = kwargs["field_type"]
            N, N_spec, _, time_axis = loudness_zwtv(s, s.fs, **loudness_args)  # noqa: N806
            # TODO(MitchellAcoustics): Add the R_spec and bark_axis back in
            # when we implement time series calcs
            # https://github.com/MitchellAcoustics/Soundscapy/issues/113
            res = _stat_calcs("N", N, res, statistics)
            if return_time_series:
                res["N_ts"] = time_axis, N

            # Prepare args specifically for sharpness_din_from_loudness
            sharpness_args = {}
            if "weighting" in kwargs:
                sharpness_args["weighting"] = kwargs["weighting"]
            # Call with filtered args
            S = sharpness_din_from_loudness(N, N_spec, **sharpness_args)  # noqa: N806
            res = _stat_calcs(label, S, res, statistics)
            if return_time_series:
                res[f"{label}_ts"] = (time_axis, S)

        elif metric == "sharpness_din_perseg":
            # Prepare args specifically for sharpness_din_perseg
            sharpness_args = {}
            if "weighting" in kwargs:
                sharpness_args["weighting"] = kwargs["weighting"]
            if "nperseg" in kwargs:
                sharpness_args["nperseg"] = kwargs["nperseg"]
            if "noverlap" in kwargs:
                sharpness_args["noverlap"] = kwargs["noverlap"]
            # Call with filtered args
            S, time_axis = sharpness_din_perseg(s, s.fs, **sharpness_args)  # noqa: N806
            res = _stat_calcs(label, S, res, statistics)
            if return_time_series:
                res[f"{label}_ts"] = (time_axis, S)

        elif metric == "sharpness_din_tv":
            # Prepare args specifically for sharpness_din_tv
            sharpness_args = {}
            if "weighting" in kwargs:
                sharpness_args["weighting"] = kwargs["weighting"]
            if "skip" in kwargs:
                sharpness_args["skip"] = kwargs["skip"]
            # Call with filtered args
            S, time_axis = sharpness_din_tv(s, s.fs, **sharpness_args)  # noqa: N806
            res = _stat_calcs(label, S, res, statistics)
            if return_time_series:
                res[f"{label}_ts"] = (time_axis, S)
        else:
            msg = f"Metric {metric} not recognized."
            logger.error(msg)
            raise ValueError(msg)  # noqa: TRY301
    except Exception as e:
        logger.error(f"Error calculating {metric}: {e!s}")
        raise

    # Return the results in the requested format
    if not as_df:
        return res

    rec = getattr(s, "recording", None)
    return pd.DataFrame(res, index=[rec])

def maad_metric_1ch(
    s: Signal | Binaural,
    metric: Literal["all_temporal_alpha_indices", "all_spectral_alpha_indices"],
    as_df: bool = False,
    func_args: dict | None = None,
) -> Any:
    """
    Run a metric from the scikit-maad library (or suite of indices) on a single channel.

    Currently only supports running all the alpha indices at once.

    Parameters
    ----------
    s
        Single channel signal to calculate the alpha indices for.
    metric
        Metric to calculate.
    as_df
        Whether to return a pandas DataFrame, by default False.
        If True, returns a MultiIndex Dataframe
        with ("Recording", "Channel") as the index.
    func_args
        Additional keyword arguments to pass to the metric function, by default {}.

    Returns
    -------
    :
        Dictionary of results if as_df is False, otherwise a pandas DataFrame.

    Raises
    ------
    ValueError
        If the signal is not single-channel or if an unrecognized metric is specified.

    See Also
    --------
    maad.features.all_spectral_alpha_indices
    maad.features.all_temporal_alpha_indices

    """
    logger.debug(f"Calculating MAAD metric: {metric}")

    if func_args is None:
        func_args = {}

    # Checks and status
    if s.channels != 1:
        msg = "Signal must be single channel."
        logger.error(msg)
        raise ValueError(msg)

    logger.debug(f"Calculating scikit-maad {metric}")

    # Start the calc
    try:
        if metric == "all_spectral_alpha_indices":
            Sxx, tn, fn, ext = spectrogram(s, s.fs, **func_args)  # noqa: N806
            res = all_spectral_alpha_indices(Sxx, tn, fn, extent=ext, **func_args)[0]
        elif metric == "all_temporal_alpha_indices":
            res = all_temporal_alpha_indices(s, s.fs, **func_args)
        else:
            msg = f"Metric {metric} not recognized."
            logger.error(msg)
            raise ValueError(msg)  # noqa: TRY301
    except Exception as e:
        logger.error(f"Error calculating {metric}: {e!s}")
        raise

    if not as_df:
        return res.to_dict("records")[0]
    try:
        res["Recording"] = s.recording
        return res.set_index(["Recording"])
    except AttributeError:
        return res

def pyacoustics_metric_1ch(  # noqa: ANN201
    s: Signal | Binaural,
    metric: str,
    statistics: list[int | str] = (
        5,
        10,
        50,
        90,
        95,
        "avg",
        "max",
        "min",
        "kurt",
        "skew",
    ),
    label: str | None = None,
    as_df: bool = False,  # noqa: FBT001, FBT002
    return_time_series: bool = False,  # noqa: FBT001, FBT002
    func_args={},  # noqa: ANN001, B006
):
    """

    !!! warning "Deprecated"
        pyacoustics is deprecated. Use `acoustics_metric_1ch` instead.
    """
    warnings.warn(
        "pyacoustics is deprecated. Use acoustics_metric_1ch instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    return acoustics_metric_1ch(
        s,
        metric,
        statistics,
        label,
        as_df,
        return_time_series,
        func_args,
    )

def acoustics_metric_1ch(
    s: Signal | Binaural,
    metric: str,
    statistics: list[int | str] = (
        5,
        10,
        50,
        90,
        95,
        "avg",
        "max",
        "min",
        "kurt",
        "skew",
    ),
    label: str | None = None,
    as_df: bool = False,
    return_time_series: bool = False,
    func_args: dict | None = None,
) -> dict | pd.DataFrame:
    """
    Run a metric from the acoustic_toolbox library on a single channel object.

    Parameters
    ----------
    s
        Single channel signal to calculate the metric for.
    metric
        The metric to run.
    statistics
        List of level statistics to calculate (e.g. L_5, L_90, etc).
    label
        Label to use for the metric in the results dictionary.
        If None, will pull from default label for that metric given in DEFAULT_LABELS.
    as_df
        Whether to return a pandas DataFrame, by default False.
        If True, returns a MultiIndex Dataframe
        with ("Recording", "Channel") as the index.
    return_time_series
        Whether to return the time series of the metric, by default False.
        Cannot return time series if as_df is True.
    func_args
        Additional keyword arguments to pass to the metric function, by default {}.

    Returns
    -------
    :
        Dictionary of the calculated statistics or a pandas DataFrame.

    Raises
    ------
    ValueError
        If the signal is not single-channel or if an unrecognized metric is specified.

    See Also
    --------
    `acoustic_toolbox`

    """
    logger.debug(f"Calculating acoustics metric: {metric}")

    if func_args is None:
        func_args = {}

    if s.channels != 1:
        msg = "Signal must be single channel"
        logger.error(msg)
        raise ValueError(msg)
    try:
        label = label or DEFAULT_LABELS[metric]
    except KeyError as e:
        msg = f"Metric {metric} not recognized."
        logger.error(msg)
        raise ValueError(msg) from e
    if as_df and return_time_series:
        logger.warning(
            "Cannot return both a dataframe and time series. Returning dataframe only."
        )

        return_time_series = False

    logger.debug(f"Calculating Acoustic Toolbox: {metric} {statistics}")

    res = {}
    try:
        if metric in {"LZeq", "Leq", "LAeq", "LCeq"}:
            if metric in {"LZeq", "Leq"}:
                weighting = "Z"
            elif metric == "LAeq":
                weighting = "A"
            elif metric == "LCeq":
                weighting = "C"
            if "avg" in statistics or "mean" in statistics:
                stat = "avg" if "avg" in statistics else "mean"
                res[f"{label}"] = s.weigh(weighting).leq()
                statistics = list(statistics)
                statistics.remove(stat)
            if len(statistics) > 0:
                res = _stat_calcs(
                    label, s.weigh(weighting).levels(**func_args)[1], res, statistics
                )

            if return_time_series:
                res[f"{label}_ts"] = s.weigh(weighting).levels(**func_args)
        elif metric == "SEL":
            res[f"{label}"] = s.sound_exposure_level()
        else:
            msg = f"Metric {metric} not recognized."
            logger.error(msg)
            raise ValueError(msg)  # noqa: TRY301
    except Exception as e:
        logger.error(f"Error calculating {metric}: {e!s}")
        raise

    if not as_df:
        return res
    try:
        rec = s.recording
        return pd.DataFrame(res, index=[rec])
    except AttributeError:
        return pd.DataFrame(res, index=[0])

def pyacoustics_metric_2ch(  # noqa: ANN201
    b: Binaural,
    metric: str,
    statistics: tuple | list = (
        5,
        10,
        50,
        90,
        95,
        "avg",
        "max",
        "min",
        "kurt",
        "skew",
    ),
    label: str | None = None,
    channel_names: tuple[str, str] = ("Left", "Right"),
    as_df: bool = False,  # noqa: FBT001, FBT002
    return_time_series: bool = False,  # noqa: FBT001, FBT002
    func_args={},  # noqa: ANN001, B006
):
    """

    !!! warning "Deprecated"
        pyacoustics is deprecated. Use `acoustics_metric_2ch` instead.
    """
    warnings.warn(
        "pyacoustics is deprecated. Use acoustics_metric_2ch instead.",
        DeprecationWarning,
        stacklevel=2,
    )
    return acoustics_metric_2ch(
        b,
        metric,
        statistics,
        label,
        channel_names,
        as_df,
        return_time_series,
        func_args,
    )

def acoustics_metric_2ch(
    b: Binaural,
    metric: str,
    statistics: tuple | list = (
        5,
        10,
        50,
        90,
        95,
        "avg",
        "max",
        "min",
        "kurt",
        "skew",
    ),
    label: str | None = None,
    channel_names: tuple[str, str] = ("Left", "Right"),
    as_df: bool = False,
    return_time_series: bool = False,
    func_args: dict | None = None,
) -> dict | pd.DataFrame:
    """
    Run a metric from the Acoustic Toolbox library on a Binaural object.

    Parameters
    ----------
    b
        Binaural signal to calculate the metric for.
    metric
        The metric to run.
    statistics
        List of level statistics to calculate (e.g. L_5, L_90, etc).
    label
        Label to use for the metric in the results dictionary.
        If None, will pull from default label for that metric given in DEFAULT_LABELS.
    channel_names
        Custom names for the channels, by default ("Left", "Right").
    as_df
        Whether to return a pandas DataFrame, by default False.
        If True, returns a MultiIndex Dataframe
        with ("Recording", "Channel") as the index.
    return_time_series
        Whether to return the time series of the metric, by default False.
        Cannot return time series if as_df is True.
    func_args
        Arguments to pass to the metric function, by default {}.

    Returns
    -------
    :
        Dictionary of results if as_df is False, otherwise a pandas DataFrame.

    Raises
    ------
    ValueError
        If the input signal is not 2-channel.

    See Also
    --------
    `acoustics_metric_1ch`

    """
    logger.debug(f"Calculating acoustics metric for 2 channels: {metric}")

    if func_args is None:
        func_args = {}

    if b.channels != 2:  # noqa: PLR2004
        msg = "Must be 2 channel signal. Use `acoustics_metric_1ch instead`."
        logger.error(msg)
        raise ValueError(msg)

    logger.debug(f"Calculating Acoustic Toolbox metrics: {metric}")

    try:
        res_l = acoustics_metric_1ch(
            b[0],
            metric,
            statistics,
            label,
            as_df=False,
            return_time_series=return_time_series,
            func_args=func_args,
        )

        res_r = acoustics_metric_1ch(
            b[1],
            metric,
            statistics,
            label,
            as_df=False,
            return_time_series=return_time_series,
            func_args=func_args,
        )

        res = {channel_names[0]: res_l, channel_names[1]: res_r}
    except Exception as e:
        logger.error(f"Error calculating {metric} for 2 channels: {e!s}")
        raise

    if not as_df:
        return res
    try:
        rec = b.recording
    except AttributeError:
        rec = 0
    else:
        data = pd.DataFrame.from_dict(res, orient="index")
        data["Recording"] = rec
        data["Channel"] = data.index
        return data.set_index(["Recording", "Channel"])
    data = pd.DataFrame.from_dict(res, orient="index")
    data["Recording"] = rec
    data["Channel"] = data.index
    return data.set_index(["Recording", "Channel"])

def mosqito_metric_2ch(
    b: Binaural,
    metric: Literal[
        "loudness_zwtv",
        "sharpness_din_from_loudness",
        "sharpness_din_perseg",
        "sharpness_din_tv",
        "roughness_dw",
    ],
    statistics: tuple | list = (
        5,
        10,
        50,
        90,
        95,
        "avg",
        "max",
        "min",
        "kurt",
        "skew",
    ),
    label: str | None = None,
    channel_names: tuple[str, str] = ("Left", "Right"),
    as_df: bool = False,
    return_time_series: bool = False,
    parallel: bool = True,
    func_args: dict | None = None,
) -> dict | pd.DataFrame:
    """
    Calculate metrics from MoSQITo for a two-channel signal with parallel processing.

    Parameters
    ----------
    b
        Binaural signal to calculate the sound quality indices for.
    metric
        Metric to calculate.
    statistics
        List of level statistics to calculate (e.g. L_5, L_90, etc.).
    label
        Label to use for the metric in the results dictionary.
        If None, will pull from default label for that metric given in DEFAULT_LABELS.
    channel_names
        Custom names for the channels, by default ("Left", "Right").
    as_df
        Whether to return a pandas DataFrame, by default False.
        If True, returns a MultiIndex Dataframe
        with ("Recording", "Channel") as the index.
    return_time_series
        Whether to return the time series of the metric, by default False.
        Only works for metrics that return a time series array.
        Cannot be returned in a dataframe.
    parallel
        Whether to process channels in parallel, by default True.
    func_args
        Additional arguments to pass to the metric function, by default {}.

    Returns
    -------
    :
        Dictionary of results if as_df is False, otherwise a pandas DataFrame.

    Raises
    ------
    ValueError
        If the input signal is not 2-channel.

    """
    logger.debug(f"Calculating MoSQITo metric for 2 channels: {metric}")

    if func_args is None:
        func_args = {}

    if b.channels != 2:  # noqa: PLR2004
        msg = "Must be 2 channel signal. Use `mosqito_metric_1ch` instead."
        logger.error(msg)
        raise ValueError(msg)

    if metric == "sharpness_din_from_loudness":
        logger.debug(
            "Calculating MoSQITo metrics: `sharpness_din` from `loudness_zwtv`"
        )
    else:
        logger.debug(f"Calculating MoSQITo metric: {metric}")

    try:
        if parallel:
            with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
                future_l = executor.submit(
                    mosqito_metric_1ch,
                    b[0],
                    metric,
                    statistics,
                    label,
                    as_df=False,
                    return_time_series=return_time_series,
                    **func_args,
                )
                future_r = executor.submit(
                    mosqito_metric_1ch,
                    b[1],
                    metric,
                    statistics,
                    label,
                    as_df=False,
                    return_time_series=return_time_series,
                    **func_args,
                )
                res_l = future_l.result()
                res_r = future_r.result()
        else:
            res_l = mosqito_metric_1ch(
                b[0],
                metric,
                statistics,
                label,
                as_df=False,
                return_time_series=return_time_series,
                **func_args,
            )
            res_r = mosqito_metric_1ch(
                b[1],
                metric,
                statistics,
                label,
                as_df=False,
                return_time_series=return_time_series,
                **func_args,
            )

        res = {channel_names[0]: res_l, channel_names[1]: res_r}
    except Exception as e:
        logger.error(f"Error calculating MoSQITo metric {metric} for 2 channels: {e!s}")
        raise

    if not as_df:
        return res
    try:
        rec = b.recording
    except AttributeError:
        rec = 0
    else:
        data = pd.DataFrame.from_dict(res, orient="index")
        data["Recording"] = rec
        data["Channel"] = data.index
        return data.set_index(["Recording", "Channel"])
    data = pd.DataFrame.from_dict(res, orient="index")
    data["Recording"] = rec
    data["Channel"] = data.index
    return data.set_index(["Recording", "Channel"])

def maad_metric_2ch(
    b: Binaural,
    metric: Literal["all_temporal_alpha_indices", "all_spectral_alpha_indices"],
    channel_names: tuple[str, str] = ("Left", "Right"),
    as_df: bool = False,
    func_args: dict | None = None,
) -> dict | pd.DataFrame:
    """
    Run a metric from scikit-maad library (or suite of indices) on a binaural signal.

    Currently only supports running all the alpha indices at once.

    Parameters
    ----------
    b
        Binaural signal to calculate the alpha indices for.
    metric
        Metric to calculate.
    channel_names
        Custom names for the channels, by default ("Left", "Right").
    as_df
        Whether to return a pandas DataFrame, by default False.
        If True, returns a MultiIndex Dataframe with ("Recording", "Channel")
        as the index.
    func_args
        Additional arguments to pass to the metric function, by default {}.

    Returns
    -------
    :
        Dictionary of results if as_df is False, otherwise a pandas DataFrame.

    Raises
    ------
    ValueError
        If the input signal is not 2-channel or if an unrecognized metric is specified.

    See Also
    --------
    `scikit-maad` library

    `maad_metric_1ch`

    """
    logger.debug(f"Calculating MAAD metric for 2 channels: {metric}")

    if func_args is None:
        func_args = {}

    if b.channels != 2:  # noqa: PLR2004
        logger.error("Must be 2 channel signal. Use `maad_metric_1ch` instead.")
        msg = "Must be 2 channel signal. Use `maad_metric_1ch` instead."
        raise ValueError(msg)

    logger.debug(f"Calculating scikit-maad {metric}")

    try:
        res_l = maad_metric_1ch(b[0], metric, as_df=False, **func_args)
        res_r = maad_metric_1ch(b[1], metric, as_df=False, **func_args)
        res = {channel_names[0]: res_l, channel_names[1]: res_r}
    except Exception as e:
        logger.error(f"Error calculating MAAD metric {metric} for 2 channels: {e!s}")
        raise

    if not as_df:
        return res
    try:
        rec = b.recording
    except AttributeError:
        rec = 0
    data = pd.DataFrame.from_dict(res, orient="index")
    data["Recording"] = rec
    data["Channel"] = data.index
    return data.set_index(["Recording", "Channel"])