tritopic 0.1.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

tritopic/core/refinement.py

@@ -0,0 +1,231 @@
+"""
+Iterative Refinement Module
+
+Implements the iterative embedding refinement process that
+improves topic quality by pulling embeddings toward topic centroids.
+"""
+
+from typing import Tuple, Optional, Callable
+import numpy as np
+from scipy import sparse
+
+
+class IterativeRefinement:
+    """
+    Iteratively refines document embeddings based on topic assignments.
+    
+    The key insight is that after initial clustering, we can use the
+    topic structure to improve the embeddings, which in turn improves
+    the topic structure.
+    """
+    
+    def __init__(
+        self,
+        max_iterations: int = 5,
+        convergence_threshold: float = 0.95,
+        refinement_strength: float = 0.15,
+        verbose: bool = True,
+    ):
+        """
+        Initialize the refinement process.
+        
+        Parameters
+        ----------
+        max_iterations : int
+            Maximum number of refinement iterations
+        convergence_threshold : float
+            ARI threshold for convergence detection
+        refinement_strength : float
+            How strongly to pull embeddings toward centroids (0-1)
+        verbose : bool
+            Print progress information
+        """
+        self.max_iterations = max_iterations
+        self.convergence_threshold = convergence_threshold
+        self.refinement_strength = refinement_strength
+        self.verbose = verbose
+        
+        self.convergence_history_ = []
+    
+    def refine(
+        self,
+        embeddings: np.ndarray,
+        initial_labels: np.ndarray,
+        graph_builder_fn: Callable[[np.ndarray], sparse.csr_matrix],
+        cluster_fn: Callable[[sparse.csr_matrix], np.ndarray],
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Iteratively refine embeddings and labels.
+        
+        Parameters
+        ----------
+        embeddings : np.ndarray
+            Initial document embeddings
+        initial_labels : np.ndarray
+            Initial topic labels
+        graph_builder_fn : callable
+            Function to build graph from embeddings
+        cluster_fn : callable
+            Function to cluster a graph
+            
+        Returns
+        -------
+        Tuple[np.ndarray, np.ndarray]
+            Refined embeddings and final labels
+        """
+        from .clustering import compute_clustering_stability
+        
+        current_embeddings = embeddings.copy()
+        current_labels = initial_labels.copy()
+        
+        self.convergence_history_ = []
+        
+        for iteration in range(self.max_iterations):
+            if self.verbose:
+                print(f"      Iteration {iteration + 1}...")
+            
+            # Refine embeddings based on current labels
+            refined_embeddings = self._refine_embeddings(
+                current_embeddings, current_labels
+            )
+            
+            # Build new graph with refined embeddings
+            new_graph = graph_builder_fn(refined_embeddings)
+            
+            # Cluster the new graph
+            new_labels = cluster_fn(new_graph)
+            
+            # Check convergence
+            if iteration > 0:
+                stability = compute_clustering_stability(current_labels, new_labels)
+                self.convergence_history_.append(stability)
+                
+                if self.verbose:
+                    print(f"         ARI vs previous: {stability:.4f}")
+                
+                if stability >= self.convergence_threshold:
+                    if self.verbose:
+                        print(f"      ✓ Converged at iteration {iteration + 1}")
+                    current_embeddings = refined_embeddings
+                    current_labels = new_labels
+                    break
+            
+            current_embeddings = refined_embeddings
+            current_labels = new_labels
+        
+        return current_embeddings, current_labels
+    
+    def _refine_embeddings(
+        self,
+        embeddings: np.ndarray,
+        labels: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Refine embeddings by pulling them toward topic centroids.
+        
+        For each document in topic t:
+        new_embedding = (1 - β) * old_embedding + β * centroid_t
+        
+        where β is the refinement_strength.
+        """
+        refined = embeddings.copy()
+        unique_labels = np.unique(labels)
+        
+        for label in unique_labels:
+            if label < 0:  # Skip outliers
+                continue
+            
+            mask = labels == label
+            if mask.sum() == 0:
+                continue
+            
+            # Compute centroid for this topic
+            centroid = embeddings[mask].mean(axis=0)
+            
+            # Pull documents toward centroid
+            refined[mask] = (
+                (1 - self.refinement_strength) * embeddings[mask] +
+                self.refinement_strength * centroid
+            )
+        
+        # Re-normalize embeddings
+        norms = np.linalg.norm(refined, axis=1, keepdims=True)
+        norms[norms == 0] = 1
+        refined = refined / norms
+        
+        return refined
+    
+    @property
+    def n_iterations(self) -> int:
+        """Number of iterations performed."""
+        return len(self.convergence_history_) + 1
+
+
+class AdaptiveRefinement(IterativeRefinement):
+    """
+    Adaptive refinement that adjusts strength based on topic coherence.
+    
+    Topics with low internal coherence get stronger refinement,
+    while already coherent topics get lighter refinement.
+    """
+    
+    def __init__(
+        self,
+        max_iterations: int = 5,
+        convergence_threshold: float = 0.95,
+        base_strength: float = 0.15,
+        min_strength: float = 0.05,
+        max_strength: float = 0.30,
+        verbose: bool = True,
+    ):
+        super().__init__(
+            max_iterations=max_iterations,
+            convergence_threshold=convergence_threshold,
+            refinement_strength=base_strength,
+            verbose=verbose,
+        )
+        self.base_strength = base_strength
+        self.min_strength = min_strength
+        self.max_strength = max_strength
+    
+    def _refine_embeddings(
+        self,
+        embeddings: np.ndarray,
+        labels: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Refine embeddings with adaptive strength per topic.
+        """
+        refined = embeddings.copy()
+        unique_labels = np.unique(labels)
+        
+        for label in unique_labels:
+            if label < 0:
+                continue
+            
+            mask = labels == label
+            if mask.sum() < 2:
+                continue
+            
+            topic_embeddings = embeddings[mask]
+            centroid = topic_embeddings.mean(axis=0)
+            
+            # Compute internal coherence (average cosine similarity to centroid)
+            centroid_norm = centroid / (np.linalg.norm(centroid) + 1e-8)
+            emb_norms = topic_embeddings / (np.linalg.norm(topic_embeddings, axis=1, keepdims=True) + 1e-8)
+            coherence = np.mean(np.dot(emb_norms, centroid_norm))
+            
+            # Adaptive strength: less coherent topics get stronger refinement
+            # coherence is in [0, 1], invert it for strength
+            strength = self.base_strength + (1 - coherence) * (self.max_strength - self.base_strength)
+            strength = np.clip(strength, self.min_strength, self.max_strength)
+            
+            # Apply refinement
+            refined[mask] = (1 - strength) * embeddings[mask] + strength * centroid
+        
+        # Re-normalize
+        norms = np.linalg.norm(refined, axis=1, keepdims=True)
+        norms[norms == 0] = 1
+        refined = refined / norms
+        
+        return refined