Loss Eigen

eigen

CLASS	DESCRIPTION
EigenvalueLoss	Eigenvalue loss.
AdaptiveEigenvalueLoss	Adaptive spectral gap loss with automatic discovery of gap index.
EigenvectorAlignmentLoss	Eigenvector alignment loss.
SpectralContrastiveLoss	Contrastive loss based on kernel or transition matrix.

FUNCTION	DESCRIPTION
eigenvalue_loss	Compute eigenvalue-based loss for spectral methods.
adaptive_eigenvalue_loss	Adaptive spectral gap loss with automatic gap selection via log-sum-exp.
eigenvector_alignment_loss	Compute the eigenvector alignment loss.
spectral_contrastive_loss	Contrastive loss based on kernel or transition matrix.
check_eigenvalues_sorted	Check if eigenvalues are sorted in descending or ascending order.

Classes

EigenvalueLoss(reduce: Literal['gap', 'gap_exp', 'sum', 'sum_pow', 'cumsum', 'entropy', 'single'] = 'gap', n_eigen: int = 0, sign: float = -1.0, eps: float = torch.finfo(torch.float32).eps, context_prefix: str | None = None, **kwargs)

Bases: Loss

Eigenvalue loss.

Computes loss based on eigenvalues to optimize specific spectral properties.

PARAMETER	DESCRIPTION
reduce	Reduction operation. "gap": \(\lambda_{n-1} - \lambda_{n}\), where \(n\) is n_eigen "gap_exp": \(-e^{\lambda_{n-1}} + e^{\lambda_{n}}\) "sum": \(\sum_{i=1}^{n} \lambda_{i}\) "sum_pow": \(\sum_{i=1}^{n} \lambda_{i}^2\) "cumsum": \(\sum_{i=1}^{n} \lambda_{i}/\sum_{i=1}^{N} \lambda_{i}\) "entropy": \(-\sum_{i=1}^{n} p_i \log(p_i)\), where \(p_i = \lambda_i / \sum_{i=1}^{n} \lambda_{i}\) "single": \(\lambda_{n}\) TYPE: str | None, optional, by default "gap" DEFAULT: 'gap'
n_eigen	Target eigenvalue index. If 0 or higher than the number of eigenvalues, uses len(eigenvalues)-1 (last valid index). Ignored when reduce="entropy". TYPE: int, optional, by default 0 DEFAULT: 0
sign	Loss sign (-1 for maximization, +1 for minimization). TYPE: float, optional, by default -1.0 DEFAULT: -1.0
eps	Small constant for numerical stability in divisions. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps
**kwargs	Additional keyword arguments for the loss function. TYPE: dict DEFAULT: {}

Examples:

>>> loss_fn = EigenvalueLoss(reduce="gap", n_eigen=1)
>>> eigenvalues = torch.tensor([2.0, 1.5, 0.5])
>>> context = {"eigenvalues": eigenvalues}
>>> loss = loss_fn(context)
>>> print(loss)
tensor(-0.5000)

METHOD	DESCRIPTION
forward	Compute eigenvalue loss.

Source code in spectre/loss/eigen.py

def __init__(
    self,
    reduce: Literal[
        "gap", "gap_exp", "sum", "sum_pow", "cumsum", "entropy", "single"
    ] = "gap",
    n_eigen: int = 0,
    sign: float = -1.0,
    eps: float = torch.finfo(torch.float32).eps,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce=reduce, sign=sign, context_prefix=context_prefix, **kwargs
    )

    if not isinstance(n_eigen, int):
        raise TypeError(
            f"Number of eigenvalues `n_eigen` must be an integer, got {type(n_eigen)}."
        )
    check_in_interval(n_eigen, "[0, inf)")
    self.n_eigen = n_eigen

    if not isinstance(eps, float):
        raise TypeError(f"eps must be a float, got {type(eps)}.")
    check_in_interval(eps, "[0, inf)")
    self.eps = eps

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute eigenvalue loss.

PARAMETER	DESCRIPTION
context	Dictionary containing the eigenvalues (required). TYPE: dict
**kwargs	Additional keyword arguments. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Eigenvalue loss.

Source code in spectre/loss/eigen.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute eigenvalue loss.

    Parameters
    ----------
    context : dict
        Dictionary containing the eigenvalues (required).

    **kwargs : dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Eigenvalue loss.
    """
    eigenvalues = self._get_context(context, "eigenvalues", None)

    if eigenvalues is None:
        raise ValueError("`EigenvalueLoss` requires 'eigenvalues' in context.")

    return eigenvalue_loss(
        eigenvalues,
        reduce=self.reduce,
        n_eigen=self.n_eigen,
        sign=self.sign,
        eps=self.eps,
    )

AdaptiveEigenvalueLoss(n_eigen: int | None = None, min_states: int = 2, max_states: int | None = None, temperature: float = 0.1, adaptive_weight: float = 1.0, anchor_weight: float = 0.0, normalize: bool = False, gap_mode: Literal['absolute', 'relative', 'log', 'timescale'] = 'absolute', sign: float = -1.0, eps: float = torch.finfo(torch.float32).eps, context_prefix: str | None = None, **kwargs)

Bases: Loss

Adaptive spectral gap loss with automatic discovery of gap index.

Uses soft-argmax to differentiably maximize the largest spectral gap, allowing the model to discover optimal number of metastable states during training rather than pre-specifying the number of states.

The loss combines two components:

Adaptive component (soft-argmax over all gaps)

• Uses temperature-scaled softmax: weights = softmax(gaps / T)
• Maximizes: sum(weights * gaps)

Anchor component (fixed gap at specified position)

• Provides stability during training
• Uses pre-training estimate as regularization

Mathematically:

\[L = \alpha \sum_i w_i g_i + \beta g_n\]

where:

• \(g_i = \lambda_{i-1} - \lambda_i\) is the spectral gap for \(i\) states
• \(w_i = \text{softmax}(g_i/T)\)
• \(\alpha\) is adaptive_weight, \(\beta\) is anchor_weight, and \(n\) is n_eigen.

PARAMETER	DESCRIPTION
n_eigen	Fixed position for anchor loss. Required if anchor_weight > 0. TYPE: int | None, optional, by default None DEFAULT: None
min_states	Minimum number of states to consider. Gaps before this threshold are ignored to prevent trivial solutions. For example, with min_states=2, the gap \(g_1 = \lambda_0 - \lambda_1\) (for 1 state) is excluded, starting the search from \(g_2 = \lambda_1 - \lambda_2\). TYPE: int, optional, by default 2 DEFAULT: 2
max_states	Maximum number of states to consider. If None, no upper limit. Restricts search space to prevent discovering too many artificial states. TYPE: int | None, optional, by default None DEFAULT: None
temperature	Softmax temperature for soft-argmax (typical range: [0.01, 1]). temperature -> 0: Focuses on single largest gap temperature -> inf: Uniform attention across all gaps TYPE: float, optional, by default 0.1 DEFAULT: 0.1
adaptive_weight	Weight for adaptive component (soft-max gap maximization). Sum of adaptive_weight and anchor_weight should equal 1. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
anchor_weight	Weight for fixed anchor component (gap at n_eigen). Set to 0 for pure adaptive mode. TYPE: float, optional, by default 0.0 DEFAULT: 0.0
normalize	If True, normalizes gaps by sum of relevant eigenvalues for scale invariance. If max_states is specified, normalizes by eigenvalues[:max_states].sum(). Otherwise, normalizes by the sum of eigenvalues up to the valid range. This prevents gaps from vanishing when many irrelevant eigenvalues exist. Helps when eigenvalue magnitudes vary significantly across training. TYPE: bool, optional, by default False DEFAULT: False
gap_mode	Mode for computing spectral gaps. "absolute": \(g_i = \lambda_{i-1} - \lambda_i\) (default) "relative": \(g_i = (\lambda_{i-1} - \lambda_i) / \lambda_{i-1}\) "log": Log-scale gaps \(g_i = \log(\lambda_{i-1}) - \log(\lambda_i)\) "timescale": Gaps in implied timescale space \(g_i = \tau_{i-1} - \tau_i\) where \(\tau_i = -1/\log(\lambda_i)\). Recommended for Markov eigenvalues in (0, 1). Handles exponential decay without the degenerate collapse that "relative" and "log" modes suffer from, because collapsed eigenvalues map to small timescales (small gaps), not large ones. TYPE: Literal["absolute", "relative", "log", "timescale"], optional, by default "absolute" DEFAULT: 'absolute'
sign	Loss sign multiplier. -1.0 for maximization (default), 1.0 for minimization. TYPE: float, optional, by default -1.0 DEFAULT: -1.0
eps	Small constant for numerical stability in divisions. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps
**kwargs	Additional keyword arguments for the loss function. DEFAULT: {}

ATTRIBUTE	DESCRIPTION
allowed_reduce	Allowed reduction operations (only ["none"] for this loss). TYPE: list[str]
loss_types	Loss function categories (SPECTRAL). TYPE: list[LossType]

Examples:

Pure adaptive mode:

>>> from spectre.loss import AdaptiveEigenvalueLoss
>>> from spectre.parametric import SpectralMap
>>>
>>> loss = AdaptiveEigenvalueLoss(
...     min_states=2,
...     max_states=10,
...     temperature=0.1,
...     adaptive_weight=1.0,
...     anchor_weight=0.0,
... )
>>>
>>> # Use with SpectralMap
>>> sm = SpectralMap(model=model, distance_fn="euclidean", loss_fn=loss)

Anchored adaptive mode:

>>> n_est = 3
>>>
>>> loss = AdaptiveEigenvalueLoss(
...     min_states=2,
...     max_states=10,
...     temperature=0.1,
...     adaptive_weight=0.7,  # 70% adaptive
...     anchor_weight=0.3,  # 30% fixed anchor
...     n_eigen=n_est,  # Anchor at estimate
... )
>>>
>>> sm = SpectralMap(model=model, kernel_fn="gaussian", loss_fn=loss)

Temperature annealing can be turned on using a PyTorch Lightning callback:

>>> from spectre.utils.callbacks import ParameterAnnealingCallback
>>>
>>> loss = AdaptiveEigenvalueLoss(
...     temperature=0.1,
...     adaptive_weight=0.8,
...     anchor_weight=0.2,
...     n_eigen=n_est,
... )

METHOD	DESCRIPTION
forward	Compute adaptive spectral gap loss.

Source code in spectre/loss/eigen.py

def __init__(
    self,
    n_eigen: int | None = None,
    min_states: int = 2,
    max_states: int | None = None,
    temperature: float = 0.1,
    adaptive_weight: float = 1.0,
    anchor_weight: float = 0.0,
    normalize: bool = False,
    gap_mode: Literal["absolute", "relative", "log", "timescale"] = "absolute",
    sign: float = -1.0,
    eps: float = torch.finfo(torch.float32).eps,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce="none", sign=sign, context_prefix=context_prefix, **kwargs
    )

    if not isinstance(min_states, int):
        raise TypeError(f"min_states must be int, got {type(min_states)}")
    check_in_interval(min_states, "[2, inf)")
    self.min_states = min_states

    if max_states is not None:
        if not isinstance(max_states, int):
            raise TypeError(
                f"max_states must be int or None, got {type(max_states)}"
            )
        if max_states < min_states:
            raise ValueError(
                f"max_states ({max_states}) must be >= min_states ({min_states})"
            )
    self.max_states = max_states

    if not isinstance(temperature, (int, float)):
        raise TypeError(f"temperature must be numeric, got {type(temperature)}")
    check_in_interval(temperature, "(0, inf)")
    self.temperature = float(temperature)

    if not isinstance(adaptive_weight, (int, float)):
        raise TypeError(
            f"adaptive_weight must be numeric, got {type(adaptive_weight)}"
        )
    check_in_interval(adaptive_weight, "[0, 1]")
    self.adaptive_weight = float(adaptive_weight)

    if not isinstance(anchor_weight, (int, float)):
        raise TypeError(f"anchor_weight must be numeric, got {type(anchor_weight)}")
    check_in_interval(anchor_weight, "[0, 1]")
    self.anchor_weight = float(anchor_weight)

    if self.adaptive_weight == 0.0 and self.anchor_weight == 0.0:
        raise ValueError("At least one weight must be > 0")

    if self.anchor_weight > 0 and n_eigen is None:
        raise ValueError("n_eigen must be provided when anchor_weight > 0")

    if n_eigen is not None:
        if not isinstance(n_eigen, int):
            raise TypeError(f"n_eigen must be int, got {type(n_eigen)}")
        check_in_interval(n_eigen, "[2, inf)")
    self.n_eigen = n_eigen

    if not isinstance(normalize, bool):
        raise TypeError(f"normalize must be bool, got {type(normalize)}")
    self.normalize = normalize

    if gap_mode not in ["absolute", "relative", "log", "timescale"]:
        raise ValueError(
            f"gap_mode must be 'absolute', 'relative', 'log', or 'timescale', "
            f"got '{gap_mode}'"
        )
    self.gap_mode = gap_mode

    if not isinstance(eps, float):
        raise TypeError(f"eps must be float, got {type(eps)}")
    check_in_interval(eps, "[0, inf)")
    self.eps = eps

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute adaptive spectral gap loss.

PARAMETER	DESCRIPTION
context	Context dictionary, must contain "eigenvalues" key. TYPE: dict
**kwargs	Additional keyword arguments. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Adaptive spectral gap loss.

Source code in spectre/loss/eigen.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute adaptive spectral gap loss.

    Parameters
    ----------
    context: dict
        Context dictionary, must contain "eigenvalues" key.

    **kwargs: dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Adaptive spectral gap loss.
    """
    eigenvalues = self._get_context(context, "eigenvalues", None)

    if eigenvalues is None:
        raise ValueError("AdaptiveEigenvalueLoss requires eigenvalues in context.")

    return adaptive_eigenvalue_loss(
        eigenvalues,
        n_eigen=self.n_eigen,
        min_states=self.min_states,
        max_states=self.max_states,
        temperature=self.temperature,
        adaptive_weight=self.adaptive_weight,
        anchor_weight=self.anchor_weight,
        normalize=self.normalize,
        gap_mode=self.gap_mode,
        sign=self.sign,
        eps=self.eps,
    )

EigenvectorAlignmentLoss(mode: Literal['reconstruction', 'cosine'] = 'reconstruction', reduce: Literal['mean', 'sum'] = 'mean', sign: float = 1.0, eps: float = torch.finfo(torch.float32).eps, context_prefix: str | None = None, **kwargs)

Bases: Loss

Eigenvector alignment loss.

Encourages embeddings to align with provided eigenvectors.

PARAMETER	DESCRIPTION
mode	Mode of alignment. "reconstruction": Minimize alignment error between Z and eigenvectors. "cosine": Maximize cosine similarity between Z and eigenvectors. TYPE: Literal["reconstruction", "cosine"], optional, by default "reconstruction" DEFAULT: 'reconstruction'
reduce	Reduction method. TYPE: Literal["mean", "sum"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
eps	Small value for numerical stability. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps
**kwargs	Additional keyword arguments for the loss function. TYPE: dict DEFAULT: {}

Examples:

>>> loss_fn = EigenvectorAlignmentLoss(mode="reconstruction")
>>> Z = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
>>> eigenvectors = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
>>> context = {"Z": Z, "eigenvectors": eigenvectors}
>>> loss = loss_fn(context)
>>> print(loss)
tensor(...)

METHOD	DESCRIPTION
forward	Compute eigenvector alignment loss.

Source code in spectre/loss/eigen.py

def __init__(
    self,
    mode: Literal["reconstruction", "cosine"] = "reconstruction",
    reduce: Literal["mean", "sum"] = "mean",
    sign: float = 1.0,
    eps: float = torch.finfo(torch.float32).eps,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce=reduce, sign=sign, context_prefix=context_prefix, **kwargs
    )

    if mode not in ["reconstruction", "cosine"]:
        raise ValueError(
            f"Unknown mode '{mode}'. Options: 'reconstruction', 'cosine'."
        )
    self.mode = mode

    if not isinstance(eps, float):
        raise TypeError(f"eps must be a float, got {type(eps)}.")
    check_in_interval(eps, "[0, inf)")
    self.eps = eps

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute eigenvector alignment loss.

PARAMETER	DESCRIPTION
context	Context dictionary containing the input data and eigenvectors. TYPE: dict
**kwargs	Additional keyword arguments. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Eigenvector alignment loss.

Source code in spectre/loss/eigen.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute eigenvector alignment loss.

    Parameters
    ----------
    context : dict
        Context dictionary containing the input data and eigenvectors.

    **kwargs
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Eigenvector alignment loss.
    """
    Z = self._get_context(context, "Z", None)
    eigenvectors = self._get_context(context, "eigenvectors", None)

    if Z is None:
        raise ValueError("`EigenvectorAlignmentLoss` requires `Z` in context.")
    if eigenvectors is None:
        raise ValueError(
            "`EigenvectorAlignmentLoss` requires `eigenvectors` in context."
        )

    return eigenvector_alignment_loss(
        Z,
        eigenvectors,
        mode=self.mode,
        reduce=self.reduce,
        sign=self.sign,
        eps=self.eps,
    )

SpectralContrastiveLoss(temperature: float = 0.1, threshold: float = 0.0, soft: bool = False, reduce: Literal['mean', 'sum'] = 'mean', sign: float = 1.0, context_prefix: str | None = None, **kwargs)

Bases: Loss

Contrastive loss based on kernel or transition matrix.

Computes a contrastive loss that encourages similar points (connected in K) to have similar embeddings in Z, while dissimilar points (not connected in K) to have dissimilar embeddings.

For binary adjacency matrices (traditional contrastive loss), connected nodes are treated as positive pairs and non-connected nodes as negative pairs.

For transition matrices (e.g., diffusion maps), a soft contrastive loss can be used where connections above a threshold are treated as positive pairs with weights proportional to the transition probabilities.

PARAMETER	DESCRIPTION
temperature	Temperature parameter for similarity scaling. TYPE: float, optional, by default 0.1 DEFAULT: 0.1
threshold	Threshold for defining positive/negative pairs when soft=True. When soft=False, this parameter is ignored and K > 0 is used. TYPE: float, optional, by default 0.0 DEFAULT: 0.0
soft	Whether to use soft contrastive loss with continuous weights. False: Binary adjacency (K > 0 vs K == 0) True: Threshold-based with probability weighting (K > threshold) TYPE: bool, optional, by default False DEFAULT: False
reduce	Reduction method. "mean": Average the loss over all samples. "sum": Sum the loss over all samples. TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
**kwargs	Additional keyword arguments for the loss function. TYPE: dict DEFAULT: {}

METHOD	DESCRIPTION
forward	Compute spectral contrastive loss.

Source code in spectre/loss/eigen.py

def __init__(
    self,
    temperature: float = 0.1,
    threshold: float = 0.0,
    soft: bool = False,
    reduce: Literal["mean", "sum"] = "mean",
    sign: float = 1.0,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce=reduce, sign=sign, context_prefix=context_prefix, **kwargs
    )

    if not isinstance(temperature, float):
        raise TypeError(f"Temperature must be a float, got {type(temperature)}.")
    check_in_interval(temperature, "(0, inf)")
    self.temperature = temperature

    if not isinstance(threshold, float):
        raise TypeError(f"Threshold must be a float, got {type(threshold)}.")
    check_in_interval(threshold, "[0, inf)")
    self.threshold = threshold

    if not isinstance(soft, bool):
        raise TypeError(f"Soft must be a boolean, got {type(soft)}.")
    self.soft = soft

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute spectral contrastive loss.

Source code in spectre/loss/eigen.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """Compute spectral contrastive loss."""
    Z = self._get_context(context, "Z", None)
    K = self._get_context(context, "K", None)

    if Z is None:
        raise ValueError("`SpectralContrastiveLoss` requires `Z` in context.")
    if K is None:
        raise ValueError("`SpectralContrastiveLoss` requires `K` in context.")

    return spectral_contrastive_loss(
        Z,
        K,
        temperature=self.temperature,
        threshold=self.threshold,
        soft=self.soft,
        reduce=self.reduce,
        sign=self.sign,
    )

Functions

eigenvalue_loss(eigenvalues: torch.Tensor, reduce: Literal['gap', 'gap_exp', 'sum', 'sum_pow', 'cumsum', 'entropy', 'single'] = 'gap', n_eigen: int = 0, sign: float = -1, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute eigenvalue-based loss for spectral methods.

Computes loss based on eigenvalues to optimize specific spectral properties.

PARAMETER	DESCRIPTION
eigenvalues	Eigenvalues in descending order TYPE: Tensor
reduce	Reduction operation. "gap": \(\lambda_{n-1} - \lambda_{n}\), where \(n\) is n_eigen "gap_exp": \(-e^{\lambda_{n-1}} + e^{\lambda_{n}}\) "sum": \(\sum_{i=1}^{n} \lambda_{i}\) "sum_pow": \(\sum_{i=1}^{n} \lambda_{i}^2\) "cumsum": \(\sum_{i=1}^{n} \lambda_{i}/\sum_{i=1}^{N} \lambda_{i}\) "entropy": \(-\sum_{i=1}^{n} p_i \log(p_i)\), where \(p_i = \lambda_i / \sum_{i=1}^{n} \lambda_{i}\) "single": \(\lambda_{n}\) TYPE: str | None, optional, by default "gap" DEFAULT: 'gap'
n_eigen	Target eigenvalue index. If 0 or higher than the number of eigenvalues, uses len(eigenvalues)-1 (last valid index) Ignored when reduce="entropy" TYPE: int, optional, by default 0 DEFAULT: 0
sign	Loss sign (-1 for maximization, +1 for minimization). TYPE: float, optional, by default -1.0 DEFAULT: -1
eps	Small constant for numerical stability in divisions. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Computed reduced eigenvalue loss.

Examples:

>>> eigenvalues = torch.tensor([2.0, 1.5, 0.5])
>>> loss = eigenvalue_loss(eigenvalues, reduce="gap", n_eigen=1)
>>> print(loss)
tensor(-0.5000)

Source code in spectre/loss/eigen.py

def eigenvalue_loss(
    eigenvalues: torch.Tensor,
    reduce: Literal[
        "gap", "gap_exp", "sum", "sum_pow", "cumsum", "entropy", "single"
    ] = "gap",
    n_eigen: int = 0,
    sign: float = -1,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute eigenvalue-based loss for spectral methods.

    Computes loss based on eigenvalues to optimize specific spectral properties.

    Parameters
    ----------
    eigenvalues : torch.Tensor
        Eigenvalues in descending order

    reduce : str | None, optional, by default "gap"
        Reduction operation.

        - "gap": $\\lambda_{n-1} - \\lambda_{n}$, where $n$ is `n_eigen`
        - "gap_exp": $-e^{\\lambda_{n-1}} + e^{\\lambda_{n}}$
        - "sum": $\\sum_{i=1}^{n} \\lambda_{i}$
        - "sum_pow": $\\sum_{i=1}^{n} \\lambda_{i}^2$
        - "cumsum": $\\sum_{i=1}^{n} \\lambda_{i}/\\sum_{i=1}^{N} \\lambda_{i}$
        - "entropy": $-\\sum_{i=1}^{n} p_i \\log(p_i)$,
            where $p_i = \\lambda_i / \\sum_{i=1}^{n} \\lambda_{i}$
        - "single": $\\lambda_{n}$

    n_eigen : int, optional, by default 0
        Target eigenvalue index.

        - If 0 or higher than the number of eigenvalues, uses
            `len(eigenvalues)-1` (last valid index)
        - Ignored when `reduce="entropy"`

    sign : float, optional, by default -1.0
        Loss sign (-1 for maximization, +1 for minimization).

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Small constant for numerical stability in divisions.

    Returns
    -------
    torch.Tensor
        Computed reduced eigenvalue loss.

    Examples
    --------
    >>> eigenvalues = torch.tensor([2.0, 1.5, 0.5])
    >>> loss = eigenvalue_loss(eigenvalues, reduce="gap", n_eigen=1)
    >>> print(loss)
    tensor(-0.5000)
    """
    check_1d(eigenvalues)

    if reduce not in ["entropy"]:
        if n_eigen == 0 or n_eigen >= len(eigenvalues):
            n_eigen = len(eigenvalues) - 1  # Last valid index.

    if reduce == "sum":
        loss = torch.sum(eigenvalues[1 : n_eigen + 1])
    elif reduce == "sum_pow":
        loss = torch.sum(torch.pow(eigenvalues[1 : n_eigen + 1], 2))
    elif reduce == "gap":
        loss = eigenvalues[n_eigen - 1] - eigenvalues[n_eigen]
    elif reduce == "gap_exp":
        loss = -torch.exp(eigenvalues[n_eigen - 1]) + torch.exp(eigenvalues[n_eigen])
    elif reduce == "cumsum":
        loss = eigenvalues[:n_eigen].sum() / (eigenvalues.sum() + eps)
    elif reduce == "entropy":
        prob = eigenvalues / (eigenvalues.sum() + eps)
        log_probs = torch.log(prob + eps)
        loss = -(prob * log_probs).sum()
    elif reduce == "single":
        loss = eigenvalues[n_eigen]
    else:
        raise ValueError(
            f"Unknown reduce operation '{reduce}'. Options: 'gap', 'gap_exp', "
            f"'sum', 'sum_pow', 'cumsum', 'entropy', 'single'."
        )

    return sign * loss

adaptive_eigenvalue_loss(eigenvalues: torch.Tensor, n_eigen: int | None = None, min_states: int = 2, max_states: int | None = None, temperature: float = 0.1, adaptive_weight: float = 1.0, anchor_weight: float = 0.0, normalize: bool = False, gap_mode: Literal['absolute', 'relative', 'log', 'timescale'] = 'absolute', sign: float = -1.0, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Adaptive spectral gap loss with automatic gap selection via log-sum-exp.

Combines adaptive gap maximization with optional fixed anchor for stability.

PARAMETER	DESCRIPTION
eigenvalues	Eigenvalues in descending order, shape (n_eigenvalues,) TYPE: Tensor
min_states	Minimum states to consider. Gaps for fewer states are ignored. TYPE: int, optional, by default 2 DEFAULT: 2
max_states	Maximum states to consider (ignore gaps after this) TYPE: int | None, optional, by default None DEFAULT: None
temperature	Temperature for log-sum-exp smooth maximum (lower = more focused). TYPE: float, optional, by default 0.1 DEFAULT: 0.1
adaptive_weight	Weight for adaptive component [0, 1] TYPE: float, optional, by default 1.0 DEFAULT: 1.0
anchor_weight	Weight for anchor component [0, 1] TYPE: float, optional, by default 0.0 DEFAULT: 0.0
n_eigen	Fixed position for anchor (required if anchor_weight > 0) TYPE: int | None, optional, by default None DEFAULT: None
normalize	Normalize gaps by sum of relevant eigenvalues. If max_states is specified, normalizes by eigenvalues[:max_states].sum(). Otherwise, normalizes by the sum of eigenvalues up to the valid range. This prevents gaps from vanishing when many irrelevant eigenvalues exist. TYPE: bool, optional, by default False DEFAULT: False
gap_mode	Mode for computing spectral gaps. "absolute": Raw gaps \(g_i = \lambda_{i-1} - \lambda_i\) (default) "relative": Gaps normalized by preceding eigenvalue \(g_i = (\lambda_{i-1} - \lambda_i) / \lambda_{i-1}\) "log": Gaps in log-space \(g_i = \log(\lambda_{i-1}) - \log(\lambda_i)\) "timescale": Gaps in implied timescale space \(g_i = \tau_{i-1} - \tau_i\) where \(\tau_i = -1/\log(\lambda_i)\). Recommended for Markov eigenvalues in (0, 1). Handles exponential decay without the degenerate collapse that "relative" and "log" modes suffer from. TYPE: Literal["absolute", "relative", "log", "timescale"], optional, by default "absolute" DEFAULT: 'absolute'
sign	Loss sign (-1 for maximization) TYPE: float, optional, by default -1.0 DEFAULT: -1.0
eps	Numerical stability constant TYPE: float DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Scalar loss value
The loss combines adaptive and anchor components:
$L = \alpha L_{\text{adaptive}} + \beta L_{\text{anchor}}$
where	TYPE: Tensor
- Spectral gaps $g_i$ for $i$ states (see `gap_mode` for formulation)
- Adaptive component via log-sum-exp (smooth max):	\(L_{\text{adaptive}} = T \log \sum_i \exp(g_i / T)\)
- Anchor component: $L_{\text{anchor}} = g_n$ where $n$ is `n_eigen`
- $T$ is temperature, $\alpha$ is `adaptive_weight`, $\beta$ is `anchor_weight`

Examples:

>>> import torch
>>> eigenvalues = torch.tensor([1.0, 0.9, 0.8, 0.3, 0.2, 0.1])
>>> loss = adaptive_eigenvalue_loss(eigenvalues, temperature=0.1)
>>> print(loss)
tensor(-0.5000)

Source code in spectre/loss/eigen.py

def adaptive_eigenvalue_loss(
    eigenvalues: torch.Tensor,
    n_eigen: int | None = None,
    min_states: int = 2,
    max_states: int | None = None,
    temperature: float = 0.1,
    adaptive_weight: float = 1.0,
    anchor_weight: float = 0.0,
    normalize: bool = False,
    gap_mode: Literal["absolute", "relative", "log", "timescale"] = "absolute",
    sign: float = -1.0,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Adaptive spectral gap loss with automatic gap selection via log-sum-exp.

    Combines adaptive gap maximization with optional fixed anchor for stability.

    Parameters
    ----------
    eigenvalues : torch.Tensor
        Eigenvalues in descending order, shape (n_eigenvalues,)

    min_states : int, optional, by default 2
        Minimum states to consider. Gaps for fewer states are ignored.

    max_states : int | None, optional, by default None
        Maximum states to consider (ignore gaps after this)

    temperature : float, optional, by default 0.1
        Temperature for log-sum-exp smooth maximum (lower = more focused).

    adaptive_weight : float, optional, by default 1.0
        Weight for adaptive component [0, 1]

    anchor_weight : float, optional, by default 0.0
        Weight for anchor component [0, 1]

    n_eigen : int | None, optional, by default None
        Fixed position for anchor (required if anchor_weight > 0)

    normalize : bool, optional, by default False
        Normalize gaps by sum of relevant eigenvalues.

        If `max_states` is specified, normalizes by `eigenvalues[:max_states].sum()`.
        Otherwise, normalizes by the sum of eigenvalues up to the valid range.
        This prevents gaps from vanishing when many irrelevant eigenvalues exist.

    gap_mode : Literal["absolute", "relative", "log", "timescale"], optional, by default "absolute"
        Mode for computing spectral gaps.

        - "absolute": Raw gaps $g_i = \\lambda_{i-1} - \\lambda_i$ (default)
        - "relative": Gaps normalized by preceding eigenvalue
          $g_i = (\\lambda_{i-1} - \\lambda_i) / \\lambda_{i-1}$
        - "log": Gaps in log-space
          $g_i = \\log(\\lambda_{i-1}) - \\log(\\lambda_i)$
        - "timescale": Gaps in implied timescale space
          $g_i = \\tau_{i-1} - \\tau_i$ where $\\tau_i = -1/\\log(\\lambda_i)$.
          Recommended for Markov eigenvalues in (0, 1). Handles exponential
          decay without the degenerate collapse that "relative" and "log"
          modes suffer from.

    sign : float, optional, by default -1.0
        Loss sign (-1 for maximization)

    eps : float, optional
        Numerical stability constant

    Returns
    -------
    torch.Tensor
        Scalar loss value

    The loss combines adaptive and anchor components:
    $L = \\alpha L_{\\text{adaptive}} + \\beta L_{\\text{anchor}}$

    where:

    - Spectral gaps $g_i$ for $i$ states (see `gap_mode` for formulation)
    - Adaptive component via log-sum-exp (smooth max):
      $L_{\\text{adaptive}} = T \\log \\sum_i \\exp(g_i / T)$
    - Anchor component: $L_{\\text{anchor}} = g_n$ where $n$ is `n_eigen`
    - $T$ is temperature, $\\alpha$ is `adaptive_weight`, $\\beta$ is `anchor_weight`

    Examples
    --------
    >>> import torch
    >>> eigenvalues = torch.tensor([1.0, 0.9, 0.8, 0.3, 0.2, 0.1])
    >>> loss = adaptive_eigenvalue_loss(eigenvalues, temperature=0.1)
    >>> print(loss)
    tensor(-0.5000)
    """
    check_1d(eigenvalues)
    check_eigenvalues_sorted(eigenvalues, descending=True, raise_error=True)

    # Validate eigenvalues for modes that require positivity
    if gap_mode == "log" and (eigenvalues <= 0).any():
        raise ValueError(
            "gap_mode='log' requires all eigenvalues > 0. "
            f"Found {(eigenvalues <= 0).sum().item()} non-positive eigenvalues."
        )

    if gap_mode == "timescale" and (eigenvalues <= 0).any():
        raise ValueError(
            "gap_mode='timescale' requires all eigenvalues > 0. "
            f"Found {(eigenvalues <= 0).sum().item()} non-positive eigenvalues."
        )

    device, dtype = eigenvalues.device, eigenvalues.dtype

    if gap_mode == "timescale":
        # Implied timescales: τ_i = -1/log(λ_i) for λ ∈ (0, 1).
        # Maps eigenvalues to a space where:
        #   - λ near 1 → large τ (slow process, metastable state)
        #   - λ near 0 → small τ (fast process, noise)
        # The first eigenvalue (λ_0 = 1 for Markov) maps to τ_0 = ∞,
        # naturally excluding the trivial "one state" gap.
        # Conservation: sum(Δτ_i) = τ_1 − τ_{n-1}, bounded and → 0
        # as eigenvalues collapse, preventing the degenerate optimum
        # that relative/log modes suffer from.
        eigvals_clamped = torch.clamp(eigenvalues, min=eps, max=1.0 - eps)
        timescales = -1.0 / torch.log(eigvals_clamped)
        gaps = timescales[:-1] - timescales[1:]
    else:
        gaps = eigenvalues[:-1] - eigenvalues[1:]

        if gap_mode == "relative":
            denominators = torch.clamp(eigenvalues[:-1], min=eps)
            gaps = gaps / denominators
        elif gap_mode == "log":
            eigvals_clamped = torch.clamp(eigenvalues, min=eps)
            gaps = torch.log(eigvals_clamped[:-1]) - torch.log(eigvals_clamped[1:])

    valid_start = max(0, min_states - 1)
    valid_end = min(len(gaps), max_states) if max_states else len(gaps)

    if valid_start >= valid_end:
        raise ValueError(
            f"Invalid state range: min_states={min_states}, max_states={max_states}, "
            f"n_gaps={len(gaps)}. Valid range would be empty."
        )

    # Normalize gaps before slicing to ensure consistency between adaptive and anchor
    if normalize:
        # Use `max_states` for normalization if specified, otherwise use `valid_end`
        # This prevents gaps from vanishing when there are many irrelevant eigenvalues
        norm_end = max_states if max_states is not None else valid_end + 1
        norm_end = min(norm_end, len(eigenvalues))
        norm_sum = eigenvalues[:norm_end].sum()
        gaps = gaps / (norm_sum + eps)

    valid_gaps = gaps[valid_start:valid_end]

    # Check that we have valid gaps after slicing
    if len(valid_gaps) == 0:
        raise ValueError(
            f"No valid gaps in range [min_states={min_states}, max_states={max_states}]. "
            f"Available gaps: {len(gaps)}"
        )

    adaptive_loss = torch.tensor(0.0, device=device, dtype=dtype)
    if adaptive_weight > 0:
        # Smooth differentiable approximation to max(gaps) via log-sum-exp.
        # Converges to max(valid_gaps) as temperature → 0.
        # Unlike softmax-weighted average, LSE gradients (dL/dg_i = -softmax_i)
        # are always non-negative, avoiding sign inversion on non-dominant gaps.
        adaptive_loss = temperature * torch.logsumexp(valid_gaps / temperature, dim=0)

    anchor_loss = torch.tensor(0.0, device=device, dtype=dtype)
    if anchor_weight > 0:
        if n_eigen is None:
            raise ValueError("n_eigen required when anchor_weight > 0")

        # Gap at anchor position: lambda_{k-1} - lambda_k where k = n_eigen
        anchor_idx = n_eigen - 1

        if anchor_idx >= len(gaps):
            raise ValueError(
                f"n_eigen={n_eigen} (index {anchor_idx}) exceeds "
                f"available gaps (length {len(gaps)})"
            )

        anchor_loss = gaps[anchor_idx]

    total_loss = adaptive_weight * adaptive_loss + anchor_weight * anchor_loss

    return sign * total_loss

eigenvector_alignment_loss(Z: torch.Tensor, eigenvectors: torch.Tensor, mode: Literal['reconstruction', 'cosine'] = 'reconstruction', reduce: Literal['mean', 'sum'] = 'mean', sign: float = 1.0, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute the eigenvector alignment loss.

Encourages embeddings to align with provided eigenvectors.

PARAMETER	DESCRIPTION
Z	Input tensor of shape (n_samples, n_features). TYPE: Tensor
eigenvectors	Eigenvectors of shape (n_samples, n_eigenvectors). If n_eigenvectors > n_features, only the first n_features eigenvectors are used. TYPE: Tensor
reduce	Reduction method. TYPE: Literal["mean", "sum"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0
mode	Mode of alignment. "reconstruction": Minimize alignment error between Z and eigenvectors. "cosine": Maximize cosine similarity between Z and eigenvectors. TYPE: Literal["reconstruction", "cosine"], optional, by default "reconstruction" DEFAULT: 'reconstruction'
eps	Small value for numerical stability, by default torch.finfo(torch.float32).eps. TYPE: float DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Computed alignment loss.

Examples:

>>> Z = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
>>> eigenvectors = torch.tensor([[0.7071, 0.7071], [0.7071, -0.7071]])
>>> loss = eigenvector_alignment_loss(Z, eigenvectors, mode="reconstruction")
>>> print(loss)
tensor(...)

Source code in spectre/loss/eigen.py

def eigenvector_alignment_loss(
    Z: torch.Tensor,
    eigenvectors: torch.Tensor,
    mode: Literal["reconstruction", "cosine"] = "reconstruction",
    reduce: Literal["mean", "sum"] = "mean",
    sign: float = 1.0,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute the eigenvector alignment loss.

    Encourages embeddings to align with provided eigenvectors.

    Parameters
    ----------
    Z : torch.Tensor
        Input tensor of shape (n_samples, n_features).

    eigenvectors : torch.Tensor
        Eigenvectors of shape (n_samples, n_eigenvectors).
        If `n_eigenvectors > n_features`, only the first `n_features` eigenvectors
        are used.

    reduce : Literal["mean", "sum"], optional, by default "mean"
        Reduction method.

    sign : float, optional, by default 1.0
        Loss sign multiplier.

    mode : Literal["reconstruction", "cosine"], optional, by default "reconstruction"
        Mode of alignment.

        - "reconstruction": Minimize alignment error between Z and eigenvectors.
        - "cosine": Maximize cosine similarity between Z and eigenvectors.

    eps : float, optional
        Small value for numerical stability, by default torch.finfo(torch.float32).eps.

    Returns
    -------
    torch.Tensor
        Computed alignment loss.

    Examples
    --------
    >>> Z = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
    >>> eigenvectors = torch.tensor([[0.7071, 0.7071], [0.7071, -0.7071]])
    >>> loss = eigenvector_alignment_loss(Z, eigenvectors, mode="reconstruction")
    >>> print(loss)
    tensor(...)
    """
    check_2d(Z)
    check_2d(eigenvectors)
    check_same_shape(Z, eigenvectors, axis=0)

    eigvec = eigenvectors[:, : Z.shape[1]]

    eigvec_norms = torch.norm(eigvec, dim=0, keepdim=True)
    eigvec_norms = torch.where(
        eigvec_norms < eps, torch.ones_like(eigvec_norms), eigvec_norms
    )
    eigvec_norm = eigvec / eigvec_norms

    Z_norms = torch.norm(Z, dim=0, keepdim=True)
    Z_norms = torch.where(Z_norms < eps, torch.ones_like(Z_norms), Z_norms)
    Z_norm = Z / Z_norms

    if mode == "reconstruction":
        # Handle sign ambiguity: align signs based on correlation
        dot_products = torch.sum(Z_norm * eigvec_norm, dim=0)
        signs = torch.sign(dot_products)
        signs = torch.where(signs == 0, torch.ones_like(signs), signs)
        eigvec_aligned = eigvec_norm * signs.unsqueeze(0)

        # Alignment error
        loss = torch.mean((Z_norm - eigvec_aligned) ** 2, dim=1)

    elif mode == "cosine":
        # Normalize row-wise for cosine similarity
        Z_row_norms = torch.norm(Z, dim=1, keepdim=True)
        eigvec_row_norms = torch.norm(eigvec, dim=1, keepdim=True)

        Z_row_norms = torch.where(
            Z_row_norms < eps, torch.ones_like(Z_row_norms), Z_row_norms
        )
        eigvec_row_norms = torch.where(
            eigvec_row_norms < eps,
            torch.ones_like(eigvec_row_norms),
            eigvec_row_norms,
        )

        Z_row_norm = Z / Z_row_norms
        eigvec_row_norm = eigvec / eigvec_row_norms

        # Cosine similarity per sample
        similarities = torch.sum(Z_row_norm * eigvec_row_norm, dim=1)
        loss = -torch.abs(similarities)
    else:
        raise ValueError(f"Unknown mode '{mode}'. Options: 'reconstruction', 'cosine'.")

    if reduce == "mean":
        loss = loss.mean()
    elif reduce == "sum":
        loss = loss.sum()
    else:
        raise ValueError(
            f"Unknown reduction method '{reduce}'. Options: 'mean', 'sum'."
        )

    return sign * loss

spectral_contrastive_loss(Z: torch.Tensor, K: torch.Tensor, temperature: float = 0.1, threshold: float = 0.0, soft: bool = False, reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0) -> torch.Tensor

Contrastive loss based on kernel or transition matrix.

Computes a contrastive loss that encourages similar points (connected in K) to have similar embeddings in input data, while dissimilar points (not connected in K) to have dissimilar embeddings.

This implements an InfoNCE-style contrastive loss where for each anchor sample, positive pairs (connected in K) should have higher similarity than all other samples (including both unconnected samples and other positives).

For binary adjacency matrices (traditional contrastive loss), connected nodes are treated as positive pairs and non-connected nodes as negative pairs.

For transition matrices (e.g., diffusion maps), a soft contrastive loss can be used where connections above a threshold are treated as positive pairs with weights proportional to the transition probabilities.

PARAMETER	DESCRIPTION
Z	Input embeddings of shape (n_samples, in_features). TYPE: Tensor
K	Kernel, adjacency, or transition matrix of shape (n_samples, n_samples). TYPE: Tensor
temperature	Temperature parameter for similarity scaling. TYPE: float, optional, by default 0.1 DEFAULT: 0.1
threshold	Threshold for defining positive/negative pairs when soft=True. When soft=False, this parameter is ignored and K > 0 is used. TYPE: float, optional, by default 0.0 DEFAULT: 0.0
soft	Whether to use soft contrastive loss with continuous weights. False: Binary adjacency (K > 0 vs K == 0) True: Threshold-based with probability weighting (K > threshold) TYPE: bool, optional, by default False DEFAULT: False
reduce	Reduction method. "mean": Average the loss over all samples. "sum": Sum the loss over all samples. TYPE: Literal["sum", "mean"], optional, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, optional, by default 1.0 DEFAULT: 1.0

RETURNS	DESCRIPTION
Tensor	Computed contrastive loss.

Notes

The loss is computed per anchor sample \(i\) as:

\(L_i = -\log \frac{\sum_{p \in P_i} \exp(s_{ip}/T)} {\sum_{j \neq i} \exp(s_{ij}/T)}\)

where \(P_i\) is the set of positive pairs for anchor \(i\), \(s_{ij}\) is the similarity between samples \(i\) and \(j\), and \(T\) is the temperature.

In soft mode, positive pairs are weighted by their values in \(K\):

\(L_i = -\log \frac{\sum_{p \in P_i} K_{ip} \exp(s_{ip}/T)} {\sum_{j \neq i} \exp(s_{ij}/T)}\)

Examples:

For binary adjacency matrices (standard)

>>> loss = spectral_contrastive_loss(
...     embeddings,
...     adjacency_matrix,
...     soft=False,
... )

For diffusion maps / transition matrices (manifold-aware)

>>> loss = spectral_contrastive_loss(
...     embeddings,
...     transition_matrix,
...     threshold=0.1,
...     soft=True,
... )

Source code in spectre/loss/eigen.py

def spectral_contrastive_loss(
    Z: torch.Tensor,
    K: torch.Tensor,
    temperature: float = 0.1,
    threshold: float = 0.0,
    soft: bool = False,
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
) -> torch.Tensor:
    """
    Contrastive loss based on kernel or transition matrix.

    Computes a contrastive loss that encourages similar points (connected in `K`)
    to have similar embeddings in input data, while dissimilar points
    (not connected in `K`) to have dissimilar embeddings.

    This implements an InfoNCE-style contrastive loss where for each anchor sample,
    positive pairs (connected in `K`) should have higher similarity than all other
    samples (including both unconnected samples and other positives).

    For binary adjacency matrices (traditional contrastive loss), connected nodes
    are treated as positive pairs and non-connected nodes as negative pairs.

    For transition matrices (e.g., diffusion maps), a soft contrastive loss can be used
    where connections above a threshold are treated as positive pairs with weights
    proportional to the transition probabilities.

    Parameters
    ----------
    Z : torch.Tensor
        Input embeddings of shape (n_samples, in_features).

    K : torch.Tensor
        Kernel, adjacency, or transition matrix of shape (n_samples, n_samples).

    temperature : float, optional, by default 0.1
        Temperature parameter for similarity scaling.

    threshold : float, optional, by default 0.0
        Threshold for defining positive/negative pairs when soft=True.
        When soft=False, this parameter is ignored and K > 0 is used.

    soft : bool, optional, by default False
        Whether to use soft contrastive loss with continuous weights.

        - False: Binary adjacency (K > 0 vs K == 0)
        - True: Threshold-based with probability weighting (K > threshold)

    reduce : Literal["sum", "mean"], optional, by default "mean"
        Reduction method.

        - "mean": Average the loss over all samples.
        - "sum": Sum the loss over all samples.

    sign : float, optional, by default 1.0
        Loss sign multiplier.

    Returns
    -------
    torch.Tensor
        Computed contrastive loss.

    Notes
    -----
    The loss is computed per anchor sample $i$ as:

    $L_i = -\\log \\frac{\\sum_{p \\in P_i} \\exp(s_{ip}/T)}
                           {\\sum_{j \\neq i} \\exp(s_{ij}/T)}$

    where $P_i$ is the set of positive pairs for anchor $i$, $s_{ij}$ is the
    similarity between samples $i$ and $j$, and $T$ is the temperature.

    In soft mode, positive pairs are weighted by their values in $K$:

    $L_i = -\\log \\frac{\\sum_{p \\in P_i} K_{ip} \\exp(s_{ip}/T)}
                           {\\sum_{j \\neq i} \\exp(s_{ij}/T)}$

    Examples
    --------
    For binary adjacency matrices (standard)

    >>> loss = spectral_contrastive_loss(
    ...     embeddings,
    ...     adjacency_matrix,
    ...     soft=False,
    ... )

    For diffusion maps / transition matrices (manifold-aware)

    >>> loss = spectral_contrastive_loss(
    ...     embeddings,
    ...     transition_matrix,
    ...     threshold=0.1,
    ...     soft=True,
    ... )
    """
    check_2d(Z)
    check_2d(K)
    check_same_len(Z, K)

    n = Z.shape[0]

    # Compute pairwise cosine similarities: (n, n)
    Z_norm = torch.nn.functional.normalize(Z, p=2, dim=1)
    similarities = torch.mm(Z_norm, Z_norm.T) / temperature

    # Define positive masks based on mode
    if soft:
        pos_mask = K > threshold
        weights = K.clone()
    else:
        pos_mask = K > 0
        weights = pos_mask.float()

    # Remove self-connections
    mask_diag = ~torch.eye(n, dtype=torch.bool, device=Z.device)
    pos_mask = pos_mask & mask_diag

    # Compute per-sample contrastive loss
    losses = []

    for i in range(n):
        # Get positive pairs for anchor i
        pos_idx = pos_mask[i]

        # Skip if no positive pairs
        if not pos_idx.any():
            continue

        # Get similarities for anchor i (excluding self)
        sim_i = similarities[i, mask_diag[i]]  # (n-1,)

        # Adjust indexing after removing diagonal
        # mask_diag[i] gives us indices [0, ..., i-1, i+1, ..., n-1]
        # We need to map pos_mask[i] (which includes diagonal) to this space
        pos_idx_no_diag = pos_mask[i, mask_diag[i]]  # Exclude diagonal

        # Extract positive similarities
        pos_sims = sim_i[pos_idx_no_diag]  # (n_pos,)

        if soft:
            # Get weights for positive pairs (excluding diagonal)
            pos_weights = weights[i, mask_diag[i]][pos_idx_no_diag]  # (n_pos,)
        else:
            pos_weights = None

        # Compute log-sum-exp for numerical stability
        # Denominator: sum over all non-self samples
        log_sum_all = torch.logsumexp(sim_i, dim=0)

        # Numerator: sum over positive samples
        if soft and pos_weights is not None:
            # Weighted positive: log(sum(w_p * exp(s_p))) = logsumexp(log(w_p) + s_p)
            log_weighted_pos = torch.logsumexp(
                torch.log(pos_weights + 1e-10) + pos_sims, dim=0
            )
            loss_i = -(log_weighted_pos - log_sum_all)
        else:
            # Unweighted: log(sum(exp(s_p))) = logsumexp(s_p)
            log_sum_pos = torch.logsumexp(pos_sims, dim=0)
            loss_i = -(log_sum_pos - log_sum_all)

        losses.append(loss_i)

    # Handle case where no samples have positive pairs
    if len(losses) == 0:
        return torch.tensor(0.0, device=Z.device, dtype=Z.dtype)

    loss_tensor = torch.stack(losses)

    if reduce == "mean":
        loss = loss_tensor.mean()
    elif reduce == "sum":
        loss = loss_tensor.sum()
    else:
        raise ValueError(
            f"Unknown reduction method '{reduce}'. Options: 'mean', 'sum'."
        )

    return sign * loss

check_eigenvalues_sorted(eigenvalues: torch.Tensor, descending: bool = True, raise_error: bool = True) -> bool

Check if eigenvalues are sorted in descending or ascending order.

PARAMETER	DESCRIPTION
eigenvalues	1D tensor of eigenvalues. TYPE: Tensor
descending	If True, check descending order; if False, check ascending order. TYPE: bool, optional, by default True DEFAULT: True
raise_error	If True, raise ValueError when not sorted; if False, return False. TYPE: bool, optional, by default True DEFAULT: True

RETURNS	DESCRIPTION
bool	True if sorted correctly, False otherwise (when raise_error=False).

RAISES	DESCRIPTION
ValueError	If eigenvalues are not sorted and raise_error=True.

Examples:

>>> eigenvalues = torch.tensor([3.0, 2.0, 1.0])
>>> check_eigenvalues_sorted(eigenvalues)
True

>>> eigenvalues = torch.tensor([1.0, 3.0, 2.0])
>>> check_eigenvalues_sorted(eigenvalues, raise_error=False)
False

Source code in spectre/loss/eigen.py

def check_eigenvalues_sorted(
    eigenvalues: torch.Tensor, descending: bool = True, raise_error: bool = True
) -> bool:
    """
    Check if eigenvalues are sorted in descending or ascending order.

    Parameters
    ----------
    eigenvalues : torch.Tensor
        1D tensor of eigenvalues.

    descending : bool, optional, by default True
        If True, check descending order; if False, check ascending order.

    raise_error : bool, optional, by default True
        If True, raise ValueError when not sorted; if False, return False.

    Returns
    -------
    bool
        True if sorted correctly, False otherwise (when raise_error=False).

    Raises
    ------
    ValueError
        If eigenvalues are not sorted and raise_error=True.

    Examples
    --------
    >>> eigenvalues = torch.tensor([3.0, 2.0, 1.0])
    >>> check_eigenvalues_sorted(eigenvalues)
    True

    >>> eigenvalues = torch.tensor([1.0, 3.0, 2.0])
    >>> check_eigenvalues_sorted(eigenvalues, raise_error=False)
    False
    """
    check_1d(eigenvalues)

    if len(eigenvalues) < 2:
        return True

    if descending:
        is_sorted = torch.all(eigenvalues[:-1] >= eigenvalues[1:]).item()
        order_str = "descending"
    else:
        is_sorted = torch.all(eigenvalues[:-1] <= eigenvalues[1:]).item()
        order_str = "ascending"

    if not is_sorted:
        if raise_error:
            raise ValueError(
                f"Eigenvalues must be sorted in {order_str} order. "
                f"Got eigenvalues with violations. Consider sorting with "
                f"`torch.sort(eigenvalues, descending={descending})[0]`."
            )
        return False

    return True