Source code for pysrw.tools

"""Generic utils for data analysis."""

import numpy as np
from copy import deepcopy

"""A collection of general purpose functions or tools for the manipulation 
of the computed data."""



def tupledeepcopy(tupleIn):
    return tuple([deepcopy(el) for el in tupleIn])




def findIdx(array, value):
    """Find the closest value out of an array.

    Args:
        array: array to search for.
        value: desired value.

    Returns:
        int: the index of the `array` item closest to `value`.
    """
    return np.argmin(np.abs(array - value))




def gaussian1p(x, sigma):
    r"""A non-normalized Gaussian with unitary amplitude and zero mean.

    .. math::
        y = e^{-\frac{x^2}{2 \sigma^2}}

    Args:
        x: coordinate axis.
        sigma: standard deviation.

    Returns:
        float: array of y values
    """
    return np.exp(-x** 2 / (2 * sigma ** 2))




def gaussian2p(x, a, sigma):
    r"""A non-normalized Gaussian with zero mean.

    .. math::
        y = a \cdot e^{-\frac{x^2}{2 \sigma^2}}

    Args:
        x: coordinate axis.
        a: amplitude.
        sigma: standard deviation.

    Returns:
        float: array of y values
    """
    return a * np.exp(-x**2 / (2 * sigma ** 2))




def gaussian3p(x, a, mu, sigma):
    r"""A non-normalized Gaussian.

    .. math::
        y = a \cdot e^{-\frac{(x-\mu)^2}{2 \sigma^2}}

    Args:
        x: coordinate axis.
        a: amplitude.
        mu: mean.
        sigma: standard deviation.

    Returns:
        float: array of values
    """
    return a * np.exp(-(x - mu) ** 2 / (2 * sigma ** 2))




def gaussian4p(x, a0, a1, mu, sigma):
    r"""A non-normalized Gaussian with a constant offset.

    .. math::
        y = a0 + a1 \cdot e^{-\frac{(x-\mu)^2}{2 \sigma^2}}

    Args:
        x: coordinate axis.
        a0: constant offset.
        a1: amplitude.
        mu: mean.
        sigma: standard deviation.

    Returns:
        float: array of values
    """
    return a0 + a1 * np.exp(-(x - mu) ** 2 / (2 * sigma ** 2))




def gaussian1pNorm(x, sigma):
    r"""The normalized version of :py:func:`gaussian1p`.

    .. math::
        y = \frac{1}{\sqrt{2 \pi} \sigma} e^{-\frac{x^2}{2 \sigma^2}}

    Args:
        x: coordinate axis.
        sigma: standard deviation.

    Returns:
        float: array of y values
    """
    return 1 / (np.sqrt(2 * np.pi) * sigma) * np.exp(-x** 2 / (2 * sigma ** 2))




def gaussian2pNorm(x, a, sigma):
    r"""The normalized version of :py:func:`gaussian2p`.

    .. math::
        y = \frac{a}{\sqrt{2 \pi} \sigma} e^{-\frac{x^2}{2 \sigma^2}}

    Args:
        x: coordinate axis.
        a: coefficient scaling the area of the distribution.
        sigma: standard deviation.

    Returns:
        float: array of y values
    """
    return a / (np.sqrt(2 * np.pi) * sigma) * np.exp(-x** 2 / (2 * sigma ** 2))




def extractCrossCuts(spatialDistr, point=[0,0]):
    """Get the cross cut from a two-dimensional distribution.

    Args:
        spatialDistr: any dictionary representing a spatial distribution with
            'xax' and 'yax' as coordinate axes and 'data' as data matrix.
        point (optional): the point where to extract the cross cut. The cross 
            cuts are extracted at the closest available position to this value.
            Defaults to [0,0].

    Returns:
        dict: dictionary with 'xax', 'yax', 'xdata' and 'ydata'.
    """

    xidx = findIdx(spatialDistr["xax"], point[0])
    yidx = findIdx(spatialDistr["yax"], point[1])

    xdata = deepcopy(spatialDistr["data"][yidx, :])
    ydata = deepcopy(spatialDistr["data"][:, xidx])
    xax = deepcopy(spatialDistr["xax"])
    yax = deepcopy(spatialDistr["yax"])

    return {"xax": xax, "yax": yax, "xdata": xdata, "ydata": ydata}




def extractProfiles(spatialDistr):
    """Get the integrated profiles from a two-dimensional distribution.

    Args:
        spatialDistr: any dictionary representing a spatial distribution with
            'xax' and 'yax' as coordinate axes and 'data' as data matrix.

    Returns:
        dict: dictionary with 'xax', 'yax', 'xdata' and 'ydata'.
    """

    xdata = np.nansum(spatialDistr["data"], axis=0)
    ydata = np.nansum(spatialDistr["data"], axis=1)
    xax = deepcopy(spatialDistr["xax"])
    yax = deepcopy(spatialDistr["yax"])

    return {"xax": xax, "yax": yax, "xdata": xdata, "ydata": ydata}