physioex.explain.prototypes.posthoc#

class physioex.explain.prototypes.posthoc.VQBottleneck(codebook_init, commitment_weight=0.25)[source]#

Bases: Module

Differentiable vector quantization with straight-through estimator.

Forward: z_q = codebook[argmin_k ||z - c_k||^2]. Backward: gradients pass through as if z_q = z (STE), but the codebook receives gradients from ||sg[z] - c_q||^2.

Parameters:

  • codebook_init (np.ndarray) – (M, d_model) initial codebook (e.g. from K-Means).

  • commitment_weight (float) – Weight β for commitment loss ||z - sg[c_q]||².

forward(z)[source]#

Quantize with STE.

Parameters:

z (Tensor) – (N, d_model) embeddings.

Returns:

(N, d_model) quantized embeddings (STE gradients to z) vq_loss: scalar, codebook_loss + β * commitment_loss

Return type:

z_q

physioex.explain.prototypes.posthoc.discover_prototypes_nmf(Z, Y, k_per_class=10, n_classes=5, random_state=42, max_iter=500)[source]#

Discover prototypes via per-class NMF decomposition.

For each class c, collects all embeddings F_c = Z[Y == c] and runs NMF: F_c ≈ E_c @ P_c, where P_c has k rows (the prototypes).

Parameters:

  • Z (ndarray) – (N, d_model) epoch embeddings (float32).

  • Y (ndarray) – (N,) integer class labels.

  • k_per_class (int) – Number of prototypes to discover per class.

  • n_classes (int) – Total number of classes.

  • random_state (int) – Seed for NMF.

  • max_iter (int) – Maximum NMF iterations.

Returns:

prototypes: (k_per_class * n_classes, d_model) proto_labels: (k_per_class * n_classes,) class label per prototype

Return type:

Tuple of

physioex.explain.prototypes.posthoc.evaluate_metrics(y_true, y_pred, class_names=None, n_classes=5)[source]#

Compute standard sleep staging metrics.

Parameters:

  • y_true (ndarray) – (N,) ground truth labels.

  • y_pred (ndarray) – (N,) predicted labels.

  • class_names (Sequence[str] | None) – Names for each class (default: W, N1, N2, N3, REM).

  • n_classes (int) – Number of classes.

Returns:

Dict with accuracy, f1_macro, kappa, per_class_f1, confusion_matrix.

Return type:

dict

physioex.explain.prototypes.posthoc.learn_codebook_kmeans(Z, n_prototypes=50, random_state=42, batch_size=4096, max_iter=300)[source]#

Learn initial codebook via K-Means clustering.

This is Stage 1 (initialization). Use train_codebook() for supervised refinement (Stage 2).

Parameters:

  • Z (ndarray) – (N, d_model) epoch embeddings (float32).

  • n_prototypes (int) – Number of codebook entries (clusters).

  • random_state (int) – Seed for K-Means.

  • batch_size (int) – Mini-batch size for MiniBatchKMeans.

  • max_iter (int) – Maximum iterations.

Returns:

(n_prototypes, d_model) cluster centroids.

Return type:

codebook

physioex.explain.prototypes.posthoc.load_epoch_embeddings(emb_dir, split='train')[source]#

Load pre-extracted epoch embeddings and labels for a split.

Supports two layouts:

  1. Per-subject (preferred):

    emb_dir/{split}/{subject_id}_embeddings.npy
emb_dir/{split}/{subject_id}_labels.npy

  2. Flat (legacy):

    emb_dir/{split}_embeddings.npy
emb_dir/{split}_labels.npy
Parameters:

  • emb_dir (str | Path) – Directory containing the embeddings.

  • split (str) – One of "train", "valid", "test".

Returns:

Tuple of (Z, Y) where Z is (N, d_model) float32 and Y is (N,) int64.

Return type:

tuple[ndarray, ndarray]

physioex.explain.prototypes.posthoc.load_epoch_embeddings_per_subject(emb_dir, split='train')[source]#

Load per-subject epoch embeddings and labels for a split.

Returns a list of (Z_subj, Y_subj) tuples, one per subject, preserving the temporal order within each subject.

Parameters:

  • emb_dir (str | Path) – Directory containing per-subject .npy files.

  • split (str) – One of "train", "valid", "test".

Returns:

List of (Z, Y) tuples where Z is (n_epochs, d_model) and Y is (n_epochs,).

Return type:

list[tuple[ndarray, ndarray]]

physioex.explain.prototypes.posthoc.nearest_prototype_classify(Z, prototypes, proto_labels, metric='cosine')[source]#

Classify embeddings via nearest prototype.

For each embedding z_i, finds the prototype with highest similarity (cosine) or lowest distance (euclidean) and assigns its class label.

Parameters:

  • Z (ndarray) – (N, d_model) embeddings to classify.

  • prototypes (ndarray) – (M, d_model) prototype vectors.

  • proto_labels (ndarray) – (M,) class label for each prototype.

  • metric (str) – "cosine" or "euclidean".

Returns:

(N,) predicted class labels.

Return type:

ndarray

physioex.explain.prototypes.posthoc.quantize_embeddings(Z, codebook)[source]#

Quantize each embedding to its nearest codebook entry (numpy).

Parameters:

  • Z (ndarray) – (N, d_model) epoch embeddings.

  • codebook (ndarray) – (M, d_model) codebook entries.

Returns:

Z_quantized: (N, d_model) quantized embeddings assignments: (N,) index of assigned codebook entry

Return type:

Tuple of

physioex.explain.prototypes.posthoc.train_codebook(train_subjects, downstream_fn, codebook_init, val_subjects=None, n_epochs=50, patience=5, batch_size=32, lr=0.0001, commitment_weight=0.25, device='cpu', sequence_length=21, save_path=None)[source]#

Supervised codebook refinement (Stage 2).

Optimizes the codebook so that quantized embeddings, when passed through the frozen downstream model, preserve classification accuracy.

Uses stride-1 sliding windows on real per-subject recordings so the sequence encoder sees coherent temporal context. Windows are fetched on-the-fly (no pre-materialization) to avoid OOM on large datasets.

Parameters:

  • train_subjects (list[tuple[ndarray, ndarray]]) – List of (Z, Y) tuples per subject. Z: (n_epochs, d_model), Y: (n_epochs,).

  • downstream_fn (Callable[[Tensor], Tensor]) – Takes (B, L, d_model) -> (B, L, n_classes) logits.

  • codebook_init (ndarray) – (M, d_model) initial codebook from K-Means.

  • val_subjects (list[tuple[ndarray, ndarray]] | None) – Optional list of (Z, Y) per subject for validation.

  • n_epochs (int) – Maximum training epochs.

  • patience (int) – Early stopping patience.

  • batch_size (int) – Number of L-length windows per optimization step.

  • lr (float) – Learning rate for codebook.

  • commitment_weight (float) – beta for commitment loss.

  • device (str) – Torch device string.

  • sequence_length (int) – L, window length for sliding windows.

  • save_path (str | None) – Save best codebook here on every val improvement.

Returns:

(M, d_model) best codebook as numpy array.

Return type:

codebook