🛠️ Add pyright (mmschlk#409)

* adds pyright to project

* fixes some typing problems

* reworks workflows (mmschlk#405)

* reworks workflows

* limit tensorflow to not install on windows

* adds device="cpu" to SentimentAnalysis test

* adds doc building workflow as a test

* reverts windows unit test runner to use python 3.12 as a fix

* makes related software an rst file

* moved copy_notebooks into scripts

* improves docs

* try re-installing python in uv to fix win errors

* removed unncessary extension

* Add beeswarm plot to SHAP-IQ (mmschlk#406)

* add beeswarm plot

* improve plot readability

* add tests for beeswarm plot

* update docs for beeswarm plot

* more informative error messages for beeswarm plot

* refactor code

* add beeswarm plot to init

* simplify colormap generation in beeswarm plot

* add tutorial for beeswarm plot

---------

Co-authored-by: Anna Przybyłowska <[email protected]>

Author: mmschlk

pyproject.toml

@@ -151,6 +151,9 @@ exclude = [
 # Exclude a variety of commonly ignored directories.
 # Note to developers: If you add a new ignore at least put a comment next to it why it is ignored
 [tool.ruff.lint.per-file-ignores]
+"shapiq/typing.py" = [
+    "A005",  # we want to be similar to other libraries that also shadow typing
+]
 "tests/*.py" = [
     # "ALL",
     "S101", # we need asserts in tests
@@ -212,6 +215,11 @@ force-wrap-aliases = true
 no-lines-before = ["future"]
 required-imports = ["from __future__ import annotations"]
 
+[tool.pyright]
+include = ["shapiq"]
+exclude = ["tests", "docs", "benchmark", "scripts", "shapiq/plot"]
+pythonVersion = "3.10"
+
 [dependency-groups]
 test = [
     "pytest>=8.3.5",
@@ -222,6 +230,7 @@ test = [
 lint = [
     "ruff>=0.11.2",
     "pre-commit>=4.2.0",
+    "pyright>=1.1.402",
 ]
 docs = [
     "sphinx>=8.0.0",

shapiq/explainer/tree/base.py

@@ -3,11 +3,15 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 from .utils import compute_empty_prediction
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
 
 @dataclass
 class TreeModel:
@@ -52,22 +56,22 @@ class TreeModel:
 
     """
 
-    children_left: np.ndarray[int]
-    children_right: np.ndarray[int]
-    features: np.ndarray[int]
-    thresholds: np.ndarray[float]
-    values: np.ndarray[float]
-    node_sample_weight: np.ndarray[float]
-    empty_prediction: float | None = None
-    leaf_mask: np.ndarray[bool] | None = None
-    n_features_in_tree: int | None = None
-    max_feature_id: int | None = None
-    feature_ids: set | None = None
-    root_node_id: int | None = None
-    n_nodes: int | None = None
-    nodes: np.ndarray[int] | None = None
-    feature_map_original_internal: dict[int, int] | None = None
-    feature_map_internal_original: dict[int, int] | None = None
+    children_left: NDArray[np.int_]
+    children_right: NDArray[np.int_]
+    features: NDArray[np.int_]
+    thresholds: NDArray[np.floating]
+    values: NDArray[np.floating]
+    node_sample_weight: NDArray[np.floating]
+    empty_prediction: float = None  # type: ignore[assignment]
+    leaf_mask: NDArray[np.bool_] = None  # type: ignore[assignment]
+    n_features_in_tree: int = None  # type: ignore[assignment]
+    max_feature_id: int = None  # type: ignore[assignment]
+    feature_ids: set = None  # type: ignore[assignment]
+    root_node_id: int = None  # type: ignore[assignment]
+    n_nodes: int = None  # type: ignore[assignment]
+    nodes: NDArray[np.int_] = None  # type: ignore[assignment]
+    feature_map_original_internal: dict[int, int] = None  # type: ignore[assignment]
+    feature_map_internal_original: dict[int, int] = None  # type: ignore[assignment]
     original_output_type: str = "raw"  # not used at the moment
 
     def compute_empty_prediction(self) -> None:
@@ -187,7 +191,7 @@ def predict_one(self, x: np.ndarray) -> float:
             else:
                 node = self.children_right[node]
             is_leaf = self.leaf_mask[node]
-        return self.values[node]
+        return float(self.values[node])
 
 
 @dataclass

shapiq/explainer/tree/conversion/edges.py

@@ -2,22 +2,27 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 from scipy.special import binom
 
 from shapiq.explainer.tree.base import EdgeTree
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
 
 def create_edge_tree(
-    children_left: np.ndarray[int],
-    children_right: np.ndarray[int],
-    features: np.ndarray[int],
-    node_sample_weight: np.ndarray[float],
-    values: np.ndarray[float],
+    children_left: NDArray[np.int_],
+    children_right: NDArray[np.int_],
+    features: NDArray[np.int_],
+    node_sample_weight: NDArray[np.floating],
+    values: NDArray[np.floating],
     n_nodes: int,
     n_features: int,
     max_interaction: int,
-    subset_updates_pos_store: dict[int, dict[int, np.ndarray[int]]],
+    subset_updates_pos_store: dict[int, dict[int, NDArray[np.int_]]],
 ) -> EdgeTree:
     """Extracts edge information recursively from the tree information.
 

shapiq/explainer/tree/utils.py

@@ -2,15 +2,21 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
 __all__ = ["compute_empty_prediction", "get_conditional_sample_weights"]
 
 
 def get_conditional_sample_weights(
-    sample_count: np.ndarray[int],
-    parent_array: np.ndarray[int],
-) -> np.ndarray[float]:
+    sample_count: NDArray[np.int_],
+    parent_array: NDArray[np.int_],
+) -> NDArray[np.floating]:
     """Get the conditional sample weights for a tree at each decision node.
 
     The conditional sample weights are the probabilities of going left or right at each decision
@@ -41,8 +47,8 @@ def get_conditional_sample_weights(
 
 
 def compute_empty_prediction(
-    leaf_values: np.ndarray[float],
-    leaf_sample_weights: np.ndarray[float],
+    leaf_values: NDArray[np.floating],
+    leaf_sample_weights: NDArray[np.floating],
 ) -> float:
     """Compute the empty prediction of a tree model.
 
@@ -56,4 +62,4 @@ def compute_empty_prediction(
         The empty prediction of the tree model.
 
     """
-    return np.sum(leaf_values * leaf_sample_weights) / np.sum(leaf_sample_weights)
+    return float(np.sum(leaf_values * leaf_sample_weights) / np.sum(leaf_sample_weights))

shapiq/games/base.py

@@ -15,6 +15,7 @@
 
 if TYPE_CHECKING:
     from shapiq.interaction_values import InteractionValues
+    from shapiq.utils.custom_types import CoalitionMatrix, GameValues
 
 
 class Game:
@@ -29,7 +30,7 @@ class Game:
         game_name: The name of the game.
 
     Attributes:
-        precompute_flag: A flag to manually override the precomputed check. If set to ``True``, the
+        _precompute_flag: A flag to manually override the precomputed check. If set to ``True``, the
             game is considered precomputed and only uses the lookup.
         value_storage: The storage for the game values without normalization applied.
         coalition_lookup: A lookup dictionary mapping from coalitions to indices in the
@@ -79,6 +80,19 @@ class Game:
 
     """
 
+    n_players: int
+    """The number of players in the game."""
+
+    game_id: str
+    """A unique identifier for the game, based on its class name and hash."""
+
+    normalization_value: float
+    """The value which is used to normalize (center) the game values such that the value for the
+    empty coalition is zero. If this is zero, the game values are not normalized."""
+
+    value_storage: GameValues
+    """The storage for the game values without normalization applied."""
+
     def __init__(
         self,
         n_players: int | None = None,
@@ -120,19 +134,19 @@ def __init__(
 
         """
         # manual flag for choosing precomputed values even if not all values might be stored
-        self.precompute_flag: bool = False  # flag to manually override the precomputed check
+        self._precompute_flag: bool = False  # flag to manually override the precomputed check
 
         # define storage variables
-        self.value_storage: np.ndarray = np.zeros(0, dtype=float)
+        self.value_storage: GameValues = np.zeros(0, dtype=float)
         self.coalition_lookup: dict[tuple[int, ...], int] = {}
-        self.n_players: int = n_players  # if path_to_values is provided, this may be overwritten
+        self.n_players = n_players
 
         if n_players is None and path_to_values is None:
             msg = "The number of players has to be provided if game is not loaded from values."
             raise ValueError(msg)
 
         # setup normalization of the game
-        self.normalization_value: float = 0.0
+        self.normalization_value = 0.0
         if normalize and path_to_values is None:
             self.normalization_value = normalization_value
             if normalization_value is None:
@@ -147,7 +161,7 @@ def __init__(
                     stacklevel=2,
                 )
 
-        game_id: str = str(hash(self))[:8]
+        game_id = str(hash(self))[:8]
         self.game_id = f"{self.get_game_name()}_{game_id}"
         if path_to_values is not None:
             self.load_values(path_to_values, precomputed=True)
@@ -177,7 +191,7 @@ def n_values_stored(self) -> int:
     @property
     def precomputed(self) -> bool:
         """Indication whether the game has been precomputed."""
-        return self.n_values_stored >= 2**self.n_players or self.precompute_flag
+        return self.n_values_stored >= 2**self.n_players or self._precompute_flag
 
     @property
     def normalize(self) -> bool:
@@ -192,13 +206,13 @@ def is_normalized(self) -> bool:
     def _check_coalitions(
         self,
         coalitions: (
-            np.ndarray
+            CoalitionMatrix
             | list[tuple[int, ...]]
             | list[tuple[str, ...]]
             | tuple[int, ...]
             | tuple[str, ...]
         ),
-    ) -> np.ndarray:
+    ) -> CoalitionMatrix:
         """Validates the coalitions and convert them to one-hot encoding.
 
         Check if the coalitions are in the correct format and convert them to one-hot encoding.
@@ -280,15 +294,15 @@ def _check_coalitions(
     def __call__(
         self,
         coalitions: (
-            np.ndarray
+            CoalitionMatrix
             | list[tuple[int, ...]]
             | list[tuple[str, ...]]
             | tuple[int, ...]
             | tuple[str, ...]
         ),
         *,
         verbose: bool = False,
-    ) -> np.ndarray:
+    ) -> GameValues:
         """Calls the game with the given coalitions.
 
         Calls the game's value function with the given coalitions and returns the output of the
@@ -305,7 +319,7 @@ def __call__(
             The values of the coalitions.
 
         """
-        coalitions: np.ndarray = self._check_coalitions(coalitions)
+        coalitions = self._check_coalitions(coalitions)
         verbose = verbose or self.verbose
         if not self.precomputed and not verbose:
             values = self.value_function(coalitions)
@@ -335,7 +349,7 @@ def _lookup_coalitions(self, coalitions: np.ndarray) -> np.ndarray:
                 raise KeyError(msg) from error
         return values
 
-    def value_function(self, coalitions: np.ndarray) -> np.ndarray:
+    def value_function(self, coalitions: CoalitionMatrix) -> GameValues:
         """Returns the value of the coalitions.
 
         The value function of the game, which models the behavior of the game. The value function
@@ -355,7 +369,7 @@ def value_function(self, coalitions: np.ndarray) -> np.ndarray:
         msg = "The value function has to be implemented in inherited classes."
         raise NotImplementedError(msg)
 
-    def precompute(self, coalitions: np.ndarray | None = None) -> None:
+    def precompute(self, coalitions: CoalitionMatrix | None = None) -> None:
         """Precompute the game values for all or a given set of coalitions.
 
         The pre-computation iterates over the powerset of all coalitions or a given set of
@@ -401,26 +415,25 @@ def precompute(self, coalitions: np.ndarray | None = None) -> None:
                 stacklevel=2,
             )
         if coalitions is None:
-            coalitions = list(powerset(range(self.n_players)))  # might be getting slow
-            coalitions_array = transform_coalitions_to_array(coalitions, self.n_players)
-            coalitions_dict = {coal: i for i, coal in enumerate(coalitions)}
+            all_coalitions = list(powerset(range(self.n_players)))  # might be getting slow
+            coalitions = transform_coalitions_to_array(all_coalitions, self.n_players)
+            coalitions_dict = {coal: i for i, coal in enumerate(all_coalitions)}
         else:
-            coalitions_array = coalitions
-            coalitions_tuple = transform_array_to_coalitions(coalitions=coalitions_array)
+            coalitions_tuple = transform_array_to_coalitions(coalitions=coalitions)
             coalitions_dict = {coal: i for i, coal in enumerate(coalitions_tuple)}
 
         # run the game for all coalitions (no normalization)
         norm_value, self.normalization_value = self.normalization_value, 0
-        game_values: np.ndarray = self(coalitions_array)  # call the game with the coalitions
+        game_values = self(coalitions)  # call the game with the coalitions
         self.normalization_value = norm_value
 
         # update the storage with the new coalitions and values
         self.value_storage = game_values.astype(float)
         self.coalition_lookup = coalitions_dict
-        self.precompute_flag = True
+        self._precompute_flag = True
 
     def compute(
-        self, coalitions: np.ndarray | None = None
+        self, coalitions: CoalitionMatrix | None = None
     ) -> tuple[np.ndarray, dict[tuple[int, ...], int], float]:
         """Compute the game values for all or a given set of coalitions.
 
@@ -443,8 +456,7 @@ def compute(
             (array([0.25, 1.5]), {(1): 0, (1, 2): 1.5}, 0.0)
 
         """
-        coalitions: np.ndarray = self._check_coalitions(coalitions)
-        game_values = self.value_function(coalitions)
+        game_values = self.value_function(self._check_coalitions(coalitions))
 
         return game_values, self.coalition_lookup, self.normalization_value
 
@@ -470,7 +482,7 @@ def save_values(self, path: Path | str) -> None:
             self.precompute()
 
         # transform the values_storage to float16 for compression
-        self.value_storage.astype(np.float16)
+        values = self.value_storage.astype(np.float16)
 
         # cast the coalitions_in_storage to bool
         coalitions_in_storage = transform_coalitions_to_array(
@@ -481,7 +493,7 @@ def save_values(self, path: Path | str) -> None:
         # save the data
         np.savez_compressed(
             path,
-            values=self.value_storage,
+            values=values,
             coalitions=coalitions_in_storage,
             n_players=self.n_players,
             normalization_value=self.normalization_value,
@@ -517,7 +529,7 @@ def load_values(self, path: Path | str, *, precomputed: bool = False) -> None:
         self.value_storage = data["values"]
         coalition_lookup: list[tuple] = transform_array_to_coalitions(data["coalitions"])
         self.coalition_lookup = {coal: i for i, coal in enumerate(coalition_lookup)}
-        self.precompute_flag = precomputed
+        self._precompute_flag = precomputed
         self.normalization_value = float(data["normalization_value"])
 
     def save(self, path: Path | str) -> None: