Parametric Spectral Map

spectral_map

CLASS	DESCRIPTION
SpectralMap	Spectral map for dimensionality reduction via eigenvalue optimization.

Classes

SpectralMap(*, model: ModelType, kernel_fn: Kernel | str = 'gaussian', kernel_kwargs: dict | None = None, distance_fn: PairwiseDistance | str = 'euclidean', distance_kwargs: dict | None = None, loss_fn: Literal['eigenvalue', 'adaptive_eigenvalue'] = 'eigenvalue', loss_kwargs: dict | None = None, loss_weights: dict[str, float] | None = None, additional_loss: list | dict | None = None, additional_loss_weights: list[float] | dict[str, float] | None = None, optimizer_fn: type[torch.optim.Optimizer] | str | None = None, optimizer_kwargs: dict | None = None, lr_scheduler_fn: Any | str | None = None, lr_scheduler_kwargs: dict | None = None, n_eigen: int = 10, symmetric_eigendecomposition: bool = True, device: Literal['cuda', 'cpu'] = 'cuda', store_loss_context: bool = False)

Bases: Parametric

Spectral map for dimensionality reduction via eigenvalue optimization.

Learns low-dimensional representations by optimizing spectral properties of kernel matrices computed in the latent space. By default, the method maximizes the spectral gap between leading eigenvalues to identify metastable states or dominant geometric structures in high-dimensional data.

The optimization objective is based on eigenvalue decomposition of a kernel matrix constructed from neural network embeddings, enabling end-to-end learning of spectral properties.

The spectral gap loss is defined as (loss_fn="eigenvalue" with reduce="gap"): \(\lambda_{n-1} - \lambda_n\), where \(\lambda\) are the ordered eigenvalues of the kernel matrix. Larger spectral gaps indicate better separation between leading eigenvalues.

PARAMETER	DESCRIPTION
model	Neural network architecture for learning spectral embeddings. Can be a single model or dictionary of models. See Parametric for supported model types. TYPE: ModelType
kernel_fn	Kernel function specification. Kernel object: Pre-configured kernel instance str: Kernel type name (see available options below) Available string options: "gaussian": GaussianKernel "t": Student's TKernel TYPE: Kernel | str, optional, by default "gaussian" DEFAULT: 'gaussian'
kernel_kwargs	Keyword arguments for kernel instantiation when kernel_fn is a string. See kernels for available options. TYPE: dict | None, optional, by default None DEFAULT: None
distance_fn	Distance metric specification. PairwiseDistance object: Pre-configured distance instance str: Distance type name (see available options below) Available string options: "euclidean": Euclidean L2 distance "mahalanobis": Mahalanobis distance "covariance": Covariance-based distance TYPE: PairwiseDistance | str, optional, by default "euclidean" DEFAULT: 'euclidean'
distance_kwargs	Keyword arguments for distance instantiation when distance_fn is a string. TYPE: dict | None, optional, by default None DEFAULT: None
loss_fn	Loss function specification. Only EigenvalueLoss (loss_fn = "eigenvalue") and AdaptiveEigenvalueLoss (loss_fn = "adaptive_eigenvalue") are supported for spectral optimization. TYPE: str, by default "eigenvalue" DEFAULT: 'eigenvalue'
loss_kwargs	Additional keyword arguments for the loss function. reduce: Literal["gap", "gap_exp", "sum", "sum_pow", "single", "none", None], optional, by default "gap" See losses for more information. sign: float, optional, by default -1.0 Loss sign multiplier. TYPE: dict, optional, by default None DEFAULT: None
loss_weights	Weights for built-in loss components. Keys must match names in DEFAULT_LOSS_WEIGHTS. Unspecified keys keep their default weight. For SpectralMap, the available key is: "loss_eigenvalue": Weight for the spectral (eigenvalue) loss (default 1.0) Example: loss_weights={"loss_eigenvalue": 0.5} TYPE: dict[str, float] | None, optional, by default None DEFAULT: None
additional_loss	Additional loss functions to combine with the primary spectral loss. list[Loss]: Unnamed losses combined with primary loss dict[str, Loss]: Named losses for interpretable logging None: Use only the primary spectral loss (default) The total loss becomes a weighted sum of primary and additional losses. Useful for adding regularization (e.g., OrthogonalLoss, DecorrelationLoss). TYPE: list[Loss] | dict[str, Loss] | None, by default None DEFAULT: None
additional_loss_weights	Weights for additional loss functions. For list[float]: provide list of floats (same length as additional_loss) For dict[str, float]: provide dict with matching keys If None: all additional losses weighted at 1.0 TYPE: list | dict | None, by default None DEFAULT: None
optimizer_fn	Optimizer specification for training. None: Uses Adam optimizer (default) str: Optimizer name from torch.optim (e.g., "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (supports third-party) TYPE: torch.optim.Optimizer | str | None, optional, by default None DEFAULT: None
optimizer_kwargs	Keyword arguments for optimizer instantiation (e.g., {"lr": 1e-3, "weight_decay": 1e-4}). TYPE: dict | None, optional, by default None DEFAULT: None
lr_scheduler_fn	Learning rate scheduler specification. None: No scheduler (default) str: Scheduler name from torch.optim.lr_scheduler (e.g., "ReduceLROnPlateau") type: Scheduler class TYPE: Any | str | None, optional, by default None DEFAULT: None
lr_scheduler_kwargs	Keyword arguments for scheduler instantiation. TYPE: dict | None, optional, by default None DEFAULT: None
n_eigen	Number of eigenvalues to compute. TYPE: int, optional, by default 10 DEFAULT: 10
symmetric_eigendecomposition	Whether to use symmetric eigendecomposition. TYPE: bool, optional, by default True DEFAULT: True
device	Computation device. Falls back to "cpu" if CUDA unavailable. TYPE: Literal["cuda", "cpu"], optional, by default "cuda" DEFAULT: 'cuda'

ATTRIBUTE	DESCRIPTION
DEFAULT_LOSS_WEIGHTS	Default weights for built-in loss components: {"loss_eigenvalue": 1.0}. TYPE: ClassVar[dict[str, float]]
kernel_fn	Kernel function instance. TYPE: Kernel
distance_fn	Distance function instance. TYPE: PairwiseDistance
loss_fn	Composite loss function containing spectral loss and any additional losses. TYPE: CompositeLoss
eigenvalues	Most recently computed eigenvalues from forward pass, shape (n_eigen,). Updated during loss() computation. TYPE: Tensor

Examples:

String-based configuration:

>>> import torch
>>> from spectre.parametric import SpectralMap
>>> from spectre.core import Model
>>>
>>> model = Model(in_features=10, out_features=2, multipliers=[0.5])
>>>
>>> sm = SpectralMap(
...     model=model,
...     kernel_fn="gaussian",
...     kernel_kwargs={"bw_method": "median"},
...     distance_fn="euclidean",
...     loss_kwargs={"n_eigen": 2, "reduce": "gap"},
...     device="cpu",
... )
>>>
>>> X = torch.randn(100, 10)
>>> loss = sm.loss(X)

Object-based configuration:

>>> from spectre.kernel import TKernel
>>> from spectre.pairwise_distance import PairwiseDistanceMahalanobis
>>>
>>> sm = SpectralMap(
...     model=model,
...     kernel_fn=TKernel(alpha=torch.tensor(1.0), beta=torch.tensor(3.0)),
...     distance_fn=PairwiseDistanceMahalanobis(),
...     loss_kwargs={"n_eigen": 2, "reduce": "gap"},
...     device="cpu",
... )

With custom loss weights and additional losses:

>>> from spectre.loss import OrthogonalLoss, DecorrelationLoss
>>>
>>> sm = SpectralMap(
...     model=model,
...     kernel_fn="gaussian",
...     loss_fn="eigenvalue",
...     loss_kwargs={"n_eigen": 2, "reduce": "gap"},
...     loss_weights={"loss_eigenvalue": 0.5},
...     additional_loss={
...         "orthogonal": OrthogonalLoss(),
...         "decorrelation": DecorrelationLoss(method="correlation"),
...     },
...     additional_loss_weights={
...         "orthogonal": 0.1,
...         "decorrelation": 0.05,
...     },
...     device="cpu",
... )
>>>
>>> # Total 0.5 * spectral_loss + 0.1 * orthogonal + 0.05 * decorrelation
>>> X = torch.randn(100, 10)
>>> loss = sm.loss(X)  # Composite loss value

METHOD	DESCRIPTION
training_step	PyTorch Lightning training step for spectral map optimization.
compute_loss_context	Compute inputs needed for spectral loss computation.
forward_step	Compute forward pass through network.
loss	Compute spectral loss based on kernel matrix eigenvalues in latent
score_step	Compute spectral loss for evaluation without training status checks.

Source code in spectre/parametric/spectral_map.py

def __init__(
    self,
    *,
    model: ModelType,
    kernel_fn: Kernel | str = "gaussian",
    kernel_kwargs: dict | None = None,
    distance_fn: PairwiseDistance | str = "euclidean",
    distance_kwargs: dict | None = None,
    loss_fn: Literal["eigenvalue", "adaptive_eigenvalue"] = "eigenvalue",
    loss_kwargs: dict | None = None,
    loss_weights: dict[str, float] | None = None,
    additional_loss: list | dict | None = None,
    additional_loss_weights: list[float] | dict[str, float] | None = None,
    optimizer_fn: type[torch.optim.Optimizer] | str | None = None,
    optimizer_kwargs: dict | None = None,
    lr_scheduler_fn: Any | str | None = None,
    lr_scheduler_kwargs: dict | None = None,
    n_eigen: int = 10,
    symmetric_eigendecomposition: bool = True,
    device: Literal["cuda", "cpu"] = "cuda",
    store_loss_context: bool = False,
) -> None:
    super().__init__(
        model=model,
        device=device,
        store_loss_context=store_loss_context,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )

    self.kernel_fn = initialize_kernel_fn(kernel_fn, kernel_kwargs or {})
    self.distance_fn = initialize_distance_fn(distance_fn, distance_kwargs or {})

    if loss_kwargs is None:
        loss_kwargs = {}

    if loss_fn == "eigenvalue":
        self.loss_fn = EigenvalueLoss(**loss_kwargs)
    elif loss_fn == "adaptive_eigenvalue":
        self.loss_fn = AdaptiveEigenvalueLoss(**loss_kwargs)
    else:
        raise ValueError(f"Loss function `{loss_fn}` not available.")

    self._loss_weights = {**self.DEFAULT_LOSS_WEIGHTS}
    if loss_weights is not None:
        if not isinstance(loss_weights, dict):
            raise TypeError(
                f"loss_weights must be dict or None, got {type(loss_weights)}."
            )
        unknown = set(loss_weights.keys()) - set(self.DEFAULT_LOSS_WEIGHTS)
        if unknown:
            raise ValueError(
                f"Unknown loss weight keys: {unknown}. "
                f"Available: {set(self.DEFAULT_LOSS_WEIGHTS.keys())}"
            )
        self._loss_weights.update(loss_weights)

    self.loss_fn = CompositeLoss(
        losses={"loss_eigenvalue": self.loss_fn},
        weights={"loss_eigenvalue": self._loss_weights["loss_eigenvalue"]},
    )

    if additional_loss is not None:
        self.loss_fn = self.loss_fn + CompositeLoss(
            losses=additional_loss, weights=additional_loss_weights
        )

    if not isinstance(n_eigen, int) or n_eigen < 2:
        raise ValueError("n_eigen must be an integer > 1.")
    self.n_eigen = n_eigen

    if not isinstance(symmetric_eigendecomposition, bool):
        raise ValueError("symmetric_eigendecomposition must be a boolean.")
    self.symmetric_eigendecomposition = symmetric_eigendecomposition

    self._compute_eigenvectors = False
    if self.loss_fn.exists("eigenvector"):
        self._compute_eigenvectors = True

    self._compute_inverse_diffusion_tensor = False
    if isinstance(self.distance_fn, PairwiseDistanceMahalanobis):
        self._compute_inverse_diffusion_tensor = True

Functions

training_step(batch: DataBatch, batch_idx: int) -> torch.Tensor

PyTorch Lightning training step for spectral map optimization.

Processes a single training batch by computing spectral loss and logging eigenvalue metrics. Automatically logs loss and eigenvalues to the trainer. When additional losses are provided, logs individual loss components.

PARAMETER	DESCRIPTION
batch	Training batch. TYPE: DataBatch
batch_idx	Index of current batch within the epoch. TYPE: int

RETURNS	DESCRIPTION
Tensor	Scalar spectral loss value for backpropagation. If additional losses are provided, returns weighted sum of all components.

Source code in spectre/parametric/spectral_map.py

def training_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    PyTorch Lightning training step for spectral map optimization.

    Processes a single training batch by computing spectral loss and
    logging eigenvalue metrics. Automatically logs loss and eigenvalues to
    the trainer. When additional losses are provided, logs individual loss
    components.

    Parameters
    ----------
    batch : DataBatch
        Training batch.

    batch_idx : int
        Index of current batch within the epoch.

    Returns
    -------
    torch.Tensor
        Scalar spectral loss value for backpropagation. If additional
        losses are provided, returns weighted sum of all components.
    """
    loss = self.loss(X=batch.data, weights=batch.weights, target=batch.target)

    self.log_training_metrics(loss, **self.loss_fn.loss_components)

    return loss

compute_loss_context(X: torch.Tensor, weights: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]

Compute inputs needed for spectral loss computation.

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: Tensor
weights	Optional sample weights. TYPE: Tensor | None DEFAULT: None

RETURNS	DESCRIPTION
tuple	(z_batch, z_kernel, eigenvalues, eigenvectors)

Source code in spectre/parametric/spectral_map.py

def compute_loss_context(
    self, X: torch.Tensor, weights: torch.Tensor | None = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """
    Compute inputs needed for spectral loss computation.

    Parameters
    ----------
    X : torch.Tensor
        Input data tensor.

    weights : torch.Tensor | None
        Optional sample weights.

    Returns
    -------
    tuple
        (z_batch, z_kernel, eigenvalues, eigenvectors)
    """
    z_batch = self.forward_step(X)

    if self._compute_inverse_diffusion_tensor:
        inv_diff = inverse_diffusion_tensor(X, self.model)
        z_pdist = self.distance_fn(z_batch, metric_inv=inv_diff)
    else:
        z_pdist = self.distance_fn(z_batch)

    z_kernel = self.kernel_fn(z_pdist, weights=weights)

    if self._compute_eigenvectors:
        eigenvalues, eigenvectors = eigvecs(
            z_kernel,
            n_eigen=self.n_eigen,
            symmetric=self.symmetric_eigendecomposition,
            sort_descending=True,
            UPLO="U",
        )
        eigenvectors = eigenvectors[:, 1 : self.out_features + 1]
    else:
        eigenvalues = eigvals(
            z_kernel,
            n_eigen=self.n_eigen,
            symmetric=self.symmetric_eigendecomposition,
            sort_descending=True,
            UPLO="U",
        )
        eigenvectors = None

    return z_batch, z_kernel, eigenvalues, eigenvectors

forward_step(X: torch.Tensor) -> torch.Tensor

Compute forward pass through network.

Applies the neural network model to input data to obtain low-dimensional representations in the latent space where kernel matrices are computed.

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: torch.Tensor of shape (n_samples, in_features)

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, out_features)	Latent space embeddings.

Source code in spectre/parametric/spectral_map.py

def forward_step(self, X: torch.Tensor) -> torch.Tensor:
    """
    Compute forward pass through network.

    Applies the neural network model to input data to obtain
    low-dimensional representations in the latent space where kernel
    matrices are computed.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data tensor.

    Returns
    -------
    torch.Tensor of shape (n_samples, out_features)
        Latent space embeddings.
    """
    X = self.model(X)

    return X

loss(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute spectral loss based on kernel matrix eigenvalues in latent space.

When additional losses are provided, computes a composite loss combining the primary spectral loss with additional regularization terms.

PARAMETER	DESCRIPTION
X	Input data tensor. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Optional sample weights for weighted kernel computation. Applied during kernel normalization if provided. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Not used for loss computation (unsupervised method). TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar loss value. If additional losses are provided, returns weighted sum of all loss components.

Source code in spectre/parametric/spectral_map.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute spectral loss based on kernel matrix eigenvalues in latent
    space.

    When additional losses are provided, computes a composite loss
    combining the primary spectral loss with additional regularization
    terms.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data tensor.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Optional sample weights for weighted kernel computation. Applied
        during kernel normalization if provided.

    target : torch.Tensor or None, by default None
        Not used for loss computation (unsupervised method).

    Returns
    -------
    torch.Tensor
        Scalar loss value. If additional losses are provided, returns
        weighted sum of all loss components.
    """
    z_batch, z_kernel, eigenvalues, eigenvectors = self.compute_loss_context(
        X=X, weights=weights
    )

    self.eigenvalues = eigenvalues
    self.eigenvectors = eigenvectors

    context = {
        "eigenvalues": eigenvalues,
        "eigenvectors": eigenvectors,
        "weights": weights,
        "K": z_kernel,
        "Z": z_batch,
    }

    if self.store_loss_context:
        self.loss_context = context

    loss = self.loss_fn(context)

    return loss

score_step(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute spectral loss for evaluation without training status checks.

Evaluates the spectral map quality by computing eigenvalue-based loss on the kernel matrix in the latent space without requiring the model to be marked as trained.

PARAMETER	DESCRIPTION
X	Input data for evaluation. TYPE: torch.Tensor of shape (n_samples, n_features)
weights	Sample weights for weighted kernel computation. TYPE: torch.Tensor (n_samples,) or None, optional, by default None DEFAULT: None
target	Target tensor for supervised scoring. Can be a tensor of reference targets, e.g., reference eigenvalues. TYPE: torch.Tensor or None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar spectral loss value (negative indicates better spectral gap).

Source code in spectre/parametric/spectral_map.py

def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute spectral loss for evaluation without training status checks.

    Evaluates the spectral map quality by computing eigenvalue-based loss
    on the kernel matrix in the latent space without requiring the model to
    be marked as trained.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, n_features)
        Input data for evaluation.

    weights : torch.Tensor (n_samples,) or None, optional, by default None
        Sample weights for weighted kernel computation.

    target : torch.Tensor or None, optional, by default None
        Target tensor for supervised scoring.

        Can be a tensor of reference targets, e.g., reference eigenvalues.

    Returns
    -------
    torch.Tensor
        Scalar spectral loss value (negative indicates better spectral
        gap).
    """
    return self.loss(X, weights=weights, target=target)

Functions