Source code for openfisca_core.commons.formulas

from typing import Any, Dict, Sequence, TypeVar

import numpy

from openfisca_core.types import ArrayLike, ArrayType

T = TypeVar("T")


[docs]def apply_thresholds( input: ArrayType[float], thresholds: ArrayLike[float], choices: ArrayLike[float], ) -> ArrayType[float]: """Makes a choice based on an input and thresholds. From a list of ``choices``, this function selects one of these values based on a list of inputs, depending on the value of each ``input`` within a list of ``thresholds``. Args: input: A list of inputs to make a choice from. thresholds: A list of thresholds to choose. choices: A list of the possible values to choose from. Returns: :obj:`numpy.ndarray` of :obj:`float`: A list of the values chosen. Raises: :exc:`AssertionError`: When the number of ``thresholds`` (t) and the number of choices (c) are not either t == c or t == c - 1. Examples: >>> input = numpy.array([4, 5, 6, 7, 8]) >>> thresholds = [5, 7] >>> choices = [10, 15, 20] >>> apply_thresholds(input, thresholds, choices) array([10, 10, 15, 15, 20]) """ condlist: Sequence[ArrayType[bool]] condlist = [input <= threshold for threshold in thresholds] if len(condlist) == len(choices) - 1: # If a choice is provided for input > highest threshold, last condition # must be true to return it. condlist += [True] assert len(condlist) == len(choices), \ " ".join([ "'apply_thresholds' must be called with the same number of", "thresholds than choices, or one more choice.", ]) return numpy.select(condlist, choices)
[docs]def concat(this: ArrayLike[str], that: ArrayLike[str]) -> ArrayType[str]: """Concatenates the values of two arrays. Args: this: An array to concatenate. that: Another array to concatenate. Returns: :obj:`numpy.ndarray` of :obj:`float`: An array with the concatenated values. Examples: >>> this = ["this", "that"] >>> that = numpy.array([1, 2.5]) >>> concat(this, that) array(['this1.0', 'that2.5']...) """ if isinstance(this, numpy.ndarray) and \ not numpy.issubdtype(this.dtype, numpy.str_): this = this.astype('str') if isinstance(that, numpy.ndarray) and \ not numpy.issubdtype(that.dtype, numpy.str_): that = that.astype('str') return numpy.char.add(this, that)
[docs]def switch( conditions: ArrayType[Any], value_by_condition: Dict[float, T], ) -> ArrayType[T]: """Mimicks a switch statement. Given an array of conditions, returns an array of the same size, replacing each condition item with the matching given value. Args: conditions: An array of conditions. value_by_condition: Values to replace for each condition. Returns: :obj:`numpy.ndarray`: An array with the replaced values. Raises: :exc:`AssertionError`: When ``value_by_condition`` is empty. Examples: >>> conditions = numpy.array([1, 1, 1, 2]) >>> value_by_condition = {1: 80, 2: 90} >>> switch(conditions, value_by_condition) array([80, 80, 80, 90]) """ assert len(value_by_condition) > 0, \ "'switch' must be called with at least one value." condlist = [ conditions == condition for condition in value_by_condition.keys() ] return numpy.select(condlist, value_by_condition.values())