nmreval/src/nmreval/models/basic.py

"""
***************
Basic functions
***************

Simple functions
"""

import numpy as np

from ..lib.utils import ArrayLike
from ..math.mittagleffler import mlf


class Constant:
    """
    A boring constant line.
    """
    type = 'Basic'
    name = 'Constant'
    equation = 'C'
    params = ['C']

    @staticmethod
    def func(x: ArrayLike, c: float) -> ArrayLike:
        """
        Constant

        .. math::
            y = c

        Args:
            x (array-like): Input values
            c (float): constant
        """
        return c*np.ones(len(x))


class Linear:
    """
    Slightly more exciting line
    """
    type = 'Basic'
    name = 'Straight line'
    equation = 'm*x + t'
    params = ['m', 't']

    @staticmethod
    def func(x: ArrayLike, m: float, t: float) -> ArrayLike:
        """
        Straight line.

        .. math::
            y = m\cdot x + t

        Args:
            x (array_like): Input values
            m (float): Slope
            t (float): Intercept

        """
        return m*x + t


class PowerLaw:
    """
    Power law

    .. math::
        y = A\cdot x^b

    Args:
        x (array_like): Input values
        A (float): Prefactor
        b (float): Exponent

    """

    type = 'Basic'
    name = 'Power law'
    equation = 'A*x^{b}'
    params = ['A', 'b']

    @staticmethod
    def func(x, a: float, b: float):
        return a * x**b


class Log:
    """
    Logarithm

    .. math::
        y = C\cdot \ln(x-x_0)

    Args:
        x (array_like): Input values
        C (float): Prefactor
        x0 (float): Offset

    """
    type = 'Basic'
    name = 'Logarithm'
    equation = 'C*ln(x-x_{0})'
    params = ['C', 'x_{0}']

    @staticmethod
    def func(x, c, x0):
        return c * np.log(x-x0)


class Parabola:
    """
    Parabola with vertex :math:`(x_0, y_0)`

    .. math::
        y = C\cdot (x-x_{0}) + y_0

    Args:
        x (array_like): Input values
        c (float): Slope
        x0 (float): x position of vertex
        y0 (float): y position of vertex

    """
    type = 'Basic'
    name = 'Parabola'
    equation = 'C*(x-x_{0})^{2} + y_{0}'
    params = ['C', 'x_{0}', 'y_{0}']

    @staticmethod
    def func(x, c, x0, y0):
        return c * (x-x0)**2 + y0


class ParabolaLog:
    """
    Parabola (on a log-scale) with vertex :math:`(x_0, y_0)`

    .. math::
        y = C\cdot (x-x_{0}) + y_0

    Args:
        x (array_like): Input values
        c (float): Slope
        x0 (float): x position of vertex
        y0 (float): y position of vertex

    """
    type = 'Basic'
    name = 'Log-Parabola'
    equation = 'exp[C*(x-x_{0})^{2} + y_{0}]'
    params = ['C', 'x_{0}', 'y_{0}']

    @staticmethod
    def func(x, c, x0, y0):
        return np.exp(c * (x-x0)**2 + y0)


class PowerLawCross:
    """
    Crossover between power laws
    """
    type = 'Basic'
    name = 'Crossing Power Laws'
    params = ['C', 'b_{1}', 'b_{2}', 'x_{0}']

    @staticmethod
    def func(x, c, b1, b2, x0):
        """
        Crossover between to power laws at position :math:`x_0`

        .. math::
            y \\propto
            \\begin{cases}
                x^{b_1}, & x \le x_0 \\\\
                x^{b_2}, & x > x_0
            \\end{cases}

        Args:
            x (array_like): Input values
            c (float): Prefactor
            b1 (float): Exponent of first power law
            b2 (float): Exponent of second power law
            x0 (float): x position of crossover
        """
        mas = np.nonzero(x > x0)
        ret_val = c * x**b1
        c2 = c * x0**(b1-b2)
        ret_val[mas] = c2 * x[mas]**b2
        return ret_val


class Sinc:
    type = 'Basic'
    name = 'Sinc'
    equation = 'C * sinc((x-x_{0})/w)'
    params = ['C', 'x_{0}', 'w']

    @staticmethod
    def func(x, c: float, x0: float, w: float):
        # numpy sinc is defined as sin(pi*x)/(pi*x)
        return c * np.sinc(((x-x0)/w)/np.pi)


class Sine:
    """
    Wavy sine function
    """
    type = 'Basic'
    name = 'Sine'
    equation = r'C*sin(a*x-\phi)'
    params = ['C', 'a', r'\phi']

    @staticmethod
    def func(x, c: float, a: float, phi: float):
        """
        Calculate sine function

        .. math::
            y = C\sin(a x - \phi)

        Args:
            x (array_like): Input values
            c (float): Prefactor
            a (float): frequency
            phi (float): shift

        """
        return c*np.sin(a*x-phi)


class ExpFunc:
    """
    Stretched exponential function

    .. math::
        y = C\exp[\pm (x\cdot x_0)^\\beta] \\text{ or } C\exp[\pm (x/x_0)^\\beta]

    Args:
        x (array_like): Input values
        C (float): Prefactor
        x0 (float): Decay/growth constant
        beta (float): Stretching parameter

    Keyword Args:
        pm (int): Sign of the number determines if growing or decaying function.
                  Positive values result in growing exponentials, negative in decaying functions.
                  Default: -1
        mode (str): Interpretation of x0 as either rate (:math:`x\cdot x_0`) or time (:math:`x/x_0`).
                    Possible values are *time* or *rate*, default is *time*.
    """
    type = 'Basic'
    name = 'Exponential Function'
    equation = r'C*exp[\pm(x_{0}x)^{\beta}] or C*exp[\pm(x/x_{0})^{\beta}]'
    params = ['C', r'x_{0}', r'\beta']
    choices = [('Sign', 'pm', {'decaying': -1, 'growing': 1}),
               ('x0 type', 'mode', {'Time (x/x0)': 'time', 'Rate (x*x0)': 'rate'})]

    @staticmethod
    def func(x, c, x0, beta, pm: int = -1, mode: str = 'time'):
        if mode == 'time':
            return c * np.exp(np.sign(pm) * (x / x0) ** beta)
        elif mode == 'rate':
            return c * np.exp(np.sign(pm) * (x * x0) ** beta)
        else:
            raise ValueError('Unknown mode %s. Use either "rate" or "time".' % str(mode))


class MittagLeffler:
    """
    Mittag-Leffler function

    .. math::
        y = C\cdot E_\\alpha[-(x/x_0)^\\alpha] \\text{ or } C\cdot E_\\alpha[-(x\cdot x_0)^\\alpha]

    where

    .. math::
        E_a(z)= \sum_{k=0}^\infty \\frac{z^k}{\Gamma(\\alpha k + 1)}

    Args:
        x (array_like): Input values
        C (float): Prefactor
        x0 (float): Decay constant
        alpha (float): Stretching parameter

    Keyword Args:
        mode (str): Interpretation of x0 as either rate (:math:`x\cdot x_0`) or time (:math:`x/x_0`).
                    Possible values are *time* or *rate*, default is *time*.
    """
    type = 'Basic'
    name = 'Mittag-Leffler'
    equation = r'C*E_{\alpha}(-(x/x_{0}), \alpha)'
    params = ['C', 'x_{0}', r'\alpha']

    @staticmethod
    def func(x, c, x0, alpha, mode: str = 'time'):
        if mode == 'time':
            return c*mlf(-(x/x0), alpha)
        elif mode == 'rate':
            return c*mlf(-(x*x0), alpha)
        else:
            raise ValueError('Unknown mode {mode}. Use either "rate" or "time".')