# This file is part of pyGPCCA.
#
# The code and documentation of the functions below origins (with some adjustments) from MSMTools.
#
# Copyright (c) 2015, 2014 Computational Molecular Biology Group, Freie Universitaet Berlin (GER)
#
# MSMTools is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
from typing import List, Tuple, Union
from functools import singledispatch
from scipy.linalg import eig, lu_solve, lu_factor
from scipy.sparse import csgraph, spmatrix, csr_matrix, isspmatrix_csr
from scipy.sparse.linalg import eigs
import numpy as np
from pygpcca._types import ArrayLike
from pygpcca.utils._docs import d
from pygpcca.utils._checks import ensure_ndarray_or_sparse
from pygpcca.utils._constants import EPS
__all__ = [
"connected_sets",
"is_transition_matrix",
"stationary_distribution",
]
@singledispatch
def connected_sets(C: Union[ArrayLike, spmatrix], directed: bool = True) -> List[ArrayLike]:
"""
Compute connected sets of microstates.
Connected components for a directed graph with edge-weights
given by the count matrix.
Parameters
----------
C
Count matrix specifying edge weights.
directed
Whether to compute connected components for a directed or undirected graph.
Returns
-------
Each entry is an array containing all vertices (states) in the corresponding connected component. The list is sorted
according to the size of the individual components. The largest connected set is the first entry in the list.
Notes
-----
Viewing the count matrix as the adjacency matrix of a (directed) graph the connected components are given by the
connected components of that graph. Connected components of a graph can be efficiently computed using
Tarjan's algorithm [1]_.
References
----------
.. [1] Tarjan, R E. 1972. Depth-first search and linear graph
algorithms. SIAM Journal on Computing 1 (2): 146-160.
Credits
-------
The code and docstring of this function origins (with some adjustments) from MSMTools,
Copyright (c) 2015, 2014 Computational Molecular Biology Group,
Freie Universitaet Berlin (GER).
"""
raise NotImplementedError(type(C))
@connected_sets.register(np.ndarray)
def _csd(C: ArrayLike, directed: bool = True) -> List[ArrayLike]:
return connected_sets(csr_matrix(C), directed=directed)
@connected_sets.register(spmatrix)
def _css(C: spmatrix, directed: bool = True) -> List[ArrayLike]:
if not isspmatrix_csr(C):
C = csr_matrix(C)
M = C.shape[0]
# compute connected components of C. nc is the number of components,
# indices contain the component labels of the states
nc, indices = csgraph.connected_components(C, directed=directed, connection="strong")
# discrete states
states = np.arange(M)
# order indices
ind = np.argsort(indices)
indices = indices[ind]
# order states
states = states[ind]
# the state index tuple is now of the following form (states, indices)=
# ([s_23, s_17,...,s_3, s_2, ...], [0, 0, ..., 1, 1, ...])
# find number of states per component
count = np.bincount(indices)
# cumulative sum of count gives start and end indices of components
csum = np.zeros(len(count) + 1, dtype=np.uint32)
csum[1:] = np.cumsum(count)
# generate list containing components, sort each component by increasing state label
cc = [np.sort(states[csum[i] : csum[i + 1]]) for i in range(nc)]
# sort by size of component - largest component first
return sorted(cc, key=lambda x: -len(x))
@singledispatch
def is_transition_matrix(T: Union[ArrayLike, spmatrix], tol: float = 1e-12) -> bool:
r"""
Check if the given matrix is a transition matrix.
Parameters
----------
T
Matrix to check.
tol
Floating point tolerance to check with.
Returns
-------
True, if ``T`` is a valid transition matrix, false otherwise.
Notes
-----
A valid transition matrix :math:`P=(p_{ij})` has non-negative elements, :math:`p_{ij} \geq 0`, and elements of each
row sum up to one, :math:`\sum_j p_{ij} = 1`. Matrices wit this property are also called stochastic matrices.
Credits
-------
The code and docstring of this function origins (with some adjustments) from MSMTools,
Copyright (c) 2015, 2014 Computational Molecular Biology Group,
Freie Universitaet Berlin (GER).
"""
raise NotImplementedError(type(T))
@is_transition_matrix.register(spmatrix)
def _itmd(T: spmatrix, tol: float = 1e-12) -> bool:
T = ensure_ndarray_or_sparse(T, ndim=2, uniform=True, kind="numeric")
if not isspmatrix_csr(T):
T = csr_matrix(T) # compressed sparse row for fast row slicing
values = T.data # non-zero entries of T
# check entry-wise positivity
is_positive: bool = np.allclose(values, np.abs(values), rtol=tol)
# check row normalization
is_normed: bool = np.allclose(T.sum(axis=1), 1.0, rtol=tol)
return is_positive and is_normed
@is_transition_matrix.register(np.ndarray)
def _itms(T: ArrayLike, tol: float = 1e-12) -> bool:
T = ensure_ndarray_or_sparse(T, ndim=2, uniform=True, kind="numeric")
dim = T.shape[0]
X = np.abs(T) - T
x = np.sum(T, axis=1)
return X.max() < 2.0 * tol and np.abs(x - np.ones(dim)).max() < dim * tol # type: ignore[no-any-return]
[docs]@singledispatch
@d.dedent
def stationary_distribution(P: Union[ArrayLike, spmatrix]) -> ArrayLike:
r"""
Compute stationary distribution of stochastic matrix `P`.
Parameters
----------
%(P)s
Returns
-------
Vector of stationary probabilities.
Notes
-----
The stationary distribution :math:`\pi` is the left eigenvector corresponding to the non-degenerate eigenvalue
:math:`\lambda=1` of a reversible transition matrix :math:`P`,
.. math:: \pi^T P =\pi^T.
Credits
-------
The code and docstring of this function origins (with some adjustments) from MSMTools,
Copyright (c) 2015, 2014 Computational Molecular Biology Group,
Freie Universitaet Berlin (GER).
"""
raise NotImplementedError(type(P))
@stationary_distribution.register(np.ndarray)
def _sdd(P: ArrayLike) -> ArrayLike:
try:
mu = stationary_distribution_from_backward_iteration(P)
if np.any(mu < 0): # numerical problem, fall back to more robust algorithm.
raise RuntimeError("Encountered negative value.")
except RuntimeError:
mu = stationary_distribution_from_eigenvector(P)
if np.any(mu < 0): # still? Then set to 0 and renormalize
mu = np.maximum(mu, 0.0)
mu /= mu.sum()
# check whether this really is a stationary distribution
_is_stationary_distribution(P, mu)
return mu
def _eigs_slepc(P: spmatrix, k: int, which: "str" = "LR", tol: float = EPS) -> Tuple[ArrayLike, ArrayLike]:
from petsc4py import PETSc
from slepc4py import SLEPc
M = PETSc.Mat().create()
if not isspmatrix_csr(P):
P = csr_matrix(P)
M.createAIJ(size=P.shape, csr=(P.indptr, P.indices, P.data))
E = SLEPc.EPS()
E.create()
E.setOperators(M)
E.setDimensions(k)
E.setTolerances(tol=tol)
if which == "LR":
E.setWhichEigenpairs(E.Which.LARGEST_REAL)
elif which == "LM":
E.setWhichEigenpairs(E.Which.LARGEST_MAGNITUDE)
else:
raise NotImplementedError(f"`which={which}` is not implemented.")
E.solve()
nconv = E.getConverged()
if nconv < k:
raise ValueError(f"Requested `{k}` eigenvalues/vectors, but only `{nconv}` converged.")
xr, _ = M.getVecs()
xi, _ = M.getVecs()
eigenvalues, eigenvectors = [], []
for i in range(k):
# Get the i-th eigenvalue as computed by solve().
eigenvalues.append(E.getEigenpair(i, xr, xi))
if eigenvalues[-1].imag != 0.0:
eigenvectors.append([complex(r, i) for r, i in zip(xr.getArray(), xi.getArray())])
else:
eigenvectors.append(list(xr.getArray()))
return np.asarray(eigenvalues), np.asarray(eigenvectors).T
@stationary_distribution.register(spmatrix)
def _sds(P: spmatrix) -> ArrayLike:
# get the top two eigenvalues and vecs so that we can check for irreducibility
try:
vals, vecs = _eigs_slepc(P.T, k=2, which="LR")
except ImportError:
vals, vecs = eigs(P.T, k=2, which="LR", ncv=None)
# check for irreducibility
if np.allclose(vals, 1, rtol=1e2 * EPS, atol=1e2 * EPS):
second_largest = np.min(vals)
raise ValueError(f"This matrix is reducible. The second largest eigenvalue is {second_largest}.")
# sort by real part and take the top one
p = np.argsort(vals.real)[::-1]
vecs = vecs[:, p]
top_vec = vecs[:, 0]
# check for imaginary component
imaginary_component = top_vec.imag
if not np.allclose(imaginary_component, 0, rtol=EPS, atol=EPS):
raise ValueError("Top eigenvector has imaginary component.")
top_vec = top_vec.real
# check the sign structure
if not (top_vec > -1e4 * EPS).all() and not (top_vec < 1e4 * EPS).all():
el_min, el_max = np.min(top_vec), np.max(top_vec)
raise ValueError(f"Top eigenvector has both positive and negative entries. It has range = [{el_min}, {el_max}]")
top_vec = np.abs(top_vec)
pi = top_vec / np.sum(top_vec)
# check whether this really is a stationary distribution
_is_stationary_distribution(P, pi)
# normalize to 1 and return
return pi
def backward_iteration(A: ArrayLike, mu: float, x0: ArrayLike, tol: float = 1e-14, maxiter: int = 100) -> ArrayLike:
"""
Find eigenvector to approximate eigenvalue via backward iteration.
Parameters
----------
A
Matrix for which eigenvector is desired.
mu
Approximate eigenvalue for desired eigenvector.
x0
Initial guess for eigenvector.
tol
Tolerance parameter for termination of iteration.
Returns
-------
Eigenvector to approximate eigenvalue ``mu``.
Credits
-------
The code and docstring of this function origins (with some adjustments) from MSMTools,
Copyright (c) 2015, 2014 Computational Molecular Biology Group,
Freie Universitaet Berlin (GER).
"""
T = A - mu * np.eye(A.shape[0])
# LU-factor of T
lupiv = lu_factor(T)
# starting iterate with ||y_0||=1
r0 = 1.0 / np.linalg.norm(x0)
y0 = x0 * r0
# local variables for inverse iteration
y = 1.0 * y0
r = 1.0 * r0
for _ in range(maxiter):
x = lu_solve(lupiv, y)
r = 1.0 / np.linalg.norm(x)
y = x * r
if r <= tol:
return y
raise RuntimeError(f"Failed to converge after `{maxiter}` iterations, residuum is `{r}`.")
@d.dedent
def stationary_distribution_from_backward_iteration(P: ArrayLike, eps: float = 1e-15) -> ArrayLike:
"""
Fast computation of the stationary vector using backward iteration.
Parameters
----------
%(P)s
eps
Perturbation parameter for the true eigenvalue.
Returns
-------
Stationary vector.
Credits
-------
The code and docstring of this function origins (with some adjustments) from MSMTools,
Copyright (c) 2015, 2014 Computational Molecular Biology Group,
Freie Universitaet Berlin (GER).
"""
A = np.transpose(P)
mu = 1.0 - eps
x0 = np.ones(P.shape[0])
y = backward_iteration(A, mu, x0)
return y / np.sum(y)
@d.dedent
def stationary_distribution_from_eigenvector(P: ArrayLike) -> ArrayLike:
r"""
Compute stationary distribution of stochastic matrix `P`.
The stationary distribution is the left eigenvector corresponding to the
non-degenerate eigenvalue :math:`\lambda=1`.
Parameters
----------
%(P)s
Returns
-------
Vector of stationary probabilities.
Credits
-------
The code and docstring of this function origins (with some adjustments) from MSMTools,
Copyright (c) 2015, 2014 Computational Molecular Biology Group,
Freie Universitaet Berlin (GER).
"""
val, L = eig(P, left=True, right=False)
# sorted eigenvalues and left and right eigenvectors
perm = np.argsort(val)[::-1]
L = L[:, perm]
# make sure that stationary distribution is non-negative and l1-normalized
nu = np.abs(L[:, 0])
return nu / np.sum(nu)
def _is_stationary_distribution(T: Union[ArrayLike, spmatrix], pi: ArrayLike) -> bool:
# check the shapes
if not T.shape[0] == T.shape[1] or not T.shape[0] == pi.shape[0]:
raise ValueError("Shape mismatch.")
# check for invariance
if not np.allclose(T.T.dot(pi), pi, rtol=1e6 * EPS, atol=1e6 * EPS):
dev = np.max(np.abs(T.T.dot(pi) - pi))
raise ValueError(
f"Stationary distribution is not invariant under the transition matrix. Maximal deviation = " f"{dev}"
)
# check for positivity
if not (pi > -1e4 * EPS).all():
dev = np.min(pi)
raise ValueError(f"Stationary distribution has negative elements. Minimal element = {dev}")
# check whether it sums to one
if not np.allclose(pi.sum(), 1, rtol=1e4 * EPS, atol=1e4 * EPS):
dev = np.abs(pi.sum() - 1)
raise ValueError(f"Stationary distribution doe not sum to one. Deviation = {dev}.")
return True