Compute Kde

kde

CLASS	DESCRIPTION
KernelDensity	Kernel Density Estimation with GPU support.

Classes

KernelDensity(bw_method: float | torch.Tensor | str = 'scott', kernel_fn: Literal['gaussian', 'epanechnikov', 'tophat'] = 'gaussian', kernel_kwargs: dict | None = None, eps: float = torch.finfo(torch.float32).eps, dtype: torch.dtype = torch.float32, device: torch.device | str | None = None)

Bases: Estimator

Kernel Density Estimation with GPU support.

This class provides GPU-accelerated kernel density estimation compatible with scipy.stats.gaussian_kde.

PARAMETER	DESCRIPTION
bw_method	Bandwidth parameter for the Gaussian kernel. float: Covariance factor torch.Tensor: Covariance factor tensor str: Bandwidth estimation method ("scott", "silverman") TYPE: float | torch.Tensor | str, optional, by default "scott" DEFAULT: 'scott'
kernel_fn	Kernel function for bandwidth estimation. TYPE: Literal["gaussian", "epanechnikov", "tophat"], by default "gaussian" DEFAULT: 'gaussian'
kernel_kwargs	Additional parameters (reserved for future extensions). TYPE: dict | None, optional, by default None DEFAULT: None
eps	Small value to avoid division by zero. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps
dtype	Data type for tensors. TYPE: torch.dtype, optional, by default torch.float32 DEFAULT: float32
device	Device to run computations on. If None, automatically selects CUDA if available, otherwise CPU. TYPE: torch.device | str | None, by default None DEFAULT: None

ATTRIBUTE	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Training sample weights of shape (n_samples,). TYPE: Tensor | None
covariance	Kernel covariance matrix: data_cov * factor^2. TYPE: Tensor
inv_cov	Inverse of kernel covariance matrix. TYPE: Tensor
factor	Bandwidth factor used. TYPE: float
neff	Effective number of samples (accounting for weights). TYPE: float
eps	Small value to avoid division by zero. TYPE: float
dtype	Data type for tensors. TYPE: dtype
device	Device to run computations on. TYPE: device

METHOD	DESCRIPTION
fit	Fit the Kernel Density model.
predict	Compute log-likelihood of samples under the model.
score_samples	Compute log-likelihood of each sample under the model.
score	Compute the total log-likelihood of the data.
optimize_bandwidth	Find optimal bandwidth using cross-validation likelihood.
sample	Generate random samples from the fitted kernel density model.

Source code in spectre/compute/kde.py

def __init__(
    self,
    bw_method: float | torch.Tensor | str = "scott",
    kernel_fn: Literal["gaussian", "epanechnikov", "tophat"] = "gaussian",
    kernel_kwargs: dict | None = None,
    eps: float = torch.finfo(torch.float32).eps,
    dtype: torch.dtype = torch.float32,
    device: torch.device | str | None = None,
) -> None:
    super().__init__()

    self.dtype = dtype

    if isinstance(bw_method, str):
        if bw_method not in ["scott", "silverman"]:
            raise ValueError(
                f"bw_method must be 'scott', 'silverman', or a numeric factor, "
                f"got '{bw_method}'"
            )
        self.bw_method = bw_method

    elif isinstance(bw_method, (int, float)):
        self.bw_method = "custom"
        self._custom_factor = float(bw_method)

    elif torch.is_tensor(bw_method):
        self.bw_method = "custom"
        self._custom_factor = float(bw_method.item())

    else:
        raise TypeError(
            f"Expected `bw_method` to be a str, float, or torch.Tensor, got "
            f"{type(bw_method)}."
        )

    if kernel_fn not in ["gaussian", "epanechnikov", "tophat"]:
        raise ValueError(
            f"kernel_fn must be 'gaussian', 'epanechnikov', or 'tophat', "
            f"got '{kernel_fn}'"
        )
    self.kernel_fn = kernel_fn

    if not isinstance(eps, (int, float)):
        raise TypeError(f"Expected `eps` to be a number, got {type(eps).__name__}.")
    check_in_interval(eps, "(0, inf)")
    self.eps = float(eps)

    if device is None:
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        self.device = torch.device(device)

    self.kernel_kwargs = kernel_kwargs or {}

    self._data: torch.Tensor | None = None
    self._weights: torch.Tensor | None = None
    self._covariance: torch.Tensor | None = None
    self._inv_cov: torch.Tensor | None = None
    self._factor: float | None = None
    self._neff: float | None = None
    self._norm_factor: float | None = None

Attributes

data: torch.Tensor | None property

Return training data.

weights: torch.Tensor | None property

Return training sample weights.

covariance: torch.Tensor | None property

Return kernel covariance matrix.

inv_cov: torch.Tensor | None property

Return inverse kernel covariance matrix.

factor: float | None property

Return bandwidth factor.

neff: float | None property

Return effective number of samples.

Functions

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

Fit the Kernel Density model.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). If provided, the KDE will be weighted, giving more importance to samples with higher weights. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values. Not used for density estimation but included for compatibility with Estimator interface. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
KernelDensity	Returns self for method chaining.

Source code in spectre/compute/kde.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> "KernelDensity":
    """
    Fit the Kernel Density model.

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

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,). If provided, the KDE will be
        weighted, giving more importance to samples with higher weights.

    target : torch.Tensor | None, optional, by default None
        Target values. Not used for density estimation but included for
        compatibility with `Estimator` interface.

    Returns
    -------
    KernelDensity
        Returns self for method chaining.
    """
    check_2d(X)

    self._data = X.to(device=self.device, dtype=self.dtype)
    n_samples, n_features = self._data.shape
    self.in_features = n_features
    self.out_features = 1

    if weights is not None:
        check_1d(weights)
        check_same_shape(X, weights, axis=0)
        check_sample_weights(weights)

        self._weights = weights.to(device=self.device, dtype=self.dtype)
        self._weights = self._weights / self._weights.sum()

        self._neff = float((self._weights.sum() ** 2) / (self._weights**2).sum())
    else:
        self._weights = (
            torch.ones(n_samples, device=self.device, dtype=self.dtype) / n_samples
        )
        self._neff = float(n_samples)

    mean = (self._data.T * self._weights).sum(dim=1)
    centered = self._data.T - mean.unsqueeze(1)

    # Weighted covariance with bias=False (divide by sum of weights, not
    # sum-1)
    self._data_cov = (centered * self._weights) @ centered.T

    if self.bw_method == "scott":
        # Scott's rule: neff^(-1/(d+4))
        self._factor = self._neff ** (-1.0 / (n_features + 4))

    elif self.bw_method == "silverman":
        # Silverman's rule: (neff*(d+2)/4)^(-1/(d+4))
        self._factor = (self._neff * (n_features + 2) / 4.0) ** (
            -1.0 / (n_features + 4)
        )

    else:
        self._factor = self._custom_factor

    # Compute kernel covariance: factor^2 * data_cov + regularization (eps)
    self._covariance = self._factor**2 * self._data_cov
    self._covariance = self._covariance + self.eps * torch.eye(
        n_features, device=self.device, dtype=self.dtype
    )

    # Compute inverse covariance for Mahalanobis distance
    self._inv_cov = torch.linalg.inv(self._covariance)

    # Cache normalization factor
    try:
        L = torch.linalg.cholesky(self._covariance)
        sqrt_det_cov = torch.prod(torch.diag(L))
    except RuntimeError:
        # Fallback if Cholesky fails
        sqrt_det_cov = torch.sqrt(torch.linalg.det(self._covariance))

    self._norm_factor = float(
        (1.0 / (((2 * math.pi) ** (n_features / 2.0)) * sqrt_det_cov)).item()
    )

    self.is_fitted = True

    return self

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

Compute log-likelihood of samples under the model.

This method is an alias for score_samples() to maintain compatibility with the Estimator interface.

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

RETURNS	DESCRIPTION
Tensor	Log-likelihood of each sample, shape (n_samples,).

Source code in spectre/compute/kde.py

def predict(self, X: torch.Tensor) -> torch.Tensor:
    """
    Compute log-likelihood of samples under the model.

    This method is an alias for `score_samples()` to maintain compatibility
    with the `Estimator` interface.

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

    Returns
    -------
    torch.Tensor
        Log-likelihood of each sample, shape (n_samples,).
    """
    return self.score_samples(X)

score_samples(X: torch.Tensor, batch_size: int | None = None) -> torch.Tensor

Compute log-likelihood of each sample under the model.

The log-likelihood is computed using the scipy-compatible formula: \(\log p(x) = \log \left( \sum_{i=1}^{n} w_i K(x, x_i) \right)\) where \(K\) is the kernel and \(w_i\) are (normalized) weights.

PARAMETER	DESCRIPTION
X	Query points of shape (n_samples, in_features). TYPE: Tensor
batch_size	Number of samples to process at once. If None, processes all samples together. TYPE: int | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Log-likelihood of each sample, shape (n_samples,).

Source code in spectre/compute/kde.py

def score_samples(
    self, X: torch.Tensor, batch_size: int | None = None
) -> torch.Tensor:
    """
    Compute log-likelihood of each sample under the model.

    The log-likelihood is computed using the scipy-compatible formula:
    $\\log p(x) = \\log \\left( \\sum_{i=1}^{n} w_i K(x, x_i) \\right)$
    where $K$ is the kernel and $w_i$ are (normalized) weights.

    Parameters
    ----------
    X : torch.Tensor
        Query points of shape (n_samples, in_features).

    batch_size : int | None, optional, by default None
        Number of samples to process at once. If None, processes all
        samples together.

    Returns
    -------
    torch.Tensor
        Log-likelihood of each sample, shape (n_samples,).
    """
    self.check_fitted()
    check_2d(X)

    if X.shape[1] != self.in_features:
        raise ValueError(
            f"X has {X.shape[1]} features, but KernelDensity was fitted with "
            f"{self.in_features} features."
        )

    X = X.to(device=self.device, dtype=self.dtype)

    if batch_size is not None and X.shape[0] > batch_size:
        scores = []
        for i in range(0, X.shape[0], batch_size):
            batch = X[i : i + batch_size]
            scores.append(self._score_samples_impl(batch))
        return torch.cat(scores)
    else:
        return self._score_samples_impl(X)

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

Compute the total log-likelihood of the data.

This method computes the mean log-likelihood over all samples.

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). If provided, computes weighted mean log-likelihood. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Not used for density estimation but included for compatibility. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Mean log-likelihood over all samples. Higher is better.

Source code in spectre/compute/kde.py

def score(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute the total log-likelihood of the data.

    This method computes the mean log-likelihood over all samples.

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

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,). If provided, computes
        weighted mean log-likelihood.

    target : torch.Tensor | None, optional, by default None
        Not used for density estimation but included for compatibility.

    Returns
    -------
    torch.Tensor
        Mean log-likelihood over all samples. Higher is better.
    """
    log_density = self.score_samples(X)

    if weights is not None:
        check_1d(weights)
        check_same_shape(X, weights, axis=0)
        check_sample_weights(weights)

        weights = weights.to(device=self.device, dtype=self.dtype)
        return (log_density * weights).sum() / weights.sum()
    else:
        return log_density.mean()

optimize_bandwidth(X: torch.Tensor, weights: torch.Tensor | None = None, cv_folds: int = 5, bandwidth_range: tuple[float, float] = (0.1, 2.0), n_trials: int = 20, random_state: int | None = None) -> float

Find optimal bandwidth using cross-validation likelihood.

Uses k-fold CV to evaluate log-likelihood for different bandwidths.

PARAMETER	DESCRIPTION
X	Training data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights of shape (n_samples,). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
cv_folds	Number of cross-validation folds. TYPE: int, optional, by default 5 DEFAULT: 5
bandwidth_range	Range of bandwidth factors to search. TYPE: tuple[float, float], optional, by default (0.1, 2.0) DEFAULT: (0.1, 2.0)
n_trials	Number of bandwidth values to try. TYPE: int, optional, by default 20 DEFAULT: 20
random_state	Random seed for reproducibility. TYPE: int | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
float	Optimal bandwidth factor.

Source code in spectre/compute/kde.py

def optimize_bandwidth(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    cv_folds: int = 5,
    bandwidth_range: tuple[float, float] = (0.1, 2.0),
    n_trials: int = 20,
    random_state: int | None = None,
) -> float:
    """
    Find optimal bandwidth using cross-validation likelihood.

    Uses k-fold CV to evaluate log-likelihood for different bandwidths.

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

    weights : torch.Tensor | None, optional, by default None
        Sample weights of shape (n_samples,).

    cv_folds : int, optional, by default 5
        Number of cross-validation folds.

    bandwidth_range : tuple[float, float], optional, by default (0.1, 2.0)
        Range of bandwidth factors to search.

    n_trials : int, optional, by default 20
        Number of bandwidth values to try.

    random_state : int | None, optional, by default None
        Random seed for reproducibility.

    Returns
    -------
    float
        Optimal bandwidth factor.
    """
    check_2d(X)

    if random_state is not None:
        check_random_state(random_state=random_state)

    n_samples = X.shape[0]
    indices = torch.randperm(n_samples)
    fold_sizes = torch.full((cv_folds,), n_samples // cv_folds, dtype=torch.long)
    fold_sizes[: n_samples % cv_folds] += 1

    best_bw = None
    best_score = -float("inf")

    bandwidths = torch.linspace(bandwidth_range[0], bandwidth_range[1], n_trials)

    for bw in bandwidths:
        scores = []
        current = 0

        for fold in range(cv_folds):
            start, stop = current, current + fold_sizes[fold]
            current = stop

            test_idx = indices[start:stop]
            train_idx = torch.cat([indices[:start], indices[stop:]])

            X_train = X[train_idx]
            X_test = X[test_idx]

            w_train = weights[train_idx] if weights is not None else None

            # Create temporary KDE with current bandwidth
            temp_kde = KernelDensity(
                bw_method=float(bw),
                kernel_fn=self.kernel_fn,
                eps=self.eps,
                dtype=self.dtype,
                device=self.device,
            )
            temp_kde.fit(X_train, weights=w_train)
            score = temp_kde.score(X_test)
            scores.append(score.item())

        avg_score = sum(scores) / len(scores)

        if avg_score > best_score:
            best_score = avg_score
            best_bw = float(bw)

    return best_bw

sample(n_samples: int = 1, random_state: int | None = None) -> torch.Tensor

Generate random samples from the fitted kernel density model.

Samples are generated by:

1. Randomly selecting training samples (weighted if applicable)
2. Adding Gaussian noise from N(0, covariance)

This matches scipy.stats.gaussian_kde.resample() behavior.

PARAMETER	DESCRIPTION
n_samples	Number of samples to generate. TYPE: int, optional, by default 1 DEFAULT: 1
random_state	Random seed for reproducibility. TYPE: int | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Generated samples of shape (n_samples, in_features).

Source code in spectre/compute/kde.py

def sample(
    self, n_samples: int = 1, random_state: int | None = None
) -> torch.Tensor:
    """
    Generate random samples from the fitted kernel density model.

    Samples are generated by:

    1. Randomly selecting training samples (weighted if applicable)
    2. Adding Gaussian noise from N(0, covariance)

    This matches scipy.stats.gaussian_kde.resample() behavior.

    Parameters
    ----------
    n_samples : int, optional, by default 1
        Number of samples to generate.

    random_state : int | None, optional, by default None
        Random seed for reproducibility.

    Returns
    -------
    torch.Tensor
        Generated samples of shape (n_samples, in_features).
    """
    self.check_fitted()

    if random_state is not None:
        check_random_state(random_state=random_state)

    idx = torch.multinomial(self._weights, n_samples, replacement=True)
    samples = self._data[idx].clone()  # ty: ignore

    # Add Gaussian noise with kernel covariance
    # Generate noise from N(0, I) and transform with Cholesky factor
    try:
        L = torch.linalg.cholesky(self._covariance)
    except RuntimeError:
        # Fallback: add small regularization to diagonal
        L = torch.linalg.cholesky(
            self._covariance
            + self.eps
            * torch.eye(self.in_features, device=self.device, dtype=self.dtype)
        )

    noise = (
        torch.randn(
            n_samples, self.in_features, device=self.device, dtype=self.dtype
        )
        @ L.T
    )

    samples = samples + noise

    return samples

Functions