Clustering Pcca

pcca

CLASS	DESCRIPTION
PCCA	Perron Cluster Cluster Analysis (PCCA) for spectral clustering.

FUNCTION	DESCRIPTION
pcca	Perform PCCA spectral clustering (functional interface).

Classes

PCCA(n_states: int, optimize: bool = True, n_iter: int = 100, tol: float = 1e-05, eps: float = torch.finfo(torch.float32).eps, lr: float = 0.1, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32)

Bases: Estimator

Perron Cluster Cluster Analysis (PCCA) for spectral clustering.

PCCA performs fuzzy clustering on eigenvectors of transition matrices to identify metastable states in Markov chains or general spectral clustering tasks. The algorithm identifies a simplex structure in eigenvector space and computes soft membership assignments.

The input features should resemble eigenvectors (excluding the constant stationary eigenvector corresponding to eigenvalue 1). A constant first column is automatically prepended internally to satisfy the partition of unity constraint required by PCCA.

PARAMETER	DESCRIPTION
n_states	Number of metastable states (clusters) to identify. Must be at least 2. TYPE: int
optimize	Whether to optimize the transition matrix after initialization. If False, uses only the initial simplex-based transformation. TYPE: bool, by default True DEFAULT: True
n_iter	Maximum iterations for optimization. TYPE: int, by default 100 DEFAULT: 100
tol	Convergence tolerance for optimization. TYPE: float, by default 1e-5 DEFAULT: 1e-05
eps	Numerical stability regularization for preventing division by zero and ensuring positive definiteness. TYPE: float, by default torch.finfo(torch.float32).eps DEFAULT: eps
lr	Learning rate for optimization. TYPE: float, by default 0.1 DEFAULT: 0.1
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
dtype	Data type to use for computations. TYPE: torch.dtype, by default torch.float32 DEFAULT: float32

Properties

proba : torch.Tensor Fuzzy membership matrix of shape (n_samples, n_states). Each row sums to 1 and represents probability of belonging to each state.

labels : torch.Tensor State assignments of shape (n_samples,). Computed as argmax of membership matrix.

transition_matrix : torch.Tensor Transition matrix of shape (n_states, n_states).

References

.. [1] Deuflhard P, Weber M., "Robust Perron Cluster Analysis in Conformation Dynamics," Linear Algebra Appl., vol 398 pp 161-184, 2005. .. [2] Kube S, Weber M., "A coarse graining method for the identification of transition rates between molecular conformations," J. Chem. Phys., vol 126 pp 24103-024113, 2007. .. [3] Kube S., "Statistical Error Estimation and Grid-free Hierarchical Refinement in Conformation Dynamics," Doctoral Thesis, 2008.

Examples:

Basic PCCA clustering on eigenvectors:

>>> pcca = PCCA(n_states=3)
>>> X_eigenvectors = torch.randn(100, 2)  # 2 non-trivial eigenvectors
>>> labels = pcca.fit_predict(X_eigenvectors)
>>> labels.shape
torch.Size([100])

Accessing memberships:

>>> pcca.fit(X_eigenvectors)
>>> proba = pcca.predict_proba(X_eigenvectors)
>>> proba.shape
torch.Size([100, 3])
>>> proba.sum(dim=1)  # Each row sums to 1
tensor([1., 1., 1., ...])

METHOD	DESCRIPTION
fit	Fit PCCA+ model to eigenvector data.
predict	Predict state labels for data.
fit_predict	Fit and predict state labels.
predict_proba	Predict probabilities for data.
score	Compute state clustering score.
get_results	Get all fitted parameters as a dictionary.

ATTRIBUTE	DESCRIPTION
proba	Membership probability matrix of shape (n_samples, n_states). TYPE: Tensor
labels	State assignments of shape (n_samples,). TYPE: Tensor
transition_matrix	Transition matrix A of shape (n_states, n_states). TYPE: Tensor

Source code in spectre/clustering/pcca.py

def __init__(
    self,
    n_states: int,
    optimize: bool = True,
    n_iter: int = 100,
    tol: float = 1e-5,
    eps: float = torch.finfo(torch.float32).eps,
    lr: float = 0.1,
    device: torch.device | str | None = None,
    dtype: torch.dtype = torch.float32,
):
    super().__init__()

    if not isinstance(n_states, int):
        raise TypeError(
            f"Expected `n_states` to be an integer, got {type(n_states).__name__}."
        )
    check_in_interval(n_states, "[2, inf)")
    self.n_states = n_states

    self.out_features = n_states
    self.n_eigenvectors = self.n_states - 1

    if not isinstance(optimize, bool):
        raise TypeError(
            f"Expected `optimize` to be a bool, got {type(optimize).__name__}."
        )
    self.optimize = optimize

    if not isinstance(n_iter, int):
        raise TypeError(
            f"Expected `n_iter` to be an integer, got {type(n_iter).__name__}."
        )
    check_in_interval(n_iter, "[1, inf)")
    self.n_iter = n_iter

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

    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 not isinstance(lr, (int, float)):
        raise TypeError(f"Expected `lr` to be a number, got {type(lr).__name__}.")
    check_in_interval(lr, "(0, inf)")
    self.lr = float(lr)

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

    self.dtype = dtype

    # Internal fitted attributes
    self._proba = None
    self._labels = None
    self._transition_matrix = None

Attributes

proba: torch.Tensor property

Membership probability matrix of shape (n_samples, n_states).

labels: torch.Tensor property

State assignments of shape (n_samples,).

transition_matrix: torch.Tensor property

Transition matrix A of shape (n_states, n_states).

Functions

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

Fit PCCA+ model to eigenvector data.

PARAMETER	DESCRIPTION
X	Data of shape (n_samples, n_eigenvectors). Typically the non-trivial dominant eigenvectors from spectral decomposition (excluding the constant first eigenvector). TYPE: Tensor
weights	Sample weights. Currently not used. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values for supervised learning. Not supported. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
**kwargs	Additional keyword arguments. DEFAULT: {}

RETURNS	DESCRIPTION
PCCA	Return self.

Source code in spectre/clustering/pcca.py

def fit(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
    **kwargs,
) -> "PCCA":
    """
    Fit PCCA+ model to eigenvector data.

    Parameters
    ----------
    X : torch.Tensor
        Data of shape (n_samples, n_eigenvectors). Typically the
        non-trivial dominant eigenvectors from spectral decomposition
        (excluding the constant first eigenvector).

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

    target : torch.Tensor | None, optional, by default None
        Target values for supervised learning. Not supported.

    **kwargs
        Additional keyword arguments.

    Returns
    -------
    PCCA
        Return self.
    """
    check_is_tensor(X)
    check_2d(X)

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

    self.in_features = X.shape[1]
    if self.n_eigenvectors > X.shape[1]:
        raise ValueError(
            f"n_eigenvectors ({self.n_eigenvectors}) cannot exceed number of "
            f"columns in X ({X.shape[1]})."
        )
    if self.n_eigenvectors < self.n_states - 1:
        raise ValueError(
            f"n_eigenvectors ({self.n_eigenvectors}) must be at least "
            f"n_states - 1 ({self.n_states - 1}) for PCCA."
        )

    # Augment with constant column (stationary eigenvector) for PCCA. PCCA
    # requires the first eigenvector to be constant (all ones) to enforce
    # partition of unity in the membership matrix.
    X_e = X[:, : self.n_eigenvectors]
    X_e = self._augment_eigenvectors(X_e)

    # X_vertices is (n_states, n_states) so A is square
    vertex_indices = self._inner_simplex_algorithm(X_e)
    X_vertices = X_e[vertex_indices]
    cond_number = torch.linalg.cond(X_vertices)

    if cond_number > 1e10:
        warnings.warn(
            f"Vertex matrix is ill-conditioned (condition number: {cond_number:.2e}). "
            "Results may be numerically unstable. Consider using fewer states "
            "or more samples.",
            UserWarning,
            stacklevel=2,
        )

    try:
        A = torch.linalg.inv(X_vertices)
    except RuntimeError:
        A = torch.linalg.pinv(X_vertices)

    # Apply PCCA feasibility constraints (Deuflhard & Weber 2005)
    A = self._constrain(A, X_e)

    if self.optimize:
        A = self._optimize(X_e, A)

    self._transition_matrix = self._constrain(A, X_e)
    self._proba = self._compute_proba(X_e, self._transition_matrix)
    self._labels = torch.argmax(self._proba, dim=1)
    self._validate_simplex_structure(X_e, vertex_indices)

    self.is_fitted = True

    return self

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

Predict state labels for data.

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

RETURNS	DESCRIPTION
Tensor	State labels of shape (n_samples,).

Source code in spectre/clustering/pcca.py

def predict(self, X: torch.Tensor) -> torch.Tensor:
    """
    Predict state labels for data.

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

    Returns
    -------
    torch.Tensor
        State labels of shape (n_samples,).
    """
    self.check_fitted()
    check_2d(X)
    check_is_tensor(X)

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

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

    return torch.argmax(self._compute_proba(X_e, self._transition_matrix), dim=1)

fit_predict(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None, **kwargs) -> torch.Tensor

Fit and predict state labels.

PARAMETER	DESCRIPTION
X	Data of shape (n_samples, in_features). TYPE: Tensor
weights	Sample weights. Currently not used. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
target	Target values. Not supported. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
**kwargs	Additional keyword arguments. DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Stete labels of shape (n_samples,).

Source code in spectre/clustering/pcca.py

def fit_predict(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
    **kwargs,
) -> torch.Tensor:
    """
    Fit and predict state labels.

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

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

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

    **kwargs
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Stete labels of shape (n_samples,).
    """
    self.fit(X, weights=weights, target=target, **kwargs)

    # Use stored labels from fit to avoid re-augmenting X which may have
    # more columns than n_eigenvectors
    return self._labels

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

Predict probabilities for data.

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

RETURNS	DESCRIPTION
Tensor	Membership probabilities of shape (n_samples, n_states).

Source code in spectre/clustering/pcca.py

def predict_proba(self, X: torch.Tensor) -> torch.Tensor:
    """
    Predict probabilities for data.

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

    Returns
    -------
    torch.Tensor
        Membership probabilities of shape (n_samples, n_states).
    """
    self.check_fitted()
    check_2d(X)
    check_is_tensor(X)

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

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

    return self._compute_proba(X_e, self._transition_matrix)

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

Compute state clustering score.

PARAMETER	DESCRIPTION
X	Data of shape (n_samples, in_features). TYPE: Tensor
weights	Weights of shape (n_samples,). TYPE: torch.Tensor | None, by default None DEFAULT: None
target	Target labels of shape (n_samples,). TYPE: torch.Tensor | None, by default None DEFAULT: None

RETURNS	DESCRIPTION
float	Score of the clustering.

Source code in spectre/clustering/pcca.py

def score(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> float:
    """
    Compute state clustering score.

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

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

    target : torch.Tensor | None, by default None
        Target labels of shape (n_samples,).

    Returns
    -------
    float
        Score of the clustering.
    """
    # TODO: return proper score
    return 0.0

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

Get all fitted parameters as a dictionary.

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary containing computed results.

Source code in spectre/clustering/pcca.py

def get_results(self) -> dict[str, torch.Tensor]:
    """
    Get all fitted parameters as a dictionary.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary containing computed results.
    """
    self.check_fitted()

    return {
        "proba": self._proba,
        "labels": self._labels,
        "transition_matrix": self._transition_matrix,
    }

Functions

pcca(X: torch.Tensor, n_states: int, optimize: bool = True, n_iter: int = 100, tol: float = 1e-05, eps: float = torch.finfo(torch.float32).eps, device: torch.device | str | None = None, dtype: torch.dtype = torch.float32) -> dict[str, torch.Tensor]

Perform PCCA spectral clustering (functional interface).

PARAMETER	DESCRIPTION
X	Data of shape (n_samples, in_features). TYPE: Tensor
n_states	Number of metastable states to identify. TYPE: int
optimize	Whether to optimize transition matrix. TYPE: bool, by default True DEFAULT: True
n_iter	Maximum optimization iterations. TYPE: int, by default 100 DEFAULT: 100
tol	Convergence tolerance. TYPE: float, by default 1e-5 DEFAULT: 1e-05
eps	Numerical stability parameter. TYPE: float, by default torch.finfo(torch.float32).eps DEFAULT: eps
device	Computation device. TYPE: torch.device | str | None, by default None DEFAULT: None
dtype	Data type. TYPE: torch.dtype, by default torch.float32 DEFAULT: float32

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary with results.

Source code in spectre/clustering/pcca.py

def pcca(
    X: torch.Tensor,
    n_states: int,
    optimize: bool = True,
    n_iter: int = 100,
    tol: float = 1e-5,
    eps: float = torch.finfo(torch.float32).eps,
    device: torch.device | str | None = None,
    dtype: torch.dtype = torch.float32,
) -> dict[str, torch.Tensor]:
    """
    Perform PCCA spectral clustering (functional interface).

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

    n_states : int
        Number of metastable states to identify.

    optimize : bool, by default True
        Whether to optimize transition matrix.

    n_iter : int, by default 100
        Maximum optimization iterations.

    tol : float, by default 1e-5
        Convergence tolerance.

    eps : float, by default torch.finfo(torch.float32).eps
        Numerical stability parameter.

    device : torch.device | str | None, by default None
        Computation device.

    dtype : torch.dtype, by default torch.float32
        Data type.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary with results.
    """
    model = PCCA(
        n_states=n_states,
        optimize=optimize,
        n_iter=n_iter,
        tol=tol,
        eps=eps,
        device=device,
        dtype=dtype,
    )

    model.fit(X, weights=None, target=None)

    return model.get_results()