amachine.am_msp
1from __future__ import annotations 2 3import copy 4from fractions import Fraction 5import gc 6import warnings 7from collections import deque 8 9import matplotlib.pyplot as plt 10 11import numpy as np 12import scipy.sparse as sp_sparse 13import scipy.sparse.linalg as sp_linalg 14 15from .am_fast.distance import jensenshannondivergence_cpu as af_jensenshannondivergence_cpu 16 17# --------- Conditionally use cupy & cupyx if available else numpy and scipy 18 19try : 20 import cupy as cp 21 import cupyx as cpx 22 CUPY_AVAILABLE = True 23except ImportError : 24 CUPY_AVAILABLE = False 25 26if CUPY_AVAILABLE : 27 import cupyx.scipy.sparse as cpx_sparse 28 import cupyx.scipy.sparse.linalg as cpx_linalg 29 from .am_fast.distance import jensenshannondivergence_gpu as af_jensenshannondivergence_device 30else : 31 cp = None 32 cpx = None 33 cpx_sparse = None 34 cpx_linalg = None 35 36def to_cpu( array ): 37 return array.get() if CUPY_AVAILABLE and hasattr( array, 'get' ) else array 38 39# --------- Conditionally use cagra if available else hnswlib 40 41CAGRA_AVAILABLE=False 42 43if CUPY_AVAILABLE : 44 try : 45 from cuvs.neighbors import cagra 46 CAGRA_AVAILABLE = True 47 except ImportError : 48 CAGRA_AVAILABLE = False 49 50if not CAGRA_AVAILABLE : 51 import hnswlib 52else : 53 hnswlib = None 54 55# ------------------------------------------------------------- 56 57from .am_causal_state import CausalState 58from .am_transition import Transition 59 60class MSP: 61 """ 62 Mixed-State Presentation (MSP) for computing 'exact' intrinsic complexity. 63 64 Computes E (excess entropy), S (synchronization information), and T (transient information) 65 using spectral decomposition of the fundamental matrix. 66 67 See 68 Exact Complexity: The Spectral Decomposition of Intrinsic Computation 69 https://arxiv.org/abs/1309.3792 70 """ 71 72 def __init__( 73 self, 74 states, 75 belief_states, 76 transitions, 77 alphabet, 78 start_state_idx: int = 0, 79 gmres_rtol: float = 1e-06, 80 gmres_atol: float = 0, 81 gmres_maxiter: int = 1500, 82 eps: float = 1e-12, 83 ): 84 self.states = states 85 self.belief_states = np.array(belief_states) 86 self.transitions = transitions 87 self.alphabet = alphabet 88 self.start_state_idx = start_state_idx 89 self.gmres_rtol = gmres_rtol 90 self.gmres_atol = gmres_atol 91 self.gmres_maxiter = gmres_maxiter 92 self.EPS = eps 93 self._cache = {} 94 95 M = len(states) 96 if self.belief_states.ndim != 2 or self.belief_states.shape[0] != M: 97 raise ValueError( 98 f"belief_states must have shape (M={M}, N_causal), " 99 f"got {self.belief_states.shape}." 100 ) 101 if not (0 <= start_state_idx < M): 102 raise ValueError( 103 f"start_state_idx={start_state_idx} is out of range for M={M} states." 104 ) 105 if gmres_rtol <= 0: 106 raise ValueError(f"gmres_rtol must be positive, got {gmres_rtol}.") 107 if eps <= 0: 108 raise ValueError(f"eps must be positive, got {eps}.") 109 110 def _build_W(self) -> sp_sparse.csr_matrix: 111 """ 112 Build and return the sparse row-stochastic transition matrix W (M × M). 113 Rows are re-normalised to absorb floating-point drift. 114 """ 115 if "W" in self._cache: 116 return self._cache["W"] 117 118 M = len(self.states) 119 120 rows = [tr.origin_state_idx for tr in self.transitions] 121 cols = [tr.target_state_idx for tr in self.transitions] 122 data = [tr.prob for tr in self.transitions] 123 124 W = sp_sparse.csr_matrix((data, (rows, cols)), shape=(M, M)) 125 126 row_sums = np.array(W.sum(axis=1)).ravel() 127 128 if np.any(row_sums < self.EPS): 129 bad = np.where(row_sums < self.EPS)[0].tolist() 130 raise ValueError( 131 f"States at indices {bad} have zero total outgoing probability. " 132 "Check that every MSP state has at least one outgoing transition." 133 ) 134 135 W = (sp_sparse.diags(1.0 / row_sums) @ W).tocsr() 136 137 self._cache["W"] = W 138 return W 139 140 def _compute_pi_W(self, W) -> np.ndarray: 141 """ 142 Solve for the stationary distribution of W via GPU GMRES, then free 143 all intermediate GPU arrays before returning the CPU result. 144 """ 145 if "pi_W" in self._cache: 146 return self._cache["pi_W"] 147 148 xp = cp if CUPY_AVAILABLE else np 149 150 xp_sparse = cpx_sparse if CUPY_AVAILABLE else sp_sparse 151 xp_linalg = cpx_linalg if CUPY_AVAILABLE else sp_linalg 152 153 W_device = cpx_sparse.csr_matrix(W) 154 M = W_device.shape[0] 155 156 I_device = xp_sparse.eye(M, format="csr", dtype=W_device.dtype) 157 A_device = (W_device.T - I_device).tocsr() 158 159 ones_row = xp_sparse.csr_matrix(xp.ones((1, M), dtype=W_device.dtype)) 160 A_aug = xp_sparse.vstack([A_device[:-1, :], ones_row]).tocsr() 161 162 b_device = xp.zeros(M, dtype=W_device.dtype) 163 b_device[-1] = 1.0 164 165 diag_vals = A_aug.diagonal() 166 diag_vals[diag_vals == 0] = 1.0 167 M_inv = xp_sparse.diags(1.0 / diag_vals) 168 169 pi_W_device, info = xp_linalg.gmres( 170 A_aug, b_device, 171 M=M_inv, 172 rtol=self.gmres_rtol, 173 atol=self.gmres_atol, 174 restart=100, 175 maxiter=2000, 176 ) 177 178 if info > 0: 179 print(f"Warning: GMRES did not converge after {info} iterations") 180 elif info < 0: 181 raise ValueError("GMRES failed due to illegal input or breakdown.") 182 183 pi_W_device = xp.clip(pi_W_device, 0.0, None) 184 s = pi_W_device.sum() 185 186 if s < self.EPS: 187 raise ValueError("Stationary distribution collapsed.") 188 189 pi_W_device /= s 190 191 # Download before freeing 192 result = to_cpu( pi_W_device ) 193 194 # -- Free every GPU object created in this method ---------------------- 195 del W_device, I_device, A_device, ones_row, A_aug, b_device, M_inv, diag_vals, pi_W_device 196 197 if CUPY_AVAILABLE : 198 cp.get_default_memory_pool().free_all_blocks() 199 cp.get_default_pinned_memory_pool().free_all_blocks() 200 201 gc.collect() 202 203 self._cache["pi_W"] = result 204 205 return result 206 207 def _fundamental( 208 self, 209 W: sp_sparse.csr_matrix, 210 pi_W: np.ndarray, 211 ) -> tuple[sp_linalg.LinearOperator, np.ndarray, np.ndarray, LinearOperator]: 212 """ 213 Compute the fundamental matrix 214 """ 215 if "fundamental" in self._cache: 216 return self._cache["fundamental"] 217 218 M = W.shape[0] 219 WT = W.T.tocsc() 220 221 e0 = np.zeros(M) 222 e0[self.start_state_idx] = 1.0 223 224 def Q_T_matvec(v: np.ndarray) -> np.ndarray: 225 return v - WT @ v + pi_W * v.sum() 226 227 Q_LO = sp_linalg.LinearOperator((M, M), matvec=Q_T_matvec, dtype=float) 228 229 W_diag = np.array(W.diagonal()) 230 M_diag = 1.0 - W_diag + pi_W 231 M_diag = np.where(np.abs(M_diag) > self.EPS, M_diag, 1.0) 232 precond = sp_linalg.LinearOperator((M, M), matvec=lambda v: v / M_diag, dtype=float) 233 234 Z_row, exit_code = sp_linalg.lgmres( 235 Q_LO, 236 e0, 237 M=precond, 238 inner_m=50, 239 outer_k=15, 240 rtol=self.gmres_rtol, 241 atol=self.gmres_atol, 242 maxiter=self.gmres_maxiter, 243 ) 244 245 if exit_code != 0: 246 247 raise RuntimeError( 248 f"lgmres failed on the fundamental matrix solve (exit code {exit_code}). " 249 "Try increasing gmres_maxiter, loosening gmres_rtol, or verify that " 250 "the transition matrix is irreducible." 251 ) 252 253 fund_row = Z_row - pi_W 254 255 self._cache["fundamental"] = (Q_LO, Z_row, fund_row, precond) 256 return Q_LO, Z_row, fund_row, precond 257 258 def _fundamental_2( 259 self, 260 W: sp_sparse.csr_matrix, 261 pi_W: np.ndarray, 262 ) -> tuple[LinearOperator, np.ndarray, np.ndarray, LinearOperator]: 263 264 """Compute the fundamental matrix row using Shifted ILU Preconditioned LGMRES.""" 265 266 if "fundamental" in self._cache: 267 return self._cache["fundamental"] 268 269 M = W.shape[0] 270 WT = W.T.tocsc() # CSC format is required for spilu 271 272 e0 = np.zeros(M, dtype=float) 273 e0[self.start_state_idx] = 1.0 274 275 # Define Matrix-Vector product for Q^T = I - W^T + pi_W * 1^T 276 def Q_T_matvec(v: np.ndarray) -> np.ndarray: 277 return v - WT @ v + pi_W * v.sum() 278 279 Q_LO = sp_linalg.LinearOperator((M, M), matvec=Q_T_matvec, dtype=float) 280 281 # Construct Shifted ILU Preconditioner 282 # A_delta = (1 + delta)*I - W^T is strictly diagonally dominant and sparse 283 delta = getattr(self, "ilu_shift", 1e-3) 284 I_sp = sp_sparse.eye(M, format="csc", dtype=float) 285 A_delta = (1.0 + delta) * I_sp - WT 286 287 try: 288 # Compute Incomplete LU decomposition 289 ilu = sp_linalg.spilu( 290 A_delta, 291 drop_tol=getattr(self, "ilu_drop_tol", 1e-3), 292 fill_factor=getattr(self, "ilu_fill_factor", 2.0), 293 ) 294 295 precond = sp_linalg.LinearOperator( 296 (M, M), matvec=ilu.solve, dtype=float 297 ) 298 299 print( "LU decomposition computed" ) 300 301 except Exception as e: 302 # Fallback to Jacobi if ILU fails due to memory constraints 303 304 print(f"Warning: ILU preconditioning failed ({e}). Falling back to Jacobi.") 305 306 W_diag = np.array(W.diagonal()) 307 M_diag = 1.0 - W_diag + pi_W 308 M_diag = np.where(np.abs(M_diag) > self.EPS, M_diag, 1.0) 309 precond = sp_linalg.LinearOperator( 310 (M, M), matvec=lambda v: v / M_diag, dtype=float 311 ) 312 313 # Solve system Q^T * Z_row = e0 314 Z_row, exit_code = sp_linalg.lgmres( 315 Q_LO, 316 e0, 317 M=precond, 318 inner_m=60, 319 outer_k=15, 320 rtol=self.gmres_rtol, 321 atol=self.gmres_atol, 322 maxiter=self.gmres_maxiter, 323 ) 324 325 if exit_code != 0: 326 raise RuntimeError( 327 f"lgmres failed on the fundamental matrix solve (exit code {exit_code}). " 328 "Consider increasing ilu_fill_factor or adjusting gmres_rtol." 329 ) 330 331 fund_row = Z_row - pi_W 332 333 self._cache["fundamental"] = (Q_LO, Z_row, fund_row, precond) 334 return Q_LO, Z_row, fund_row, precond 335 336 def _get_H_WA(self) -> np.ndarray: 337 """ 338 Per-state Shannon entropy of the output symbol distribution, 339 marginalised over target states. Rows are normalised before use. 340 """ 341 if "H_WA" in self._cache: 342 return self._cache["H_WA"] 343 344 M, A = len(self.states), len(self.alphabet) 345 state_sym = np.zeros((M, A)) 346 for tr in self.transitions: 347 state_sym[tr.origin_state_idx, tr.symbol_idx] += tr.prob 348 349 sym_row_sums = state_sym.sum(axis=1, keepdims=True) 350 sym_row_sums = np.where(sym_row_sums > self.EPS, sym_row_sums, 1.0) 351 state_sym /= sym_row_sums 352 353 safe = state_sym > self.EPS 354 log_vals = np.where(safe, np.log2(np.where(safe, state_sym, 1.0)), 0.0) 355 H_WA = -np.sum(state_sym * log_vals, axis=1) 356 357 self._cache["H_WA"] = H_WA 358 359 return H_WA 360 361 def _get_H_eta(self) -> np.ndarray: 362 """ 363 Per-state Shannon entropy of the mixed-state (belief) vector. 364 Rows are normalised before entropy computation. 365 """ 366 if "H_eta" in self._cache: 367 return self._cache["H_eta"] 368 369 msv = self.belief_states.copy().astype(float) 370 row_sums = msv.sum(axis=1, keepdims=True) 371 row_sums = np.where(row_sums > self.EPS, row_sums, 1.0) 372 msv /= row_sums 373 374 safe = msv > self.EPS 375 log_vals = np.where(safe, np.log2(np.where(safe, msv, 1.0)), 0.0) 376 H_eta = -np.sum(msv * log_vals, axis=1) 377 378 self._cache["H_eta"] = H_eta 379 return H_eta 380 381 def get_E_S_T(self) -> tuple[float, float, float]: 382 383 try: 384 385 print( f"Performing Spectral Decomposition on MSP" ) 386 387 print( f"Computing W ..." ) 388 389 W = self._build_W() 390 391 print( f"Computing Pi_W ..." ) 392 393 pi_W = self._compute_pi_W(W) # GPU memory freed inside _compute_pi_W 394 395 print( f"Computing H_WA ..." ) 396 397 H_WA = self._get_H_WA() 398 399 print( f"Computing H_eta ..." ) 400 401 H_eta = self._get_H_eta() 402 403 print( f"Computing fundamental ..." ) 404 405 Q_LO, Z_row, fund_row, precond = self._fundamental_2( W, pi_W) 406 407 E = float(fund_row @ H_WA) 408 S = float(fund_row @ H_eta) 409 410 print( f"Computing ZZ_row ..." ) 411 412 ZZ_row, exit_code = sp_linalg.lgmres( 413 Q_LO, 414 Z_row, 415 x0=Z_row, 416 M=precond, 417 inner_m=50, 418 outer_k=15, 419 rtol=self.gmres_rtol, 420 atol=self.gmres_atol, 421 maxiter=self.gmres_maxiter, 422 ) 423 424 if exit_code != 0: 425 raise RuntimeError( 426 f"GMRES failed on the Z² solve for T (exit code {exit_code}). " 427 "Try increasing gmres_maxiter or loosening gmres_rtol." 428 ) 429 430 t_row = ZZ_row - pi_W 431 T = float(t_row @ H_WA) 432 433 print( f"Done." ) 434 435 return E, S, T 436 437 finally: 438 439 if CUPY_AVAILABLE : 440 cp.get_default_memory_pool().free_all_blocks() 441 cp.get_default_pinned_memory_pool().free_all_blocks() 442 443 gc.collect() 444 445 def clear_cache(self) -> None: 446 self._cache.clear() 447 448def compute_msp_exact( 449 T_x : list[ list[ list[ Fraction ] ] ], 450 pi : tuple[Fraction], 451 n_states : int, 452 alphabet : tuple, 453 exact_state_cap: int = 1000, 454 verbose: bool = True, 455): 456 457 n_symbols = len(alphabet) 458 459 def _apply(mu: tuple[Fraction, ...], x: int) : 460 461 Tx = T_x[x] 462 mu_Tx = [Fraction(0)] * n_states 463 for i, w in enumerate(mu): 464 if w == Fraction(0): 465 continue 466 row = Tx[i] 467 for j in range(n_states): 468 if row[j]: 469 mu_Tx[j] += w * row[j] 470 471 prob = sum(mu_Tx) 472 if prob == Fraction(0): 473 return None, Fraction(0) 474 475 mu_next = tuple(v / prob for v in mu_Tx) 476 return mu_next, prob 477 478 belief_states: list[tuple[Fraction, ...]] = [] 479 seen_states: dict[tuple, int] = {} 480 msp_transitions: list = [] 481 482 def _register(mu: tuple[Fraction, ...]) -> int: 483 idx = len(belief_states) 484 belief_states.append(mu) 485 seen_states[mu] = idx 486 return idx 487 488 start_mu = tuple(pi) 489 _register(start_mu) 490 frontier = deque([0]) 491 492 while frontier and len(belief_states) < exact_state_cap: 493 494 idx = frontier.popleft() 495 mu = belief_states[idx] 496 497 for x in range(n_symbols): 498 499 mu_next, prob = _apply(mu, x) 500 501 if mu_next is None: 502 continue 503 504 if mu_next not in seen_states: 505 next_idx = _register(mu_next) 506 frontier.append(next_idx) 507 else: 508 next_idx = seen_states[mu_next] 509 510 msp_transitions.append( 511 Transition( 512 origin_state_idx = idx, 513 target_state_idx = next_idx, 514 prob = float(prob), 515 symbol_idx = x, 516 pq = prob 517 ) 518 ) 519 520 if frontier: 521 raise RuntimeError( f"Exact state cap {exact_state_cap} exceded." ) 522 523 if verbose: 524 print(f"Exact MSP found: {len(belief_states)} states.") 525 526 msp = MSP( 527 states = [CausalState(name=f"MS_{i}") for i in range(len(belief_states))], 528 belief_states = belief_states, 529 transitions = msp_transitions, 530 alphabet = alphabet, 531 ) 532 533 return msp 534 535def compute_msp( 536 T_x : list[np.ndarray], 537 pi : np.ndarray, 538 n_states : int, 539 alphabet : tuple, 540 exact_state_cap: int = 175_000, 541 jsd_eps: float = 1e-7, 542 k_ann: int = 50, 543 EPS : float = 1e-12, 544 verbose = True, 545) : 546 547 xp = cp if CUPY_AVAILABLE else np 548 549 n_input_states = T_x[0].shape[0] 550 n_symbols = len(alphabet) 551 552 T_flat_device = xp.asarray( 553 np.stack(T_x).transpose(1, 0, 2).reshape(n_input_states, n_symbols * n_input_states) 554 ) 555 556 # -- Sparse Belief Vectors ------------------------------------------- 557 558 sparse_belief_vector_indices: list[np.ndarray] = [] 559 sparse_belief_vector_values: list[np.ndarray] = [] 560 561 def _to_sparse(v: np.ndarray): 562 563 non_zero = np.where( np.abs( v ) > 1e-20 )[ 0 ] 564 return non_zero.astype(np.int32), v[ non_zero ] 565 566 def _to_dense( sparse_indices: np.ndarray, sparse_values: np.ndarray ): 567 dense = np.zeros( n_input_states ) 568 dense[ sparse_indices ] = sparse_values 569 return dense 570 571 def _batch_to_dense( indicies: list[int] ) -> np.ndarray: 572 573 out = np.zeros((len(indicies), n_input_states), dtype=np.float64) 574 for row, i in enumerate(indicies): 575 out[row, sparse_belief_vector_indices[i]] = sparse_belief_vector_values[i] 576 577 return out 578 579 # -- Transitions ---------------------------------------------- 580 581 t_origin: list[int] = [] 582 t_target: list[int] = [] 583 t_prob: list[float] = [] 584 t_symbol: list[int] = [] 585 586 # -- Sparse dual bucket coarse hash ---------------------------------- 587 588 # Since we use dual bins 589 # Some x, y in p, q needs to differ by more than 1/(2*coarse_scale) to not share a bin 590 # Thus minimal total variation between non-bin sharing distributions is 591 # To ensure jsd(p,q) < jsd_eps -> p,q share a bin, we use Pinsker's inequality 592 # 1 /(2*coarse_scale) <= 2*sqrt(2*jsd_eps) 593 # 1 /( 4*sqrt(2*jsd_eps) ) <= coarse_scale 594 595 coarse_scale = 1 / ( 4*np.sqrt( 2 * jsd_eps ) ) 596 597 def _coarse_keys( sparse_indices: np.ndarray, sparse_values: np.ndarray) -> list[bytes]: 598 scaled = _to_dense( sparse_indices, sparse_values*coarse_scale ) 599 return [ 600 np.round(scaled).astype(np.int32).tobytes(), 601 np.round(scaled + 0.5).astype(np.int32).tobytes() 602 ] 603 604 buckets: dict[bytes, list[int]] = {} 605 606 # Jenson Shannon Divergence for sparse vectors 607 def _jsd_sparse( 608 ai: np.ndarray, av: np.ndarray, 609 bi: np.ndarray, bv: np.ndarray, 610 ) -> float: 611 612 # get sorted union of the [non-zero] indices over both vectors 613 all_i = np.union1d( ai, bi ) 614 615 # form dense vectors 616 p = np.zeros( len( all_i ), dtype=np.float64 ) 617 q = np.zeros( len( all_i ), dtype=np.float64 ) 618 619 # np.searchsorted( all_i, ai ) finds position in all_i where each ai is 620 p[ np.searchsorted( all_i, ai ) ] = av 621 q[ np.searchsorted( all_i, bi ) ] = bv 622 623 # Jenson Shannon divergence (JDS) 624 return af_jensenshannondivergence_cpu( p, q ) 625 626 # Check if an existing belief vector is within jsd_eps 627 def _lookup( sparse_indices: np.ndarray, sparse_values: np.ndarray ) -> int: 628 629 keys = _coarse_keys( sparse_indices, sparse_values ) 630 candidates = list( { idx for key in keys for idx in buckets.get( key, [] ) } ) 631 632 if not candidates : 633 return -1 634 635 dense = _to_dense( sparse_indices, sparse_values ).reshape( 1, n_input_states ) 636 candidate_vectors = _batch_to_dense( candidates ) 637 638 distances = af_jensenshannondivergence_cpu( candidate_vectors, dense, axis=1 ) 639 min_idx = np.argmin( distances ) 640 min_dist = distances[ min_idx ] 641 642 if min_dist < jsd_eps : 643 return candidates[ min_idx ] 644 645 return -1 646 647 # Map belief vector to coarse hash buckets for subsequent fast lookup 648 def _register( sparse_indices: np.ndarray, sparse_values: np.ndarray, idx: int): 649 for key in _coarse_keys( sparse_indices, sparse_values ): 650 buckets.setdefault( key, [] ).append( idx ) 651 652 pi_indicies, pi_values = _to_sparse( pi ) 653 sparse_belief_vector_indices.append( pi_indicies ) 654 sparse_belief_vector_values.append( pi_values ) 655 n_belief_states = 1 656 657 _register( pi_indicies, pi_values, 0 ) 658 659 frontier = deque([0]) 660 661 FRONTIER_CHUNK = 256 662 663 try: 664 while frontier and n_belief_states < exact_state_cap: 665 666 batch_indices = [] 667 while frontier and len(batch_indices) < FRONTIER_CHUNK: 668 batch_indices.append(frontier.popleft()) 669 670 B = len(batch_indices) 671 672 mus_device = xp.asarray(_batch_to_dense( batch_indices ) ) 673 flat_device = mus_device @ T_flat_device 674 all_mu_Tx_device = flat_device.reshape( B, n_symbols, n_input_states ) 675 probs_device = all_mu_Tx_device.sum(axis=2) 676 677 vb_device, vx_device = xp.where(probs_device > EPS) 678 679 if vb_device.size == 0: 680 del mus_device, flat_device, all_mu_Tx_device, probs_device, vb_device, vx_device 681 continue 682 683 vp_device = probs_device[vb_device, vx_device] 684 mu_next_device = all_mu_Tx_device[vb_device, vx_device] / vp_device[:, None] 685 686 mu_next_cpu = to_cpu( mu_next_device ) 687 688 vp_cpu = to_cpu( vp_device ) 689 vb_cpu = to_cpu( vb_device ) 690 vx_cpu = to_cpu( vx_device ) 691 692 # Free GPU arrays for this batch immediately after download 693 del mus_device, flat_device, all_mu_Tx_device, probs_device 694 del vb_device, vx_device, vp_device, mu_next_device 695 696 if not np.isfinite(mu_next_cpu).all(): 697 raise ValueError("Non-finite belief vectors in BFS batch") 698 699 for q in range( len( vb_cpu ) ): 700 701 indicies, values = _to_sparse(mu_next_cpu[q]) 702 orig = int( batch_indices[ int( vb_cpu[ q ] ) ] ) 703 idx = _lookup( indicies, values ) 704 705 if idx == -1: 706 707 idx = n_belief_states 708 sparse_belief_vector_indices.append( indicies ) 709 sparse_belief_vector_values.append( values ) 710 _register( indicies, values, idx ) 711 712 n_belief_states += 1 713 frontier.append(idx) 714 715 t_origin.append(orig) 716 t_target.append(idx) 717 t_prob.append(float(vp_cpu[q])) 718 t_symbol.append(int(vx_cpu[q])) 719 720 # -- Close the graph by mapping each dangling state to the nearest closed state 721 722 if frontier: 723 724 if verbose: 725 print(f"Max exact states ({exact_state_cap}) exceded ({n_belief_states}). Closing graph...") 726 727 dense_all = np.zeros((n_belief_states, n_input_states), dtype=np.float32) 728 729 for i in range(n_belief_states): 730 dense_all[i, sparse_belief_vector_indices[i]] = sparse_belief_vector_values[i].astype(np.float32) 731 732 ixp = xp if CAGRA_AVAILABLE else np 733 734 if CAGRA_AVAILABLE: 735 736 bv_device_f32 = xp.asarray(dense_all) 737 del dense_all 738 739 if verbose : 740 print("Building CAGRA index...") 741 742 index = cagra.build( 743 cagra.IndexParams(graph_degree=32, metric='sqeuclidean'), 744 bv_device_f32, 745 ) 746 747 else : 748 749 if verbose : 750 print("Building hnswlib index...") 751 752 T_flat_device = to_cpu( T_flat_device ) 753 bv_device_f32 = dense_all 754 index = hnswlib.Index(space='l2', dim=bv_device_f32[0].size) 755 index.init_index(max_elements=len(bv_device_f32), ef_construction=500, M=50) 756 index.set_ef(50) 757 index.add_items(bv_device_f32) 758 759 if verbose: 760 print("Index built.") 761 762 closure_buckets: dict[bytes, list[tuple]] = {} 763 max_error = 0.0 764 avg_error = 0.0 765 n_mapped = 0 766 767 def _closure_lookup( indicies: np.ndarray, values: np.ndarray) -> int: 768 seen = set() 769 for key in _coarse_keys( indicies, values): 770 for entry in closure_buckets.get(key, []): 771 eid = id(entry) 772 if eid not in seen: 773 seen.add(eid) 774 s_idx, s_val, target = entry 775 if _jsd_sparse(indicies, values, s_idx, s_val) < jsd_eps: 776 return target 777 return -1 778 779 def _closure_register(indicies: np.ndarray, values: np.ndarray, target: int): 780 entry = (indicies, values, target) 781 for key in _coarse_keys(indicies, values): 782 closure_buckets.setdefault(key, []).append(entry) 783 784 CHUNK = 10_000 785 frontier_list = list(frontier) 786 787 for start in range(0, len(frontier_list), CHUNK): 788 789 chunk_idxs = frontier_list[start : start + CHUNK] 790 C = len(chunk_idxs) 791 792 if CAGRA_AVAILABLE : 793 chunk_device = cp.asarray(_batch_to_dense(chunk_idxs)) 794 else : 795 chunk_device = _batch_to_dense(chunk_idxs) 796 797 flat_device = chunk_device @ T_flat_device 798 799 mu_Tx_device = flat_device.reshape(C, n_symbols, n_input_states) 800 801 probs_device = mu_Tx_device.sum(axis=2) 802 803 # Free intermediates we no longer need 804 del chunk_device, flat_device 805 806 vf_device, vx_device = ixp.where(probs_device > EPS) 807 808 if vf_device.size == 0: 809 del mu_Tx_device, probs_device, vf_device, vx_device 810 continue 811 812 vp_device = probs_device[vf_device, vx_device] 813 mu_v_device = mu_Tx_device[vf_device, vx_device] / vp_device[:, None] 814 815 # Free now-consumed GPU arrays 816 del mu_Tx_device, probs_device 817 818 mu_v_cpu = to_cpu( mu_v_device ) 819 vp_cpu = to_cpu( vp_device ) 820 vf_cpu = to_cpu( vf_device ) 821 vx_cpu = to_cpu( vx_device ) 822 823 del vf_device, vx_device, vp_device 824 825 # float32 copy for CAGRA; free float64 immediately after cast 826 unknown_device = mu_v_device.astype(ixp.float32) 827 del mu_v_device 828 829 N = len(vf_cpu) 830 orig_arr = np.array([chunk_idxs[f] for f in vf_cpu], dtype=np.int32) 831 next_idx = np.full(N, -1, dtype=np.int32) 832 unknown_mask = np.ones(N, dtype=bool) 833 834 sparse_cache = [_to_sparse(mu_v_cpu[q]) for q in range(N)] 835 836 for q in range(N): 837 indicies, values = sparse_cache[q] 838 target = _closure_lookup(indicies, values) 839 if target != -1: 840 next_idx[q] = target 841 unknown_mask[q] = False 842 843 if unknown_mask.any(): 844 845 unknown_rows = unknown_device[ixp.asarray(unknown_mask)] 846 847 if CAGRA_AVAILABLE : 848 849 if verbose: 850 print(f"CARGA search for {unknown_mask.sum()} unknown states...") 851 852 _, indices_device = cagra.search( 853 cagra.SearchParams(itopk_size=128), 854 index, unknown_rows, k_ann, 855 ) 856 857 else : 858 if verbose : 859 print(f"hnswlib search for {unknown_mask.sum()} unknown states...") 860 861 indices_device, _ = index.knn_query( to_cpu( unknown_rows ), k=k_ann ) 862 863 output_indices = ixp.asarray( indices_device ) 864 865 min_dist = ixp.full(unknown_rows.shape[0], ixp.inf, dtype=ixp.float32) 866 min_ks = ixp.zeros(unknown_rows.shape[0], dtype=ixp.int32) 867 868 for idx_k in range(output_indices.shape[1]): 869 sv = bv_device_f32[output_indices[:, idx_k], :] 870 871 if CAGRA_AVAILABLE : 872 d = af_jensenshannondivergence_device(unknown_rows, sv, axis=1) 873 else : 874 d = af_jensenshannondivergence_cpu(unknown_rows, sv, axis=1) 875 876 mask = d < min_dist 877 min_dist = ixp.where(mask, d, min_dist) 878 min_ks = ixp.where(mask, idx_k, min_ks) 879 880 max_dist = to_cpu( ixp.max(min_dist) ) 881 mean_dist = to_cpu( ixp.mean(min_dist) ) 882 883 n_mapped += N 884 avg_error += mean_dist 885 max_error = max(max_dist, max_error) 886 887 resolved = to_cpu( output_indices[ixp.arange(unknown_rows.shape[0]), min_ks] ) 888 unknown_qs = np.where(unknown_mask)[0] 889 890 for i, q in enumerate(unknown_qs): 891 nz_idx, nz_val = sparse_cache[q] 892 next_idx[q] = int(resolved[i]) 893 _closure_register(nz_idx, nz_val, int(resolved[i])) 894 895 # Free ANN-related GPU arrays 896 del unknown_rows, indices_device, output_indices 897 del min_dist, min_ks 898 899 del unknown_device 900 901 for q in range(N): 902 t_origin.append(int(orig_arr[q])) 903 t_target.append(int(next_idx[q])) 904 t_prob.append(float(vp_cpu[q])) 905 t_symbol.append(int(vx_cpu[q])) 906 907 # -- Free CAGRA index and float32 belief matrix -------------------- 908 909 # CAGRA memory lives outside CuPy's pool and is not released by free_all_blocks() 910 del index 911 del bv_device_f32 912 913 if CUPY_AVAILABLE : 914 915 cp.get_default_memory_pool().free_all_blocks() 916 cp.get_default_pinned_memory_pool().free_all_blocks() 917 918 gc.collect() 919 920 if verbose: 921 print(f"Closed: {n_mapped} dangling states") 922 print(f"Avg JSD state closure error: {(1.0 / np.log(2)) * avg_error / n_mapped} bits") 923 print(f"Max JSD state closure error: {(1.0 / np.log(2)) * max_error} bits") 924 925 # ------------------ Create the MSP --------------------# 926 927 msp_transitions = [ 928 Transition( 929 origin_state_idx = t_origin[i], 930 target_state_idx = t_target[i], 931 prob = t_prob[i], 932 symbol_idx = t_symbol[i], 933 ) 934 for i in range(len(t_origin)) 935 ] 936 937 belief_states = list(_batch_to_dense(list(range(n_belief_states)))) 938 939 940 msp = MSP( 941 states = [CausalState(name=f"MS_{i}") for i in range(n_belief_states)], 942 belief_states = belief_states, 943 transitions = msp_transitions, 944 alphabet = alphabet, 945 ) 946 947 if verbose: 948 print("Done.\n") 949 950 return msp 951 952 finally: 953 954 # ------------------ Cleanup --------------------# 955 956 if 'T_flat_device' in dir(): 957 del T_flat_device 958 959 if 'index' in locals(): 960 del index 961 962 if 'bv_device_f32' in locals(): 963 del bv_device_f32 964 965 if CUPY_AVAILABLE : 966 967 cp.get_default_memory_pool().free_all_blocks() 968 cp.get_default_pinned_memory_pool().free_all_blocks() 969 970 gc.collect()
def
to_cpu(array):
CAGRA_AVAILABLE =
True
class
MSP:
61class MSP: 62 """ 63 Mixed-State Presentation (MSP) for computing 'exact' intrinsic complexity. 64 65 Computes E (excess entropy), S (synchronization information), and T (transient information) 66 using spectral decomposition of the fundamental matrix. 67 68 See 69 Exact Complexity: The Spectral Decomposition of Intrinsic Computation 70 https://arxiv.org/abs/1309.3792 71 """ 72 73 def __init__( 74 self, 75 states, 76 belief_states, 77 transitions, 78 alphabet, 79 start_state_idx: int = 0, 80 gmres_rtol: float = 1e-06, 81 gmres_atol: float = 0, 82 gmres_maxiter: int = 1500, 83 eps: float = 1e-12, 84 ): 85 self.states = states 86 self.belief_states = np.array(belief_states) 87 self.transitions = transitions 88 self.alphabet = alphabet 89 self.start_state_idx = start_state_idx 90 self.gmres_rtol = gmres_rtol 91 self.gmres_atol = gmres_atol 92 self.gmres_maxiter = gmres_maxiter 93 self.EPS = eps 94 self._cache = {} 95 96 M = len(states) 97 if self.belief_states.ndim != 2 or self.belief_states.shape[0] != M: 98 raise ValueError( 99 f"belief_states must have shape (M={M}, N_causal), " 100 f"got {self.belief_states.shape}." 101 ) 102 if not (0 <= start_state_idx < M): 103 raise ValueError( 104 f"start_state_idx={start_state_idx} is out of range for M={M} states." 105 ) 106 if gmres_rtol <= 0: 107 raise ValueError(f"gmres_rtol must be positive, got {gmres_rtol}.") 108 if eps <= 0: 109 raise ValueError(f"eps must be positive, got {eps}.") 110 111 def _build_W(self) -> sp_sparse.csr_matrix: 112 """ 113 Build and return the sparse row-stochastic transition matrix W (M × M). 114 Rows are re-normalised to absorb floating-point drift. 115 """ 116 if "W" in self._cache: 117 return self._cache["W"] 118 119 M = len(self.states) 120 121 rows = [tr.origin_state_idx for tr in self.transitions] 122 cols = [tr.target_state_idx for tr in self.transitions] 123 data = [tr.prob for tr in self.transitions] 124 125 W = sp_sparse.csr_matrix((data, (rows, cols)), shape=(M, M)) 126 127 row_sums = np.array(W.sum(axis=1)).ravel() 128 129 if np.any(row_sums < self.EPS): 130 bad = np.where(row_sums < self.EPS)[0].tolist() 131 raise ValueError( 132 f"States at indices {bad} have zero total outgoing probability. " 133 "Check that every MSP state has at least one outgoing transition." 134 ) 135 136 W = (sp_sparse.diags(1.0 / row_sums) @ W).tocsr() 137 138 self._cache["W"] = W 139 return W 140 141 def _compute_pi_W(self, W) -> np.ndarray: 142 """ 143 Solve for the stationary distribution of W via GPU GMRES, then free 144 all intermediate GPU arrays before returning the CPU result. 145 """ 146 if "pi_W" in self._cache: 147 return self._cache["pi_W"] 148 149 xp = cp if CUPY_AVAILABLE else np 150 151 xp_sparse = cpx_sparse if CUPY_AVAILABLE else sp_sparse 152 xp_linalg = cpx_linalg if CUPY_AVAILABLE else sp_linalg 153 154 W_device = cpx_sparse.csr_matrix(W) 155 M = W_device.shape[0] 156 157 I_device = xp_sparse.eye(M, format="csr", dtype=W_device.dtype) 158 A_device = (W_device.T - I_device).tocsr() 159 160 ones_row = xp_sparse.csr_matrix(xp.ones((1, M), dtype=W_device.dtype)) 161 A_aug = xp_sparse.vstack([A_device[:-1, :], ones_row]).tocsr() 162 163 b_device = xp.zeros(M, dtype=W_device.dtype) 164 b_device[-1] = 1.0 165 166 diag_vals = A_aug.diagonal() 167 diag_vals[diag_vals == 0] = 1.0 168 M_inv = xp_sparse.diags(1.0 / diag_vals) 169 170 pi_W_device, info = xp_linalg.gmres( 171 A_aug, b_device, 172 M=M_inv, 173 rtol=self.gmres_rtol, 174 atol=self.gmres_atol, 175 restart=100, 176 maxiter=2000, 177 ) 178 179 if info > 0: 180 print(f"Warning: GMRES did not converge after {info} iterations") 181 elif info < 0: 182 raise ValueError("GMRES failed due to illegal input or breakdown.") 183 184 pi_W_device = xp.clip(pi_W_device, 0.0, None) 185 s = pi_W_device.sum() 186 187 if s < self.EPS: 188 raise ValueError("Stationary distribution collapsed.") 189 190 pi_W_device /= s 191 192 # Download before freeing 193 result = to_cpu( pi_W_device ) 194 195 # -- Free every GPU object created in this method ---------------------- 196 del W_device, I_device, A_device, ones_row, A_aug, b_device, M_inv, diag_vals, pi_W_device 197 198 if CUPY_AVAILABLE : 199 cp.get_default_memory_pool().free_all_blocks() 200 cp.get_default_pinned_memory_pool().free_all_blocks() 201 202 gc.collect() 203 204 self._cache["pi_W"] = result 205 206 return result 207 208 def _fundamental( 209 self, 210 W: sp_sparse.csr_matrix, 211 pi_W: np.ndarray, 212 ) -> tuple[sp_linalg.LinearOperator, np.ndarray, np.ndarray, LinearOperator]: 213 """ 214 Compute the fundamental matrix 215 """ 216 if "fundamental" in self._cache: 217 return self._cache["fundamental"] 218 219 M = W.shape[0] 220 WT = W.T.tocsc() 221 222 e0 = np.zeros(M) 223 e0[self.start_state_idx] = 1.0 224 225 def Q_T_matvec(v: np.ndarray) -> np.ndarray: 226 return v - WT @ v + pi_W * v.sum() 227 228 Q_LO = sp_linalg.LinearOperator((M, M), matvec=Q_T_matvec, dtype=float) 229 230 W_diag = np.array(W.diagonal()) 231 M_diag = 1.0 - W_diag + pi_W 232 M_diag = np.where(np.abs(M_diag) > self.EPS, M_diag, 1.0) 233 precond = sp_linalg.LinearOperator((M, M), matvec=lambda v: v / M_diag, dtype=float) 234 235 Z_row, exit_code = sp_linalg.lgmres( 236 Q_LO, 237 e0, 238 M=precond, 239 inner_m=50, 240 outer_k=15, 241 rtol=self.gmres_rtol, 242 atol=self.gmres_atol, 243 maxiter=self.gmres_maxiter, 244 ) 245 246 if exit_code != 0: 247 248 raise RuntimeError( 249 f"lgmres failed on the fundamental matrix solve (exit code {exit_code}). " 250 "Try increasing gmres_maxiter, loosening gmres_rtol, or verify that " 251 "the transition matrix is irreducible." 252 ) 253 254 fund_row = Z_row - pi_W 255 256 self._cache["fundamental"] = (Q_LO, Z_row, fund_row, precond) 257 return Q_LO, Z_row, fund_row, precond 258 259 def _fundamental_2( 260 self, 261 W: sp_sparse.csr_matrix, 262 pi_W: np.ndarray, 263 ) -> tuple[LinearOperator, np.ndarray, np.ndarray, LinearOperator]: 264 265 """Compute the fundamental matrix row using Shifted ILU Preconditioned LGMRES.""" 266 267 if "fundamental" in self._cache: 268 return self._cache["fundamental"] 269 270 M = W.shape[0] 271 WT = W.T.tocsc() # CSC format is required for spilu 272 273 e0 = np.zeros(M, dtype=float) 274 e0[self.start_state_idx] = 1.0 275 276 # Define Matrix-Vector product for Q^T = I - W^T + pi_W * 1^T 277 def Q_T_matvec(v: np.ndarray) -> np.ndarray: 278 return v - WT @ v + pi_W * v.sum() 279 280 Q_LO = sp_linalg.LinearOperator((M, M), matvec=Q_T_matvec, dtype=float) 281 282 # Construct Shifted ILU Preconditioner 283 # A_delta = (1 + delta)*I - W^T is strictly diagonally dominant and sparse 284 delta = getattr(self, "ilu_shift", 1e-3) 285 I_sp = sp_sparse.eye(M, format="csc", dtype=float) 286 A_delta = (1.0 + delta) * I_sp - WT 287 288 try: 289 # Compute Incomplete LU decomposition 290 ilu = sp_linalg.spilu( 291 A_delta, 292 drop_tol=getattr(self, "ilu_drop_tol", 1e-3), 293 fill_factor=getattr(self, "ilu_fill_factor", 2.0), 294 ) 295 296 precond = sp_linalg.LinearOperator( 297 (M, M), matvec=ilu.solve, dtype=float 298 ) 299 300 print( "LU decomposition computed" ) 301 302 except Exception as e: 303 # Fallback to Jacobi if ILU fails due to memory constraints 304 305 print(f"Warning: ILU preconditioning failed ({e}). Falling back to Jacobi.") 306 307 W_diag = np.array(W.diagonal()) 308 M_diag = 1.0 - W_diag + pi_W 309 M_diag = np.where(np.abs(M_diag) > self.EPS, M_diag, 1.0) 310 precond = sp_linalg.LinearOperator( 311 (M, M), matvec=lambda v: v / M_diag, dtype=float 312 ) 313 314 # Solve system Q^T * Z_row = e0 315 Z_row, exit_code = sp_linalg.lgmres( 316 Q_LO, 317 e0, 318 M=precond, 319 inner_m=60, 320 outer_k=15, 321 rtol=self.gmres_rtol, 322 atol=self.gmres_atol, 323 maxiter=self.gmres_maxiter, 324 ) 325 326 if exit_code != 0: 327 raise RuntimeError( 328 f"lgmres failed on the fundamental matrix solve (exit code {exit_code}). " 329 "Consider increasing ilu_fill_factor or adjusting gmres_rtol." 330 ) 331 332 fund_row = Z_row - pi_W 333 334 self._cache["fundamental"] = (Q_LO, Z_row, fund_row, precond) 335 return Q_LO, Z_row, fund_row, precond 336 337 def _get_H_WA(self) -> np.ndarray: 338 """ 339 Per-state Shannon entropy of the output symbol distribution, 340 marginalised over target states. Rows are normalised before use. 341 """ 342 if "H_WA" in self._cache: 343 return self._cache["H_WA"] 344 345 M, A = len(self.states), len(self.alphabet) 346 state_sym = np.zeros((M, A)) 347 for tr in self.transitions: 348 state_sym[tr.origin_state_idx, tr.symbol_idx] += tr.prob 349 350 sym_row_sums = state_sym.sum(axis=1, keepdims=True) 351 sym_row_sums = np.where(sym_row_sums > self.EPS, sym_row_sums, 1.0) 352 state_sym /= sym_row_sums 353 354 safe = state_sym > self.EPS 355 log_vals = np.where(safe, np.log2(np.where(safe, state_sym, 1.0)), 0.0) 356 H_WA = -np.sum(state_sym * log_vals, axis=1) 357 358 self._cache["H_WA"] = H_WA 359 360 return H_WA 361 362 def _get_H_eta(self) -> np.ndarray: 363 """ 364 Per-state Shannon entropy of the mixed-state (belief) vector. 365 Rows are normalised before entropy computation. 366 """ 367 if "H_eta" in self._cache: 368 return self._cache["H_eta"] 369 370 msv = self.belief_states.copy().astype(float) 371 row_sums = msv.sum(axis=1, keepdims=True) 372 row_sums = np.where(row_sums > self.EPS, row_sums, 1.0) 373 msv /= row_sums 374 375 safe = msv > self.EPS 376 log_vals = np.where(safe, np.log2(np.where(safe, msv, 1.0)), 0.0) 377 H_eta = -np.sum(msv * log_vals, axis=1) 378 379 self._cache["H_eta"] = H_eta 380 return H_eta 381 382 def get_E_S_T(self) -> tuple[float, float, float]: 383 384 try: 385 386 print( f"Performing Spectral Decomposition on MSP" ) 387 388 print( f"Computing W ..." ) 389 390 W = self._build_W() 391 392 print( f"Computing Pi_W ..." ) 393 394 pi_W = self._compute_pi_W(W) # GPU memory freed inside _compute_pi_W 395 396 print( f"Computing H_WA ..." ) 397 398 H_WA = self._get_H_WA() 399 400 print( f"Computing H_eta ..." ) 401 402 H_eta = self._get_H_eta() 403 404 print( f"Computing fundamental ..." ) 405 406 Q_LO, Z_row, fund_row, precond = self._fundamental_2( W, pi_W) 407 408 E = float(fund_row @ H_WA) 409 S = float(fund_row @ H_eta) 410 411 print( f"Computing ZZ_row ..." ) 412 413 ZZ_row, exit_code = sp_linalg.lgmres( 414 Q_LO, 415 Z_row, 416 x0=Z_row, 417 M=precond, 418 inner_m=50, 419 outer_k=15, 420 rtol=self.gmres_rtol, 421 atol=self.gmres_atol, 422 maxiter=self.gmres_maxiter, 423 ) 424 425 if exit_code != 0: 426 raise RuntimeError( 427 f"GMRES failed on the Z² solve for T (exit code {exit_code}). " 428 "Try increasing gmres_maxiter or loosening gmres_rtol." 429 ) 430 431 t_row = ZZ_row - pi_W 432 T = float(t_row @ H_WA) 433 434 print( f"Done." ) 435 436 return E, S, T 437 438 finally: 439 440 if CUPY_AVAILABLE : 441 cp.get_default_memory_pool().free_all_blocks() 442 cp.get_default_pinned_memory_pool().free_all_blocks() 443 444 gc.collect() 445 446 def clear_cache(self) -> None: 447 self._cache.clear()
Mixed-State Presentation (MSP) for computing 'exact' intrinsic complexity.
Computes E (excess entropy), S (synchronization information), and T (transient information) using spectral decomposition of the fundamental matrix.
See Exact Complexity: The Spectral Decomposition of Intrinsic Computation https://arxiv.org/abs/1309.3792
MSP( states, belief_states, transitions, alphabet, start_state_idx: int = 0, gmres_rtol: float = 1e-06, gmres_atol: float = 0, gmres_maxiter: int = 1500, eps: float = 1e-12)
73 def __init__( 74 self, 75 states, 76 belief_states, 77 transitions, 78 alphabet, 79 start_state_idx: int = 0, 80 gmres_rtol: float = 1e-06, 81 gmres_atol: float = 0, 82 gmres_maxiter: int = 1500, 83 eps: float = 1e-12, 84 ): 85 self.states = states 86 self.belief_states = np.array(belief_states) 87 self.transitions = transitions 88 self.alphabet = alphabet 89 self.start_state_idx = start_state_idx 90 self.gmres_rtol = gmres_rtol 91 self.gmres_atol = gmres_atol 92 self.gmres_maxiter = gmres_maxiter 93 self.EPS = eps 94 self._cache = {} 95 96 M = len(states) 97 if self.belief_states.ndim != 2 or self.belief_states.shape[0] != M: 98 raise ValueError( 99 f"belief_states must have shape (M={M}, N_causal), " 100 f"got {self.belief_states.shape}." 101 ) 102 if not (0 <= start_state_idx < M): 103 raise ValueError( 104 f"start_state_idx={start_state_idx} is out of range for M={M} states." 105 ) 106 if gmres_rtol <= 0: 107 raise ValueError(f"gmres_rtol must be positive, got {gmres_rtol}.") 108 if eps <= 0: 109 raise ValueError(f"eps must be positive, got {eps}.")
def
get_E_S_T(self) -> tuple[float, float, float]:
382 def get_E_S_T(self) -> tuple[float, float, float]: 383 384 try: 385 386 print( f"Performing Spectral Decomposition on MSP" ) 387 388 print( f"Computing W ..." ) 389 390 W = self._build_W() 391 392 print( f"Computing Pi_W ..." ) 393 394 pi_W = self._compute_pi_W(W) # GPU memory freed inside _compute_pi_W 395 396 print( f"Computing H_WA ..." ) 397 398 H_WA = self._get_H_WA() 399 400 print( f"Computing H_eta ..." ) 401 402 H_eta = self._get_H_eta() 403 404 print( f"Computing fundamental ..." ) 405 406 Q_LO, Z_row, fund_row, precond = self._fundamental_2( W, pi_W) 407 408 E = float(fund_row @ H_WA) 409 S = float(fund_row @ H_eta) 410 411 print( f"Computing ZZ_row ..." ) 412 413 ZZ_row, exit_code = sp_linalg.lgmres( 414 Q_LO, 415 Z_row, 416 x0=Z_row, 417 M=precond, 418 inner_m=50, 419 outer_k=15, 420 rtol=self.gmres_rtol, 421 atol=self.gmres_atol, 422 maxiter=self.gmres_maxiter, 423 ) 424 425 if exit_code != 0: 426 raise RuntimeError( 427 f"GMRES failed on the Z² solve for T (exit code {exit_code}). " 428 "Try increasing gmres_maxiter or loosening gmres_rtol." 429 ) 430 431 t_row = ZZ_row - pi_W 432 T = float(t_row @ H_WA) 433 434 print( f"Done." ) 435 436 return E, S, T 437 438 finally: 439 440 if CUPY_AVAILABLE : 441 cp.get_default_memory_pool().free_all_blocks() 442 cp.get_default_pinned_memory_pool().free_all_blocks() 443 444 gc.collect()
def
compute_msp_exact( T_x: list[list[list[fractions.Fraction]]], pi: tuple[fractions.Fraction], n_states: int, alphabet: tuple, exact_state_cap: int = 1000, verbose: bool = True):
449def compute_msp_exact( 450 T_x : list[ list[ list[ Fraction ] ] ], 451 pi : tuple[Fraction], 452 n_states : int, 453 alphabet : tuple, 454 exact_state_cap: int = 1000, 455 verbose: bool = True, 456): 457 458 n_symbols = len(alphabet) 459 460 def _apply(mu: tuple[Fraction, ...], x: int) : 461 462 Tx = T_x[x] 463 mu_Tx = [Fraction(0)] * n_states 464 for i, w in enumerate(mu): 465 if w == Fraction(0): 466 continue 467 row = Tx[i] 468 for j in range(n_states): 469 if row[j]: 470 mu_Tx[j] += w * row[j] 471 472 prob = sum(mu_Tx) 473 if prob == Fraction(0): 474 return None, Fraction(0) 475 476 mu_next = tuple(v / prob for v in mu_Tx) 477 return mu_next, prob 478 479 belief_states: list[tuple[Fraction, ...]] = [] 480 seen_states: dict[tuple, int] = {} 481 msp_transitions: list = [] 482 483 def _register(mu: tuple[Fraction, ...]) -> int: 484 idx = len(belief_states) 485 belief_states.append(mu) 486 seen_states[mu] = idx 487 return idx 488 489 start_mu = tuple(pi) 490 _register(start_mu) 491 frontier = deque([0]) 492 493 while frontier and len(belief_states) < exact_state_cap: 494 495 idx = frontier.popleft() 496 mu = belief_states[idx] 497 498 for x in range(n_symbols): 499 500 mu_next, prob = _apply(mu, x) 501 502 if mu_next is None: 503 continue 504 505 if mu_next not in seen_states: 506 next_idx = _register(mu_next) 507 frontier.append(next_idx) 508 else: 509 next_idx = seen_states[mu_next] 510 511 msp_transitions.append( 512 Transition( 513 origin_state_idx = idx, 514 target_state_idx = next_idx, 515 prob = float(prob), 516 symbol_idx = x, 517 pq = prob 518 ) 519 ) 520 521 if frontier: 522 raise RuntimeError( f"Exact state cap {exact_state_cap} exceded." ) 523 524 if verbose: 525 print(f"Exact MSP found: {len(belief_states)} states.") 526 527 msp = MSP( 528 states = [CausalState(name=f"MS_{i}") for i in range(len(belief_states))], 529 belief_states = belief_states, 530 transitions = msp_transitions, 531 alphabet = alphabet, 532 ) 533 534 return msp
def
compute_msp( T_x: list[numpy.ndarray], pi: numpy.ndarray, n_states: int, alphabet: tuple, exact_state_cap: int = 175000, jsd_eps: float = 1e-07, k_ann: int = 50, EPS: float = 1e-12, verbose=True):
536def compute_msp( 537 T_x : list[np.ndarray], 538 pi : np.ndarray, 539 n_states : int, 540 alphabet : tuple, 541 exact_state_cap: int = 175_000, 542 jsd_eps: float = 1e-7, 543 k_ann: int = 50, 544 EPS : float = 1e-12, 545 verbose = True, 546) : 547 548 xp = cp if CUPY_AVAILABLE else np 549 550 n_input_states = T_x[0].shape[0] 551 n_symbols = len(alphabet) 552 553 T_flat_device = xp.asarray( 554 np.stack(T_x).transpose(1, 0, 2).reshape(n_input_states, n_symbols * n_input_states) 555 ) 556 557 # -- Sparse Belief Vectors ------------------------------------------- 558 559 sparse_belief_vector_indices: list[np.ndarray] = [] 560 sparse_belief_vector_values: list[np.ndarray] = [] 561 562 def _to_sparse(v: np.ndarray): 563 564 non_zero = np.where( np.abs( v ) > 1e-20 )[ 0 ] 565 return non_zero.astype(np.int32), v[ non_zero ] 566 567 def _to_dense( sparse_indices: np.ndarray, sparse_values: np.ndarray ): 568 dense = np.zeros( n_input_states ) 569 dense[ sparse_indices ] = sparse_values 570 return dense 571 572 def _batch_to_dense( indicies: list[int] ) -> np.ndarray: 573 574 out = np.zeros((len(indicies), n_input_states), dtype=np.float64) 575 for row, i in enumerate(indicies): 576 out[row, sparse_belief_vector_indices[i]] = sparse_belief_vector_values[i] 577 578 return out 579 580 # -- Transitions ---------------------------------------------- 581 582 t_origin: list[int] = [] 583 t_target: list[int] = [] 584 t_prob: list[float] = [] 585 t_symbol: list[int] = [] 586 587 # -- Sparse dual bucket coarse hash ---------------------------------- 588 589 # Since we use dual bins 590 # Some x, y in p, q needs to differ by more than 1/(2*coarse_scale) to not share a bin 591 # Thus minimal total variation between non-bin sharing distributions is 592 # To ensure jsd(p,q) < jsd_eps -> p,q share a bin, we use Pinsker's inequality 593 # 1 /(2*coarse_scale) <= 2*sqrt(2*jsd_eps) 594 # 1 /( 4*sqrt(2*jsd_eps) ) <= coarse_scale 595 596 coarse_scale = 1 / ( 4*np.sqrt( 2 * jsd_eps ) ) 597 598 def _coarse_keys( sparse_indices: np.ndarray, sparse_values: np.ndarray) -> list[bytes]: 599 scaled = _to_dense( sparse_indices, sparse_values*coarse_scale ) 600 return [ 601 np.round(scaled).astype(np.int32).tobytes(), 602 np.round(scaled + 0.5).astype(np.int32).tobytes() 603 ] 604 605 buckets: dict[bytes, list[int]] = {} 606 607 # Jenson Shannon Divergence for sparse vectors 608 def _jsd_sparse( 609 ai: np.ndarray, av: np.ndarray, 610 bi: np.ndarray, bv: np.ndarray, 611 ) -> float: 612 613 # get sorted union of the [non-zero] indices over both vectors 614 all_i = np.union1d( ai, bi ) 615 616 # form dense vectors 617 p = np.zeros( len( all_i ), dtype=np.float64 ) 618 q = np.zeros( len( all_i ), dtype=np.float64 ) 619 620 # np.searchsorted( all_i, ai ) finds position in all_i where each ai is 621 p[ np.searchsorted( all_i, ai ) ] = av 622 q[ np.searchsorted( all_i, bi ) ] = bv 623 624 # Jenson Shannon divergence (JDS) 625 return af_jensenshannondivergence_cpu( p, q ) 626 627 # Check if an existing belief vector is within jsd_eps 628 def _lookup( sparse_indices: np.ndarray, sparse_values: np.ndarray ) -> int: 629 630 keys = _coarse_keys( sparse_indices, sparse_values ) 631 candidates = list( { idx for key in keys for idx in buckets.get( key, [] ) } ) 632 633 if not candidates : 634 return -1 635 636 dense = _to_dense( sparse_indices, sparse_values ).reshape( 1, n_input_states ) 637 candidate_vectors = _batch_to_dense( candidates ) 638 639 distances = af_jensenshannondivergence_cpu( candidate_vectors, dense, axis=1 ) 640 min_idx = np.argmin( distances ) 641 min_dist = distances[ min_idx ] 642 643 if min_dist < jsd_eps : 644 return candidates[ min_idx ] 645 646 return -1 647 648 # Map belief vector to coarse hash buckets for subsequent fast lookup 649 def _register( sparse_indices: np.ndarray, sparse_values: np.ndarray, idx: int): 650 for key in _coarse_keys( sparse_indices, sparse_values ): 651 buckets.setdefault( key, [] ).append( idx ) 652 653 pi_indicies, pi_values = _to_sparse( pi ) 654 sparse_belief_vector_indices.append( pi_indicies ) 655 sparse_belief_vector_values.append( pi_values ) 656 n_belief_states = 1 657 658 _register( pi_indicies, pi_values, 0 ) 659 660 frontier = deque([0]) 661 662 FRONTIER_CHUNK = 256 663 664 try: 665 while frontier and n_belief_states < exact_state_cap: 666 667 batch_indices = [] 668 while frontier and len(batch_indices) < FRONTIER_CHUNK: 669 batch_indices.append(frontier.popleft()) 670 671 B = len(batch_indices) 672 673 mus_device = xp.asarray(_batch_to_dense( batch_indices ) ) 674 flat_device = mus_device @ T_flat_device 675 all_mu_Tx_device = flat_device.reshape( B, n_symbols, n_input_states ) 676 probs_device = all_mu_Tx_device.sum(axis=2) 677 678 vb_device, vx_device = xp.where(probs_device > EPS) 679 680 if vb_device.size == 0: 681 del mus_device, flat_device, all_mu_Tx_device, probs_device, vb_device, vx_device 682 continue 683 684 vp_device = probs_device[vb_device, vx_device] 685 mu_next_device = all_mu_Tx_device[vb_device, vx_device] / vp_device[:, None] 686 687 mu_next_cpu = to_cpu( mu_next_device ) 688 689 vp_cpu = to_cpu( vp_device ) 690 vb_cpu = to_cpu( vb_device ) 691 vx_cpu = to_cpu( vx_device ) 692 693 # Free GPU arrays for this batch immediately after download 694 del mus_device, flat_device, all_mu_Tx_device, probs_device 695 del vb_device, vx_device, vp_device, mu_next_device 696 697 if not np.isfinite(mu_next_cpu).all(): 698 raise ValueError("Non-finite belief vectors in BFS batch") 699 700 for q in range( len( vb_cpu ) ): 701 702 indicies, values = _to_sparse(mu_next_cpu[q]) 703 orig = int( batch_indices[ int( vb_cpu[ q ] ) ] ) 704 idx = _lookup( indicies, values ) 705 706 if idx == -1: 707 708 idx = n_belief_states 709 sparse_belief_vector_indices.append( indicies ) 710 sparse_belief_vector_values.append( values ) 711 _register( indicies, values, idx ) 712 713 n_belief_states += 1 714 frontier.append(idx) 715 716 t_origin.append(orig) 717 t_target.append(idx) 718 t_prob.append(float(vp_cpu[q])) 719 t_symbol.append(int(vx_cpu[q])) 720 721 # -- Close the graph by mapping each dangling state to the nearest closed state 722 723 if frontier: 724 725 if verbose: 726 print(f"Max exact states ({exact_state_cap}) exceded ({n_belief_states}). Closing graph...") 727 728 dense_all = np.zeros((n_belief_states, n_input_states), dtype=np.float32) 729 730 for i in range(n_belief_states): 731 dense_all[i, sparse_belief_vector_indices[i]] = sparse_belief_vector_values[i].astype(np.float32) 732 733 ixp = xp if CAGRA_AVAILABLE else np 734 735 if CAGRA_AVAILABLE: 736 737 bv_device_f32 = xp.asarray(dense_all) 738 del dense_all 739 740 if verbose : 741 print("Building CAGRA index...") 742 743 index = cagra.build( 744 cagra.IndexParams(graph_degree=32, metric='sqeuclidean'), 745 bv_device_f32, 746 ) 747 748 else : 749 750 if verbose : 751 print("Building hnswlib index...") 752 753 T_flat_device = to_cpu( T_flat_device ) 754 bv_device_f32 = dense_all 755 index = hnswlib.Index(space='l2', dim=bv_device_f32[0].size) 756 index.init_index(max_elements=len(bv_device_f32), ef_construction=500, M=50) 757 index.set_ef(50) 758 index.add_items(bv_device_f32) 759 760 if verbose: 761 print("Index built.") 762 763 closure_buckets: dict[bytes, list[tuple]] = {} 764 max_error = 0.0 765 avg_error = 0.0 766 n_mapped = 0 767 768 def _closure_lookup( indicies: np.ndarray, values: np.ndarray) -> int: 769 seen = set() 770 for key in _coarse_keys( indicies, values): 771 for entry in closure_buckets.get(key, []): 772 eid = id(entry) 773 if eid not in seen: 774 seen.add(eid) 775 s_idx, s_val, target = entry 776 if _jsd_sparse(indicies, values, s_idx, s_val) < jsd_eps: 777 return target 778 return -1 779 780 def _closure_register(indicies: np.ndarray, values: np.ndarray, target: int): 781 entry = (indicies, values, target) 782 for key in _coarse_keys(indicies, values): 783 closure_buckets.setdefault(key, []).append(entry) 784 785 CHUNK = 10_000 786 frontier_list = list(frontier) 787 788 for start in range(0, len(frontier_list), CHUNK): 789 790 chunk_idxs = frontier_list[start : start + CHUNK] 791 C = len(chunk_idxs) 792 793 if CAGRA_AVAILABLE : 794 chunk_device = cp.asarray(_batch_to_dense(chunk_idxs)) 795 else : 796 chunk_device = _batch_to_dense(chunk_idxs) 797 798 flat_device = chunk_device @ T_flat_device 799 800 mu_Tx_device = flat_device.reshape(C, n_symbols, n_input_states) 801 802 probs_device = mu_Tx_device.sum(axis=2) 803 804 # Free intermediates we no longer need 805 del chunk_device, flat_device 806 807 vf_device, vx_device = ixp.where(probs_device > EPS) 808 809 if vf_device.size == 0: 810 del mu_Tx_device, probs_device, vf_device, vx_device 811 continue 812 813 vp_device = probs_device[vf_device, vx_device] 814 mu_v_device = mu_Tx_device[vf_device, vx_device] / vp_device[:, None] 815 816 # Free now-consumed GPU arrays 817 del mu_Tx_device, probs_device 818 819 mu_v_cpu = to_cpu( mu_v_device ) 820 vp_cpu = to_cpu( vp_device ) 821 vf_cpu = to_cpu( vf_device ) 822 vx_cpu = to_cpu( vx_device ) 823 824 del vf_device, vx_device, vp_device 825 826 # float32 copy for CAGRA; free float64 immediately after cast 827 unknown_device = mu_v_device.astype(ixp.float32) 828 del mu_v_device 829 830 N = len(vf_cpu) 831 orig_arr = np.array([chunk_idxs[f] for f in vf_cpu], dtype=np.int32) 832 next_idx = np.full(N, -1, dtype=np.int32) 833 unknown_mask = np.ones(N, dtype=bool) 834 835 sparse_cache = [_to_sparse(mu_v_cpu[q]) for q in range(N)] 836 837 for q in range(N): 838 indicies, values = sparse_cache[q] 839 target = _closure_lookup(indicies, values) 840 if target != -1: 841 next_idx[q] = target 842 unknown_mask[q] = False 843 844 if unknown_mask.any(): 845 846 unknown_rows = unknown_device[ixp.asarray(unknown_mask)] 847 848 if CAGRA_AVAILABLE : 849 850 if verbose: 851 print(f"CARGA search for {unknown_mask.sum()} unknown states...") 852 853 _, indices_device = cagra.search( 854 cagra.SearchParams(itopk_size=128), 855 index, unknown_rows, k_ann, 856 ) 857 858 else : 859 if verbose : 860 print(f"hnswlib search for {unknown_mask.sum()} unknown states...") 861 862 indices_device, _ = index.knn_query( to_cpu( unknown_rows ), k=k_ann ) 863 864 output_indices = ixp.asarray( indices_device ) 865 866 min_dist = ixp.full(unknown_rows.shape[0], ixp.inf, dtype=ixp.float32) 867 min_ks = ixp.zeros(unknown_rows.shape[0], dtype=ixp.int32) 868 869 for idx_k in range(output_indices.shape[1]): 870 sv = bv_device_f32[output_indices[:, idx_k], :] 871 872 if CAGRA_AVAILABLE : 873 d = af_jensenshannondivergence_device(unknown_rows, sv, axis=1) 874 else : 875 d = af_jensenshannondivergence_cpu(unknown_rows, sv, axis=1) 876 877 mask = d < min_dist 878 min_dist = ixp.where(mask, d, min_dist) 879 min_ks = ixp.where(mask, idx_k, min_ks) 880 881 max_dist = to_cpu( ixp.max(min_dist) ) 882 mean_dist = to_cpu( ixp.mean(min_dist) ) 883 884 n_mapped += N 885 avg_error += mean_dist 886 max_error = max(max_dist, max_error) 887 888 resolved = to_cpu( output_indices[ixp.arange(unknown_rows.shape[0]), min_ks] ) 889 unknown_qs = np.where(unknown_mask)[0] 890 891 for i, q in enumerate(unknown_qs): 892 nz_idx, nz_val = sparse_cache[q] 893 next_idx[q] = int(resolved[i]) 894 _closure_register(nz_idx, nz_val, int(resolved[i])) 895 896 # Free ANN-related GPU arrays 897 del unknown_rows, indices_device, output_indices 898 del min_dist, min_ks 899 900 del unknown_device 901 902 for q in range(N): 903 t_origin.append(int(orig_arr[q])) 904 t_target.append(int(next_idx[q])) 905 t_prob.append(float(vp_cpu[q])) 906 t_symbol.append(int(vx_cpu[q])) 907 908 # -- Free CAGRA index and float32 belief matrix -------------------- 909 910 # CAGRA memory lives outside CuPy's pool and is not released by free_all_blocks() 911 del index 912 del bv_device_f32 913 914 if CUPY_AVAILABLE : 915 916 cp.get_default_memory_pool().free_all_blocks() 917 cp.get_default_pinned_memory_pool().free_all_blocks() 918 919 gc.collect() 920 921 if verbose: 922 print(f"Closed: {n_mapped} dangling states") 923 print(f"Avg JSD state closure error: {(1.0 / np.log(2)) * avg_error / n_mapped} bits") 924 print(f"Max JSD state closure error: {(1.0 / np.log(2)) * max_error} bits") 925 926 # ------------------ Create the MSP --------------------# 927 928 msp_transitions = [ 929 Transition( 930 origin_state_idx = t_origin[i], 931 target_state_idx = t_target[i], 932 prob = t_prob[i], 933 symbol_idx = t_symbol[i], 934 ) 935 for i in range(len(t_origin)) 936 ] 937 938 belief_states = list(_batch_to_dense(list(range(n_belief_states)))) 939 940 941 msp = MSP( 942 states = [CausalState(name=f"MS_{i}") for i in range(n_belief_states)], 943 belief_states = belief_states, 944 transitions = msp_transitions, 945 alphabet = alphabet, 946 ) 947 948 if verbose: 949 print("Done.\n") 950 951 return msp 952 953 finally: 954 955 # ------------------ Cleanup --------------------# 956 957 if 'T_flat_device' in dir(): 958 del T_flat_device 959 960 if 'index' in locals(): 961 del index 962 963 if 'bv_device_f32' in locals(): 964 del bv_device_f32 965 966 if CUPY_AVAILABLE : 967 968 cp.get_default_memory_pool().free_all_blocks() 969 cp.get_default_pinned_memory_pool().free_all_blocks() 970 971 gc.collect()