Source code for openfisca_core.commons.formulas

from __future__ import annotations

from collections.abc import Mapping

import numpy

from . import types as t


[docs] def apply_thresholds( input: t.Array[numpy.float32], thresholds: t.ArrayLike[float], choices: t.ArrayLike[float], ) -> t.Array[numpy.float32]: """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: ndarray[float32]: A list of the values chosen. 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: list[t.Array[numpy.bool_] | 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] msg = ( "'apply_thresholds' must be called with the same number of thresholds " "than choices, or one more choice." ) assert len(condlist) == len(choices), msg return numpy.select(condlist, choices)
[docs] def concat( this: t.Array[numpy.str_] | t.ArrayLike[object], that: t.Array[numpy.str_] | t.ArrayLike[object], ) -> t.Array[numpy.str_]: """Concatenate the values of two arrays. Args: this: An array to concatenate. that: Another array to concatenate. Returns: ndarray[str_]: 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 not isinstance(this, numpy.ndarray): this = numpy.array(this) if not numpy.issubdtype(this.dtype, numpy.str_): this = this.astype("str") if not isinstance(that, numpy.ndarray): that = numpy.array(that) if not numpy.issubdtype(that.dtype, numpy.str_): that = that.astype("str") return numpy.char.add(this, that)
[docs] def switch( conditions: t.Array[numpy.float32] | t.ArrayLike[float], value_by_condition: Mapping[float, float], ) -> t.Array[numpy.float32]: """Mimick 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: ndarray[float32]: An array with the replaced values. 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] return numpy.select(condlist, tuple(value_by_condition.values()))
__all__ = ["apply_thresholds", "concat", "switch"]