Spaces:
Runtime error
Runtime error
| """ | |
| Wrapper for sklearn's KDTree and BallTree. | |
| """ | |
| from typing import TYPE_CHECKING, Any | |
| import numpy as np | |
| from numpy.typing import NDArray | |
| from .base import SurprisalTree | |
| if TYPE_CHECKING: | |
| from sklearn.neighbors import ( | |
| BallTree, # pyright: ignore[reportUnknownVariableType] | |
| KDTree, # pyright: ignore[reportUnknownVariableType] | |
| ) | |
| class SklearnTreeWrapper(SurprisalTree): | |
| """ | |
| Wrapper for sklearn's KDTree and BallTree with surprisal computation. | |
| Uses density estimation via k-nearest neighbors. | |
| """ | |
| tree_type: str | |
| k: int | |
| points: list[NDArray[np.floating[Any]]] | |
| tree: "KDTree | BallTree | None" | |
| total_points: int | |
| def __init__( | |
| self, tree_type: str = "kd", k: int = 5, max_leaf_size: int = 10 | |
| ) -> None: | |
| super().__init__(max_leaf_size) | |
| self.tree_type = tree_type | |
| self.k = k | |
| self.points = [] | |
| self.tree = None | |
| def insert(self, point: np.ndarray) -> None: | |
| self.points.append(point) | |
| self.total_points += 1 | |
| self._rebuild_tree() | |
| def batch_insert(self, points: np.ndarray) -> None: | |
| """More efficient batch insertion.""" | |
| self.points.extend(points) | |
| self.total_points += len(points) | |
| self._rebuild_tree() | |
| def _rebuild_tree(self) -> None: | |
| if len(self.points) == 0: | |
| return | |
| from sklearn.neighbors import ( | |
| BallTree, # pyright: ignore[reportUnknownVariableType] | |
| KDTree, # pyright: ignore[reportUnknownVariableType] | |
| ) | |
| points_array: NDArray[np.floating[Any]] = np.array(self.points) | |
| if self.tree_type == "kd": | |
| self.tree = KDTree(points_array) | |
| elif self.tree_type == "ball": | |
| self.tree = BallTree(points_array) | |
| else: | |
| raise ValueError(f"Unknown tree type: {self.tree_type}") | |
| def surprisal(self, point: np.ndarray) -> float: | |
| """ | |
| Compute surprisal using k-NN density estimation. | |
| S(e) ≈ log(V_k(e)) where V_k is the volume of k-ball | |
| """ | |
| if self.tree is None or len(self.points) < self.k: # pyright: ignore[reportUnknownMemberType] | |
| return float("inf") | |
| k_actual = min(self.k, len(self.points)) | |
| distances, _indices = self.tree.query([point], k=k_actual) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] | |
| avg_distance: float = float(np.mean(distances[0])) # pyright: ignore[reportUnknownArgumentType] | |
| dim = point.shape[0] | |
| surprisal_value: float = dim * np.log(avg_distance + 1e-10) | |
| return surprisal_value | |