reppi
reppi — Representation Learning Algorithms
A Python library implementing classical sparse representation and dictionary learning algorithms.
Modules
sparse Sparse coding (OMP, Batch-OMP). dictionary Dictionary learning (K-SVD, LC-KSVD1, LC-KSVD2).
1""" 2reppi — Representation Learning Algorithms 3========================================== 4 5A Python library implementing classical sparse representation and 6dictionary learning algorithms. 7 8Modules 9------- 10sparse 11 Sparse coding (OMP, Batch-OMP). 12dictionary 13 Dictionary learning (K-SVD, LC-KSVD1, LC-KSVD2). 14""" 15 16from reppi.dictionary import ( 17 FDDL, 18 KSVD, 19 LCKSVD, 20 FrozenDictionaryLearner, 21 IncrementalFrozenDictionary, 22) 23from reppi.sparse import FISTA, OMP, fista_core 24 25__all__ = ["FDDL", "FISTA", "KSVD", "LCKSVD", "OMP", "FrozenDictionaryLearner", "IncrementalFrozenDictionary", "fista_core"] 26__version__ = "1.0.0"
178class FDDL: 179 """ 180 Fisher Discrimination Dictionary Learning. 181 182 Parameters 183 ---------- 184 n_components : int or sequence of int 185 Total number of dictionary atoms, split evenly across classes 186 (remainder to the last class), or an explicit per-class atom 187 count. The paper usually sets all pi equal (Sec. 6.1). 188 lambda1 : float 189 L1 sparsity weight (Eq. 6). 190 lambda2 : float 191 Weight of the Fisher discriminative coefficient term (Eq. 6). 192 eta : float 193 Elastic term weight in f(X) (Eq. 5). The paper fixes eta=1, 194 which is sufficient for fi(Xi) to be strictly convex whenever 195 eta > 1 - ni/n for every class (Appendix A); eta=1 satisfies 196 this for any class with more than one training sample. 197 n_iter : int 198 Maximum number of outer alternating iterations (D/X updates). 199 tol : float or None 200 Outer-loop relative convergence tolerance on J(D,X) (Eq. 6) 201 between consecutive iterations (Table 1, step 4: "return to 202 step 2 until the values of J(D,X) in adjacent iterations are 203 close enough"). None disables early stopping. 204 coding_max_iter, coding_tol : controls for the Eq. (7) solve. 205 dict_max_iter, dict_tol : controls for the Eq. (8) BCD atom update 206 (passed through to ``bcd_dictionary_update``). 207 coding_chunk_size : int 208 Maximum number of sample columns streamed to the GPU at once 209 while solving each class's Eq. (7) FISTA problem (Step 2). 210 Bounds Step 2's peak GPU memory independent of class size; 211 lower it if Step 2 itself OOMs (see ``solve_class_codes_chunked``). 212 This is typically the step that needs the smallest chunk size, 213 since it needs several full-chunk-size buffers alive at once 214 (FISTA's own iterate bookkeeping plus gradient temporaries). 215 dict_update_chunk_size : int 216 Maximum number of sample columns streamed to the GPU at once 217 while accumulating the Eq. (8) sufficient statistics (Step 3). 218 Bounds Step 3's peak GPU memory independent of dataset size; 219 lower it if Step 3 itself OOMs (see ``build_di_update_system_streaming``). 220 pin_memory : bool 221 Whether the CPU-resident training data / coefficients 222 (``A_list``/``X_list``) are allocated as pinned memory, for 223 faster host<->device transfers. Set False to save that memory 224 overhead if transfer speed isn't a concern. 225 random_state : int or None 226 verbose : bool 227 228 Attributes 229 ---------- 230 D_list_ : list of np.ndarray 231 Learned per-class sub-dictionaries, D_list_[i] has shape 232 (n_features, p_i). Converted to numpy at the end of `fit` 233 (computed on GPU internally). 234 D_ : np.ndarray, shape (n_features, n_components) 235 Learned dictionary, horizontally stacked D_list_. 236 X_list_ : list of np.ndarray 237 Learned per-class coding coefficients (full n_components rows). 238 Converted to numpy at the end of `fit`. 239 atom_boundaries_ : dict[int, tuple[int, int]] 240 Atom row-range owned by each class within D_/X_list_[*]. 241 classes_ : np.ndarray 242 Class labels seen during ``fit``, in the internal class-index 243 order (index i corresponds to D_list_[i]). 244 sample_order_ : np.ndarray 245 Indices into the original training X that produce the 246 class-grouped column order used internally (useful for 247 re-aligning ``X_list_``/errors with the original sample order). 248 objective_history_ : list of float 249 J(D,X) (Eq. 6) after every outer iteration. 250 """ 251 252 def __init__( 253 self, 254 n_components: int | list[int], 255 lambda1: float = 0.005, 256 lambda2: float = 0.005, 257 eta: float = 1.0, 258 n_iter: int = 15, 259 tol: float | None = 1e-4, 260 coding_max_iter: int = 200, 261 coding_tol: float | None = 1e-6, 262 dict_max_iter: int = 1, 263 dict_tol: float = 1e-6, 264 coding_chunk_size: int = 8192, 265 dict_update_chunk_size: int = 8192, 266 pin_memory: bool = True, 267 random_state: int | None = None, 268 verbose: bool = False, 269 ) -> None: 270 if lambda1 < 0 or lambda2 < 0: 271 raise ValueError("lambda1 and lambda2 must be >= 0.") 272 if eta <= 0: 273 raise ValueError("eta must be > 0.") 274 if coding_chunk_size < 1: 275 raise ValueError("coding_chunk_size must be >= 1.") 276 if dict_update_chunk_size < 1: 277 raise ValueError("dict_update_chunk_size must be >= 1.") 278 279 self.n_components = n_components 280 self.lambda1 = lambda1 281 self.lambda2 = lambda2 282 self.eta = eta 283 self.n_iter = n_iter 284 self.tol = tol 285 self.coding_max_iter = coding_max_iter 286 self.coding_tol = coding_tol 287 self.dict_max_iter = dict_max_iter 288 self.dict_tol = dict_tol 289 self.coding_chunk_size = coding_chunk_size 290 self.dict_update_chunk_size = dict_update_chunk_size 291 self.pin_memory = pin_memory 292 self.random_state = random_state 293 self.verbose = verbose 294 295 self.D_list_: list[np.ndarray] | None = None 296 self.X_list_: list[np.ndarray] | None = None 297 self.atom_boundaries_: dict[int, tuple[int, int]] | None = None 298 self.classes_: np.ndarray | None = None 299 self.sample_order_: np.ndarray | None = None 300 self.objective_history_: list[float] = [] 301 302 # ------------------------------------------------------------------ 303 # Public API 304 # ------------------------------------------------------------------ 305 306 def fit( 307 self, 308 X: np.ndarray, 309 y: np.ndarray, 310 D_init: list[np.ndarray] | None = None, 311 checkpoint_dir: str | None = None, 312 resume: bool = True, 313 ) -> FDDL: 314 """ 315 Learn a Fisher discriminative dictionary. 316 317 Parameters 318 ---------- 319 X : np.ndarray, shape (n_features, n_samples) 320 Training signals. 321 y : np.ndarray, shape (n_samples,) 322 Integer or hashable class labels, one per column of X. 323 D_init : list of np.ndarray or None 324 Optional initial per-class sub-dictionaries (unit-norm 325 columns). If None, atoms are initialized as random 326 unit-norm vectors (Table 1, step 1). Ignored when resuming. 327 checkpoint_dir : str or None 328 If given, a checkpoint of the outer loop is written to 329 ``<checkpoint_dir>/fddl_checkpoint.npz`` after every outer 330 iteration, overwriting the previous one. 331 resume : bool 332 If True (default) and a checkpoint exists, resume from it. 333 334 Returns 335 ------- 336 self 337 """ 338 device = _require_gpu_device() 339 340 X = np.asarray(X, dtype=np.float32) 341 y = np.asarray(y) 342 n_features, n_samples = X.shape 343 if y.shape[0] != n_samples: 344 raise DictionaryLearningError( 345 f"y has {y.shape[0]} labels but X has {n_samples} samples." 346 ) 347 348 classes, y_idx = np.unique(y, return_inverse=True) 349 n_classes = len(classes) 350 if n_classes < 2: 351 raise DictionaryLearningError("FDDL requires at least 2 classes.") 352 353 # Group columns by class (Table 1 operates on per-class blocks A_i). 354 sample_order = np.argsort(y_idx, kind="stable") 355 y_sorted = y_idx[sample_order] 356 sizes = [int(np.sum(y_sorted == i)) for i in range(n_classes)] 357 if any(s < 2 for s in sizes): 358 raise DictionaryLearningError( 359 "Every class needs >= 2 samples (class means/scatter are undefined otherwise)." 360 ) 361 sample_boundaries = block_boundaries(sizes) 362 363 # A_list / X_list are the CPU-resident "source of truth" storage 364 # (see module docstring): the compute loop below streams only 365 # the slice it currently needs to `device`, so this is the only 366 # place the full dataset is copied at once, and it stays on CPU. 367 X_grouped = X[:, sample_order] 368 A_list = [ 369 _to_cpu_storage(X_grouped[:, s:e], self.pin_memory) 370 for _, (s, e) in sample_boundaries.items() 371 ] 372 del X_grouped 373 374 atoms_per_class = resolve_atoms_per_class(self.n_components, n_classes) 375 atom_boundaries = block_boundaries(atoms_per_class) 376 n_atoms = sum(atoms_per_class) 377 378 rng = np.random.RandomState(self.random_state) 379 380 checkpoint_path = None 381 start_iter = 0 382 D_list = None 383 X_list = None 384 self.objective_history_ = [] 385 386 if checkpoint_dir is not None: 387 os.makedirs(checkpoint_dir, exist_ok=True) 388 checkpoint_path = os.path.join(checkpoint_dir, _CHECKPOINT_FILENAME) 389 if resume and os.path.exists(checkpoint_path): 390 D_list, X_list, self.objective_history_, start_iter = self._load_checkpoint( 391 checkpoint_path, n_classes 392 ) 393 # Checkpoints are numpy on disk. Dictionaries move to the 394 # GPU (small, live there for the whole optimization); 395 # coefficients stay on the CPU as resident storage. 396 D_list = [_to_device(Di, device) for Di in D_list] 397 X_list = [_to_cpu_storage(Xi, self.pin_memory) for Xi in X_list] 398 if self.verbose: 399 print( 400 f"Resuming from checkpoint at outer iteration " 401 f"{start_iter}/{self.n_iter} ({checkpoint_path})" 402 ) 403 404 if D_list is None: 405 if D_init is not None: 406 if len(D_init) != n_classes: 407 raise DictionaryLearningError( 408 f"D_init has {len(D_init)} sub-dictionaries but there are " 409 f"{n_classes} classes." 410 ) 411 D_list = [normalize_columns(_to_device(Di, device)) for Di in D_init] 412 for Di in D_list: 413 _check_dict_normalized(Di) 414 else: 415 # Table 1, step 1: random unit-norm atoms. 416 D_list = [ 417 normalize_columns(_to_device(rng.randn(n_features, p), device)) 418 for p in atoms_per_class 419 ] 420 X_list = [ 421 _to_cpu_storage( 422 np.zeros((n_atoms, sizes[i]), dtype=np.float32), self.pin_memory 423 ) 424 for i in range(n_classes) 425 ] 426 427 D_full = torch.hstack(D_list) 428 429 for it in range(start_iter, self.n_iter): 430 # ---- Step 2 (Eq. 7): update X class-by-class, D fixed ---- 431 # Coefficient means are tiny (one n_atoms-length vector per 432 # class); kept GPU-resident throughout the sweep so they match 433 # the device of each class's Xi while it's being solved there, 434 # even though X_list itself lives on the CPU. 435 tracker = GlobalMeanTracker(X_list, sizes, device=device) 436 for i in range(n_classes): 437 logger.info(f"Updating class {i} codes") 438 stats = tracker.exclude(i) 439 # Xi0/Ai stay CPU-resident throughout -- solve_class_codes_chunked 440 # streams `coding_chunk_size`-sized column chunks to `device` 441 # internally; no full-class GPU tensor is ever created here. 442 Xi_new_cpu, _, _ = solve_class_codes_chunked( 443 X_list[i], 444 i, 445 D_list, 446 D_full, 447 A_list[i], 448 atom_boundaries, 449 stats, 450 self.lambda1, 451 self.lambda2, 452 self.eta, 453 device, 454 self.coding_chunk_size, 455 max_iter=self.coding_max_iter, 456 tol=self.coding_tol, 457 ) 458 tracker.update(i, Xi_new_cpu) # mean computed on CPU, moved to device internally 459 X_list[i] = _to_cpu_storage(Xi_new_cpu.numpy(), self.pin_memory) 460 del Xi_new_cpu 461 _empty_cache(device) 462 463 # ---- Step 3 (Eq. 8): update D class-by-class, X fixed ---- 464 # Sufficient statistics are streamed in `dict_update_chunk_size` 465 # column chunks (see build_di_update_system_streaming) instead 466 # of materializing a full-dataset-width tensor on the GPU. 467 for i in range(n_classes): 468 logger.info(f"Updating class {i} dictionary") 469 A_stat, B_stat = build_di_update_system_streaming( 470 i, 471 D_list, 472 A_list, 473 X_list, 474 atom_boundaries, 475 device, 476 self.dict_update_chunk_size, 477 ) 478 # bcd_dictionary_update mutates its D argument in place 479 # and returns that same tensor object. 480 Di_updated = bcd_dictionary_update( 481 D_list[i], A_stat, B_stat, 0, self.dict_max_iter, self.dict_tol 482 ) 483 del A_stat, B_stat 484 D_list[i] = normalize_columns(Di_updated) 485 _empty_cache(device) 486 487 D_full = torch.hstack(D_list) 488 489 # ---- Objective (Eq. 6) for convergence tracking ---- 490 # global_fisher_value and the L1 term run directly on the 491 # CPU-resident X_list -- both are plain reductions with no 492 # matmuls, so there's no benefit to moving them to the GPU 493 # and every reason not to for a class this large. Only the 494 # fidelity term needs D (GPU-resident), so only it streams. 495 logger.info("Computing global fischer value") 496 obj = self.lambda2 * global_fisher_value(X_list, self.eta) 497 for i in range(n_classes): 498 logger.info(f"Computing fidelity value for class {i}") 499 obj += fidelity_value_chunked( 500 X_list[i], i, D_list, A_list[i], atom_boundaries, 501 device, self.dict_update_chunk_size, 502 ) 503 obj += self.lambda1 * float(torch.sum(torch.abs(X_list[i]))) 504 _empty_cache(device) 505 self.objective_history_.append(obj) 506 507 if self.verbose: 508 print(f"[FDDL] Iter {it + 1}/{self.n_iter} J={obj:.6f}") 509 510 if checkpoint_path is not None: 511 self._save_checkpoint(checkpoint_path, D_list, X_list, self.objective_history_, it + 1) 512 513 if ( 514 self.tol is not None 515 and len(self.objective_history_) >= 2 516 and abs(self.objective_history_[-2] - obj) < self.tol * abs(self.objective_history_[-2]) 517 ): 518 if self.verbose: 519 print(f"[FDDL] Converged at iteration {it + 1}.") 520 break 521 522 # Public attributes: converted back to numpy here. 523 self.D_list_ = [Di.detach().cpu().numpy() for Di in D_list] 524 self.X_list_ = [Xi.detach().cpu().numpy() for Xi in X_list] # already CPU 525 self.atom_boundaries_ = atom_boundaries 526 self.classes_ = classes 527 self.sample_order_ = sample_order 528 return self 529 530 @property 531 def D_(self) -> np.ndarray: 532 if self.D_list_ is None: 533 raise DictionaryLearningError("Call fit() before accessing D_.") 534 return np.hstack(self.D_list_) 535 536 # ------------------------------------------------------------------ 537 # Checkpointing 538 # ------------------------------------------------------------------ 539 540 def _save_checkpoint(self, path, D_list, X_list, history, iteration) -> None: 541 tmp_path = path + ".tmp.npz" 542 payload = {"n_classes": len(D_list), "history": np.array(history), "iteration": iteration} 543 for i, (Di, Xi) in enumerate(zip(D_list, X_list)): 544 payload[f"D_{i}"] = Di.detach().cpu().numpy() 545 payload[f"X_{i}"] = Xi.detach().cpu().numpy() 546 np.savez(tmp_path, **payload) 547 os.replace(tmp_path, path) 548 549 def _load_checkpoint(self, path, n_classes): 550 data = np.load(path) 551 if int(data["n_classes"]) != n_classes: 552 raise DictionaryLearningError( 553 "Checkpoint's class count does not match the current training data." 554 ) 555 D_list = [data[f"D_{i}"] for i in range(n_classes)] 556 X_list = [data[f"X_{i}"] for i in range(n_classes)] 557 history = list(data["history"]) 558 iteration = int(data["iteration"]) 559 return D_list, X_list, history, iteration
Fisher Discrimination Dictionary Learning.
Parameters
n_components : int or sequence of int
Total number of dictionary atoms, split evenly across classes
(remainder to the last class), or an explicit per-class atom
count. The paper usually sets all pi equal (Sec. 6.1).
lambda1 : float
L1 sparsity weight (Eq. 6).
lambda2 : float
Weight of the Fisher discriminative coefficient term (Eq. 6).
eta : float
Elastic term weight in f(X) (Eq. 5). The paper fixes eta=1,
which is sufficient for fi(Xi) to be strictly convex whenever
eta > 1 - ni/n for every class (Appendix A); eta=1 satisfies
this for any class with more than one training sample.
n_iter : int
Maximum number of outer alternating iterations (D/X updates).
tol : float or None
Outer-loop relative convergence tolerance on J(D,X) (Eq. 6)
between consecutive iterations (Table 1, step 4: "return to
step 2 until the values of J(D,X) in adjacent iterations are
close enough"). None disables early stopping.
coding_max_iter, coding_tol : controls for the Eq. (7) solve.
dict_max_iter, dict_tol : controls for the Eq. (8) BCD atom update
(passed through to bcd_dictionary_update).
coding_chunk_size : int
Maximum number of sample columns streamed to the GPU at once
while solving each class's Eq. (7) FISTA problem (Step 2).
Bounds Step 2's peak GPU memory independent of class size;
lower it if Step 2 itself OOMs (see solve_class_codes_chunked).
This is typically the step that needs the smallest chunk size,
since it needs several full-chunk-size buffers alive at once
(FISTA's own iterate bookkeeping plus gradient temporaries).
dict_update_chunk_size : int
Maximum number of sample columns streamed to the GPU at once
while accumulating the Eq. (8) sufficient statistics (Step 3).
Bounds Step 3's peak GPU memory independent of dataset size;
lower it if Step 3 itself OOMs (see build_di_update_system_streaming).
pin_memory : bool
Whether the CPU-resident training data / coefficients
(A_list/X_list) are allocated as pinned memory, for
faster host<->device transfers. Set False to save that memory
overhead if transfer speed isn't a concern.
random_state : int or None
verbose : bool
Attributes
D_list_ : list of np.ndarray
Learned per-class sub-dictionaries, D_list_[i] has shape
(n_features, p_i). Converted to numpy at the end of fit
(computed on GPU internally).
D_ : np.ndarray, shape (n_features, n_components)
Learned dictionary, horizontally stacked D_list_.
X_list_ : list of np.ndarray
Learned per-class coding coefficients (full n_components rows).
Converted to numpy at the end of fit.
atom_boundaries_ : dict[int, tuple[int, int]]
Atom row-range owned by each class within D_/X_list_[*].
classes_ : np.ndarray
Class labels seen during fit, in the internal class-index
order (index i corresponds to D_list_[i]).
sample_order_ : np.ndarray
Indices into the original training X that produce the
class-grouped column order used internally (useful for
re-aligning X_list_/errors with the original sample order).
objective_history_ : list of float
J(D,X) (Eq. 6) after every outer iteration.
252 def __init__( 253 self, 254 n_components: int | list[int], 255 lambda1: float = 0.005, 256 lambda2: float = 0.005, 257 eta: float = 1.0, 258 n_iter: int = 15, 259 tol: float | None = 1e-4, 260 coding_max_iter: int = 200, 261 coding_tol: float | None = 1e-6, 262 dict_max_iter: int = 1, 263 dict_tol: float = 1e-6, 264 coding_chunk_size: int = 8192, 265 dict_update_chunk_size: int = 8192, 266 pin_memory: bool = True, 267 random_state: int | None = None, 268 verbose: bool = False, 269 ) -> None: 270 if lambda1 < 0 or lambda2 < 0: 271 raise ValueError("lambda1 and lambda2 must be >= 0.") 272 if eta <= 0: 273 raise ValueError("eta must be > 0.") 274 if coding_chunk_size < 1: 275 raise ValueError("coding_chunk_size must be >= 1.") 276 if dict_update_chunk_size < 1: 277 raise ValueError("dict_update_chunk_size must be >= 1.") 278 279 self.n_components = n_components 280 self.lambda1 = lambda1 281 self.lambda2 = lambda2 282 self.eta = eta 283 self.n_iter = n_iter 284 self.tol = tol 285 self.coding_max_iter = coding_max_iter 286 self.coding_tol = coding_tol 287 self.dict_max_iter = dict_max_iter 288 self.dict_tol = dict_tol 289 self.coding_chunk_size = coding_chunk_size 290 self.dict_update_chunk_size = dict_update_chunk_size 291 self.pin_memory = pin_memory 292 self.random_state = random_state 293 self.verbose = verbose 294 295 self.D_list_: list[np.ndarray] | None = None 296 self.X_list_: list[np.ndarray] | None = None 297 self.atom_boundaries_: dict[int, tuple[int, int]] | None = None 298 self.classes_: np.ndarray | None = None 299 self.sample_order_: np.ndarray | None = None 300 self.objective_history_: list[float] = []
306 def fit( 307 self, 308 X: np.ndarray, 309 y: np.ndarray, 310 D_init: list[np.ndarray] | None = None, 311 checkpoint_dir: str | None = None, 312 resume: bool = True, 313 ) -> FDDL: 314 """ 315 Learn a Fisher discriminative dictionary. 316 317 Parameters 318 ---------- 319 X : np.ndarray, shape (n_features, n_samples) 320 Training signals. 321 y : np.ndarray, shape (n_samples,) 322 Integer or hashable class labels, one per column of X. 323 D_init : list of np.ndarray or None 324 Optional initial per-class sub-dictionaries (unit-norm 325 columns). If None, atoms are initialized as random 326 unit-norm vectors (Table 1, step 1). Ignored when resuming. 327 checkpoint_dir : str or None 328 If given, a checkpoint of the outer loop is written to 329 ``<checkpoint_dir>/fddl_checkpoint.npz`` after every outer 330 iteration, overwriting the previous one. 331 resume : bool 332 If True (default) and a checkpoint exists, resume from it. 333 334 Returns 335 ------- 336 self 337 """ 338 device = _require_gpu_device() 339 340 X = np.asarray(X, dtype=np.float32) 341 y = np.asarray(y) 342 n_features, n_samples = X.shape 343 if y.shape[0] != n_samples: 344 raise DictionaryLearningError( 345 f"y has {y.shape[0]} labels but X has {n_samples} samples." 346 ) 347 348 classes, y_idx = np.unique(y, return_inverse=True) 349 n_classes = len(classes) 350 if n_classes < 2: 351 raise DictionaryLearningError("FDDL requires at least 2 classes.") 352 353 # Group columns by class (Table 1 operates on per-class blocks A_i). 354 sample_order = np.argsort(y_idx, kind="stable") 355 y_sorted = y_idx[sample_order] 356 sizes = [int(np.sum(y_sorted == i)) for i in range(n_classes)] 357 if any(s < 2 for s in sizes): 358 raise DictionaryLearningError( 359 "Every class needs >= 2 samples (class means/scatter are undefined otherwise)." 360 ) 361 sample_boundaries = block_boundaries(sizes) 362 363 # A_list / X_list are the CPU-resident "source of truth" storage 364 # (see module docstring): the compute loop below streams only 365 # the slice it currently needs to `device`, so this is the only 366 # place the full dataset is copied at once, and it stays on CPU. 367 X_grouped = X[:, sample_order] 368 A_list = [ 369 _to_cpu_storage(X_grouped[:, s:e], self.pin_memory) 370 for _, (s, e) in sample_boundaries.items() 371 ] 372 del X_grouped 373 374 atoms_per_class = resolve_atoms_per_class(self.n_components, n_classes) 375 atom_boundaries = block_boundaries(atoms_per_class) 376 n_atoms = sum(atoms_per_class) 377 378 rng = np.random.RandomState(self.random_state) 379 380 checkpoint_path = None 381 start_iter = 0 382 D_list = None 383 X_list = None 384 self.objective_history_ = [] 385 386 if checkpoint_dir is not None: 387 os.makedirs(checkpoint_dir, exist_ok=True) 388 checkpoint_path = os.path.join(checkpoint_dir, _CHECKPOINT_FILENAME) 389 if resume and os.path.exists(checkpoint_path): 390 D_list, X_list, self.objective_history_, start_iter = self._load_checkpoint( 391 checkpoint_path, n_classes 392 ) 393 # Checkpoints are numpy on disk. Dictionaries move to the 394 # GPU (small, live there for the whole optimization); 395 # coefficients stay on the CPU as resident storage. 396 D_list = [_to_device(Di, device) for Di in D_list] 397 X_list = [_to_cpu_storage(Xi, self.pin_memory) for Xi in X_list] 398 if self.verbose: 399 print( 400 f"Resuming from checkpoint at outer iteration " 401 f"{start_iter}/{self.n_iter} ({checkpoint_path})" 402 ) 403 404 if D_list is None: 405 if D_init is not None: 406 if len(D_init) != n_classes: 407 raise DictionaryLearningError( 408 f"D_init has {len(D_init)} sub-dictionaries but there are " 409 f"{n_classes} classes." 410 ) 411 D_list = [normalize_columns(_to_device(Di, device)) for Di in D_init] 412 for Di in D_list: 413 _check_dict_normalized(Di) 414 else: 415 # Table 1, step 1: random unit-norm atoms. 416 D_list = [ 417 normalize_columns(_to_device(rng.randn(n_features, p), device)) 418 for p in atoms_per_class 419 ] 420 X_list = [ 421 _to_cpu_storage( 422 np.zeros((n_atoms, sizes[i]), dtype=np.float32), self.pin_memory 423 ) 424 for i in range(n_classes) 425 ] 426 427 D_full = torch.hstack(D_list) 428 429 for it in range(start_iter, self.n_iter): 430 # ---- Step 2 (Eq. 7): update X class-by-class, D fixed ---- 431 # Coefficient means are tiny (one n_atoms-length vector per 432 # class); kept GPU-resident throughout the sweep so they match 433 # the device of each class's Xi while it's being solved there, 434 # even though X_list itself lives on the CPU. 435 tracker = GlobalMeanTracker(X_list, sizes, device=device) 436 for i in range(n_classes): 437 logger.info(f"Updating class {i} codes") 438 stats = tracker.exclude(i) 439 # Xi0/Ai stay CPU-resident throughout -- solve_class_codes_chunked 440 # streams `coding_chunk_size`-sized column chunks to `device` 441 # internally; no full-class GPU tensor is ever created here. 442 Xi_new_cpu, _, _ = solve_class_codes_chunked( 443 X_list[i], 444 i, 445 D_list, 446 D_full, 447 A_list[i], 448 atom_boundaries, 449 stats, 450 self.lambda1, 451 self.lambda2, 452 self.eta, 453 device, 454 self.coding_chunk_size, 455 max_iter=self.coding_max_iter, 456 tol=self.coding_tol, 457 ) 458 tracker.update(i, Xi_new_cpu) # mean computed on CPU, moved to device internally 459 X_list[i] = _to_cpu_storage(Xi_new_cpu.numpy(), self.pin_memory) 460 del Xi_new_cpu 461 _empty_cache(device) 462 463 # ---- Step 3 (Eq. 8): update D class-by-class, X fixed ---- 464 # Sufficient statistics are streamed in `dict_update_chunk_size` 465 # column chunks (see build_di_update_system_streaming) instead 466 # of materializing a full-dataset-width tensor on the GPU. 467 for i in range(n_classes): 468 logger.info(f"Updating class {i} dictionary") 469 A_stat, B_stat = build_di_update_system_streaming( 470 i, 471 D_list, 472 A_list, 473 X_list, 474 atom_boundaries, 475 device, 476 self.dict_update_chunk_size, 477 ) 478 # bcd_dictionary_update mutates its D argument in place 479 # and returns that same tensor object. 480 Di_updated = bcd_dictionary_update( 481 D_list[i], A_stat, B_stat, 0, self.dict_max_iter, self.dict_tol 482 ) 483 del A_stat, B_stat 484 D_list[i] = normalize_columns(Di_updated) 485 _empty_cache(device) 486 487 D_full = torch.hstack(D_list) 488 489 # ---- Objective (Eq. 6) for convergence tracking ---- 490 # global_fisher_value and the L1 term run directly on the 491 # CPU-resident X_list -- both are plain reductions with no 492 # matmuls, so there's no benefit to moving them to the GPU 493 # and every reason not to for a class this large. Only the 494 # fidelity term needs D (GPU-resident), so only it streams. 495 logger.info("Computing global fischer value") 496 obj = self.lambda2 * global_fisher_value(X_list, self.eta) 497 for i in range(n_classes): 498 logger.info(f"Computing fidelity value for class {i}") 499 obj += fidelity_value_chunked( 500 X_list[i], i, D_list, A_list[i], atom_boundaries, 501 device, self.dict_update_chunk_size, 502 ) 503 obj += self.lambda1 * float(torch.sum(torch.abs(X_list[i]))) 504 _empty_cache(device) 505 self.objective_history_.append(obj) 506 507 if self.verbose: 508 print(f"[FDDL] Iter {it + 1}/{self.n_iter} J={obj:.6f}") 509 510 if checkpoint_path is not None: 511 self._save_checkpoint(checkpoint_path, D_list, X_list, self.objective_history_, it + 1) 512 513 if ( 514 self.tol is not None 515 and len(self.objective_history_) >= 2 516 and abs(self.objective_history_[-2] - obj) < self.tol * abs(self.objective_history_[-2]) 517 ): 518 if self.verbose: 519 print(f"[FDDL] Converged at iteration {it + 1}.") 520 break 521 522 # Public attributes: converted back to numpy here. 523 self.D_list_ = [Di.detach().cpu().numpy() for Di in D_list] 524 self.X_list_ = [Xi.detach().cpu().numpy() for Xi in X_list] # already CPU 525 self.atom_boundaries_ = atom_boundaries 526 self.classes_ = classes 527 self.sample_order_ = sample_order 528 return self
Learn a Fisher discriminative dictionary.
Parameters
X : np.ndarray, shape (n_features, n_samples)
Training signals.
y : np.ndarray, shape (n_samples,)
Integer or hashable class labels, one per column of X.
D_init : list of np.ndarray or None
Optional initial per-class sub-dictionaries (unit-norm
columns). If None, atoms are initialized as random
unit-norm vectors (Table 1, step 1). Ignored when resuming.
checkpoint_dir : str or None
If given, a checkpoint of the outer loop is written to
<checkpoint_dir>/fddl_checkpoint.npz after every outer
iteration, overwriting the previous one.
resume : bool
If True (default) and a checkpoint exists, resume from it.
Returns
self
34class FISTA(BaseSparseCoder): 35 """ 36 FISTA sparse coder for L1-regularized least squares. 37 38 Parameters 39 ---------- 40 alpha : float 41 L1 regularization weight (lambda in the paper). Must be >= 0. 42 mode : {'constant', 'backtracking'} 43 Stepsize strategy (Section 4 of the paper). 44 'constant' — requires (or computes) the Lipschitz constant 45 L(f) = 2 * ||D||_2^2 up front. 46 'backtracking' — no Lipschitz constant needed; adapts L via a 47 line search, useful when D is unknown/expensive 48 to bound tightly. 49 max_iter : int 50 Maximum number of FISTA iterations. 51 tol : float or None 52 Relative convergence tolerance on the iterate (practical stopping 53 rule, not from the paper). None disables early stopping. 54 L0 : float 55 Initial Lipschitz estimate for backtracking mode. 56 eta : float 57 Backtracking growth factor (> 1). 58 track_objective : bool 59 If True, evaluate and store F(x_k) at every iteration (adds 60 overhead beyond what 'backtracking' mode already requires). 61 check_dict : bool 62 Whether to verify that dictionary atoms are unit-norm (default True). 63 64 Attributes 65 ---------- 66 n_iter_ : int 67 Number of iterations run in the last call to `encode`. 68 objective_history_ : list of float 69 F(x_k) per iteration from the last call to `encode` (populated 70 whenever mode='backtracking' or track_objective=True). 71 """ 72 73 def __init__( 74 self, 75 alpha: float, 76 mode: str = "backtracking", 77 max_iter: int = 500, 78 tol: float | None = 1e-8, 79 L0: float = 1.0, 80 eta: float = 2.0, 81 track_objective: bool = False, 82 check_dict: bool = True, 83 ) -> None: 84 if alpha < 0: 85 raise ValueError("alpha must be >= 0.") 86 if mode not in ("constant", "backtracking"): 87 raise ValueError("mode must be 'constant' or 'backtracking'.") 88 self.alpha = alpha 89 self.mode = mode 90 self.max_iter = max_iter 91 self.tol = tol 92 self.L0 = L0 93 self.eta = eta 94 self.track_objective = track_objective 95 self.check_dict = check_dict 96 97 def encode( 98 self, 99 X: np.ndarray, 100 D: np.ndarray, 101 L: float | None = None, 102 x0: np.ndarray | None = None, 103 ) -> np.ndarray: 104 """ 105 Compute sparse codes for each column of X. 106 107 Parameters 108 ---------- 109 X : np.ndarray, shape (n_features, n_samples) 110 D : np.ndarray, shape (n_features, n_atoms) 111 L : float, optional 112 Lipschitz constant of grad f. If not given, computed as 113 2 * ||D||_2^2 (Example 2.2). Used directly in 'constant' 114 mode, and as the initial L0 estimate in 'backtracking' mode 115 (overriding the constructor's L0). 116 x0 : np.ndarray, optional 117 Initial point, shape (n_atoms, n_samples). Defaults to zeros. 118 119 Returns 120 ------- 121 Gamma : np.ndarray, shape (n_atoms, n_samples) 122 """ 123 X = np.asarray(X, dtype=np.float32) 124 D = np.asarray(D, dtype=np.float32) 125 126 if X.ndim == 1: 127 X = X[:, np.newaxis] 128 129 if self.check_dict: 130 _check_dict_normalized(D) 131 132 n_atoms = D.shape[1] 133 n_samples = X.shape[1] 134 135 if self.mode == "constant": 136 L_est = lipschitz_constant_lsq(D) if L is None else np.float32(L) 137 L0_init = L_est 138 else: # backtracking 139 L_est = None if L is None else np.float32(L) 140 L0_init = L_est if L_est is not None else self.L0 141 142 def grad_f(Z: np.ndarray) -> np.ndarray: 143 return 2.0 * (D.T @ (D @ Z - X)) 144 145 need_fg = self.mode == "backtracking" or self.track_objective 146 f = (lambda Z: np.float32(np.sum((D @ Z - X) ** 2))) if need_fg else None 147 g = (lambda Z: np.float32(self.alpha * np.sum(np.abs(Z)))) if need_fg else None 148 149 def prox_g(V: np.ndarray, t: float) -> np.ndarray: 150 return soft_threshold(V, self.alpha * t) 151 152 gamma0 = ( 153 np.zeros((n_atoms, n_samples), dtype=np.float32) if x0 is None else np.asarray(x0, dtype=np.float32) 154 ) 155 156 result = fista_core( 157 grad_f=grad_f, 158 prox_g=prox_g, 159 x0=gamma0, 160 f=f, 161 g=g, 162 L=L_est, 163 mode=self.mode, 164 L0=L0_init, 165 eta=self.eta, 166 max_iter=self.max_iter, 167 tol=self.tol, 168 ) 169 170 self.n_iter_ = result.n_iter 171 self.objective_history_ = result.objective_history 172 return result.x
FISTA sparse coder for L1-regularized least squares.
Parameters
alpha : float L1 regularization weight (lambda in the paper). Must be >= 0. mode : {'constant', 'backtracking'} Stepsize strategy (Section 4 of the paper). 'constant' — requires (or computes) the Lipschitz constant L(f) = 2 * ||D||_2^2 up front. 'backtracking' — no Lipschitz constant needed; adapts L via a line search, useful when D is unknown/expensive to bound tightly. max_iter : int Maximum number of FISTA iterations. tol : float or None Relative convergence tolerance on the iterate (practical stopping rule, not from the paper). None disables early stopping. L0 : float Initial Lipschitz estimate for backtracking mode. eta : float Backtracking growth factor (> 1). track_objective : bool If True, evaluate and store F(x_k) at every iteration (adds overhead beyond what 'backtracking' mode already requires). check_dict : bool Whether to verify that dictionary atoms are unit-norm (default True).
Attributes
n_iter_ : int
Number of iterations run in the last call to encode.
objective_history_ : list of float
F(x_k) per iteration from the last call to encode (populated
whenever mode='backtracking' or track_objective=True).
73 def __init__( 74 self, 75 alpha: float, 76 mode: str = "backtracking", 77 max_iter: int = 500, 78 tol: float | None = 1e-8, 79 L0: float = 1.0, 80 eta: float = 2.0, 81 track_objective: bool = False, 82 check_dict: bool = True, 83 ) -> None: 84 if alpha < 0: 85 raise ValueError("alpha must be >= 0.") 86 if mode not in ("constant", "backtracking"): 87 raise ValueError("mode must be 'constant' or 'backtracking'.") 88 self.alpha = alpha 89 self.mode = mode 90 self.max_iter = max_iter 91 self.tol = tol 92 self.L0 = L0 93 self.eta = eta 94 self.track_objective = track_objective 95 self.check_dict = check_dict
97 def encode( 98 self, 99 X: np.ndarray, 100 D: np.ndarray, 101 L: float | None = None, 102 x0: np.ndarray | None = None, 103 ) -> np.ndarray: 104 """ 105 Compute sparse codes for each column of X. 106 107 Parameters 108 ---------- 109 X : np.ndarray, shape (n_features, n_samples) 110 D : np.ndarray, shape (n_features, n_atoms) 111 L : float, optional 112 Lipschitz constant of grad f. If not given, computed as 113 2 * ||D||_2^2 (Example 2.2). Used directly in 'constant' 114 mode, and as the initial L0 estimate in 'backtracking' mode 115 (overriding the constructor's L0). 116 x0 : np.ndarray, optional 117 Initial point, shape (n_atoms, n_samples). Defaults to zeros. 118 119 Returns 120 ------- 121 Gamma : np.ndarray, shape (n_atoms, n_samples) 122 """ 123 X = np.asarray(X, dtype=np.float32) 124 D = np.asarray(D, dtype=np.float32) 125 126 if X.ndim == 1: 127 X = X[:, np.newaxis] 128 129 if self.check_dict: 130 _check_dict_normalized(D) 131 132 n_atoms = D.shape[1] 133 n_samples = X.shape[1] 134 135 if self.mode == "constant": 136 L_est = lipschitz_constant_lsq(D) if L is None else np.float32(L) 137 L0_init = L_est 138 else: # backtracking 139 L_est = None if L is None else np.float32(L) 140 L0_init = L_est if L_est is not None else self.L0 141 142 def grad_f(Z: np.ndarray) -> np.ndarray: 143 return 2.0 * (D.T @ (D @ Z - X)) 144 145 need_fg = self.mode == "backtracking" or self.track_objective 146 f = (lambda Z: np.float32(np.sum((D @ Z - X) ** 2))) if need_fg else None 147 g = (lambda Z: np.float32(self.alpha * np.sum(np.abs(Z)))) if need_fg else None 148 149 def prox_g(V: np.ndarray, t: float) -> np.ndarray: 150 return soft_threshold(V, self.alpha * t) 151 152 gamma0 = ( 153 np.zeros((n_atoms, n_samples), dtype=np.float32) if x0 is None else np.asarray(x0, dtype=np.float32) 154 ) 155 156 result = fista_core( 157 grad_f=grad_f, 158 prox_g=prox_g, 159 x0=gamma0, 160 f=f, 161 g=g, 162 L=L_est, 163 mode=self.mode, 164 L0=L0_init, 165 eta=self.eta, 166 max_iter=self.max_iter, 167 tol=self.tol, 168 ) 169 170 self.n_iter_ = result.n_iter 171 self.objective_history_ = result.objective_history 172 return result.x
Compute sparse codes for each column of X.
Parameters
X : np.ndarray, shape (n_features, n_samples) D : np.ndarray, shape (n_features, n_atoms) L : float, optional Lipschitz constant of grad f. If not given, computed as 2 * ||D||_2^2 (Example 2.2). Used directly in 'constant' mode, and as the initial L0 estimate in 'backtracking' mode (overriding the constructor's L0). x0 : np.ndarray, optional Initial point, shape (n_atoms, n_samples). Defaults to zeros.
Returns
Gamma : np.ndarray, shape (n_atoms, n_samples)
29class KSVD(BaseDictionaryLearner): 30 """ 31 K-SVD dictionary learner. 32 33 Alternates between: 34 1. Sparse coding — encode each training signal over the current D. 35 2. Dictionary update — update each atom (and its coefficients) via a 36 rank-1 approximation of the residual matrix. 37 38 Parameters 39 ---------- 40 n_components : int 41 Number of dictionary atoms this instance learns. When ``D_frozen`` 42 is supplied to ``fit()``, this is the count of *new* atoms only — 43 the frozen atoms are additional, held-constant columns and are not 44 counted here. 45 n_nonzero_coefs : int 46 Sparsity target T: each signal is represented with at most T atoms. 47 n_iter : int 48 Number of K-SVD iterations (default 10). 49 exact_svd : bool 50 If True, use full SVD for the atom update (exact K-SVD). 51 If False (default), use the faster approximate update. 52 mu_thresh : float 53 Mutual-incoherence threshold in (0, 1]. Atoms whose pairwise 54 correlation exceeds this value are replaced. Set to 1.0 to 55 disable (default 0.99). 56 mem_usage : str 57 One of 'high', 'normal' (default), 'low'. 58 Controls whether G = D'D (and DtX = D'X) are precomputed. 59 random_state : int or None 60 Seed for reproducible atom initialisation. 61 verbose : bool 62 Print iteration progress (default False). 63 64 Attributes 65 ---------- 66 D_ : np.ndarray, shape (n_features, n_frozen + n_components) 67 Learned dictionary (set after fit()). Includes any ``D_frozen`` 68 columns passed to ``fit()``, unchanged, as its leading columns. 69 Gamma_ : np.ndarray, shape (n_frozen + n_components, n_samples) 70 Sparse codes for the training data from the final iteration 71 (set after fit()). Exposed so callers that need the training 72 codes (e.g. LC-KSVD, which reuses this class on an augmented 73 system) don't have to re-run sparse coding. 74 errors_ : list of float 75 Per-iteration RMSE on the training data. 76 """ 77 78 def __init__( 79 self, 80 n_components: int, 81 n_nonzero_coefs: int, 82 n_iter: int = 10, 83 exact_svd: bool = False, 84 mu_thresh: float = 0.99, 85 mem_usage: str = "normal", 86 random_state: int | None = None, 87 verbose: bool = True, 88 ) -> None: 89 if mem_usage not in ("high", "normal", "low"): 90 raise ValueError("mem_usage must be 'high', 'normal', or 'low'.") 91 self.n_components = n_components 92 self.n_nonzero_coefs = n_nonzero_coefs 93 self.n_iter = n_iter 94 self.exact_svd = exact_svd 95 self.mu_thresh = mu_thresh 96 self.mem_usage = mem_usage 97 self.random_state = random_state 98 self.verbose = verbose 99 100 # Set after fit 101 self.D_: np.ndarray | None = None 102 self.Gamma_: np.ndarray | None = None 103 self.errors_: list[float] = [] 104 105 106 def fit( 107 self, 108 X: np.ndarray, 109 D_init: np.ndarray | None = None, 110 D_frozen: np.ndarray | None = None, 111 checkpoint_dir: str | None = None, 112 resume: bool = True, 113 ) -> KSVD: 114 """ 115 Learn a dictionary from training signals. 116 117 Parameters 118 ---------- 119 X : np.ndarray, shape (n_features, n_samples) 120 D_init : np.ndarray or None, shape (n_features, n_components) 121 Optional initial dictionary for the *new* (non-frozen) atoms 122 only. If None, random training signals are chosen as initial 123 atoms. Ignored when resuming from an existing checkpoint. 124 D_frozen : np.ndarray or None, shape (n_features, n_frozen_atoms) 125 Optional pre-trained atoms to prepend to the dictionary and 126 hold constant through every iteration. Must already have 127 unit-norm columns (validated). Signals are still sparse-coded 128 jointly over the full ``[D_frozen | D_active]`` dictionary at 129 every iteration; only the ``n_components`` new atoms are ever 130 updated by the atom-update step or by incoherence-based 131 replacement in ``_clear_dict`` — ``D_frozen`` is never 132 modified, not even at initialisation (it is concatenated 133 as-is, never passed through ``normalize_columns``). See 134 Carroll et al. 2017, Sec. III-A ("Frozen K-SVD"). 135 checkpoint_dir : str or None 136 If given, a checkpoint is written to 137 ``<checkpoint_dir>/ksvd_checkpoint.npz`` after every iteration, 138 overwriting the previous one. The directory is created if it 139 does not exist. 140 resume : bool 141 If True (default) and a checkpoint is found, training resumes from it. 142 If False, any existing checkpoint in ``checkpoint_dir`` is ignored and 143 overwritten. 144 145 Returns 146 ------- 147 self 148 """ 149 X = np.asarray(X, dtype=np.float32) 150 rng = np.random.RandomState(self.random_state) 151 152 if D_frozen is not None: 153 D_frozen = np.asarray(D_frozen, dtype=np.float32) 154 if D_frozen.shape[0] != X.shape[0]: 155 raise DictionaryLearningError( 156 f"D_frozen has {D_frozen.shape[0]} features, but X has " 157 f"{X.shape[0]} features." 158 ) 159 _check_dict_normalized(D_frozen) 160 n_frozen = 0 if D_frozen is None else D_frozen.shape[1] 161 162 if X.shape[1] < self.n_components: 163 raise DictionaryLearningError( 164 f"n_samples={X.shape[1]} is less than n_components={self.n_components}. " 165 ) 166 167 checkpoint_path = None 168 start_iter = 0 169 D = None 170 Gamma = None 171 self.errors_ = [] 172 173 if checkpoint_dir is not None: 174 os.makedirs(checkpoint_dir, exist_ok=True) 175 checkpoint_path = os.path.join(checkpoint_dir, _CHECKPOINT_FILENAME) 176 177 if resume and os.path.exists(checkpoint_path): 178 ( 179 D, 180 Gamma, 181 unused, 182 replaced, 183 self.errors_, 184 start_iter, 185 ) = load_checkpoint(self, checkpoint_path, X, n_frozen, D_frozen) 186 if self.verbose: 187 logger.info( 188 f"Resuming from checkpoint at iteration {start_iter}/" 189 f"{self.n_iter} ({checkpoint_path})" 190 ) 191 192 if D is None: 193 D_active = self._init_dict(X, D_init, rng) 194 D = np.hstack([D_frozen, D_active]) if D_frozen is not None else D_active 195 196 n_total = D.shape[1] 197 if n_frozen > 0: 198 D_frozen_cols = D[:, :n_frozen] 199 G_ff = D_frozen_cols.T @ D_frozen_cols 200 DtX_frozen = D_frozen_cols.T @ X 201 else: 202 G_ff = None 203 DtX_frozen = None 204 205 unused = np.arange(X.shape[1]) 206 replaced = np.zeros(n_total, dtype=bool) 207 208 for it in range(start_iter, self.n_iter): 209 if self.mem_usage in ("high", "normal"): 210 D_active = D[:, n_frozen:] 211 DtX_active = D_active.T @ X 212 G_aa = D_active.T @ D_active 213 214 G = np.empty((n_total, n_total), dtype=np.float32) 215 DtX = np.empty((n_total, X.shape[1]), dtype=np.float32) 216 if n_frozen > 0: 217 G_fa = D_frozen_cols.T @ D_active 218 G[:n_frozen, :n_frozen] = G_ff 219 G[:n_frozen, n_frozen:] = G_fa 220 G[n_frozen:, :n_frozen] = G_fa.T 221 DtX[:n_frozen, :] = DtX_frozen 222 G[n_frozen:, n_frozen:] = G_aa 223 DtX[n_frozen:, :] = DtX_active 224 else: 225 G = None 226 DtX = None 227 228 Gamma = self._sparse_code(X, D, G, DtX) 229 R = X - D @ Gamma 230 231 unused = np.arange(X.shape[1]) 232 replaced = np.zeros(n_total, dtype=bool) 233 234 for j in tqdm(range(n_frozen, n_total), desc="Updating atoms"): 235 D[:, j], gamma_j, idx, unused, replaced = _optimize_atom( 236 X, D, R, j, Gamma, unused, replaced, self.exact_svd 237 ) 238 Gamma[j, idx] = gamma_j 239 240 err = np.float32(np.sqrt((R ** 2).sum() / X.size)) 241 self.errors_.append(err) 242 logger.info(f"Iter {it + 1}/{self.n_iter} RMSE={err:.6f}") 243 244 D, _ = _clear_dict( 245 D, Gamma, X, R, self.mu_thresh, unused, replaced, frozen_atoms=n_frozen 246 ) 247 248 if checkpoint_path is not None: 249 save_checkpoint( 250 self, checkpoint_path, X, D, Gamma, unused, replaced, it + 1, n_frozen 251 ) 252 253 self.D_ = D 254 self.Gamma_ = Gamma 255 return self 256 257 def transform(self, X: np.ndarray) -> np.ndarray: 258 """Encode X using the learned dictionary.""" 259 if self.D_ is None: 260 raise DictionaryLearningError("Call fit() before transform().") 261 logger.info("Sparse Coding %d signals with learned dictionary of shape %s", X.shape[1], self.D_.shape) 262 coder = OMP(self.n_nonzero_coefs, mode="batch", check_dict=False) 263 return coder.encode(X, self.D_) 264 265 266 def _init_dict( 267 self, 268 X: np.ndarray, 269 D_init: np.ndarray | None, 270 rng: np.random.RandomState, 271 ) -> np.ndarray: 272 """ 273 Build the initial dictionary for this instance's *own* atoms. 274 275 Note: this only ever constructs ``self.n_components`` columns — 276 any frozen atoms are handled separately by ``fit()`` and 277 concatenated afterward, never passed through this method (so they 278 are never renormalised or otherwise touched at initialisation). 279 """ 280 n_features, _n_samples = X.shape 281 k = self.n_components 282 283 if D_init is not None: 284 D = np.asarray(D_init, dtype=np.float32) 285 if D.shape != (n_features, k): 286 raise DictionaryLearningError( 287 f"D_init shape {D.shape} does not match " 288 f"(n_features={n_features}, n_components={k})." 289 ) 290 return normalize_columns(D) 291 292 valid = np.where(col_norms_squared(X) > 1e-6)[0] 293 if len(valid) < k: 294 raise DictionaryLearningError( 295 "Not enough non-zero training signals to initialise the dictionary." 296 ) 297 298 # Greedily pick columns, rejecting candidates too coherent with atoms 299 # already chosen (same incoherence threshold used by _clear_dict). 300 candidates = rng.permutation(valid) 301 X_norm = normalize_columns(X[:, candidates].copy()) 302 303 chosen_idx: list[int] = [] 304 for pos in range(len(candidates)): 305 if len(chosen_idx) == k: 306 break 307 cand_col = X_norm[:, pos] 308 if chosen_idx: 309 chosen_cols = X_norm[:, chosen_idx] 310 max_coh = np.abs(chosen_cols.T @ cand_col).max() 311 if max_coh > self.mu_thresh: 312 continue 313 chosen_idx.append(pos) 314 315 if len(chosen_idx) < k: 316 logger.warning( 317 "_init_dict: only found %d/%d mutually-incoherent candidate " 318 "atoms (mu_thresh=%.3f) among %d valid signals; falling back " 319 "to filling remaining atoms without the incoherence check. " 320 "This suggests significant redundancy/near-duplication in the " 321 "training data for this class.", 322 len(chosen_idx), k, self.mu_thresh, len(valid), 323 ) 324 remaining_needed = k - len(chosen_idx) 325 remaining_pool = [p for p in range(len(candidates)) if p not in chosen_idx] 326 chosen_idx.extend(remaining_pool[:remaining_needed]) 327 328 D = X[:, candidates[chosen_idx]].copy() 329 return normalize_columns(D) 330 331 332 def _sparse_code( 333 self, 334 X: np.ndarray, 335 D: np.ndarray, 336 G: np.ndarray | None, 337 DtX: np.ndarray | None = None, 338 ) -> np.ndarray: 339 logger.info("Sparse coding %d signals with dictionary of shape %s", X.shape[1], D.shape) 340 coder = OMP(self.n_nonzero_coefs, mode="batch", check_dict=False) 341 return coder.encode(X, D, G=G, DtX=DtX)
K-SVD dictionary learner.
Alternates between:
- Sparse coding — encode each training signal over the current D.
- Dictionary update — update each atom (and its coefficients) via a rank-1 approximation of the residual matrix.
Parameters
n_components : int
Number of dictionary atoms this instance learns. When D_frozen
is supplied to fit(), this is the count of new atoms only —
the frozen atoms are additional, held-constant columns and are not
counted here.
n_nonzero_coefs : int
Sparsity target T: each signal is represented with at most T atoms.
n_iter : int
Number of K-SVD iterations (default 10).
exact_svd : bool
If True, use full SVD for the atom update (exact K-SVD).
If False (default), use the faster approximate update.
mu_thresh : float
Mutual-incoherence threshold in (0, 1]. Atoms whose pairwise
correlation exceeds this value are replaced. Set to 1.0 to
disable (default 0.99).
mem_usage : str
One of 'high', 'normal' (default), 'low'.
Controls whether G = D'D (and DtX = D'X) are precomputed.
random_state : int or None
Seed for reproducible atom initialisation.
verbose : bool
Print iteration progress (default False).
Attributes
D_ : np.ndarray, shape (n_features, n_frozen + n_components)
Learned dictionary (set after fit()). Includes any D_frozen
columns passed to fit(), unchanged, as its leading columns.
Gamma_ : np.ndarray, shape (n_frozen + n_components, n_samples)
Sparse codes for the training data from the final iteration
(set after fit()). Exposed so callers that need the training
codes (e.g. LC-KSVD, which reuses this class on an augmented
system) don't have to re-run sparse coding.
errors_ : list of float
Per-iteration RMSE on the training data.
78 def __init__( 79 self, 80 n_components: int, 81 n_nonzero_coefs: int, 82 n_iter: int = 10, 83 exact_svd: bool = False, 84 mu_thresh: float = 0.99, 85 mem_usage: str = "normal", 86 random_state: int | None = None, 87 verbose: bool = True, 88 ) -> None: 89 if mem_usage not in ("high", "normal", "low"): 90 raise ValueError("mem_usage must be 'high', 'normal', or 'low'.") 91 self.n_components = n_components 92 self.n_nonzero_coefs = n_nonzero_coefs 93 self.n_iter = n_iter 94 self.exact_svd = exact_svd 95 self.mu_thresh = mu_thresh 96 self.mem_usage = mem_usage 97 self.random_state = random_state 98 self.verbose = verbose 99 100 # Set after fit 101 self.D_: np.ndarray | None = None 102 self.Gamma_: np.ndarray | None = None 103 self.errors_: list[float] = []
106 def fit( 107 self, 108 X: np.ndarray, 109 D_init: np.ndarray | None = None, 110 D_frozen: np.ndarray | None = None, 111 checkpoint_dir: str | None = None, 112 resume: bool = True, 113 ) -> KSVD: 114 """ 115 Learn a dictionary from training signals. 116 117 Parameters 118 ---------- 119 X : np.ndarray, shape (n_features, n_samples) 120 D_init : np.ndarray or None, shape (n_features, n_components) 121 Optional initial dictionary for the *new* (non-frozen) atoms 122 only. If None, random training signals are chosen as initial 123 atoms. Ignored when resuming from an existing checkpoint. 124 D_frozen : np.ndarray or None, shape (n_features, n_frozen_atoms) 125 Optional pre-trained atoms to prepend to the dictionary and 126 hold constant through every iteration. Must already have 127 unit-norm columns (validated). Signals are still sparse-coded 128 jointly over the full ``[D_frozen | D_active]`` dictionary at 129 every iteration; only the ``n_components`` new atoms are ever 130 updated by the atom-update step or by incoherence-based 131 replacement in ``_clear_dict`` — ``D_frozen`` is never 132 modified, not even at initialisation (it is concatenated 133 as-is, never passed through ``normalize_columns``). See 134 Carroll et al. 2017, Sec. III-A ("Frozen K-SVD"). 135 checkpoint_dir : str or None 136 If given, a checkpoint is written to 137 ``<checkpoint_dir>/ksvd_checkpoint.npz`` after every iteration, 138 overwriting the previous one. The directory is created if it 139 does not exist. 140 resume : bool 141 If True (default) and a checkpoint is found, training resumes from it. 142 If False, any existing checkpoint in ``checkpoint_dir`` is ignored and 143 overwritten. 144 145 Returns 146 ------- 147 self 148 """ 149 X = np.asarray(X, dtype=np.float32) 150 rng = np.random.RandomState(self.random_state) 151 152 if D_frozen is not None: 153 D_frozen = np.asarray(D_frozen, dtype=np.float32) 154 if D_frozen.shape[0] != X.shape[0]: 155 raise DictionaryLearningError( 156 f"D_frozen has {D_frozen.shape[0]} features, but X has " 157 f"{X.shape[0]} features." 158 ) 159 _check_dict_normalized(D_frozen) 160 n_frozen = 0 if D_frozen is None else D_frozen.shape[1] 161 162 if X.shape[1] < self.n_components: 163 raise DictionaryLearningError( 164 f"n_samples={X.shape[1]} is less than n_components={self.n_components}. " 165 ) 166 167 checkpoint_path = None 168 start_iter = 0 169 D = None 170 Gamma = None 171 self.errors_ = [] 172 173 if checkpoint_dir is not None: 174 os.makedirs(checkpoint_dir, exist_ok=True) 175 checkpoint_path = os.path.join(checkpoint_dir, _CHECKPOINT_FILENAME) 176 177 if resume and os.path.exists(checkpoint_path): 178 ( 179 D, 180 Gamma, 181 unused, 182 replaced, 183 self.errors_, 184 start_iter, 185 ) = load_checkpoint(self, checkpoint_path, X, n_frozen, D_frozen) 186 if self.verbose: 187 logger.info( 188 f"Resuming from checkpoint at iteration {start_iter}/" 189 f"{self.n_iter} ({checkpoint_path})" 190 ) 191 192 if D is None: 193 D_active = self._init_dict(X, D_init, rng) 194 D = np.hstack([D_frozen, D_active]) if D_frozen is not None else D_active 195 196 n_total = D.shape[1] 197 if n_frozen > 0: 198 D_frozen_cols = D[:, :n_frozen] 199 G_ff = D_frozen_cols.T @ D_frozen_cols 200 DtX_frozen = D_frozen_cols.T @ X 201 else: 202 G_ff = None 203 DtX_frozen = None 204 205 unused = np.arange(X.shape[1]) 206 replaced = np.zeros(n_total, dtype=bool) 207 208 for it in range(start_iter, self.n_iter): 209 if self.mem_usage in ("high", "normal"): 210 D_active = D[:, n_frozen:] 211 DtX_active = D_active.T @ X 212 G_aa = D_active.T @ D_active 213 214 G = np.empty((n_total, n_total), dtype=np.float32) 215 DtX = np.empty((n_total, X.shape[1]), dtype=np.float32) 216 if n_frozen > 0: 217 G_fa = D_frozen_cols.T @ D_active 218 G[:n_frozen, :n_frozen] = G_ff 219 G[:n_frozen, n_frozen:] = G_fa 220 G[n_frozen:, :n_frozen] = G_fa.T 221 DtX[:n_frozen, :] = DtX_frozen 222 G[n_frozen:, n_frozen:] = G_aa 223 DtX[n_frozen:, :] = DtX_active 224 else: 225 G = None 226 DtX = None 227 228 Gamma = self._sparse_code(X, D, G, DtX) 229 R = X - D @ Gamma 230 231 unused = np.arange(X.shape[1]) 232 replaced = np.zeros(n_total, dtype=bool) 233 234 for j in tqdm(range(n_frozen, n_total), desc="Updating atoms"): 235 D[:, j], gamma_j, idx, unused, replaced = _optimize_atom( 236 X, D, R, j, Gamma, unused, replaced, self.exact_svd 237 ) 238 Gamma[j, idx] = gamma_j 239 240 err = np.float32(np.sqrt((R ** 2).sum() / X.size)) 241 self.errors_.append(err) 242 logger.info(f"Iter {it + 1}/{self.n_iter} RMSE={err:.6f}") 243 244 D, _ = _clear_dict( 245 D, Gamma, X, R, self.mu_thresh, unused, replaced, frozen_atoms=n_frozen 246 ) 247 248 if checkpoint_path is not None: 249 save_checkpoint( 250 self, checkpoint_path, X, D, Gamma, unused, replaced, it + 1, n_frozen 251 ) 252 253 self.D_ = D 254 self.Gamma_ = Gamma 255 return self
Learn a dictionary from training signals.
Parameters
X : np.ndarray, shape (n_features, n_samples)
D_init : np.ndarray or None, shape (n_features, n_components)
Optional initial dictionary for the new (non-frozen) atoms
only. If None, random training signals are chosen as initial
atoms. Ignored when resuming from an existing checkpoint.
D_frozen : np.ndarray or None, shape (n_features, n_frozen_atoms)
Optional pre-trained atoms to prepend to the dictionary and
hold constant through every iteration. Must already have
unit-norm columns (validated). Signals are still sparse-coded
jointly over the full [D_frozen | D_active] dictionary at
every iteration; only the n_components new atoms are ever
updated by the atom-update step or by incoherence-based
replacement in _clear_dict — D_frozen is never
modified, not even at initialisation (it is concatenated
as-is, never passed through normalize_columns). See
Carroll et al. 2017, Sec. III-A ("Frozen K-SVD").
checkpoint_dir : str or None
If given, a checkpoint is written to
<checkpoint_dir>/ksvd_checkpoint.npz after every iteration,
overwriting the previous one. The directory is created if it
does not exist.
resume : bool
If True (default) and a checkpoint is found, training resumes from it.
If False, any existing checkpoint in checkpoint_dir is ignored and
overwritten.
Returns
self
257 def transform(self, X: np.ndarray) -> np.ndarray: 258 """Encode X using the learned dictionary.""" 259 if self.D_ is None: 260 raise DictionaryLearningError("Call fit() before transform().") 261 logger.info("Sparse Coding %d signals with learned dictionary of shape %s", X.shape[1], self.D_.shape) 262 coder = OMP(self.n_nonzero_coefs, mode="batch", check_dict=False) 263 return coder.encode(X, self.D_)
Encode X using the learned dictionary.
53class LCKSVD(BaseDiscriminativeDictionaryLearner): 54 """ 55 Label Consistent K-SVD dictionary learner (LC-KSVD1 and LC-KSVD2). 56 57 Parameters 58 ---------- 59 n_components : int 60 Number of dictionary atoms. 61 n_nonzero_coefs : int 62 Sparsity level T. 63 alpha : float 64 Weight for the label-consistency term (this is sqrt_alpha in the 65 paper's Eq. (11) — see the module docstring for how it relates to 66 the paper's reported alpha values). 67 beta : float 68 Weight for the classifier term (sqrt_beta; LC-KSVD2 only). 69 variant : {'lcksvd1', 'lcksvd2'} 70 Which variant to train. 71 n_iter : int 72 Number of LC-KSVD iterations (default 50). 73 n_iter_init : int 74 K-SVD iterations for the initialisation phase (default 20). 75 exact_svd : bool 76 Use exact SVD in the atom-update step (slower but slightly better). 77 mu_thresh : float 78 Mutual-incoherence threshold (default 0.99). See the module 79 docstring's "Known limitation" section for how this interacts 80 with LC-KSVD's fixed atom-class labelling assumption. 81 lambda1 : float 82 Ridge weight for the classifier, Eq. (17) (default 1e-5). Used 83 both for W^(0) (LC-KSVD2's warm start) and, for LC-KSVD1, as the 84 default `RidgeClassifier`'s regularisation unless a `classifier` 85 override supplies its own. 86 lambda2 : float 87 Ridge weight for A^(0), Eq. (16) (default 1e-5). 88 classifier : object or None 89 Only valid when ``variant="lcksvd1"``. Any object exposing 90 ``fit(Gamma, H)`` / ``predict(Gamma)`` (see ``RidgeClassifier``), 91 trained once, after the dictionary has converged, on sparse codes 92 from the final D — per the paper's Sec. 3.2 ("the classifier W 93 for LC-KSVD1 is trained separately ... after D, A, and X are 94 computed"). If None, defaults to ``RidgeClassifier(lambda1)``. 95 Passing this for ``variant="lcksvd2"`` raises ``ValueError``, 96 since LC-KSVD2's classifier is trained jointly with the 97 dictionary as part of the augmented system and cannot be 98 substituted independently. 99 random_state : int or None 100 verbose : bool 101 102 Attributes 103 ---------- 104 D_ : np.ndarray, shape (n_features, n_components) 105 Learned dictionary. 106 W_ : np.ndarray or None, shape (n_classes, n_components) 107 Linear classifier weights. For LC-KSVD2, these are learned 108 jointly with the dictionary. For LC-KSVD1, these mirror 109 ``classifier_.W_`` when the fitted classifier exposes a linear 110 ``W_`` attribute (e.g. the default ``RidgeClassifier``); ``None`` 111 if a custom classifier without a ``W_`` attribute was supplied — 112 use ``classifier_.predict()`` / ``predict()`` in that case. 113 classifier_ : object or None 114 The fitted classifier instance for LC-KSVD1 (``None`` for 115 LC-KSVD2, which uses ``W_`` directly). 116 A_ : np.ndarray, shape (n_components, n_components) 117 Learned label-consistency transform. 118 errors_ : list of float 119 Per-iteration RMSE on training data (measured on the original, 120 un-augmented X). 121 """ 122 123 def __init__( 124 self, 125 n_components: int, 126 n_nonzero_coefs: int, 127 alpha: float = 4.0, 128 beta: float = 2.0, 129 variant: str = "lcksvd2", 130 n_iter: int = 50, 131 n_iter_init: int = 20, 132 exact_svd: bool = False, 133 mu_thresh: float = 0.99, 134 lambda1: float = 1e-5, 135 lambda2: float = 1e-5, 136 classifier: object | None = None, 137 random_state: int | None = None, 138 verbose: bool = False, 139 ) -> None: 140 if variant not in ("lcksvd1", "lcksvd2"): 141 raise ValueError("variant must be 'lcksvd1' or 'lcksvd2'.") 142 if classifier is not None and variant != "lcksvd1": 143 raise ValueError( 144 "classifier= is only valid for variant='lcksvd1'. " 145 "LC-KSVD2's classifier is trained jointly with the " 146 "dictionary and cannot be substituted independently." 147 ) 148 self.n_components = n_components 149 self.n_nonzero_coefs = n_nonzero_coefs 150 self.alpha = alpha 151 self.beta = beta 152 self.variant = variant 153 self.n_iter = n_iter 154 self.n_iter_init = n_iter_init 155 self.exact_svd = exact_svd 156 self.mu_thresh = mu_thresh 157 self.lambda1 = lambda1 158 self.lambda2 = lambda2 159 self.classifier = classifier 160 self.random_state = random_state 161 self.verbose = verbose 162 163 self.D_: np.ndarray | None = None 164 self.W_: np.ndarray | None = None 165 self.A_: np.ndarray | None = None 166 self.classifier_: object | None = None 167 self.errors_: list[float] = [] 168 self.class_boundaries_: dict[int, tuple[int, int]] | None = None 169 170 def fit( 171 self, 172 X: np.ndarray, 173 H: np.ndarray, 174 D_init: np.ndarray | None = None, 175 A_init: np.ndarray | None = None, 176 W_init: np.ndarray | None = None, 177 Q: np.ndarray | None = None, 178 checkpoint_dir: str | None = None, 179 resume: bool = True, 180 ) -> LCKSVD: 181 """ 182 Learn a discriminative dictionary from labelled training data. 183 184 Parameters 185 ---------- 186 X : np.ndarray, shape (n_features, n_samples) 187 Training signals. 188 H : np.ndarray, shape (n_classes, n_samples) 189 One-hot label matrix. 190 D_init : np.ndarray or None 191 Initial dictionary. If None, a K-SVD initialisation is run. 192 Ignored when resuming from an existing checkpoint. 193 A_init : np.ndarray or None 194 Initial label-consistency transform. Ignored when resuming. 195 W_init : np.ndarray or None 196 Initial classifier weights (used to warm-start LC-KSVD2's 197 joint optimisation only; LC-KSVD1 trains its classifier 198 separately at the end regardless of this argument). Ignored 199 when resuming. 200 Q : np.ndarray or None 201 Label-consistent target matrix. Computed from H if None. 202 Ignored when resuming. 203 checkpoint_dir : str or None 204 If given, a checkpoint of the *outer* LC-KSVD loop is written 205 to ``<checkpoint_dir>/lc_ksvd_checkpoint.npz`` after every 206 outer iteration, overwriting the previous one. The directory 207 is created if it does not exist. This is independent of, and 208 not passed down to, the inner per-iteration KSVD instance. 209 resume : bool 210 If True (default) and a checkpoint is found in 211 ``checkpoint_dir``, training resumes from it. If False, any 212 existing checkpoint in ``checkpoint_dir`` is ignored and 213 overwritten. 214 """ 215 X = np.asarray(X, dtype=np.float32) 216 H = np.asarray(H, dtype=np.float32) 217 n_features, n_samples = X.shape 218 n_classes = H.shape[0] 219 220 if n_samples < self.n_components: 221 raise DictionaryLearningError( 222 f"n_samples={X.shape[1]} is less than n_components={self.n_components}. " 223 ) 224 225 checkpoint_path = None 226 start_iter = 0 227 D = A = W = Q_loaded = None 228 self.errors_ = [] 229 230 if checkpoint_dir is not None: 231 os.makedirs(checkpoint_dir, exist_ok=True) 232 checkpoint_path = os.path.join(checkpoint_dir, _CHECKPOINT_FILENAME) 233 234 if resume and os.path.exists(checkpoint_path): 235 ( 236 D, A, W, Q_loaded, self.errors_, start_iter, 237 ) = load_checkpoint(self, checkpoint_path, X, H) 238 if self.verbose: 239 print( 240 f"Resuming from checkpoint at outer iteration " 241 f"{start_iter}/{self.n_iter} ({checkpoint_path})" 242 ) 243 244 # ---- Initialisation (skipped if resuming from a checkpoint) ---- 245 if D is None: 246 if D_init is None or A_init is None or W_init is None or Q is None: 247 if self.verbose: 248 print("Running initialisation K-SVD...") 249 D_init, A_init, W_init, Q = initialization4lcksvd( 250 X, H, 251 self.n_components, 252 self.n_iter_init, 253 self.n_nonzero_coefs, 254 random_state=self.random_state, 255 verbose=self.verbose, 256 lambda1=self.lambda1, 257 lambda2=self.lambda2, 258 ) 259 260 D = normalize_columns(D_init.copy()) 261 del D_init 262 A = A_init.copy() 263 del A_init 264 W = W_init.copy() 265 del W_init 266 else: 267 Q = Q_loaded 268 269 sqrt_alpha = self.alpha 270 sqrt_beta = self.beta 271 272 use_classifier_term = (self.variant == "lcksvd2") 273 274 # Build augmented training data # 275 # Y_aug = [X ; sqrt_alpha*Q ; sqrt_beta*H] (LC-KSVD2) 276 # Y_aug = [X ; sqrt_alpha*Q] (LC-KSVD1) 277 H_aug = H if use_classifier_term else None 278 X_aug, _, _ = _augment_data(X, Q, H_aug, sqrt_alpha, sqrt_beta) 279 280 atom_updater = KSVD( 281 n_components=self.n_components, 282 n_nonzero_coefs=self.n_nonzero_coefs, 283 n_iter=1, 284 exact_svd=self.exact_svd, 285 mu_thresh=self.mu_thresh, 286 mem_usage="normal", 287 random_state=self.random_state, 288 verbose=False, 289 ) 290 291 Gamma = None 292 for it in tqdm(range(start_iter, self.n_iter), desc="LC-KSVD iterations"): 293 294 # Build augmented dictionary # 295 D_aug = self._build_aug_dict(D, A, W, sqrt_alpha, sqrt_beta, use_classifier_term) 296 D_aug_norm = normalize_columns(D_aug) 297 del D_aug 298 299 # Sparse code + single atom-update pass on the augmented system 300 atom_updater.fit(X_aug, D_init=D_aug_norm, checkpoint_dir=None) 301 del D_aug_norm 302 D_aug_updated = atom_updater.D_ 303 Gamma = atom_updater.Gamma_ 304 atom_updater.D_ = None 305 atom_updater.Gamma_ = None 306 307 # De-augment: extract D, A, W from the updated augmented dict 308 D, A, W = self._split_aug_dict( 309 D_aug_updated, n_features, n_classes, sqrt_alpha, sqrt_beta, use_classifier_term 310 ) 311 del D_aug_updated 312 D = normalize_columns(D) 313 314 # Track RMSE on original X 315 err = np.float32(np.sqrt(rep_error_squared(X, D, Gamma).sum() / X.size)) 316 self.errors_.append(err) 317 318 if self.verbose: 319 print(f"[{self.variant.upper()}] Iter {it + 1}/{self.n_iter} RMSE={err:.6f}") 320 321 if checkpoint_path is not None: 322 save_checkpoint( 323 self, checkpoint_path, X, H, D, A, W, Q, it + 1 324 ) 325 326 self.D_ = D 327 self.A_ = A 328 del X_aug, Gamma, Q 329 330 # ---- Classifier ---- 331 if self.variant == "lcksvd1": 332 # Paper Sec. 3.2: LC-KSVD1's classifier is trained separately, 333 # after D, A, X are computed. Re-encode with the final, converged D. 334 clf = self.classifier if self.classifier is not None else RidgeClassifier( 335 lambda1=self.lambda1 336 ) 337 Gamma_final = self.transform(X) 338 clf.fit(Gamma_final, H) 339 self.classifier_ = clf 340 # Mirror a linear W_ for convenience/back-compat when the 341 # fitted classifier exposes one. None for classifiers that don't. 342 self.W_ = getattr(clf, "W_", None) 343 else: 344 self.W_ = W 345 self.classifier_ = None 346 347 # Record per-class atom ranges matching _build_label_consistent_target 348 atoms_per_class = self.n_components // n_classes 349 boundaries: dict[int, tuple[int, int]] = {} 350 for c in range(n_classes): 351 start = c * atoms_per_class 352 end = start + atoms_per_class if c < n_classes - 1 else self.n_components 353 boundaries[c] = (start, end) 354 self.class_boundaries_ = boundaries 355 356 return self 357 358 359 def transform(self, X: np.ndarray) -> np.ndarray: 360 """ 361 Encode X using the learned dictionary D. 362 """ 363 self._check_fitted() 364 coder = OMP(self.n_nonzero_coefs, mode="batch", check_dict=False) 365 return coder.encode(X, self.D_) 366 367 def predict(self, X: np.ndarray) -> np.ndarray: 368 """ 369 Classify test signals using the learned classifier. 370 371 For LC-KSVD2, this is the argmax of ``W_ @ gamma``. For LC-KSVD1, 372 this delegates to the fitted ``classifier_`` (default 373 ``RidgeClassifier``, or a user-supplied ``classifier``). 374 """ 375 self._check_fitted() 376 Gamma = self.transform(X) 377 378 if self.variant == "lcksvd1": 379 if self.classifier_ is None: 380 raise DictionaryLearningError( 381 "No classifier is available. This should not happen " 382 "for a successfully fit LC-KSVD1 model." 383 ) 384 return self.classifier_.predict(Gamma) 385 386 if self.W_ is None: 387 raise DictionaryLearningError( 388 "Classifier W is not available. " 389 "Access sparse codes via transform() instead." 390 ) 391 scores = self.W_ @ Gamma # (n_classes, n_samples) 392 return np.argmax(scores, axis=0) 393 394 def score(self, X: np.ndarray, H: np.ndarray) -> float: 395 """ 396 Classification accuracy on (X, H). 397 398 Parameters 399 ---------- 400 X : np.ndarray, shape (n_features, n_samples) 401 H : np.ndarray, shape (n_classes, n_samples) — one-hot labels 402 403 Returns 404 ------- 405 accuracy : float in [0, 1] 406 """ 407 true_labels = np.argmax(H, axis=0) 408 pred_labels = self.predict(X) 409 return np.float32(np.mean(pred_labels == true_labels), dtype=np.float32) 410 411 @staticmethod 412 def _build_aug_dict( 413 D: np.ndarray, 414 A: np.ndarray, 415 W: np.ndarray, 416 sqrt_alpha: float, 417 sqrt_beta: float, 418 use_classifier: bool, 419 ) -> np.ndarray: 420 """Stack [D ; sqrt_alpha*A ; (sqrt_beta*W)].""" 421 parts = [D, sqrt_alpha * A] 422 if use_classifier: 423 parts.append(sqrt_beta * W) 424 return np.vstack(parts) 425 426 @staticmethod 427 def _split_aug_dict( 428 D_aug: np.ndarray, 429 n_features: int, 430 n_classes: int, 431 sqrt_alpha: float, 432 sqrt_beta: float, 433 use_classifier: bool, 434 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: 435 """ 436 Recover (D, A, W) from the augmented dictionary D_aug. 437 438 D_aug rows are: n_features | n_components | (n_classes if lcksvd2). 439 440 D_aug's columns are unit-norm as a whole (across all stacked 441 blocks), not block-by-block. To recover a properly normalised D 442 together with A/W on a consistent per-atom scale, each atom's D 443 sub-block norm is computed and used to rescale D, A, and W 444 (in addition to removing the sqrt_alpha / sqrt_beta weighting 445 from A / W). 446 """ 447 n_components = D_aug.shape[1] 448 D = D_aug[:n_features, :] 449 A_rows = n_components 450 A = D_aug[n_features: n_features + A_rows, :] 451 if use_classifier: 452 W = D_aug[n_features + A_rows:, :] 453 else: 454 W = np.zeros((n_classes, n_components), dtype=np.float32) 455 456 l2norms = np.linalg.norm(D, axis=0) 457 l2norms = np.where(l2norms > 1e-14, l2norms, 1e-14) 458 459 D = D / l2norms 460 A = A / l2norms / max(sqrt_alpha, 1e-14) 461 if use_classifier: 462 W = W / l2norms / max(sqrt_beta, 1e-14) 463 464 return D, A, W
Label Consistent K-SVD dictionary learner (LC-KSVD1 and LC-KSVD2).
Parameters
n_components : int
Number of dictionary atoms.
n_nonzero_coefs : int
Sparsity level T.
alpha : float
Weight for the label-consistency term (this is sqrt_alpha in the
paper's Eq. (11) — see the module docstring for how it relates to
the paper's reported alpha values).
beta : float
Weight for the classifier term (sqrt_beta; LC-KSVD2 only).
variant : {'lcksvd1', 'lcksvd2'}
Which variant to train.
n_iter : int
Number of LC-KSVD iterations (default 50).
n_iter_init : int
K-SVD iterations for the initialisation phase (default 20).
exact_svd : bool
Use exact SVD in the atom-update step (slower but slightly better).
mu_thresh : float
Mutual-incoherence threshold (default 0.99). See the module
docstring's "Known limitation" section for how this interacts
with LC-KSVD's fixed atom-class labelling assumption.
lambda1 : float
Ridge weight for the classifier, Eq. (17) (default 1e-5). Used
both for W^(0) (LC-KSVD2's warm start) and, for LC-KSVD1, as the
default RidgeClassifier's regularisation unless a classifier
override supplies its own.
lambda2 : float
Ridge weight for A^(0), Eq. (16) (default 1e-5).
classifier : object or None
Only valid when variant="lcksvd1". Any object exposing
fit(Gamma, H) / predict(Gamma) (see RidgeClassifier),
trained once, after the dictionary has converged, on sparse codes
from the final D — per the paper's Sec. 3.2 ("the classifier W
for LC-KSVD1 is trained separately ... after D, A, and X are
computed"). If None, defaults to RidgeClassifier(lambda1).
Passing this for variant="lcksvd2" raises ValueError,
since LC-KSVD2's classifier is trained jointly with the
dictionary as part of the augmented system and cannot be
substituted independently.
random_state : int or None
verbose : bool
Attributes
D_ : np.ndarray, shape (n_features, n_components)
Learned dictionary.
W_ : np.ndarray or None, shape (n_classes, n_components)
Linear classifier weights. For LC-KSVD2, these are learned
jointly with the dictionary. For LC-KSVD1, these mirror
classifier_.W_ when the fitted classifier exposes a linear
W_ attribute (e.g. the default RidgeClassifier); None
if a custom classifier without a W_ attribute was supplied —
use classifier_.predict() / predict() in that case.
classifier_ : object or None
The fitted classifier instance for LC-KSVD1 (None for
LC-KSVD2, which uses W_ directly).
A_ : np.ndarray, shape (n_components, n_components)
Learned label-consistency transform.
errors_ : list of float
Per-iteration RMSE on training data (measured on the original,
un-augmented X).
123 def __init__( 124 self, 125 n_components: int, 126 n_nonzero_coefs: int, 127 alpha: float = 4.0, 128 beta: float = 2.0, 129 variant: str = "lcksvd2", 130 n_iter: int = 50, 131 n_iter_init: int = 20, 132 exact_svd: bool = False, 133 mu_thresh: float = 0.99, 134 lambda1: float = 1e-5, 135 lambda2: float = 1e-5, 136 classifier: object | None = None, 137 random_state: int | None = None, 138 verbose: bool = False, 139 ) -> None: 140 if variant not in ("lcksvd1", "lcksvd2"): 141 raise ValueError("variant must be 'lcksvd1' or 'lcksvd2'.") 142 if classifier is not None and variant != "lcksvd1": 143 raise ValueError( 144 "classifier= is only valid for variant='lcksvd1'. " 145 "LC-KSVD2's classifier is trained jointly with the " 146 "dictionary and cannot be substituted independently." 147 ) 148 self.n_components = n_components 149 self.n_nonzero_coefs = n_nonzero_coefs 150 self.alpha = alpha 151 self.beta = beta 152 self.variant = variant 153 self.n_iter = n_iter 154 self.n_iter_init = n_iter_init 155 self.exact_svd = exact_svd 156 self.mu_thresh = mu_thresh 157 self.lambda1 = lambda1 158 self.lambda2 = lambda2 159 self.classifier = classifier 160 self.random_state = random_state 161 self.verbose = verbose 162 163 self.D_: np.ndarray | None = None 164 self.W_: np.ndarray | None = None 165 self.A_: np.ndarray | None = None 166 self.classifier_: object | None = None 167 self.errors_: list[float] = [] 168 self.class_boundaries_: dict[int, tuple[int, int]] | None = None
170 def fit( 171 self, 172 X: np.ndarray, 173 H: np.ndarray, 174 D_init: np.ndarray | None = None, 175 A_init: np.ndarray | None = None, 176 W_init: np.ndarray | None = None, 177 Q: np.ndarray | None = None, 178 checkpoint_dir: str | None = None, 179 resume: bool = True, 180 ) -> LCKSVD: 181 """ 182 Learn a discriminative dictionary from labelled training data. 183 184 Parameters 185 ---------- 186 X : np.ndarray, shape (n_features, n_samples) 187 Training signals. 188 H : np.ndarray, shape (n_classes, n_samples) 189 One-hot label matrix. 190 D_init : np.ndarray or None 191 Initial dictionary. If None, a K-SVD initialisation is run. 192 Ignored when resuming from an existing checkpoint. 193 A_init : np.ndarray or None 194 Initial label-consistency transform. Ignored when resuming. 195 W_init : np.ndarray or None 196 Initial classifier weights (used to warm-start LC-KSVD2's 197 joint optimisation only; LC-KSVD1 trains its classifier 198 separately at the end regardless of this argument). Ignored 199 when resuming. 200 Q : np.ndarray or None 201 Label-consistent target matrix. Computed from H if None. 202 Ignored when resuming. 203 checkpoint_dir : str or None 204 If given, a checkpoint of the *outer* LC-KSVD loop is written 205 to ``<checkpoint_dir>/lc_ksvd_checkpoint.npz`` after every 206 outer iteration, overwriting the previous one. The directory 207 is created if it does not exist. This is independent of, and 208 not passed down to, the inner per-iteration KSVD instance. 209 resume : bool 210 If True (default) and a checkpoint is found in 211 ``checkpoint_dir``, training resumes from it. If False, any 212 existing checkpoint in ``checkpoint_dir`` is ignored and 213 overwritten. 214 """ 215 X = np.asarray(X, dtype=np.float32) 216 H = np.asarray(H, dtype=np.float32) 217 n_features, n_samples = X.shape 218 n_classes = H.shape[0] 219 220 if n_samples < self.n_components: 221 raise DictionaryLearningError( 222 f"n_samples={X.shape[1]} is less than n_components={self.n_components}. " 223 ) 224 225 checkpoint_path = None 226 start_iter = 0 227 D = A = W = Q_loaded = None 228 self.errors_ = [] 229 230 if checkpoint_dir is not None: 231 os.makedirs(checkpoint_dir, exist_ok=True) 232 checkpoint_path = os.path.join(checkpoint_dir, _CHECKPOINT_FILENAME) 233 234 if resume and os.path.exists(checkpoint_path): 235 ( 236 D, A, W, Q_loaded, self.errors_, start_iter, 237 ) = load_checkpoint(self, checkpoint_path, X, H) 238 if self.verbose: 239 print( 240 f"Resuming from checkpoint at outer iteration " 241 f"{start_iter}/{self.n_iter} ({checkpoint_path})" 242 ) 243 244 # ---- Initialisation (skipped if resuming from a checkpoint) ---- 245 if D is None: 246 if D_init is None or A_init is None or W_init is None or Q is None: 247 if self.verbose: 248 print("Running initialisation K-SVD...") 249 D_init, A_init, W_init, Q = initialization4lcksvd( 250 X, H, 251 self.n_components, 252 self.n_iter_init, 253 self.n_nonzero_coefs, 254 random_state=self.random_state, 255 verbose=self.verbose, 256 lambda1=self.lambda1, 257 lambda2=self.lambda2, 258 ) 259 260 D = normalize_columns(D_init.copy()) 261 del D_init 262 A = A_init.copy() 263 del A_init 264 W = W_init.copy() 265 del W_init 266 else: 267 Q = Q_loaded 268 269 sqrt_alpha = self.alpha 270 sqrt_beta = self.beta 271 272 use_classifier_term = (self.variant == "lcksvd2") 273 274 # Build augmented training data # 275 # Y_aug = [X ; sqrt_alpha*Q ; sqrt_beta*H] (LC-KSVD2) 276 # Y_aug = [X ; sqrt_alpha*Q] (LC-KSVD1) 277 H_aug = H if use_classifier_term else None 278 X_aug, _, _ = _augment_data(X, Q, H_aug, sqrt_alpha, sqrt_beta) 279 280 atom_updater = KSVD( 281 n_components=self.n_components, 282 n_nonzero_coefs=self.n_nonzero_coefs, 283 n_iter=1, 284 exact_svd=self.exact_svd, 285 mu_thresh=self.mu_thresh, 286 mem_usage="normal", 287 random_state=self.random_state, 288 verbose=False, 289 ) 290 291 Gamma = None 292 for it in tqdm(range(start_iter, self.n_iter), desc="LC-KSVD iterations"): 293 294 # Build augmented dictionary # 295 D_aug = self._build_aug_dict(D, A, W, sqrt_alpha, sqrt_beta, use_classifier_term) 296 D_aug_norm = normalize_columns(D_aug) 297 del D_aug 298 299 # Sparse code + single atom-update pass on the augmented system 300 atom_updater.fit(X_aug, D_init=D_aug_norm, checkpoint_dir=None) 301 del D_aug_norm 302 D_aug_updated = atom_updater.D_ 303 Gamma = atom_updater.Gamma_ 304 atom_updater.D_ = None 305 atom_updater.Gamma_ = None 306 307 # De-augment: extract D, A, W from the updated augmented dict 308 D, A, W = self._split_aug_dict( 309 D_aug_updated, n_features, n_classes, sqrt_alpha, sqrt_beta, use_classifier_term 310 ) 311 del D_aug_updated 312 D = normalize_columns(D) 313 314 # Track RMSE on original X 315 err = np.float32(np.sqrt(rep_error_squared(X, D, Gamma).sum() / X.size)) 316 self.errors_.append(err) 317 318 if self.verbose: 319 print(f"[{self.variant.upper()}] Iter {it + 1}/{self.n_iter} RMSE={err:.6f}") 320 321 if checkpoint_path is not None: 322 save_checkpoint( 323 self, checkpoint_path, X, H, D, A, W, Q, it + 1 324 ) 325 326 self.D_ = D 327 self.A_ = A 328 del X_aug, Gamma, Q 329 330 # ---- Classifier ---- 331 if self.variant == "lcksvd1": 332 # Paper Sec. 3.2: LC-KSVD1's classifier is trained separately, 333 # after D, A, X are computed. Re-encode with the final, converged D. 334 clf = self.classifier if self.classifier is not None else RidgeClassifier( 335 lambda1=self.lambda1 336 ) 337 Gamma_final = self.transform(X) 338 clf.fit(Gamma_final, H) 339 self.classifier_ = clf 340 # Mirror a linear W_ for convenience/back-compat when the 341 # fitted classifier exposes one. None for classifiers that don't. 342 self.W_ = getattr(clf, "W_", None) 343 else: 344 self.W_ = W 345 self.classifier_ = None 346 347 # Record per-class atom ranges matching _build_label_consistent_target 348 atoms_per_class = self.n_components // n_classes 349 boundaries: dict[int, tuple[int, int]] = {} 350 for c in range(n_classes): 351 start = c * atoms_per_class 352 end = start + atoms_per_class if c < n_classes - 1 else self.n_components 353 boundaries[c] = (start, end) 354 self.class_boundaries_ = boundaries 355 356 return self
Learn a discriminative dictionary from labelled training data.
Parameters
X : np.ndarray, shape (n_features, n_samples)
Training signals.
H : np.ndarray, shape (n_classes, n_samples)
One-hot label matrix.
D_init : np.ndarray or None
Initial dictionary. If None, a K-SVD initialisation is run.
Ignored when resuming from an existing checkpoint.
A_init : np.ndarray or None
Initial label-consistency transform. Ignored when resuming.
W_init : np.ndarray or None
Initial classifier weights (used to warm-start LC-KSVD2's
joint optimisation only; LC-KSVD1 trains its classifier
separately at the end regardless of this argument). Ignored
when resuming.
Q : np.ndarray or None
Label-consistent target matrix. Computed from H if None.
Ignored when resuming.
checkpoint_dir : str or None
If given, a checkpoint of the outer LC-KSVD loop is written
to <checkpoint_dir>/lc_ksvd_checkpoint.npz after every
outer iteration, overwriting the previous one. The directory
is created if it does not exist. This is independent of, and
not passed down to, the inner per-iteration KSVD instance.
resume : bool
If True (default) and a checkpoint is found in
checkpoint_dir, training resumes from it. If False, any
existing checkpoint in checkpoint_dir is ignored and
overwritten.
359 def transform(self, X: np.ndarray) -> np.ndarray: 360 """ 361 Encode X using the learned dictionary D. 362 """ 363 self._check_fitted() 364 coder = OMP(self.n_nonzero_coefs, mode="batch", check_dict=False) 365 return coder.encode(X, self.D_)
Encode X using the learned dictionary D.
367 def predict(self, X: np.ndarray) -> np.ndarray: 368 """ 369 Classify test signals using the learned classifier. 370 371 For LC-KSVD2, this is the argmax of ``W_ @ gamma``. For LC-KSVD1, 372 this delegates to the fitted ``classifier_`` (default 373 ``RidgeClassifier``, or a user-supplied ``classifier``). 374 """ 375 self._check_fitted() 376 Gamma = self.transform(X) 377 378 if self.variant == "lcksvd1": 379 if self.classifier_ is None: 380 raise DictionaryLearningError( 381 "No classifier is available. This should not happen " 382 "for a successfully fit LC-KSVD1 model." 383 ) 384 return self.classifier_.predict(Gamma) 385 386 if self.W_ is None: 387 raise DictionaryLearningError( 388 "Classifier W is not available. " 389 "Access sparse codes via transform() instead." 390 ) 391 scores = self.W_ @ Gamma # (n_classes, n_samples) 392 return np.argmax(scores, axis=0)
Classify test signals using the learned classifier.
For LC-KSVD2, this is the argmax of W_ @ gamma. For LC-KSVD1,
this delegates to the fitted classifier_ (default
RidgeClassifier, or a user-supplied classifier).
394 def score(self, X: np.ndarray, H: np.ndarray) -> float: 395 """ 396 Classification accuracy on (X, H). 397 398 Parameters 399 ---------- 400 X : np.ndarray, shape (n_features, n_samples) 401 H : np.ndarray, shape (n_classes, n_samples) — one-hot labels 402 403 Returns 404 ------- 405 accuracy : float in [0, 1] 406 """ 407 true_labels = np.argmax(H, axis=0) 408 pred_labels = self.predict(X) 409 return np.float32(np.mean(pred_labels == true_labels), dtype=np.float32)
Classification accuracy on (X, H).
Parameters
X : np.ndarray, shape (n_features, n_samples) H : np.ndarray, shape (n_classes, n_samples) — one-hot labels
Returns
accuracy : float in [0, 1]
22class OMP(BaseSparseCoder): 23 """ 24 Orthogonal Matching Pursuit sparse coder. 25 26 Parameters 27 ---------- 28 n_nonzero_coefs : int 29 Target sparsity — maximum number of non-zero coefficients per signal. 30 mode : {'batch', 'cholesky'} 31 Implementation variant. 32 'batch' — Batch-OMP; requires the full Gram matrix G = D'D. 33 Fastest when encoding many signals at once. 34 'cholesky' — Single-signal OMP-Cholesky; lower memory footprint. 35 check_dict : bool 36 Whether to verify that dictionary atoms are unit-norm (default True). 37 """ 38 39 def __init__( 40 self, 41 n_nonzero_coefs: int, 42 mode: str = "batch", 43 check_dict: bool = True, 44 ) -> None: 45 if n_nonzero_coefs < 1: 46 raise ValueError("n_nonzero_coefs must be >= 1.") 47 if mode not in ("batch", "cholesky"): 48 raise ValueError("mode must be 'batch' or 'cholesky'.") 49 self.n_nonzero_coefs = n_nonzero_coefs 50 self.mode = mode 51 self.check_dict = check_dict 52 53 def encode( 54 self, 55 X: np.ndarray, 56 D: np.ndarray, 57 G: np.ndarray | None = None, 58 DtX: np.ndarray | None = None, 59 ) -> np.ndarray: 60 """ 61 Compute sparse codes for each column of X. 62 """ 63 X = np.asarray(X, dtype=np.float32) 64 D = np.asarray(D, dtype=np.float32) 65 66 if X.ndim == 1: 67 X = X[:, np.newaxis] 68 69 if self.check_dict: 70 _check_dict_normalized(D) 71 72 T = self.n_nonzero_coefs 73 74 if self.mode == "batch": 75 if G is None: 76 G = D.T @ D 77 if DtX is None: 78 DtX = D.T @ X 79 logger.info("Encoding %d signals with dictionary of shape %s using Batch-OMP", X.shape[1], D.shape) 80 return batch_omp(DtX, G, T) 81 82 # cholesky mode — signal by signal 83 n_atoms = D.shape[1] 84 n_samples = X.shape[1] 85 Gamma = np.zeros((n_atoms, n_samples), dtype=np.float32) 86 for i in range(n_samples): 87 logger.info("Encoding signal %d/%d with dictionary of shape %s using OMP-Cholesky", i + 1, n_samples, D.shape) 88 Gamma[:, i] = omp_cholesky(D, X[:, i], T) 89 90 return Gamma
Orthogonal Matching Pursuit sparse coder.
Parameters
n_nonzero_coefs : int Target sparsity — maximum number of non-zero coefficients per signal. mode : {'batch', 'cholesky'} Implementation variant. 'batch' — Batch-OMP; requires the full Gram matrix G = D'D. Fastest when encoding many signals at once. 'cholesky' — Single-signal OMP-Cholesky; lower memory footprint. check_dict : bool Whether to verify that dictionary atoms are unit-norm (default True).
39 def __init__( 40 self, 41 n_nonzero_coefs: int, 42 mode: str = "batch", 43 check_dict: bool = True, 44 ) -> None: 45 if n_nonzero_coefs < 1: 46 raise ValueError("n_nonzero_coefs must be >= 1.") 47 if mode not in ("batch", "cholesky"): 48 raise ValueError("mode must be 'batch' or 'cholesky'.") 49 self.n_nonzero_coefs = n_nonzero_coefs 50 self.mode = mode 51 self.check_dict = check_dict
53 def encode( 54 self, 55 X: np.ndarray, 56 D: np.ndarray, 57 G: np.ndarray | None = None, 58 DtX: np.ndarray | None = None, 59 ) -> np.ndarray: 60 """ 61 Compute sparse codes for each column of X. 62 """ 63 X = np.asarray(X, dtype=np.float32) 64 D = np.asarray(D, dtype=np.float32) 65 66 if X.ndim == 1: 67 X = X[:, np.newaxis] 68 69 if self.check_dict: 70 _check_dict_normalized(D) 71 72 T = self.n_nonzero_coefs 73 74 if self.mode == "batch": 75 if G is None: 76 G = D.T @ D 77 if DtX is None: 78 DtX = D.T @ X 79 logger.info("Encoding %d signals with dictionary of shape %s using Batch-OMP", X.shape[1], D.shape) 80 return batch_omp(DtX, G, T) 81 82 # cholesky mode — signal by signal 83 n_atoms = D.shape[1] 84 n_samples = X.shape[1] 85 Gamma = np.zeros((n_atoms, n_samples), dtype=np.float32) 86 for i in range(n_samples): 87 logger.info("Encoding signal %d/%d with dictionary of shape %s using OMP-Cholesky", i + 1, n_samples, D.shape) 88 Gamma[:, i] = omp_cholesky(D, X[:, i], T) 89 90 return Gamma
Compute sparse codes for each column of X.
11class FrozenDictionaryLearner: 12 """ 13 Learn a residual dictionary given a frozen dictionary D_frozen, using 14 any unsupervised ``BaseDictionaryLearner`` that supports the frozen- 15 dictionary contract (``D_frozen`` / ``checkpoint_dir`` / ``resume`` on 16 ``fit()``, and a ``D_`` attribute afterward) — e.g. ``KSVD``. 17 The wrapped learner jointly sparse-codes over ``[D_frozen | D_active]`` 18 every iteration and only ever updates the new atoms. 19 20 For the full sequential pipeline, see ``IncrementalFrozenDictionary``. 21 22 Parameters 23 ---------- 24 D_frozen : np.ndarray, shape (n_features, n_frozen_atoms) 25 Pre-trained frozen dictionary. Never modified. 26 learner_class : type[BaseDictionaryLearner] 27 Unsupervised dictionary learning class to use for the residual, 28 e.g. ``KSVD``. Must accept ``D_frozen`` in its ``fit()`` and 29 expose ``D_`` afterward (see the base class's frozen-dictionary 30 contract). Its own ``n_components`` (or equivalent) should be the 31 count of *new* atoms only — D_frozen is additional, not counted 32 there. 33 learner_kwargs : dict 34 Init kwargs for ``learner_class``. 35 n_nonzero_coefs : int 36 Sparsity level associated with this stage. Not used internally by 37 this class anymore, but kept as a constructor parameter since it 38 describes the stage's intended encoding sparsity for any 39 downstream sparse coding against ``D_combined_``. 40 41 Attributes 42 ---------- 43 D_combined_ : np.ndarray, shape (n_features, n_frozen + n_active) 44 learner_ : fitted instance of ``learner_class`` 45 n_frozen_ : int number of frozen atoms 46 n_active_ : int number of active (residual) atoms 47 class_boundaries_ : dict[int, tuple[int, int]] 48 ``frozen_class_boundaries`` merged with a single new entry 49 ``{class_label: (n_frozen, n_frozen + n_active)}``. 50 """ 51 52 def __init__( 53 self, 54 D_frozen: np.ndarray, 55 learner_class: type[BaseDictionaryLearner], 56 learner_kwargs: dict, 57 n_nonzero_coefs: int, 58 ) -> None: 59 self.D_frozen = np.asarray(D_frozen, dtype=np.float32) 60 self.learner_class = learner_class 61 self.learner_kwargs = learner_kwargs 62 self.n_nonzero_coefs = n_nonzero_coefs 63 64 self.D_combined_: np.ndarray | None = None 65 self.learner_: BaseDictionaryLearner | None = None 66 self.n_frozen_: int = self.D_frozen.shape[1] 67 self.n_active_: int = 0 68 self.class_boundaries_: dict[int, tuple[int, int]] | None = None 69 70 def fit( 71 self, 72 X: np.ndarray, 73 class_label: int, 74 frozen_class_boundaries: dict[int, tuple[int, int]] | None = None, 75 checkpoint_dir: str | None = None, 76 resume: bool = True, 77 ) -> FrozenDictionaryLearner: 78 """ 79 Fit the residual dictionary on X. 80 81 Parameters 82 ---------- 83 X : np.ndarray, shape (n_features, n_samples) 84 Training signals for this class. Presented unsupervised to 85 the wrapped learner — it never sees labels. 86 class_label : int 87 The class these new atoms are assigned to in 88 ``class_boundaries_``. 89 frozen_class_boundaries : dict or None 90 ``class_boundaries_`` from earlier frozen stages, merged 91 unchanged into this call's own single new entry. 92 checkpoint_dir : str or None 93 Forwarded to the inner ``learner_class.fit()`` call, if 94 supported. 95 resume : bool 96 Forwarded to the inner learner's ``fit()`` alongside 97 ``checkpoint_dir``. Default True. 98 """ 99 X = np.asarray(X, dtype=np.float32) 100 101 learner = self.learner_class(**self.learner_kwargs) 102 103 loaded = ( 104 _completed_ksvd_checkpoint( 105 checkpoint_dir, 106 X, 107 n_components=self.learner_kwargs.get("n_components"), 108 n_nonzero_coefs=self.learner_kwargs.get("n_nonzero_coefs"), 109 n_iter=self.learner_kwargs.get("n_iter"), 110 n_frozen=self.n_frozen_, 111 D_frozen=self.D_frozen, 112 ) 113 if resume 114 else None 115 ) 116 if loaded is not None: 117 learner.D_, learner.Gamma_, learner.errors_ = loaded 118 else: 119 learner.fit( 120 X, 121 D_frozen=self.D_frozen, 122 checkpoint_dir=checkpoint_dir, 123 resume=resume, 124 ) 125 126 self.learner_ = learner 127 128 if learner.D_.shape[1] <= self.n_frozen_: 129 raise DictionaryLearningError( 130 f"{self.learner_class.__name__}.fit() returned a combined " 131 f"dictionary with {learner.D_.shape[1]} columns, which is " 132 f"not larger than n_frozen={self.n_frozen_}. Learners must " 133 "return the full [D_frozen | D_new] dictionary as D_, per " 134 "the frozen-dictionary contract." 135 ) 136 if not np.allclose(learner.D_[:, : self.n_frozen_], self.D_frozen, atol=1e-6): 137 raise DictionaryLearningError( 138 f"{self.learner_class.__name__}.fit() did not preserve " 139 "D_frozen unchanged in its leading columns of D_." 140 ) 141 142 self.D_combined_ = learner.D_ 143 self.n_active_ = self.D_combined_.shape[1] - self.n_frozen_ 144 145 boundaries = dict(frozen_class_boundaries) if frozen_class_boundaries else {} 146 if class_label in boundaries: 147 raise ValueError( 148 f"class_label {class_label} already present in " 149 "frozen_class_boundaries." 150 ) 151 boundaries[class_label] = (self.n_frozen_, self.n_frozen_ + self.n_active_) 152 self.class_boundaries_ = boundaries 153 154 return self
Learn a residual dictionary given a frozen dictionary D_frozen, using
any unsupervised BaseDictionaryLearner that supports the frozen-
dictionary contract (D_frozen / checkpoint_dir / resume on
fit(), and a D_ attribute afterward) — e.g. KSVD.
The wrapped learner jointly sparse-codes over [D_frozen | D_active]
every iteration and only ever updates the new atoms.
For the full sequential pipeline, see IncrementalFrozenDictionary.
Parameters
D_frozen : np.ndarray, shape (n_features, n_frozen_atoms)
Pre-trained frozen dictionary. Never modified.
learner_class : type[BaseDictionaryLearner]
Unsupervised dictionary learning class to use for the residual,
e.g. KSVD. Must accept D_frozen in its fit() and
expose D_ afterward (see the base class's frozen-dictionary
contract). Its own n_components (or equivalent) should be the
count of new atoms only — D_frozen is additional, not counted
there.
learner_kwargs : dict
Init kwargs for learner_class.
n_nonzero_coefs : int
Sparsity level associated with this stage. Not used internally by
this class anymore, but kept as a constructor parameter since it
describes the stage's intended encoding sparsity for any
downstream sparse coding against D_combined_.
Attributes
D_combined_ : np.ndarray, shape (n_features, n_frozen + n_active)
learner_ : fitted instance of learner_class
n_frozen_ : int number of frozen atoms
n_active_ : int number of active (residual) atoms
class_boundaries_ : dict[int, tuple[int, int]]
frozen_class_boundaries merged with a single new entry
{class_label: (n_frozen, n_frozen + n_active)}.
52 def __init__( 53 self, 54 D_frozen: np.ndarray, 55 learner_class: type[BaseDictionaryLearner], 56 learner_kwargs: dict, 57 n_nonzero_coefs: int, 58 ) -> None: 59 self.D_frozen = np.asarray(D_frozen, dtype=np.float32) 60 self.learner_class = learner_class 61 self.learner_kwargs = learner_kwargs 62 self.n_nonzero_coefs = n_nonzero_coefs 63 64 self.D_combined_: np.ndarray | None = None 65 self.learner_: BaseDictionaryLearner | None = None 66 self.n_frozen_: int = self.D_frozen.shape[1] 67 self.n_active_: int = 0 68 self.class_boundaries_: dict[int, tuple[int, int]] | None = None
70 def fit( 71 self, 72 X: np.ndarray, 73 class_label: int, 74 frozen_class_boundaries: dict[int, tuple[int, int]] | None = None, 75 checkpoint_dir: str | None = None, 76 resume: bool = True, 77 ) -> FrozenDictionaryLearner: 78 """ 79 Fit the residual dictionary on X. 80 81 Parameters 82 ---------- 83 X : np.ndarray, shape (n_features, n_samples) 84 Training signals for this class. Presented unsupervised to 85 the wrapped learner — it never sees labels. 86 class_label : int 87 The class these new atoms are assigned to in 88 ``class_boundaries_``. 89 frozen_class_boundaries : dict or None 90 ``class_boundaries_`` from earlier frozen stages, merged 91 unchanged into this call's own single new entry. 92 checkpoint_dir : str or None 93 Forwarded to the inner ``learner_class.fit()`` call, if 94 supported. 95 resume : bool 96 Forwarded to the inner learner's ``fit()`` alongside 97 ``checkpoint_dir``. Default True. 98 """ 99 X = np.asarray(X, dtype=np.float32) 100 101 learner = self.learner_class(**self.learner_kwargs) 102 103 loaded = ( 104 _completed_ksvd_checkpoint( 105 checkpoint_dir, 106 X, 107 n_components=self.learner_kwargs.get("n_components"), 108 n_nonzero_coefs=self.learner_kwargs.get("n_nonzero_coefs"), 109 n_iter=self.learner_kwargs.get("n_iter"), 110 n_frozen=self.n_frozen_, 111 D_frozen=self.D_frozen, 112 ) 113 if resume 114 else None 115 ) 116 if loaded is not None: 117 learner.D_, learner.Gamma_, learner.errors_ = loaded 118 else: 119 learner.fit( 120 X, 121 D_frozen=self.D_frozen, 122 checkpoint_dir=checkpoint_dir, 123 resume=resume, 124 ) 125 126 self.learner_ = learner 127 128 if learner.D_.shape[1] <= self.n_frozen_: 129 raise DictionaryLearningError( 130 f"{self.learner_class.__name__}.fit() returned a combined " 131 f"dictionary with {learner.D_.shape[1]} columns, which is " 132 f"not larger than n_frozen={self.n_frozen_}. Learners must " 133 "return the full [D_frozen | D_new] dictionary as D_, per " 134 "the frozen-dictionary contract." 135 ) 136 if not np.allclose(learner.D_[:, : self.n_frozen_], self.D_frozen, atol=1e-6): 137 raise DictionaryLearningError( 138 f"{self.learner_class.__name__}.fit() did not preserve " 139 "D_frozen unchanged in its leading columns of D_." 140 ) 141 142 self.D_combined_ = learner.D_ 143 self.n_active_ = self.D_combined_.shape[1] - self.n_frozen_ 144 145 boundaries = dict(frozen_class_boundaries) if frozen_class_boundaries else {} 146 if class_label in boundaries: 147 raise ValueError( 148 f"class_label {class_label} already present in " 149 "frozen_class_boundaries." 150 ) 151 boundaries[class_label] = (self.n_frozen_, self.n_frozen_ + self.n_active_) 152 self.class_boundaries_ = boundaries 153 154 return self
Fit the residual dictionary on X.
Parameters
X : np.ndarray, shape (n_features, n_samples)
Training signals for this class. Presented unsupervised to
the wrapped learner — it never sees labels.
class_label : int
The class these new atoms are assigned to in
class_boundaries_.
frozen_class_boundaries : dict or None
class_boundaries_ from earlier frozen stages, merged
unchanged into this call's own single new entry.
checkpoint_dir : str or None
Forwarded to the inner learner_class.fit() call, if
supported.
resume : bool
Forwarded to the inner learner's fit() alongside
checkpoint_dir. Default True.
17class IncrementalFrozenDictionary: 18 """ 19 Incrementally learn class-specific residual dictionaries, freezing all 20 previously learned atoms before training the next class. 21 22 The underlying dictionary learner 23 (``base_learner_class`` / ``residual_learner_class``, e.g. ``KSVD``) 24 is unsupervised. This class handles only that unsupervised, 25 incremental dictionary-learning pipeline and its per-class bookkeeping (``class_boundaries_``). 26 27 Pipeline 28 -------- 29 1. ``fit_base(X)`` 30 Learn a base dictionary D_n from normal/background data using 31 ``base_learner_class``. This dictionary is frozen for all 32 subsequent steps. 33 34 2. ``add_class(X, class_label)`` 35 Learn a residual dictionary D_a for the new class on top of 36 the currently frozen dictionary [ D_n | D_a_1 | … ], via 37 ``FrozenDictionaryLearner``. The underlying learner trains its new 38 atoms jointly alongside the frozen ones every iteration. 39 40 The end product of this pipeline is ``D_``, the full combined 41 dictionary. Sparse coding and classification on top of it happen 42 outside this class. 43 44 Checkpointing 45 ------------- 46 Each call to ``fit_base`` or ``add_class`` represents one *stage* of 47 the incremental pipeline, and each stage trains its own, differently 48 shaped, inner dictionary-learner instance. If a ``checkpoint_dir`` is 49 given, this class therefore does NOT hand every stage the same path. 50 Instead, each stage gets its own subdirectory, so that interrupting 51 and resuming an individual stage's training does 52 not collide with any other stage's saved state. 53 54 Parameters 55 ---------- 56 base_learner_class : type[BaseDictionaryLearner] 57 Unsupervised learner used for the initial base dictionary, e.g. 58 ``KSVD``. Must accept ``D_frozen`` (default None) in its 59 ``fit()``, per the frozen-dictionary contract. 60 base_learner_kwargs : dict 61 Init kwargs for ``base_learner_class``. 62 residual_learner_class : type[BaseDictionaryLearner] 63 Learner used for each residual dictionary, via 64 ``FrozenDictionaryLearner``. Can be the same as or different from 65 ``base_learner_class``. Must support the frozen-dictionary 66 contract (accept and honour ``D_frozen``). 67 residual_learner_kwargs : dict 68 Init kwargs for ``residual_learner_class``. Applied identically 69 for every ``add_class`` call; override per-call via 70 ``add_class(..., learner_kwargs_override=...)`` — e.g. to give 71 each class a different number of atoms. 72 n_nonzero_coefs : int 73 Sparsity level passed down to ``FrozenDictionaryLearner`` during 74 ``add_class``. Also the natural default sparsity for any sparse 75 coding you do yourself downstream against ``D_``, though this 76 class does not perform that coding. 77 78 Attributes 79 ---------- 80 D_ : np.ndarray full combined dictionary after all steps 81 class_labels_ : list[int] class labels added via add_class, in order 82 class_boundaries_ : dict[int, tuple[int, int]] 83 Per-class atom ranges in the full combined D_ (includes the base 84 stage's class label too). 85 stage_learners_ : list 86 Fitted learner (base stage) or FrozenDictionaryLearner (each 87 add_class stage) from each stage, in order (index 0 = base stage). 88 errors_ : dict[int, list[float]] 89 Per-stage training RMSE curves keyed by class_label 90 (key -1 for the base stage, regardless of its class_label). 91 """ 92 93 def __init__( 94 self, 95 base_learner_class: type[BaseDictionaryLearner], 96 base_learner_kwargs: dict, 97 residual_learner_class: type[BaseDictionaryLearner], 98 residual_learner_kwargs: dict, 99 n_nonzero_coefs: int, 100 ) -> None: 101 self.base_learner_class = base_learner_class 102 self.base_learner_kwargs = base_learner_kwargs 103 self.residual_learner_class = residual_learner_class 104 self.residual_learner_kwargs = residual_learner_kwargs 105 self.n_nonzero_coefs = n_nonzero_coefs 106 107 # State built incrementally 108 self.D_: np.ndarray | None = None 109 self.base_class_label_: int | None = None 110 self.class_labels_: list[int] = [] 111 self.class_boundaries_: dict[int, tuple[int, int]] = {} 112 self.stage_learners_: list = [] 113 self.errors_: dict[int, list[float]] = {} 114 115 def fit_base( 116 self, 117 X: np.ndarray, 118 class_label: int = 0, 119 checkpoint_dir: str | None = None, 120 resume: bool = True, 121 ) -> IncrementalFrozenDictionary: 122 """ 123 Learn the base dictionary from normal / background data. 124 125 Parameters 126 ---------- 127 X : np.ndarray, shape (n_features, n_samples) 128 class_label : int 129 The class label this base dictionary's atoms are assigned to 130 in ``class_boundaries_`` (default 0). Must be distinct from 131 every ``class_label`` later passed to ``add_class``. 132 checkpoint_dir : str or None 133 If given, a ``base`` subdirectory under this path is passed to 134 the base learner's own ``fit(..., checkpoint_dir=...)``, if it 135 supports one. See the class docstring for why this is a 136 dedicated subdirectory rather than shared across stages. 137 resume : bool 138 Forwarded to the base learner's ``fit()`` alongside the 139 ``base`` checkpoint subdirectory. Default True. 140 """ 141 X = np.asarray(X, dtype=np.float32) 142 143 learner = self.base_learner_class(**self.base_learner_kwargs) 144 logger.info("Fitting base dictionary with %d samples and %d features using %s", X.shape[1], X.shape[0], self.base_learner_class.__name__) 145 146 fit_kwargs = {} 147 base_checkpoint_dir = None 148 if checkpoint_dir is not None: 149 base_checkpoint_dir = os.path.join(checkpoint_dir, "base") 150 fit_kwargs["checkpoint_dir"] = base_checkpoint_dir 151 fit_kwargs["resume"] = resume 152 153 loaded = ( 154 _completed_ksvd_checkpoint( 155 base_checkpoint_dir, 156 X, 157 n_components=self.base_learner_kwargs.get("n_components"), 158 n_nonzero_coefs=self.base_learner_kwargs.get("n_nonzero_coefs"), 159 n_iter=self.base_learner_kwargs.get("n_iter"), 160 n_frozen=0, 161 D_frozen=None, 162 ) 163 if resume 164 else None 165 ) 166 if loaded is not None: 167 learner.D_, learner.Gamma_, learner.errors_ = loaded 168 else: 169 # Unsupervised: no D_frozen (nothing to freeze yet). 170 learner.fit(X, **fit_kwargs) 171 172 self.D_ = learner.D_ 173 self.base_class_label_ = class_label 174 self.class_boundaries_ = {class_label: (0, learner.D_.shape[1])} 175 self.stage_learners_.append(learner) 176 self.errors_[-1] = list(getattr(learner, "errors_", [])) 177 178 # self.D_ already holds its own reference to the learned array; 179 # dropping the learner's copies frees nothing-else-reads-them 180 # state rather than holding it for the lifetime of the pipeline. 181 # get_stage_dict(0) slices self.D_ via class_boundaries_, not 182 # this attribute, so it stays correct after this. 183 learner.D_ = None 184 if hasattr(learner, "Gamma_"): 185 learner.Gamma_ = None 186 187 return self 188 189 def add_class( 190 self, 191 X: np.ndarray, 192 class_label: int, 193 learner_kwargs_override: dict | None = None, 194 checkpoint_dir: str | None = None, 195 resume: bool = True, 196 ) -> IncrementalFrozenDictionary: 197 """ 198 Learn a residual dictionary for a new class and extend D_. 199 200 Parameters 201 ---------- 202 X : np.ndarray, shape (n_features, n_samples) 203 Training signals for this class only. Presented unsupervised 204 to the residual learner — it never sees labels. 205 class_label : int 206 Integer label for this class. Must not have been added 207 before, and must differ from the base stage's class_label. 208 learner_kwargs_override : dict or None 209 If supplied, overrides ``residual_learner_kwargs`` for this 210 call only. Useful for giving this class a different number 211 of atoms (e.g. ``{"n_components": 5}``). 212 checkpoint_dir : str or None 213 If given, a ``class_<class_label>`` subdirectory under this 214 path is passed down to ``FrozenDictionaryLearner.fit()`` (and 215 from there to the residual learner's own ``fit()``), if 216 supported. Kept distinct per class_label, and distinct from 217 the base stage's ``base`` subdirectory, so that resuming one 218 stage never collides with another's saved state — see the 219 class docstring. 220 resume : bool 221 Forwarded down to the residual learner's ``fit()`` alongside 222 the per-class checkpoint subdirectory. Default True. 223 """ 224 if self.D_ is None: 225 raise DictionaryLearningError( 226 "Call fit_base() before add_class()." 227 ) 228 if class_label in self.class_labels_ or class_label == self.base_class_label_: 229 raise ValueError( 230 f"class_label {class_label} has already been used." 231 ) 232 233 X = np.asarray(X, dtype=np.float32) 234 235 kwargs = {**self.residual_learner_kwargs, **(learner_kwargs_override or {})} 236 logger.info("Adding class %d with %d samples and %d features using %s", class_label, X.shape[1], X.shape[0], self.residual_learner_class.__name__) 237 238 frozen_step = FrozenDictionaryLearner( 239 D_frozen=self.D_, 240 learner_class=self.residual_learner_class, 241 learner_kwargs=kwargs, 242 n_nonzero_coefs=self.n_nonzero_coefs, 243 ) 244 stage_checkpoint_dir = ( 245 os.path.join(checkpoint_dir, f"class_{class_label}") 246 if checkpoint_dir is not None 247 else None 248 ) 249 frozen_step.fit( 250 X, 251 class_label=class_label, 252 frozen_class_boundaries=dict(self.class_boundaries_), 253 checkpoint_dir=stage_checkpoint_dir, 254 resume=resume, 255 ) 256 257 self.D_ = frozen_step.D_combined_ 258 self.class_boundaries_ = dict(frozen_step.class_boundaries_) 259 self.class_labels_.append(class_label) 260 self.stage_learners_.append(frozen_step) 261 self.errors_[class_label] = list( 262 getattr(frozen_step.learner_, "errors_", []) 263 ) 264 265 # self.D_ already holds its own reference to this array — free the 266 # now-superseded snapshots nothing downstream ever reads again 267 # (get_stage_dict/get_class_dict always slice the top-level 268 # self.D_, never these). 269 frozen_step.D_frozen = None 270 frozen_step.D_combined_ = None 271 frozen_step.learner_.D_ = None 272 if hasattr(frozen_step.learner_, "Gamma_"): 273 frozen_step.learner_.Gamma_ = None 274 275 return self
Incrementally learn class-specific residual dictionaries, freezing all previously learned atoms before training the next class.
The underlying dictionary learner
(base_learner_class / residual_learner_class, e.g. KSVD)
is unsupervised. This class handles only that unsupervised,
incremental dictionary-learning pipeline and its per-class bookkeeping (class_boundaries_).
Pipeline
fit_base(X)Learn a base dictionary D_n from normal/background data usingbase_learner_class. This dictionary is frozen for all subsequent steps.add_class(X, class_label)Learn a residual dictionary D_a for the new class on top of the currently frozen dictionary [ D_n | D_a_1 | … ], viaFrozenDictionaryLearner. The underlying learner trains its new atoms jointly alongside the frozen ones every iteration.
The end product of this pipeline is D_, the full combined
dictionary. Sparse coding and classification on top of it happen
outside this class.
Checkpointing
Each call to fit_base or add_class represents one stage of
the incremental pipeline, and each stage trains its own, differently
shaped, inner dictionary-learner instance. If a checkpoint_dir is
given, this class therefore does NOT hand every stage the same path.
Instead, each stage gets its own subdirectory, so that interrupting
and resuming an individual stage's training does
not collide with any other stage's saved state.
Parameters
base_learner_class : type[BaseDictionaryLearner]
Unsupervised learner used for the initial base dictionary, e.g.
KSVD. Must accept D_frozen (default None) in its
fit(), per the frozen-dictionary contract.
base_learner_kwargs : dict
Init kwargs for base_learner_class.
residual_learner_class : type[BaseDictionaryLearner]
Learner used for each residual dictionary, via
FrozenDictionaryLearner. Can be the same as or different from
base_learner_class. Must support the frozen-dictionary
contract (accept and honour D_frozen).
residual_learner_kwargs : dict
Init kwargs for residual_learner_class. Applied identically
for every add_class call; override per-call via
add_class(..., learner_kwargs_override=...) — e.g. to give
each class a different number of atoms.
n_nonzero_coefs : int
Sparsity level passed down to FrozenDictionaryLearner during
add_class. Also the natural default sparsity for any sparse
coding you do yourself downstream against D_, though this
class does not perform that coding.
Attributes
D_ : np.ndarray full combined dictionary after all steps class_labels_ : list[int] class labels added via add_class, in order class_boundaries_ : dict[int, tuple[int, int]] Per-class atom ranges in the full combined D_ (includes the base stage's class label too). stage_learners_ : list Fitted learner (base stage) or FrozenDictionaryLearner (each add_class stage) from each stage, in order (index 0 = base stage). errors_ : dict[int, list[float]] Per-stage training RMSE curves keyed by class_label (key -1 for the base stage, regardless of its class_label).
93 def __init__( 94 self, 95 base_learner_class: type[BaseDictionaryLearner], 96 base_learner_kwargs: dict, 97 residual_learner_class: type[BaseDictionaryLearner], 98 residual_learner_kwargs: dict, 99 n_nonzero_coefs: int, 100 ) -> None: 101 self.base_learner_class = base_learner_class 102 self.base_learner_kwargs = base_learner_kwargs 103 self.residual_learner_class = residual_learner_class 104 self.residual_learner_kwargs = residual_learner_kwargs 105 self.n_nonzero_coefs = n_nonzero_coefs 106 107 # State built incrementally 108 self.D_: np.ndarray | None = None 109 self.base_class_label_: int | None = None 110 self.class_labels_: list[int] = [] 111 self.class_boundaries_: dict[int, tuple[int, int]] = {} 112 self.stage_learners_: list = [] 113 self.errors_: dict[int, list[float]] = {}
115 def fit_base( 116 self, 117 X: np.ndarray, 118 class_label: int = 0, 119 checkpoint_dir: str | None = None, 120 resume: bool = True, 121 ) -> IncrementalFrozenDictionary: 122 """ 123 Learn the base dictionary from normal / background data. 124 125 Parameters 126 ---------- 127 X : np.ndarray, shape (n_features, n_samples) 128 class_label : int 129 The class label this base dictionary's atoms are assigned to 130 in ``class_boundaries_`` (default 0). Must be distinct from 131 every ``class_label`` later passed to ``add_class``. 132 checkpoint_dir : str or None 133 If given, a ``base`` subdirectory under this path is passed to 134 the base learner's own ``fit(..., checkpoint_dir=...)``, if it 135 supports one. See the class docstring for why this is a 136 dedicated subdirectory rather than shared across stages. 137 resume : bool 138 Forwarded to the base learner's ``fit()`` alongside the 139 ``base`` checkpoint subdirectory. Default True. 140 """ 141 X = np.asarray(X, dtype=np.float32) 142 143 learner = self.base_learner_class(**self.base_learner_kwargs) 144 logger.info("Fitting base dictionary with %d samples and %d features using %s", X.shape[1], X.shape[0], self.base_learner_class.__name__) 145 146 fit_kwargs = {} 147 base_checkpoint_dir = None 148 if checkpoint_dir is not None: 149 base_checkpoint_dir = os.path.join(checkpoint_dir, "base") 150 fit_kwargs["checkpoint_dir"] = base_checkpoint_dir 151 fit_kwargs["resume"] = resume 152 153 loaded = ( 154 _completed_ksvd_checkpoint( 155 base_checkpoint_dir, 156 X, 157 n_components=self.base_learner_kwargs.get("n_components"), 158 n_nonzero_coefs=self.base_learner_kwargs.get("n_nonzero_coefs"), 159 n_iter=self.base_learner_kwargs.get("n_iter"), 160 n_frozen=0, 161 D_frozen=None, 162 ) 163 if resume 164 else None 165 ) 166 if loaded is not None: 167 learner.D_, learner.Gamma_, learner.errors_ = loaded 168 else: 169 # Unsupervised: no D_frozen (nothing to freeze yet). 170 learner.fit(X, **fit_kwargs) 171 172 self.D_ = learner.D_ 173 self.base_class_label_ = class_label 174 self.class_boundaries_ = {class_label: (0, learner.D_.shape[1])} 175 self.stage_learners_.append(learner) 176 self.errors_[-1] = list(getattr(learner, "errors_", [])) 177 178 # self.D_ already holds its own reference to the learned array; 179 # dropping the learner's copies frees nothing-else-reads-them 180 # state rather than holding it for the lifetime of the pipeline. 181 # get_stage_dict(0) slices self.D_ via class_boundaries_, not 182 # this attribute, so it stays correct after this. 183 learner.D_ = None 184 if hasattr(learner, "Gamma_"): 185 learner.Gamma_ = None 186 187 return self
Learn the base dictionary from normal / background data.
Parameters
X : np.ndarray, shape (n_features, n_samples)
class_label : int
The class label this base dictionary's atoms are assigned to
in class_boundaries_ (default 0). Must be distinct from
every class_label later passed to add_class.
checkpoint_dir : str or None
If given, a base subdirectory under this path is passed to
the base learner's own fit(..., checkpoint_dir=...), if it
supports one. See the class docstring for why this is a
dedicated subdirectory rather than shared across stages.
resume : bool
Forwarded to the base learner's fit() alongside the
base checkpoint subdirectory. Default True.
189 def add_class( 190 self, 191 X: np.ndarray, 192 class_label: int, 193 learner_kwargs_override: dict | None = None, 194 checkpoint_dir: str | None = None, 195 resume: bool = True, 196 ) -> IncrementalFrozenDictionary: 197 """ 198 Learn a residual dictionary for a new class and extend D_. 199 200 Parameters 201 ---------- 202 X : np.ndarray, shape (n_features, n_samples) 203 Training signals for this class only. Presented unsupervised 204 to the residual learner — it never sees labels. 205 class_label : int 206 Integer label for this class. Must not have been added 207 before, and must differ from the base stage's class_label. 208 learner_kwargs_override : dict or None 209 If supplied, overrides ``residual_learner_kwargs`` for this 210 call only. Useful for giving this class a different number 211 of atoms (e.g. ``{"n_components": 5}``). 212 checkpoint_dir : str or None 213 If given, a ``class_<class_label>`` subdirectory under this 214 path is passed down to ``FrozenDictionaryLearner.fit()`` (and 215 from there to the residual learner's own ``fit()``), if 216 supported. Kept distinct per class_label, and distinct from 217 the base stage's ``base`` subdirectory, so that resuming one 218 stage never collides with another's saved state — see the 219 class docstring. 220 resume : bool 221 Forwarded down to the residual learner's ``fit()`` alongside 222 the per-class checkpoint subdirectory. Default True. 223 """ 224 if self.D_ is None: 225 raise DictionaryLearningError( 226 "Call fit_base() before add_class()." 227 ) 228 if class_label in self.class_labels_ or class_label == self.base_class_label_: 229 raise ValueError( 230 f"class_label {class_label} has already been used." 231 ) 232 233 X = np.asarray(X, dtype=np.float32) 234 235 kwargs = {**self.residual_learner_kwargs, **(learner_kwargs_override or {})} 236 logger.info("Adding class %d with %d samples and %d features using %s", class_label, X.shape[1], X.shape[0], self.residual_learner_class.__name__) 237 238 frozen_step = FrozenDictionaryLearner( 239 D_frozen=self.D_, 240 learner_class=self.residual_learner_class, 241 learner_kwargs=kwargs, 242 n_nonzero_coefs=self.n_nonzero_coefs, 243 ) 244 stage_checkpoint_dir = ( 245 os.path.join(checkpoint_dir, f"class_{class_label}") 246 if checkpoint_dir is not None 247 else None 248 ) 249 frozen_step.fit( 250 X, 251 class_label=class_label, 252 frozen_class_boundaries=dict(self.class_boundaries_), 253 checkpoint_dir=stage_checkpoint_dir, 254 resume=resume, 255 ) 256 257 self.D_ = frozen_step.D_combined_ 258 self.class_boundaries_ = dict(frozen_step.class_boundaries_) 259 self.class_labels_.append(class_label) 260 self.stage_learners_.append(frozen_step) 261 self.errors_[class_label] = list( 262 getattr(frozen_step.learner_, "errors_", []) 263 ) 264 265 # self.D_ already holds its own reference to this array — free the 266 # now-superseded snapshots nothing downstream ever reads again 267 # (get_stage_dict/get_class_dict always slice the top-level 268 # self.D_, never these). 269 frozen_step.D_frozen = None 270 frozen_step.D_combined_ = None 271 frozen_step.learner_.D_ = None 272 if hasattr(frozen_step.learner_, "Gamma_"): 273 frozen_step.learner_.Gamma_ = None 274 275 return self
Learn a residual dictionary for a new class and extend D_.
Parameters
X : np.ndarray, shape (n_features, n_samples)
Training signals for this class only. Presented unsupervised
to the residual learner — it never sees labels.
class_label : int
Integer label for this class. Must not have been added
before, and must differ from the base stage's class_label.
learner_kwargs_override : dict or None
If supplied, overrides residual_learner_kwargs for this
call only. Useful for giving this class a different number
of atoms (e.g. {"n_components": 5}).
checkpoint_dir : str or None
If given, a class_<class_label> subdirectory under this
path is passed down to FrozenDictionaryLearner.fit() (and
from there to the residual learner's own fit()), if
supported. Kept distinct per class_label, and distinct from
the base stage's base subdirectory, so that resuming one
stage never collides with another's saved state — see the
class docstring.
resume : bool
Forwarded down to the residual learner's fit() alongside
the per-class checkpoint subdirectory. Default True.
122def fista_core( 123 grad_f: ArrayFunc, 124 prox_g: ProxFunc, 125 x0: Array, 126 f: float | None = None, 127 g: float | None = None, 128 L: float | None = None, 129 mode: str = "backtracking", 130 L0: float = 1.0, 131 eta: float = 2.0, 132 max_iter: int = 500, 133 tol: float | None = 1e-8, 134 max_backtrack_iter: int = 100, 135) -> FISTAResult: 136 """ 137 Run FISTA (Beck & Teboulle, 2009) to minimize F(x) = f(x) + g(x). 138 139 Parameters 140 ---------- 141 grad_f : callable 142 Gradient of the smooth part, grad_f(x) -> array (same shape/type as x). 143 prox_g : callable 144 Proximal operator of g: prox_g(v, t) -> argmin_x g(x) + ||x-v||^2/(2t). 145 x0 : np.ndarray or torch.Tensor 146 Initial point. Any shape is supported (vector, matrix, ...); the 147 Frobenius inner product is used internally, so the analysis holds 148 verbatim in this more general Hilbert-space setting (Remark 2.1). 149 Pass a torch.Tensor already on a GPU device to run the whole solve 150 on GPU (grad_f/prox_g/f/g must then also operate on that device). 151 f, g : callable, optional 152 Value of the smooth / nonsmooth parts. Required when 153 mode='backtracking' (needed to check the descent condition 154 (3.2)/eq. before (4.1)). If supplied, also used to record 155 objective_history. 156 L : float, optional 157 Lipschitz constant of grad_f. Required when mode='constant'. 158 mode : {'constant', 'backtracking'} 159 Stepsize strategy ("FISTA with constant stepsize" / 160 "FISTA with backtracking" in Section 4 of the paper). 161 L0 : float 162 Initial Lipschitz estimate for backtracking mode (ignored for 163 mode='constant'). 164 eta : float 165 Backtracking growth factor, eta > 1. 166 max_iter : int 167 Maximum number of iterations. 168 tol : float or None 169 Relative stopping tolerance on ||x_k - x_{k-1}|| / max(1, ||x_{k-1}||). 170 This is a practical stopping rule, not part of the original paper 171 (which only bounds F(x_k) - F(x*)); pass tol=None to always run 172 max_iter iterations. 173 max_backtrack_iter : int 174 Safety cap on the number of backtracking growth steps per outer 175 iteration. 176 177 Returns 178 ------- 179 FISTAResult 180 """ 181 if mode not in ("constant", "backtracking"): 182 raise ValueError("mode must be 'constant' or 'backtracking'.") 183 if mode == "constant" and (L is None or L <= 0): 184 raise ValueError("mode='constant' requires a positive Lipschitz constant L.") 185 if mode == "backtracking" and (f is None or g is None): 186 raise ValueError( 187 "mode='backtracking' requires both f and g to evaluate the descent condition." 188 ) 189 if eta <= 1: 190 raise ValueError("eta must be > 1.") 191 if max_iter < 1: 192 raise ValueError("max_iter must be >= 1.") 193 194 x_prev = _prepare_x0(x0) # x_0 195 y = _copy(x_prev) # y_1 = x_0 196 t = 1.0 # t_1 = 1 197 Lk = float(L) if mode == "constant" else float(L0) 198 199 obj_history: list[float] = [] 200 converged = False 201 n_iter = 0 202 x_k = x_prev 203 204 for k in range(1, max_iter + 1): 205 n_iter = k 206 grad_y = grad_f(y) 207 f_xk: float | None = None 208 g_xk: float | None = None 209 210 if mode == "constant": 211 x_k = _p_L(y, grad_y, Lk, prox_g) 212 else: 213 f_y = f(y) 214 L_bar = Lk 215 for _ in range(max_backtrack_iter): 216 x_k = _p_L(y, grad_y, L_bar, prox_g) 217 g_xk = g(x_k) 218 f_xk = f(x_k) 219 if f_xk <= _q(x_k, y, f_y, grad_y, L_bar, g_xk): 220 break 221 L_bar *= eta 222 else: 223 raise RuntimeError( 224 "Backtracking line search failed to satisfy the descent " 225 "condition; check f/grad_f/L0/eta." 226 ) 227 Lk = L_bar 228 229 # eqs. (4.2)-(4.3): momentum update and extrapolation point. 230 t_next = (1.0 + np.sqrt(1.0 + 4.0 * t * t)) / 2.0 231 x_diff = x_k - x_prev 232 y = x_k + ((t - 1.0) / t_next) * x_diff 233 234 if f is not None and g is not None: 235 if f_xk is None or g_xk is None: # mode == "constant" 236 f_xk = f(x_k) 237 g_xk = g(x_k) 238 obj_history.append(f_xk + g_xk) 239 240 if tol is not None: 241 denom = max(1.0, _norm(x_prev)) 242 if _norm(x_diff) / denom < tol: 243 x_prev = x_k 244 converged = True 245 break 246 247 x_prev = x_k 248 t = t_next 249 250 if k % 50 == 0 or k == max_iter: 251 logger.info("FISTA iteration %d/%d (L=%.4g)", k, max_iter, Lk) 252 253 return FISTAResult( 254 x=x_prev, 255 n_iter=n_iter, 256 converged=converged, 257 L=Lk, 258 objective_history=obj_history, 259 )
Run FISTA (Beck & Teboulle, 2009) to minimize F(x) = f(x) + g(x).
Parameters
grad_f : callable Gradient of the smooth part, grad_f(x) -> array (same shape/type as x). prox_g : callable Proximal operator of g: prox_g(v, t) -> argmin_x g(x) + ||x-v||^2/(2t). x0 : np.ndarray or torch.Tensor Initial point. Any shape is supported (vector, matrix, ...); the Frobenius inner product is used internally, so the analysis holds verbatim in this more general Hilbert-space setting (Remark 2.1). Pass a torch.Tensor already on a GPU device to run the whole solve on GPU (grad_f/prox_g/f/g must then also operate on that device). f, g : callable, optional Value of the smooth / nonsmooth parts. Required when mode='backtracking' (needed to check the descent condition (3.2)/eq. before (4.1)). If supplied, also used to record objective_history. L : float, optional Lipschitz constant of grad_f. Required when mode='constant'. mode : {'constant', 'backtracking'} Stepsize strategy ("FISTA with constant stepsize" / "FISTA with backtracking" in Section 4 of the paper). L0 : float Initial Lipschitz estimate for backtracking mode (ignored for mode='constant'). eta : float Backtracking growth factor, eta > 1. max_iter : int Maximum number of iterations. tol : float or None Relative stopping tolerance on ||x_k - x_{k-1}|| / max(1, ||x_{k-1}||). This is a practical stopping rule, not part of the original paper (which only bounds F(x_k) - F(x*)); pass tol=None to always run max_iter iterations. max_backtrack_iter : int Safety cap on the number of backtracking growth steps per outer iteration.
Returns
FISTAResult