Kernel Base

base

CLASS	DESCRIPTION
Kernel	Abstract base class for kernel functions used in spectral methods such as diffusion

FUNCTION	DESCRIPTION
initialize_kernel_fn	Initialize a kernel from an instance or registry name.

Classes

Kernel(learnable: bool = False, learnable_log: bool = True, normalization: KernelNormalizer | NormalizationType | dict | None = None, dtype: torch.dtype = torch.float32)

Bases: Module, ABC

Abstract base class for kernel functions used in spectral methods such as diffusion maps or spectral maps.

Provides a unified interface for computing similarity matrices from distance matrices, with support for learnable parameters, automatic differentiation, and various normalization schemes. All kernel implementations should inherit from this class.

PARAMETER	DESCRIPTION
learnable	Whether kernel parameters should be learnable. TYPE: bool, optional, by default False DEFAULT: False
learnable_log	Whether to parameterize learnable parameters in log space for numerical stability and to enforce positivity constraints. TYPE: bool, optional, by default True DEFAULT: True
normalization	Kernel normalization specification. Multiple formats supported: KernelNormalizer: Custom normalization instance for advanced control str: Normalization type shortcut. Creates KernelNormalizer(norm_type=str) Valid values: [NormalizationType][spectre.kernel.NormalizationType] dict: Dictionary of KernelNormalizer parameters. Creates KernelNormalizer(**dict) None: No normalization applied (returns raw kernel matrix) TYPE: KernelNormalizer | NormalizationType | dict | None, by default None DEFAULT: None
dtype	Data type for internal tensor computations. TYPE: torch.dtype, optional, by default torch.float32 DEFAULT: float32

Examples:

Creating a custom kernel by inheriting from Kernel:

>>> import torch
>>> from spectre.kernel import Kernel
>>> from spectre.pairwise_distance import pairwise_distance_euclidean
>>>
>>> class LinearKernel(Kernel):
...     def get_kernel_params(self):
...         return {}
...
...     def forward_step(self, D, weights=None):
...         return 1.0 - D / D.max()  # Simple linear decay
>>>
>>> X = torch.randn(10, 3)
>>> D = pairwise_distance_euclidean(X)
>>> kernel = LinearKernel()
>>> K = kernel(D)

Using existing kernel implementations:

>>> from spectre.kernel import GaussianKernel, TKernel
>>> gaussian_kernel = GaussianKernel(bw_method=torch.tensor(1.0))
>>> t_kernel = TKernel(alpha=torch.tensor(1.0), beta=torch.tensor(2.0))

Notes

• Learnable parameters enable end-to-end optimization in neural networks
• The normalize_alpha parameter controls anisotropic scaling behavior

METHOD	DESCRIPTION
get_kernel_params	Return current kernel parameters as dictionary.
get_param_summary	Return parameter summary for logging.
freeze_params	Freeze kernel parameters to prevent recomputation.
forward_step	Compute the kernel matrix.
forward	Compute the kernel matrix and apply optional normalization to create

ATTRIBUTE	DESCRIPTION
normalization_type	Get kernel normalization type. TYPE: str
param_info	Return detailed parameter information. TYPE: dict[str, dict]

Source code in spectre/kernel/base.py

def __init__(
    self,
    learnable: bool = False,
    learnable_log: bool = True,
    normalization: KernelNormalizer | NormalizationType | dict | None = None,
    dtype: torch.dtype = torch.float32,
) -> None:
    super().__init__()

    if not isinstance(learnable, bool):
        raise TypeError("Parameter `learnable` should be a boolean.")
    self.learnable = learnable

    if not isinstance(learnable_log, bool):
        raise TypeError("Parameter `learnable_log` should be a boolean.")
    self.learnable_log = learnable_log
    self._get_param = (
        lambda param: torch.exp(param) if self.learnable_log else param
    )

    self.normalization = self._parse_normalization(normalization)

    self.dtype = dtype

Attributes

normalization_type: str property

Get kernel normalization type.

param_info: dict[str, dict] property

Return detailed parameter information.

RETURNS	DESCRIPTION
dict[str, dict]	Dictionary with parameter names as keys and metadata dicts with: value: current parameter value (torch.Tensor) learnable: bool indicating if parameter is learnable shape: tuple of parameter shape numel: number of elements

Examples:

>>> kernel = GaussianKernel(bw_method=torch.tensor(1.0), learnable=True)
>>> info = kernel.param_info
>>> info["bandwidth"]
{'value': tensor([1.0]), 'learnable': True, 'shape': (1,), 'numel': 1}

Functions

get_kernel_params() -> dict[str, torch.Tensor] abstractmethod

Return current kernel parameters as dictionary.

Returns dictionary with parameter names as keys and current values (after applying _get_param transformation for learnable parameters).

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary of parameter names to current values.

Examples:

For GaussianKernel:

>>> kernel = GaussianKernel(bw_method=torch.tensor(1.0))
>>> params = kernel.get_kernel_params()
>>> params
{'bandwidth': tensor([1.0])}

For TKernel:

>>> kernel = TKernel(alpha=torch.tensor(1.0), beta=torch.tensor(2.0))
>>> params = kernel.get_kernel_params()
>>> params
{'alpha': tensor([1.0]), 'beta': tensor([2.0])}

Source code in spectre/kernel/base.py

@abstractmethod
def get_kernel_params(self) -> dict[str, torch.Tensor]:
    """
    Return current kernel parameters as dictionary.

    Returns dictionary with parameter names as keys and current values
    (after applying `_get_param` transformation for learnable parameters).

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary of parameter names to current values.

    Examples
    --------
    For GaussianKernel:

    >>> kernel = GaussianKernel(bw_method=torch.tensor(1.0))
    >>> params = kernel.get_kernel_params()
    >>> params
    {'bandwidth': tensor([1.0])}

    For TKernel:

    >>> kernel = TKernel(alpha=torch.tensor(1.0), beta=torch.tensor(2.0))
    >>> params = kernel.get_kernel_params()
    >>> params
    {'alpha': tensor([1.0]), 'beta': tensor([2.0])}
    """
    raise NotImplementedError()

get_param_summary() -> dict[str, float | dict]

Return parameter summary for logging.

Returns scalar values as-is, vector/matrix as statistics.

RETURNS	DESCRIPTION
dict[str, float | dict]	For scalars: {param_name: float_value} For vectors/matrices: {param_name: {mean, std, min, max, shape}}

Examples:

Scalar parameter:

>>> kernel = GaussianKernel(bw_method=torch.tensor(1.5))
>>> summary = kernel.get_param_summary()
>>> summary
{'bandwidth': 1.5}

Vector parameter:

>>> kernel = TKernel(alpha=torch.ones(10, 1), beta=torch.tensor(2.0))
>>> summary = kernel.get_param_summary()
>>> summary
{'alpha': {'mean': 1.0, 'std': 0.0, 'min': 1.0, 'max': 1.0, 'shape': (10, 1)},
 'beta': 2.0}

Source code in spectre/kernel/base.py

def get_param_summary(self) -> dict[str, float | dict]:
    """
    Return parameter summary for logging.

    Returns scalar values as-is, vector/matrix as statistics.

    Returns
    -------
    dict[str, float | dict]
        For scalars: `{param_name: float_value}`
        For vectors/matrices: `{param_name: {mean, std, min, max, shape}}`

    Examples
    --------
    Scalar parameter:

    >>> kernel = GaussianKernel(bw_method=torch.tensor(1.5))
    >>> summary = kernel.get_param_summary()
    >>> summary
    {'bandwidth': 1.5}

    Vector parameter:

    >>> kernel = TKernel(alpha=torch.ones(10, 1), beta=torch.tensor(2.0))
    >>> summary = kernel.get_param_summary()
    >>> summary
    {'alpha': {'mean': 1.0, 'std': 0.0, 'min': 1.0, 'max': 1.0, 'shape': (10, 1)},
     'beta': 2.0}
    """
    params = self.get_kernel_params()
    summary = {}

    for name, value in params.items():
        if value.numel() == 1:
            summary[name] = value.item()
        else:
            summary[name] = {
                "mean": value.mean().item(),
                "std": value.std().item(),
                "min": value.min().item(),
                "max": value.max().item(),
                "shape": tuple(value.shape),
            }

    return summary

freeze_params() -> Kernel

Freeze kernel parameters to prevent recomputation.

Converts dynamic parameter computation (e.g., bandwidth methods like "median") to constant parameters using current values.

The base implementation is a no-op. Subclasses with dynamic parameters (e.g., GaussianKernel) override this method to freeze their specific parameters using cached values from previous forward passes.

RETURNS	DESCRIPTION
Kernel	Returns self for method chaining.

Source code in spectre/kernel/base.py

def freeze_params(self) -> "Kernel":
    """
    Freeze kernel parameters to prevent recomputation.

    Converts dynamic parameter computation (e.g., bandwidth methods like "median")
    to constant parameters using current values.

    The base implementation is a no-op. Subclasses with dynamic parameters
    (e.g., `GaussianKernel`) override this method to freeze their specific
    parameters using cached values from previous forward passes.

    Returns
    -------
    Kernel
        Returns self for method chaining.
    """
    return self

forward_step(D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs) -> torch.Tensor abstractmethod

Compute the kernel matrix.

PARAMETER	DESCRIPTION
D	Pairwise distance matrix of shape (n_samples, n_samples). TYPE: Tensor
weights	Weights for each sample. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
**kwargs	Additional keyword arguments passed to the kernel. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Kernel matrix of shape (n_samples, n_samples).

Source code in spectre/kernel/base.py

@abstractmethod
def forward_step(
    self, D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs
) -> torch.Tensor:
    """
    Compute the kernel matrix.

    Parameters
    ----------
    D : torch.Tensor
        Pairwise distance matrix of shape (n_samples, n_samples).

    weights : torch.Tensor | None, optional, by default None
        Weights for each sample.

    **kwargs : dict
        Additional keyword arguments passed to the kernel.

    Returns
    -------
    torch.Tensor
        Kernel matrix of shape (n_samples, n_samples).
    """
    raise NotImplementedError()

forward(D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs) -> torch.Tensor

Compute the kernel matrix and apply optional normalization to create transition matrices for spectral methods.

PARAMETER	DESCRIPTION
D	Pairwise distance matrix of shape (n_samples, n_samples). TYPE: Tensor
weights	Weights for each sample. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
**kwargs	Additional keyword arguments passed to the kernel. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Kernel matrix of shape (n_samples, n_samples).

Source code in spectre/kernel/base.py

def forward(
    self, D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs
) -> torch.Tensor:
    """
    Compute the kernel matrix and apply optional normalization to create
    transition matrices for spectral methods.

    Parameters
    ----------
    D : torch.Tensor
        Pairwise distance matrix of shape (n_samples, n_samples).

    weights : torch.Tensor | None, optional, by default None
        Weights for each sample.

    **kwargs : dict
        Additional keyword arguments passed to the kernel.

    Returns
    -------
    torch.Tensor
        Kernel matrix of shape (n_samples, n_samples).
    """
    K = self.forward_step(D, weights=weights, **kwargs)

    if self.normalization is not None:
        return self.normalization(K, weights=weights)
    else:
        return K

KernelRegistry

Bases: Registry

Registry for kernel classes.

Functions

initialize_kernel_fn(kernel_fn: Kernel | str | None, kernel_kwargs: dict[str, Any] | None = None) -> Kernel

Initialize a kernel from an instance or registry name.

Arguments

kernel_fn : Kernel | str | None Kernel instance or name of registered kernel. kernel_kwargs : dict[str, Any] | None Keyword arguments passed to kernel constructor (only for string names).

RETURNS	DESCRIPTION
Kernel	Initialized kernel instance.

RAISES	DESCRIPTION
ValueError	If kernel_fn is None.
TypeError	If kernel_fn is not a Kernel instance or string.

Source code in spectre/kernel/base.py

def initialize_kernel_fn(
    kernel_fn: Kernel | str | None, kernel_kwargs: dict[str, Any] | None = None
) -> Kernel:
    """
    Initialize a kernel from an instance or registry name.

    Arguments
    ---------
    kernel_fn : Kernel | str | None
        Kernel instance or name of registered kernel.
    kernel_kwargs : dict[str, Any] | None
        Keyword arguments passed to kernel constructor (only for string names).

    Returns
    -------
    Kernel
        Initialized kernel instance.

    Raises
    ------
    ValueError
        If `kernel_fn` is None.
    TypeError
        If `kernel_fn` is not a Kernel instance or string.
    """
    return _initialize_from_registry(
        KernelRegistry, kernel_fn, kernel_kwargs, param_name="kernel_fn"
    )