Source code for shapley_numba.game_templates

"""Game templates module.

Provides templates for even easier game authoring.
"""

from typing import Any, Iterable, Mapping, cast

import numba
import numpy as np
from numba.types import DictType

from shapley_numba.core import numba_game

__all__ = [
    'ParameterChangeExplanation',
    'parameter_change_explanation_spec',
    'TableGame',
]

table_game_spec = [
    ('values_dict', DictType(numba.int32, numba.float64)),
    ('default_value', numba.float64),
    ('vacuum_value', numba.float64),
]


@numba.njit
def _subset_to_key(subset: np.ndarray) -> np.int32:
    key = numba.int32(0)
    for i in range(len(subset)):
        key += numba.int32(2**i) * subset[i]
    return np.int32(key)


def _prepare_table_game(
    game_table: Mapping[Iterable[int], float],
) -> DictType:
    """Convert a dictionary to a numba Dict to use in `TableGame`."""
    values_dict = numba.typed.Dict.empty(numba.int32, numba.float64)
    for key, value in game_table.items():
        values_dict[_subset_to_key(np.array(key))] = np.float64(value)
    return values_dict


[docs] @numba_game(table_game_spec) class TableGame: """Table Game Template. A game where values are stored in a lookup table (dictionary or array). For best results use `prepare_table_game` to create a numba dictionary from a mapping of tuples to floats. Without compilation, supply a mapping of ints to floats, where ints represent subsets as binary bits. >>> import numpy as np >>> from shapley_numba.game_templates import TableGame >>> game = TableGame({(0, 1): 1.0, (1, 1): 2.0}) >>> game.value(np.array([1, 0])) np.float64(0.0) >>> game.value(np.array([0, 1])) np.float64(1.0) >>> game.value(np.array([0, 0])) np.float64(0.0) >>> game.value(np.array([1, 1])) np.float64(2.0) """ @staticmethod def __prepare_jit__( values_dict: Any, default_value: float = 0.0, vacuum_value: float = 0.0, ) -> tuple[tuple[Any, ...], dict[str, Any]]: """Convert a tuple-keyed dict to a numba typed dict for jit compilation.""" prepared: Any = values_dict if isinstance(prepared, dict): d = cast(dict[Any, Any], prepared) if isinstance(next(iter(d), None), (tuple, list)): prepared = cast(Any, _prepare_table_game(d)) return (cast(Any, prepared), default_value, vacuum_value), {} def __init__( self, values_dict: DictType | dict[int, float], default_value: float = 0.0, vacuum_value: float = 0.0, ): """Initialize the TableGame. Parameters ---------- values_dict : Mapping[Iterable[int], float | np.float64 | int]|Mapping[int, float] The lookup table default_value : float, optional, default 0.0 The value of a coalition not in the table. vacuum_value: float, optional, default 0.0 The value of the empty set (canonically should be 0. but shapley-numba allows non-zero values) """ self.values_dict = values_dict self.default_value = np.float64(default_value) self.vacuum_value = np.float64(vacuum_value) def value(self, subset: np.ndarray) -> np.float64: """Compute the value of the coalition.""" if subset.sum() == 0: return self.vacuum_value key = _subset_to_key(subset) if key in self.values_dict: return np.float64(self.values_dict[key]) return self.default_value
parameter_change_explanation_spec = [ ('old_parameters', numba.float64[:]), ('new_parameters', numba.float64[:]), ]
[docs] class ParameterChangeExplanation: """Parameter Change Explanation Template. The "game" template for computing attribution due to a change in a multi-parameter model. Provides an attribution to each parameter change. See BlackScholesCallGame for an example of implementation. """
[docs] def __init__(self, old_parameters, new_parameters): """Initialize the ParameterChangeExplanation.""" self.old_parameters = old_parameters self.new_parameters = new_parameters
[docs] def model_evaluate(self, parameters): """Evaluate the model with the given parameters.""" raise NotImplementedError
[docs] def value(self, subset): """Compute the value of a given subset of players.""" parameters = np.where(subset == 1, self.new_parameters, self.old_parameters) return self.model_evaluate(parameters)