Ensemble Embedding

embedding

CLASS	DESCRIPTION
EmbeddingEnsemble	Ensemble of multiple embeddings with configurable aggregation.

Attributes

Classes

EmbeddingEnsemble(estimators: list[ModelType] | dict[str, ModelType], reduce: str = 'mean', estimator_weights: list[float] | dict[str, float] | None = None, alignment: str = 'procrustes', fit_manager: FitManager | None = None)

Bases: Ensemble

Ensemble of multiple embeddings with configurable aggregation.

Combines embeddings from multiple diffusion maps, PCAs, or other methods using various reduction strategies. Provides robust embeddings that are less sensitive to hyperparameter choices.

PARAMETER	DESCRIPTION
estimators	Models to construct ensemble. TYPE: list[ModelType] | dict[str, ModelType]
reduce	Reduction method for combining embeddings: "mean": Average embeddings across models "median": Element-wise median "weighted_mean": Weighted average (requires weights parameter) TYPE: str, by default "mean" DEFAULT: 'mean'
estimator_weights	Estimator weights for "weighted_mean" reduce. Must match estimators structure. If None and reduce is "weighted_mean", raises ValueError. TYPE: list[float] | dict[str, float] | None, by default None DEFAULT: None
alignment	Whether to align embeddings before reduction. "procrustes": Align via Procrustes analysis (rotation + scaling) "none": No alignment (assumes embeddings in same coordinates) TYPE: str, by default "procrustes" DEFAULT: 'procrustes'
fit_manager	Manager for unified training of both Estimator and Parametric models. Required for using Parametric (PyTorch Lightning) models in embedding ensemble. See FitManager for configuration options. TYPE: FitManager | None DEFAULT: None

ATTRIBUTE	DESCRIPTION
estimators	Dictionary of fitted models with names as keys. TYPE: dict[str, ModelType]
estimator_weights	Normalized estimator weights for reduction. TYPE: dict[str, float]
embeddings	Individual embeddings from each model (stored during fit). TYPE: list[Tensor]

Examples:

Ensemble diffusion maps with different kernels:

>>> from spectre.decomposition import DiffusionMap
>>> from spectre.kernel import GaussianKernel, TKernel
>>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
>>> from spectre.ensemble import EmbeddingEnsemble
>>>
>>> # Create estimators with different kernels
>>> dm_gaussian = DiffusionMap(
...     kernel_fn=GaussianKernel(bw_method=torch.tensor(1.0)),
...     distance_fn=PairwiseDistanceEuclidean(),
...     out_features=5,
... )
>>> dm_t = DiffusionMap(
...     kernel_fn=TKernel(
...         alpha=torch.tensor(1.0),
...         beta=torch.tensor(3.0),
...     ),
...     distance_fn=PairwiseDistanceEuclidean(),
...     out_features=5,
... )
>>>
>>> # Create ensemble
>>> ensemble = EmbeddingEnsemble(
...     estimators=[dm_gaussian, dm_t],
...     reduce="mean",
...     alignment="procrustes",
... )
>>>
>>> X = torch.randn(200, 10)
>>> ensemble.fit(X)
>>> X_embedded = ensemble.predict(X)

Named estimators with custom weights:

>>> ensemble = EmbeddingEnsemble(
...     estimators={"gaussian": dm_gaussian, "t_kernel": dm_t},
...     reduce="weighted_mean",
...     estimator_weights={"gaussian": 0.7, "t_kernel": 0.3},
... )
>>> ensemble.fit(X)

Access individual estimator embeddings:

>>> ensemble.fit(X)
>>> for name, emb in zip(["gaussian", "t_kernel"], ensemble.embeddings):
...     print(f"{name}: {emb.shape}")

METHOD	DESCRIPTION
fit	Fit all estimators independently on the data.
predict	Combine predictions from all estimators.
score	Reduce scores from all estimators.
get_stability_score	Compute stability across ensemble members.

Source code in spectre/ensemble/embedding.py

def __init__(
    self,
    estimators: list[ModelType] | dict[str, ModelType],
    reduce: str = "mean",
    estimator_weights: list[float] | dict[str, float] | None = None,
    alignment: str = "procrustes",
    fit_manager: FitManager | None = None,
):
    super().__init__(estimators, fit_manager=fit_manager)

    allowed_reduce = ["mean", "median", "weighted_mean"]
    if reduce not in allowed_reduce:
        raise ValueError(f"Unknown reduce '{reduce}'. Options: {allowed_reduce}.")
    self.reduce = reduce

    if reduce == "weighted_mean" and estimator_weights is None:
        raise ValueError("Weights must be provided for 'weighted_mean' reduction.")

    if estimator_weights is not None:
        self._estimator_weights = self.normalize_estimator_weights(
            estimator_weights
        )
    else:
        self._estimator_weights = self.get_equal_estimator_weights()

    allowed_alignments = ["procrustes", "none"]
    if alignment not in allowed_alignments:
        raise ValueError(
            f"Unknown alignment '{alignment}'. Options: {allowed_alignments}."
        )
    self.alignment = alignment

    # Storage for embeddings
    self._embeddings = []
    self._reference_embedding = None

Attributes

embeddings: list[torch.Tensor] property

Return the embeddings of the fitted estimators.

estimator_weights: dict[str, float] property

Return the weights of the fitted estimators.

Functions

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

Fit all estimators independently on the data.

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
target	Target values for supervised methods. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
EmbeddingEnsemble	Returns self for method chaining.

Source code in spectre/ensemble/embedding.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> "EmbeddingEnsemble":
    """
    Fit all estimators independently on the data.

    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, ).

    target : torch.Tensor | None, optional, by default None
        Target values for supervised methods.

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

    self.in_features = X.shape[1]

    # Fit each estimator and collect embeddings
    self._embeddings = []

    for name, estimator in self._estimators.items():
        # Fit estimator
        self._fit_estimator(
            estimator, X, weights=weights, target=target, estimator_name=name
        )

        # Get embedding (predict on training data)
        emb = estimator.predict(X)
        self._embeddings.append(emb)

    # Set output features from first embedding
    first_shape = self._embeddings[0].shape
    self.out_features = first_shape[1]

    # Validate all embeddings have same shape
    if not all(emb.shape == first_shape for emb in self._embeddings):
        shapes = [emb.shape for emb in self._embeddings]
        raise ValueError(
            f"All embeddings must have the same shape. Got shapes: {shapes}."
        )

    # Store reference embedding for alignment (use first estimator)
    self._reference_embedding = self._embeddings[0]

    self.is_fitted = True

    return self

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

Combine predictions from all estimators.

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

RETURNS	DESCRIPTION
Tensor	Reduced embedding of shape (n_samples, out_features).

Source code in spectre/ensemble/embedding.py

def predict(self, X: torch.Tensor) -> torch.Tensor:
    """
    Combine predictions from all estimators.

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

    Returns
    -------
    torch.Tensor
        Reduced embedding of shape (n_samples, out_features).
    """
    self.check_fitted()

    check_2d(X)

    if X.shape[1] != self.in_features:
        raise ValueError(f"Expected {self.in_features} features, got {X.shape[1]}.")

    # Get predictions from each estimator
    predictions = list(self.predict_ensemble(X).values())

    # Align predictions if requested
    if self.alignment == "procrustes":
        predictions = self._align_embeddings(predictions)

    # Reduce predictions using base class method
    return self.reduce_predictions(
        predictions, reduce=self.reduce, estimator_weights=self._estimator_weights
    )

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

Reduce scores from all estimators.

Computes the mean score across all ensemble members.

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
float	Mean score across estimators.

Source code in spectre/ensemble/embedding.py

def score(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> float:
    """
    Reduce scores from all estimators.

    Computes the mean score across all ensemble members.

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

    weights : torch.Tensor | None, optional, by default None
        Sample weights.

    target : torch.Tensor | None, optional, by default None
        Target values.

    Returns
    -------
    float
        Mean score across estimators.
    """
    return self.reduce_scores(X, weights=weights, target=target, reduce="mean")

get_stability_score(X: torch.Tensor) -> float

Compute stability across ensemble members.

Measures pairwise agreement between estimator embeddings using Procrustes error. Lower values indicate higher agreement.

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

RETURNS	DESCRIPTION
float	Mean Procrustes error between all pairs of embeddings.

Source code in spectre/ensemble/embedding.py

def get_stability_score(self, X: torch.Tensor) -> float:
    """
    Compute stability across ensemble members.

    Measures pairwise agreement between estimator embeddings using
    Procrustes error. Lower values indicate higher agreement.

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

    Returns
    -------
    float
        Mean Procrustes error between all pairs of embeddings.
    """
    self.check_fitted()

    # Get all predictions
    predictions = list(self.predict_ensemble(X).values())

    # Compute pairwise Procrustes errors
    n_estimators = len(predictions)
    errors = []

    for i in range(n_estimators):
        for j in range(i + 1, n_estimators):
            # Use metrics module function
            error = procrustes_error(
                predictions[i], predictions[j], scaling=True, eps=1e-8
            )
            errors.append(error)

    if not errors:
        return 0.0

    return float(torch.tensor(errors).mean().item())

Functions