Source code for ultrasound_metrics.metrics.snr

"""
Signal-to-Noise Ratio (SNR) metric for ultrasound, commonly used for speckle.

This module provides functions to calculate the Signal-to-Noise Ratio (SNR)
for ultrasound speckle patterns. The SNR is calculated as mean/standard-deviation.

References
----------
.. [1] C. B. Burckhardt, "Speckle in Ultrasound B-mode Scans," in IEEE Transactions
       on Sonics and Ultrasonics, vol. 25, no. 1, pp. 1-6, Jan. 1978,
       doi: 10.1109/T-SU.1978.30978.
.. [2] D. Perdios, M. Vonlanthen, F. Martinez, M. Arditi and J. -P. Thiran,
       "CNN-Based Image Reconstruction Method for Ultrafast Ultrasound Imaging,"
       in IEEE Transactions on Ultrasonics, Ferroelectrics, and Frequency Control,
       vol. 69, no. 4, pp. 1154-1168, April 2022, doi: 10.1109/TUFFC.2021.3131383.
"""

import warnings

from array_api_compat import array_namespace
from beartype import beartype as typechecker
from jaxtyping import Num, jaxtyped

from .._utils.array_api import ArrayAPIObj


@jaxtyped(typechecker=typechecker)


def snr(
    values: Num[ArrayAPIObj, " *pixels"],
    *,
    db: bool = False,
) -> float:
    """
    Calculate Signal-to-Noise Ratio (SNR) of an image.

    This function computes the Signal-to-Noise Ratio (SNR) for ultrasound speckle
    patterns by calculating mean/standard-deviation.
    This metric provides a single value per ROI/image.

    .. math::

        \\text{SNR} = \\frac{\\mathbb{E}[X]}{\\sqrt{\\text{Var}(X)}}

    where:

    - :math:`\\mathbb{E}[X]` is the mean (average pixel value) across the (post-envelope) image
    - :math:`\\text{Var}(X)` is the variance of pixel values across the (post-envelope) image

    In ultrasound, the statistics of fully developed speckle signals
    (post envelope detection) follow a Rayleigh distribution, and
    the SNR of a Rayleigh distribution is equal to:

    .. math::

        \\text{SNR}_{\\text{Rayleigh}} = \\sqrt{\\frac{\\pi}{4 - \\pi}} \\approx 1.91

    Parameters
    ----------
    values : array
        Input values.
        - spatial dimensions are flattened
        - Can be real or complex floating-point data

    db : bool, optional
        If True, return the SNR in decibels (dB). Default is False.

    Returns
    -------
    float
        Signal-to-Noise Ratio (SNR) of the values. If ``db=True``, the SNR is returned in decibels (dB).

    Warns
    -----
    UserWarning
        If variance is zero (would cause division by zero).

    Notes
    -----
    - This implementation is an amplitude ratio, not a power ratio.
    - Complex data is handled by computing the magnitude (absolute value) before processing
    - For Rayleigh-distributed speckle, the theoretical SNR is approximately 1.91
    - If ROI is used, it should be applied to the input values before calling this function
    """
    xp = array_namespace(values)

    if values.ndim < 1:
        raise ValueError(f"Values must have at least 1 dimension, got {values.ndim}")

    # SNR operates on the envelope of the values
    values = xp.abs(values)

    mean = xp.mean(values)
    std = xp.std(values)

    if bool(std == 0):
        warnings.warn(
            "Standard deviation is zero. This will cause division by zero. Check if your values have sufficient variation.",
            UserWarning,
        )

    snr_value = mean / std

    if db:
        snr_value = 20.0 * xp.log10(snr_value)

    return float(snr_value)