survivalist.functions

  1# This program is free software: you can redistribute it and/or modify
  2# it under the terms of the GNU General Public License as published by
  3# the Free Software Foundation, either version 3 of the License, or
  4# (at your option) any later version.
  5#
  6# This program is distributed in the hope that it will be useful,
  7# but WITHOUT ANY WARRANTY; without even the implied warranty of
  8# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  9# GNU General Public License for more details.
 10#
 11# You should have received a copy of the GNU General Public License
 12# along with this program.  If not, see <http://www.gnu.org/licenses/>.
 13
 14import numpy as np
 15from sklearn.utils import check_consistent_length
 16
 17__all__ = ["StepFunction"]
 18
 19
 20class StepFunction:
 21    """Callable step function.
 22
 23    .. math::
 24
 25        f(z) = a * y_i + b,
 26        x_i \\leq z < x_{i + 1}
 27
 28    Parameters
 29    ----------
 30    x : ndarray, shape = (n_points,)
 31        Values on the x axis in ascending order.
 32
 33    y : ndarray, shape = (n_points,)
 34        Corresponding values on the y axis.
 35
 36    a : float, optional, default: 1.0
 37        Constant to multiply by.
 38
 39    b : float, optional, default: 0.0
 40        Constant offset term.
 41
 42    domain : tuple, optional
 43        A tuple with two entries that sets the limits of the
 44        domain of the step function.
 45        If entry is `None`, use the first/last value of `x` as limit.
 46    """
 47
 48    def __init__(self, x, y, *, a=1.0, b=0.0, domain=(0, None)):
 49        check_consistent_length(x, y)
 50        self.x = x
 51        self.y = y
 52        self.a = a
 53        self.b = b
 54        domain_lower = self.x[0] if domain[0] is None else domain[0]
 55        domain_upper = self.x[-1] if domain[1] is None else domain[1]
 56        self._domain = (float(domain_lower), float(domain_upper))
 57
 58    @property
 59    def domain(self):
 60        """Returns the domain of the function, that means
 61        the range of values that the function accepts.
 62
 63        Returns
 64        -------
 65        lower_limit : float
 66            Lower limit of domain.
 67
 68        upper_limit : float
 69            Upper limit of domain.
 70        """
 71        return self._domain
 72
 73    def __call__(self, x):
 74        """Evaluate step function.
 75
 76        Values outside the interval specified by `self.domain`
 77        will raise an exception.
 78        Values in `x` that are in the interval `[self.domain[0]; self.x[0]]`
 79        get mapped to `self.y[0]`.
 80
 81        Parameters
 82        ----------
 83        x : float|array-like, shape=(n_values,)
 84            Values to evaluate step function at.
 85
 86        Returns
 87        -------
 88        y : float|array-like, shape=(n_values,)
 89            Values of step function at `x`.
 90        """
 91        x = np.atleast_1d(x)
 92        if not np.isfinite(x).all():
 93            raise ValueError("x must be finite")
 94        if np.min(x) < self._domain[0] or np.max(x) > self.domain[1]:
 95            raise ValueError(
 96                f"x must be within [{self.domain[0]:f}; {self.domain[1]:f}]")
 97
 98        # x is within the domain, but we need to account for self.domain[0] <= x < self.x[0]
 99        x = np.clip(x, a_min=self.x[0], a_max=None)
100
101        i = np.searchsorted(self.x, x, side="left")
102        not_exact = self.x[i] != x
103        i[not_exact] -= 1
104        value = self.a * self.y[i] + self.b
105        if value.shape[0] == 1:
106            return value[0]
107        return value
108
109    def __repr__(self):
110        return f"StepFunction(x={self.x!r}, y={self.y!r}, a={self.a!r}, b={self.b!r})"
111
112    def __eq__(self, other):
113        if isinstance(other, type(self)):
114            return all(self.x == other.x) and all(self.y == other.y) and self.a == other.a and self.b == other.b
115        return False