115 lines
3.3 KiB
Python
115 lines
3.3 KiB
Python
"""
|
|
Cover Tree implementation.
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
import numpy as np
|
|
|
|
from .base import SurprisalTree, TreeNode
|
|
|
|
|
|
@dataclass
|
|
class CoverNode(TreeNode):
|
|
"""Node for Cover Tree with point, scale, and children."""
|
|
|
|
point: np.ndarray | None = None
|
|
scale: float = 0.0
|
|
children: list["CoverNode"] = field(default_factory=list)
|
|
|
|
|
|
class CoverTree(SurprisalTree):
|
|
"""
|
|
Cover Tree implementation with surprisal computation.
|
|
Organizes points hierarchically by scale.
|
|
"""
|
|
|
|
base: float
|
|
root: CoverNode | None
|
|
total_points: int
|
|
|
|
def __init__(self, base: float = 2.0, max_leaf_size: int = 10) -> None:
|
|
super().__init__(max_leaf_size)
|
|
self.base = base
|
|
self.root = None
|
|
|
|
def insert(self, point: np.ndarray) -> None:
|
|
if self.root is None:
|
|
self.root = CoverNode(point=point, scale=0.0, count=1)
|
|
else:
|
|
self._insert_recursive(self.root, point, self.root.scale)
|
|
self.total_points += 1
|
|
|
|
def _insert_recursive(
|
|
self, node: CoverNode, point: np.ndarray, scale: float
|
|
) -> None:
|
|
"""
|
|
Insert point into cover tree recursively.
|
|
Fixed to ensure proper tree structure and varied paths.
|
|
"""
|
|
node.count += 1
|
|
dist = float(np.linalg.norm(node.point - point))
|
|
|
|
cover_radius = self.base**scale
|
|
if dist <= cover_radius:
|
|
for child in node.children:
|
|
child_dist = float(np.linalg.norm(child.point - point))
|
|
child_radius = self.base ** (scale - 1)
|
|
if child_dist <= child_radius:
|
|
self._insert_recursive(child, point, scale - 1)
|
|
return
|
|
|
|
new_child = CoverNode(point=point, scale=scale - 1, count=1)
|
|
node.children.append(new_child)
|
|
else:
|
|
new_scale = scale + 1
|
|
|
|
new_root = CoverNode(
|
|
point=node.point, scale=new_scale, count=node.count + 1
|
|
)
|
|
new_root.children = [node]
|
|
|
|
new_sibling = CoverNode(point=point, scale=scale, count=1)
|
|
new_root.children.append(new_sibling)
|
|
|
|
self.root = new_root
|
|
|
|
def surprisal(self, point: np.ndarray) -> float:
|
|
"""
|
|
Compute surprisal based on path through cover tree.
|
|
Uses combination of branch probabilities and distance to final node.
|
|
"""
|
|
if self.root is None:
|
|
return float("inf")
|
|
|
|
surprisal_value = 0.0
|
|
node: CoverNode = self.root
|
|
|
|
while node.children:
|
|
best_child: CoverNode | None = None
|
|
best_dist = float("inf")
|
|
|
|
for child in node.children:
|
|
dist = float(np.linalg.norm(child.point - point))
|
|
if dist < best_dist:
|
|
best_dist = dist
|
|
best_child = child
|
|
|
|
if best_child is None:
|
|
break
|
|
|
|
parent_count = node.count
|
|
child_count = best_child.count
|
|
p_branch = child_count / parent_count
|
|
|
|
surprisal_value += -np.log(p_branch + 1e-10)
|
|
|
|
node = best_child
|
|
|
|
dist_to_rep = float(np.linalg.norm(node.point - point))
|
|
|
|
dim = len(point)
|
|
distance_surprisal = dim * np.log(dist_to_rep + 0.01)
|
|
|
|
return float(surprisal_value + distance_surprisal)
|