Spaces:
Runtime error
Runtime error
| """ | |
| Graph-theoretic surprisal using k-NN graph and random walk. | |
| """ | |
| from __future__ import annotations | |
| 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 NearestNeighbors | |
| def _knn_indices( | |
| points: NDArray[np.floating[Any]], n_neighbors: int | |
| ) -> NDArray[np.intp]: | |
| """Get k-nearest neighbor indices for each point.""" | |
| from sklearn.neighbors import NearestNeighbors | |
| knn: NearestNeighbors = NearestNeighbors(n_neighbors=n_neighbors, algorithm="auto") | |
| knn.fit(points) # pyright: ignore[reportUnknownMemberType] | |
| _distances, indices = knn.kneighbors(points) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] | |
| return indices # pyright: ignore[reportUnknownVariableType] | |
| def _nearest_index(points: NDArray[np.floating[Any]], query: np.ndarray) -> int: | |
| """Find index of nearest point to query.""" | |
| from sklearn.neighbors import NearestNeighbors | |
| knn: NearestNeighbors = NearestNeighbors(n_neighbors=1) | |
| knn.fit(points) # pyright: ignore[reportUnknownMemberType] | |
| _distances, indices = knn.kneighbors([query]) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] | |
| return int(indices[0, 0]) # pyright: ignore[reportUnknownArgumentType] | |
| class GraphSurprisal(SurprisalTree): | |
| """ | |
| Graph-theoretic surprisal using k-NN graph and random walk. | |
| Surprisal based on stationary distribution of random walk. | |
| """ | |
| k: int | |
| max_iter: int | |
| points: list[NDArray[np.floating[Any]]] | |
| stationary_dist: NDArray[np.floating[Any]] | None | |
| graph_built: bool | |
| total_points: int | |
| def __init__( | |
| self, k: int = 5, max_iter: int = 100, max_leaf_size: int = 10 | |
| ) -> None: | |
| super().__init__(max_leaf_size) | |
| self.k = k | |
| self.max_iter = max_iter | |
| self.points = [] | |
| self.stationary_dist = None | |
| self.graph_built = False | |
| def insert(self, point: np.ndarray) -> None: | |
| self.points.append(point) | |
| self.total_points += 1 | |
| self.graph_built = False | |
| def batch_insert(self, points: np.ndarray) -> None: | |
| """More efficient batch insertion.""" | |
| self.points.extend(points) | |
| self.total_points += len(points) | |
| self.graph_built = False | |
| def _build_graph_and_compute_stationary(self) -> None: | |
| """Build k-NN graph and compute stationary distribution.""" | |
| if len(self.points) < 2: | |
| return | |
| points_array: NDArray[np.floating[Any]] = np.array(self.points) | |
| k_actual = min(self.k, len(self.points) - 1) | |
| indices = _knn_indices(points_array, k_actual + 1) | |
| n = len(self.points) | |
| transition = np.zeros((n, n)) | |
| for i in range(n): | |
| neighbors: NDArray[np.intp] = indices[i, 1:] | |
| for j_idx in neighbors: | |
| j: int = int(j_idx) | |
| transition[i, j] = 1.0 | |
| row_sums = transition.sum(axis=1, keepdims=True) | |
| # Add self-loops for isolated nodes to maintain stochasticity | |
| for i in range(n): | |
| if row_sums[i, 0] == 0: | |
| transition[i, i] = 1.0 | |
| row_sums[i, 0] = 1.0 | |
| transition = transition / row_sums | |
| stationary: NDArray[np.floating[Any]] = np.ones(n) / n | |
| for _ in range(self.max_iter): | |
| new_stationary: NDArray[np.floating[Any]] = transition.T @ stationary | |
| if np.allclose(new_stationary, stationary, atol=1e-6): | |
| break | |
| stationary = new_stationary | |
| self.stationary_dist = stationary / stationary.sum() | |
| self.graph_built = True | |
| def surprisal(self, point: np.ndarray) -> float: | |
| """ | |
| Compute surprisal based on stationary distribution. | |
| Well-connected central facts = high probability = low surprisal | |
| Peripheral isolated facts = low probability = high surprisal | |
| """ | |
| if not self.graph_built: | |
| self._build_graph_and_compute_stationary() | |
| if self.stationary_dist is None or len(self.points) == 0: | |
| return float("inf") | |
| points_array: NDArray[np.floating[Any]] = np.array(self.points) | |
| nearest_idx = _nearest_index(points_array, point) | |
| prob = self.stationary_dist[nearest_idx] | |
| return float(-np.log(prob + 1e-10)) | |