def estimate_P(C, reversible=True, fixed_statdist=None, maxiter=1000000, maxerr=1e-8, mincount_connectivity=0):
""" Estimates full transition matrix for general connectivity structure
Parameters
----------
C : ndarray
count matrix
reversible : bool
estimate reversible?
fixed_statdist : ndarray or None
estimate with given stationary distribution
maxiter : int
Maximum number of reversible iterations.
maxerr : float
Stopping criterion for reversible iteration: Will stop when infinity
norm of difference vector of two subsequent equilibrium distributions
is below maxerr.
mincount_connectivity : float
Minimum count which counts as a connection.
"""
import deeptime.markov.tools.estimation as msmest
n = np.shape(C)[0]
# output matrix. Set initially to Identity matrix in order to handle empty states
P = np.eye(n, dtype=np.float64)
# decide if we need to proceed by weakly or strongly connected sets
if reversible and fixed_statdist is None: # reversible to unknown eq. dist. - use strongly connected sets.
S = compute_connected_sets(C, connectivity_threshold=mincount_connectivity, directed=True)
for s in S:
mask = np.zeros(n, dtype=bool)
mask[s] = True
if C[np.ix_(mask, ~mask)].sum() > np.finfo(C.dtype).eps: # outgoing transitions - use partial rev algo.
transition_matrix_partial_rev(C, P, mask, maxiter=maxiter, maxerr=maxerr)
else: # closed set - use standard estimator
indices = np.ix_(mask, mask)
if s.size > 1: # leave diagonal 1 if single closed state.
P[indices] = msmest.transition_matrix(C[indices], reversible=True, warn_not_converged=False,
maxiter=maxiter, maxerr=maxerr)
else: # nonreversible or given equilibrium distribution - weakly connected sets
S = compute_connected_sets(C, connectivity_threshold=mincount_connectivity, directed=False)
for s in S:
indices = np.ix_(s, s)
if not reversible:
Csub = C[indices]
# any zero rows? must set Cii = 1 to avoid dividing by zero
zero_rows = np.where(Csub.sum(axis=1) == 0)[0]
Csub[zero_rows, zero_rows] = 1.0
P[indices] = msmest.transition_matrix(Csub, reversible=False)
elif reversible and fixed_statdist is not None:
P[indices] = msmest.transition_matrix(C[indices], reversible=True, fixed_statdist=fixed_statdist,
maxiter=maxiter, maxerr=maxerr)
else: # unknown case
raise NotImplementedError('Transition estimation for the case reversible=' + str(reversible) +
' fixed_statdist=' + str(fixed_statdist is not None) + ' not implemented.')
# done
return P
def test_dense(self):
# non-reversible
P = transition_matrix(self.C)
assert_allclose(P, self.P_nonrev)
# reversible maximum likelihood
P = transition_matrix(self.C, reversible=True)
assert_allclose(P, self.P_rev_ml)
P = transition_matrix(self.C, reversible=True, rev_pisym=False)
assert_allclose(P, self.P_rev_ml)
# reversible, pi symmetrization
P = transition_matrix(self.C, reversible=True, rev_pisym=True)
assert_allclose(P, self.P_rev_pirev)
def setUp(self):
"""Store state of the rng"""
self.state = np.random.mtrand.get_state()
"""Reseed the rng to enforce 'deterministic' behavior"""
np.random.mtrand.seed(42)
"""Meta-stable birth-death chain"""
b = 2
q = np.zeros(7)
p = np.zeros(7)
q[1:] = 0.5
p[0:-1] = 0.5
q[2] = 1.0 - 10**(-b)
q[4] = 10**(-b)
p[2] = 10**(-b)
p[4] = 1.0 - 10**(-b)
bdc = BirthDeathChain(q, p)
self.dtraj = bdc.msm.simulate(10000, start=0)
self.tau = 1
"""Estimate MSM"""
self.C_MSM = count_matrix(self.dtraj, self.tau, sliding=True)
self.lcc_MSM = largest_connected_set(self.C_MSM)
self.Ccc_MSM = largest_connected_submatrix(self.C_MSM,
lcc=self.lcc_MSM)
self.mle_rev_max_err = 1E-8
self.P_MSM = transition_matrix(self.Ccc_MSM,
reversible=True,
maxerr=self.mle_rev_max_err)
self.mu_MSM = stationary_distribution(self.P_MSM)
self.k = 3
self.ts = timescales(self.P_MSM, k=self.k, tau=self.tau)
def test_simulate_recover_transition_matrix(msm):
# test if transition matrix can be reconstructed
N = 5000
trajs = msm.simulate(N, seed=42)
# trajs = msmgen.generate_traj(self.P, N, random_state=self.random_state)
C = count_matrix(trajs, 1, sparse_return=False)
T = transition_matrix(C)
np.testing.assert_allclose(T, msm.transition_matrix, atol=.01)
def test_transition_matrix(self):
"""Non-reversible"""
T = transition_matrix(self.C1).toarray()
assert_allclose(T, self.T1.toarray())
T = transition_matrix(self.C2).toarray()
assert_allclose(T, self.T2.toarray())
"""Reversible"""
T = transition_matrix(self.C1, reversible=True).toarray()
assert_allclose(T, self.T1.toarray())
T = transition_matrix(self.C2, reversible=True).toarray()
assert_allclose(T, self.T2.toarray())
"""Reversible with fixed pi"""
T = transition_matrix(self.C1, reversible=True, mu=self.pi1).toarray()
assert_allclose(T, self.T1.toarray())
"""Reversible with fixed pi"""
T = transition_matrix(self.C1,
reversible=True,
mu=self.pi1,
method='sparse').toarray()
assert_allclose(T, self.T1.toarray())
T = transition_matrix(self.C2, reversible=True, mu=self.pi2).toarray()
assert_allclose(T, self.T2.toarray())
def from_data(dtrajs, n_hidden_states, reversible):
r""" Makes an initial guess :class:`HMM <HiddenMarkovModel>` with Gaussian output model.
To this end, a Gaussian mixture model is estimated using `scikit-learn <https://scikit-learn.org/>`_.
Parameters
----------
dtrajs : array_like or list of array_like
Trajectories which are used for making the initial guess.
n_hidden_states : int
Number of hidden states.
reversible : bool
Whether the hidden transition matrix is estimated so that it is reversible.
Returns
-------
hmm_init : HiddenMarkovModel
An initial guess for the HMM
See Also
--------
deeptime.markov.hmm.GaussianOutputModel : The type of output model this heuristic uses.
deeptime.markov.hmm.init.discrete.metastable_from_data
deeptime.markov.hmm.init.discrete.metastable_from_msm
"""
from deeptime.markov.hmm import HiddenMarkovModel, GaussianOutputModel
from sklearn.mixture import GaussianMixture
import deeptime.markov.tools.estimation as msmest
import deeptime.markov.tools.analysis as msmana
from deeptime.util.types import ensure_timeseries_data
dtrajs = ensure_timeseries_data(dtrajs)
collected_observations = np.concatenate(dtrajs)
gmm = GaussianMixture(n_components=n_hidden_states)
gmm.fit(collected_observations[:, None])
output_model = GaussianOutputModel(n_hidden_states, means=gmm.means_[:, 0], sigmas=np.sqrt(gmm.covariances_[:, 0]))
# Compute fractional state memberships.
Nij = np.zeros((n_hidden_states, n_hidden_states))
for o_t in dtrajs:
# length of trajectory
T = o_t.shape[0]
# output probability
pobs = output_model.to_state_probability_trajectory(o_t)
# normalize
pobs /= pobs.sum(axis=1)[:, None]
# Accumulate fractional transition counts from this trajectory.
for t in range(T - 1):
Nij += np.outer(pobs[t, :], pobs[t + 1, :])
# Compute transition matrix maximum likelihood estimate.
transition_matrix = msmest.transition_matrix(Nij, reversible=reversible)
initial_distribution = msmana.stationary_distribution(transition_matrix)
return HiddenMarkovModel(transition_model=transition_matrix, output_model=output_model,
initial_distribution=initial_distribution)
def test_simulate(msm):
N = 1000
traj = msm.simulate(n_steps=N, start=0, seed=42)
# test shapes and sizes
assert traj.size == N
assert traj.min() >= 0
assert traj.max() <= 1
# test statistics of transition matrix
C = count_matrix(traj, 1)
Pest = transition_matrix(C)
assert np.max(np.abs(Pest - msm.transition_matrix)) < 0.025
def test_sparse(self):
# non-reversible
P = transition_matrix(self.Cs).toarray()
assert_allclose(P, self.P_nonrev)
# non-rev return pi
P, pi = transition_matrix(self.Cs, return_statdist=True)
assert_allclose(P.T.dot(pi), pi)
P, pi = transition_matrix(self.Cs,
method='sparse',
return_statdist=True)
assert_allclose(P.T.dot(pi), pi)
# reversible maximum likelihood
P = transition_matrix(self.Cs, reversible=True).toarray()
assert_allclose(P, self.P_rev_ml)
P = transition_matrix(self.Cs, reversible=True,
rev_pisym=False).toarray()
assert_allclose(P, self.P_rev_ml)
P, pi = transition_matrix(self.Cs,
reversible=True,
rev_pisym=False,
return_statdist=True,
method='sparse')
assert_allclose(P.T.dot(pi), pi)
# reversible, pi symmetrization
P = transition_matrix(self.Cs, reversible=True,
rev_pisym=True).toarray()
assert_allclose(P, self.P_rev_pirev)
P, pi = transition_matrix(self.Cs,
reversible=True,
rev_pisym=True,
return_statdist=True)
assert_allclose(P.T.dot(pi), pi)
P, pi = transition_matrix(self.Cs,
reversible=True,
rev_pisym=True,
return_statdist=True,
method='sparse')
assert_allclose(P.T.dot(pi), pi)
def _update_transition_matrix(model: _SampleStorage,
transition_matrix_prior,
initial_distribution_prior,
reversible: bool = True,
stationary: bool = False,
n_sampling_steps: int = 1000):
""" Updates the hidden-state transition matrix and the initial distribution """
C = _bd_util.count_matrix(model.hidden_trajs, 1,
model.transition_matrix.shape[0])
model.counts[...] = C
C = C + transition_matrix_prior
# check if we work with these options
if reversible and not is_connected(C, directed=True):
raise NotImplementedError(
'Encountered disconnected count matrix with sampling option reversible:\n '
f'{C}\nUse prior to ensure connectivity or use reversible=False.'
)
# ensure consistent sparsity pattern (P0 might have additional zeros because of underflows)
# TODO: these steps work around a bug in msmtools. Should be fixed there
P0 = transition_matrix(C,
reversible=reversible,
maxiter=10000,
warn_not_converged=False)
zeros = np.where(P0 + P0.T == 0)
C[zeros] = 0
# run sampler
Tij = sample_tmatrix(C,
nsample=1,
nsteps=n_sampling_steps,
reversible=reversible)
# INITIAL DISTRIBUTION
if stationary: # p0 is consistent with P
p0 = stationary_distribution(Tij, C=C)
else:
n0 = BayesianHMM._count_init(model.hidden_trajs,
model.transition_matrix.shape[0])
first_timestep_counts_with_prior = n0 + initial_distribution_prior
positive = first_timestep_counts_with_prior > 0
p0 = np.zeros(n0.shape, dtype=model.transition_matrix.dtype)
p0[positive] = np.random.dirichlet(
first_timestep_counts_with_prior[positive]
) # sample p0 from posterior
# update HMM with new sample
model.transition_matrix[...] = Tij
model.stationary_distribution[...] = p0
def test_birth_death_chain(fixed_seed, sparse):
"""Meta-stable birth-death chain"""
b = 2
q = np.zeros(7)
p = np.zeros(7)
q[1:] = 0.5
p[0:-1] = 0.5
q[2] = 1.0 - 10**(-b)
q[4] = 10**(-b)
p[2] = 10**(-b)
p[4] = 1.0 - 10**(-b)
bdc = deeptime.data.birth_death_chain(q, p)
dtraj = bdc.msm.simulate(10000, start=0)
tau = 1
reference_count_matrix = msmest.count_matrix(dtraj, tau, sliding=True)
reference_largest_connected_component = msmest.largest_connected_set(
reference_count_matrix)
reference_lcs = msmest.largest_connected_submatrix(
reference_count_matrix, lcc=reference_largest_connected_component)
reference_msm = msmest.transition_matrix(reference_lcs,
reversible=True,
maxerr=1e-8)
reference_statdist = msmana.stationary_distribution(reference_msm)
k = 3
reference_timescales = msmana.timescales(reference_msm, k=k, tau=tau)
msm = estimate_markov_model(dtraj, tau, sparse=sparse)
assert_equal(tau, msm.count_model.lagtime)
assert_array_equal(reference_largest_connected_component,
msm.count_model.connected_sets()[0])
assert_(scipy.sparse.issparse(msm.count_model.count_matrix) == sparse)
assert_(scipy.sparse.issparse(msm.transition_matrix) == sparse)
if sparse:
count_matrix = msm.count_model.count_matrix.toarray()
transition_matrix = msm.transition_matrix.toarray()
else:
count_matrix = msm.count_model.count_matrix
transition_matrix = msm.transition_matrix
assert_array_almost_equal(reference_lcs.toarray(), count_matrix)
assert_array_almost_equal(reference_count_matrix.toarray(), count_matrix)
assert_array_almost_equal(reference_msm.toarray(), transition_matrix)
assert_array_almost_equal(reference_statdist, msm.stationary_distribution)
assert_array_almost_equal(reference_timescales[1:], msm.timescales(k - 1))
def setUpClass(cls):
n_states = 50
traj_length = 10000
dtraj = np.zeros(traj_length, dtype=int)
dtraj[::2] = np.random.randint(1,
n_states,
size=(traj_length - 1) // 2 + 1)
c = count_matrix(dtraj, lag=1)
while not is_connected(c, directed=True):
dtraj = np.zeros(traj_length, dtype=int)
dtraj[::2] = np.random.randint(1,
n_states,
size=(traj_length - 1) // 2 + 1)
c = count_matrix(dtraj, lag=1)
#state_counts = np.bincount(dtraj)[:,np.newaxis]
ttraj = np.zeros(traj_length, dtype=int)
btraj = np.zeros((traj_length, 1))
cls.tram_trajs = ([ttraj], [dtraj], [btraj])
cls.T_ref = transition_matrix(c, reversible=True).toarray()
def _estimate(self, dtrajs):
if self.E is None or self.w is None or self.m is None:
raise ValueError("E, w or m was not specified. Stopping.")
# get trajectory counts. This sets _C_full and _nstates_full
dtrajstats = self._get_dtraj_stats(dtrajs)
self._C_full = dtrajstats.count_matrix() # full count matrix
self._nstates_full = self._C_full.shape[0] # number of states
# set active set. This is at the same time a mapping from active to full
if self.connectivity == 'largest':
# statdist not given - full connectivity on all states
self.active_set = dtrajstats.largest_connected_set
else:
# for 'None' and 'all' all visited states are active
self.active_set = dtrajstats.visited_set
# FIXME: setting is_estimated before so that we can start using the parameters just set, but this is not clean!
# is estimated
self._is_estimated = True
# if active set is empty, we can't do anything.
if _np.size(self.active_set) == 0:
raise RuntimeError('Active set is empty. Cannot estimate AMM.')
# active count matrix and number of states
self._C_active = dtrajstats.count_matrix(subset=self.active_set)
self._nstates = self._C_active.shape[0]
# computed derived quantities
# back-mapping from full to lcs
self._full2active = -1 * _np.ones(dtrajstats.nstates, dtype=int)
self._full2active[self.active_set] = _np.arange(len(self.active_set))
# slice out active states from E matrix
_dset = list(set(_np.concatenate(self._dtrajs_full)))
_rras = [_dset.index(s) for s in self.active_set]
self.E_active = self.E[_rras]
if not self.sparse:
self._C_active = self._C_active.toarray()
self._C_full = self._C_full.toarray()
# reversibly counted
self._C2 = 0.5 * (self._C_active + self._C_active.T)
self._nz = _np.nonzero(self._C2)
self._csum = _np.sum(self._C_active, axis=1) # row sums C
# get ranges of Markov model expectation values
if self.support_ci == 1:
self.E_min = _np.min(self.E_active, axis=0)
self.E_max = _np.max(self.E_active, axis=0)
else:
# PyEMMA confidence interval calculation fails sometimes with conf=1.0
self.E_min, self.E_max = _ci(self.E_active, conf=self.support_ci)
# dimensions of E matrix
self.n_mstates_active, self.n_exp_active = _np.shape(self.E_active)
assert self.n_exp_active == len(self.w)
assert self.n_exp_active == len(self.m)
self.count_outside = []
self.count_inside = []
self._lls = []
i = 0
# Determine which experimental values are outside the support as defined by the Confidence interval
for emi, ema, mm, mw in zip(self.E_min, self.E_max, self.m, self.w):
if mm < emi or ema < mm:
self.logger.info(
"Experimental value %f is outside the support (%f,%f)" %
(mm, emi, ema))
self.count_outside.append(i)
else:
self.count_inside.append(i)
i = i + 1
self.logger.info(
"Total experimental constraints outside support %d of %d" %
(len(self.count_outside), len(self.E_min)))
# A number of initializations
self.P, self.pi = transition_matrix(self._C_active,
reversible=True,
return_statdist=True)
self.lagrange = _np.zeros(self.m.shape)
self._pihat = self.pi.copy()
self._update_mhat()
self._dmhat = 1e-1 * _np.ones(_np.shape(self.mhat))
# Determine number of slices of R-tensors computable at once with the given cache size
self._slicesz = _np.floor(self.max_cache /
(self.P.nbytes / 1.e6)).astype(int)
# compute first bundle of slices
self._update_Rslices(0)
self._ll_old = self._log_likelihood_biased(self._C_active, self.P,
self.m, self.mhat, self.w)
self._lls = [self._ll_old]
# make sure everything is initialized
self._update_pihat()
self._update_mhat()
self._update_Q()
self._update_X_and_pi()
self._ll_old = self._log_likelihood_biased(self._C_active, self.P,
self.m, self.mhat, self.w)
self._update_G()
#
# Main estimation algorithm
# 2-step algorithm, lagrange multipliers and pihat have different convergence criteria
# when the lagrange multipliers have converged, pihat is updated until the log-likelihood has converged (changes are smaller than 1e-3).
# These do not always converge together, but usually within a few steps of each other.
# A better heuristic for the latter may be necessary. For realistic cases (the two ubiquitin examples in [1])
# this yielded results very similar to those with more stringent convergence criteria (changes smaller than 1e-9) with convergence times
# which are seconds instead of tens of minutes.
#
converged = False # Convergence flag for lagrange multipliers
i = 0
die = False
while i <= self.maxiter:
pihat_old = self._pihat.copy()
self._update_pihat()
if not _np.all(self._pihat > 0):
self._pihat = pihat_old.copy()
die = True
self.logger.warning(
"pihat does not have a finite probability for all states, terminating"
)
self._update_mhat()
self._update_Q()
if i > 1:
X_old = self.X.copy()
self._update_X_and_pi()
if _np.any(self.X[self._nz] < 0) and i > 0:
die = True
self.logger.warning(
"Warning: new X is not proportional to C... reverting to previous step and terminating"
)
self.X = X_old.copy()
if not converged:
self._newton_lagrange()
else: # once Lagrange multipliers are converged compute likelihood here
P = self.X / self.pi[:, None]
_ll_new = self._log_likelihood_biased(self._C_active, P,
self.m, self.mhat,
self.w)
self._lls.append(_ll_new)
# General case fixed-point iteration
if len(self.count_outside) > 0:
if i > 1 and _np.all(
(_np.abs(self._dmhat) /
self.sigmas) < self.eps) and not converged:
self.logger.info(
"Converged Lagrange multipliers after %i steps..." % i)
converged = True
# Special case
else:
if _np.abs(self._lls[-2] - self._lls[-1]) < 1e-8:
self.logger.info(
"Converged Lagrange multipliers after %i steps..." % i)
converged = True
# if Lagrange multipliers are converged, check whether log-likelihood has converged
if converged and _np.abs(self._lls[-2] - self._lls[-1]) < 1e-8:
self.logger.info("Converged pihat after %i steps..." % i)
die = True
if die:
break
if i == self.maxiter:
self.logger.info("Failed to converge within %i iterations. "
"Consider increasing max_iter(now=%i)" %
(i, self.max_iter))
i += 1
_P = transition_matrix(self._C_active, reversible=True, mu=self._pihat)
self._connected_sets = connected_sets(self._C_full)
self.set_model_params(P=_P,
pi=self._pihat,
reversible=True,
dt_model=self.timestep_traj.get_scaled(self.lag))
return self
def fit(self, data, *args, **kw):
r""" Fits an AMM.
Parameters
----------
data : TransitionCountModel or (N, N) ndarray
Count matrix over data.
*args
scikit-learn compatibility argument
**kw
scikit-learn compatibility argument
Returns
-------
self : AugmentedMSMEstimator
Reference to self.
"""
if not isinstance(data, (TransitionCountModel, np.ndarray)):
raise ValueError("Can only fit on a TransitionCountModel or a count matrix directly.")
if isinstance(data, np.ndarray):
if data.ndim != 2 or data.shape[0] != data.shape[1] or np.any(data < 0.):
raise ValueError("If fitting a count matrix directly, only non-negative square matrices can be used.")
count_model = TransitionCountModel(data)
else:
count_model = data
if len(self.experimental_measurement_weights) != self.expectations_by_state.shape[1]:
raise ValueError("Experimental weights must span full observable space.")
if len(self.experimental_measurements) != self.expectations_by_state.shape[1]:
raise ValueError("Experimental measurements must span full observable state space.")
count_matrix = count_model.count_matrix
if issparse(count_matrix):
count_matrix = count_matrix.toarray()
# slice out active states from E matrix
expectations_selected = self.expectations_by_state[count_model.state_symbols]
count_matrix_symmetric = 0.5 * (count_matrix + count_matrix.T)
nonzero_counts = np.nonzero(count_matrix_symmetric)
counts_row_sums = np.sum(count_matrix, axis=1)
expectations_confidence_interval = confidence_interval(expectations_selected, conf=self.support_confidence)
measurements = self.experimental_measurements
measurement_weights = self.experimental_measurement_weights
count_outside = []
count_inside = []
i = 0
# Determine which experimental values are outside the support as defined by the Confidence interval
for confidence_lower, confidence_upper, measurement, weight in zip(
expectations_confidence_interval[0], expectations_confidence_interval[1],
measurements, measurement_weights):
if measurement < confidence_lower or confidence_upper < measurement:
self._log.info(f"Experimental value {measurement} is outside the "
f"support ({confidence_lower, confidence_upper})")
count_outside.append(i)
else:
count_inside.append(i)
i = i + 1
# A number of initializations
transition_matrix, stationary_distribution = msmest.transition_matrix(count_matrix, reversible=True,
return_statdist=True)
if issparse(transition_matrix):
transition_matrix = transition_matrix.toarray()
# Determine number of slices of R-tensors computable at once with the given cache size
slices_z = np.floor(self.max_cache / (transition_matrix.nbytes / 1.e6)).astype(int)
# Optimizer state
state = AMMOptimizerState(expectations_selected, measurements, measurement_weights,
stationary_distribution, slices_z, count_matrix_symmetric, counts_row_sums)
ll_old = state.log_likelihood_biased(count_matrix, transition_matrix)
state.log_likelihoods.append(ll_old)
# make sure everything is initialized
state.update_pi_hat()
state.update_m_hat()
state.update_Q()
state.update_X_and_pi()
ll_old = state.log_likelihood_biased(count_matrix, transition_matrix)
state.log_likelihood_prev = ll_old
state.update_G()
#
# Main estimation algorithm
# 2-step algorithm, lagrange multipliers and pihat have different convergence criteria
# when the lagrange multipliers have converged, pihat is updated until the log-likelihood has converged
# (changes are smaller than 1e-3).
# These do not always converge together, but usually within a few steps of each other.
# A better heuristic for the latter may be necessary. For realistic cases (the two ubiquitin examples in [1])
# this yielded results very similar to those with more stringent convergence criteria
# (changes smaller than 1e-9) with convergence times
# which are seconds instead of tens of minutes.
#
converged = False # Convergence flag for lagrange multipliers
i = 0
die = False
while i <= self.maxiter:
pi_hat_old = state.pi_hat.copy()
state.update_pi_hat()
if not np.all(state.pi_hat > 0):
state.pi_hat = pi_hat_old.copy()
die = True
self._log.warning("pihat does not have a finite probability for all states, terminating")
state.update_m_hat()
state.update_Q()
if i > 1:
X_old = np.copy(state.X)
state.update_X_and_pi()
if np.any(state.X[nonzero_counts] < 0) and i > 0:
die = True
self._log.warning(
"Warning: new X is not proportional to C... reverting to previous step and terminating")
state.X = X_old
if not converged:
self._newton_lagrange(state, count_matrix)
else: # once Lagrange multipliers are converged compute likelihood here
transition_matrix = state.X / state.pi[:, None]
_ll_new = state.log_likelihood_biased(count_matrix, transition_matrix)
state.log_likelihoods.append(_ll_new)
# General case fixed-point iteration
if len(count_outside) > 0:
if i > 1 and np.all((np.abs(state.delta_m_hat) / self.uncertainties) < self.convergence_criterion_lagrange)\
and not converged:
self._log.info(f"Converged Lagrange multipliers after {i} steps...")
converged = True
# Special case
else:
if np.abs(state.log_likelihoods[-2] - state.log_likelihoods[-1]) < 1e-8:
self._log.info(f"Converged Lagrange multipliers after {i} steps...")
converged = True
# if Lagrange multipliers are converged, check whether log-likelihood has converged
if converged and np.abs(state.log_likelihoods[-2] - state.log_likelihoods[-1]) < 1e-8:
self._log.info(f"Converged pihat after {i} steps...")
die = True
if die:
break
if i == self.maxiter:
ll_diff = np.abs(state.log_likelihoods[-2] - state.log_likelihoods[-1])
self._log.info(f"Failed to converge within {i} iterations. Log-likelihoods lastly changed by {ll_diff}."
f" Consider increasing max_iter(now={self.max_iter})")
i += 1
transition_matrix = msmest.transition_matrix(count_matrix, reversible=True, mu=state.pi_hat)
self._model = AugmentedMSM(transition_matrix=transition_matrix, stationary_distribution=state.pi_hat,
reversible=True, count_model=count_model, amm_optimizer_state=state)
return self
def _fit_connected(self, counts):
from .. import _transition_matrix as tmat
if isinstance(counts, np.ndarray):
if not is_square_matrix(counts) or np.any(counts < 0.):
raise ValueError(
"If fitting a count matrix directly, only non-negative square matrices can be used."
)
count_model = TransitionCountModel(counts)
elif isinstance(counts, TransitionCountModel):
count_model = counts
else:
raise ValueError(
f"Unknown type of counts {counts}, only n x n ndarray, TransitionCountModel,"
f" or TransitionCountEstimators with a count model are supported."
)
if self.stationary_distribution_constraint is not None:
if len(self.stationary_distribution_constraint
) != count_model.n_states_full:
raise ValueError(
f"Stationary distribution constraint must be defined over full "
f"set of states ({count_model.n_states_full}), but contained "
f"{len(self.stationary_distribution_constraint)} elements."
)
if np.any(self.stationary_distribution_constraint[
count_model.state_symbols]) == 0.:
raise ValueError(
"The count matrix contains symbols that have no probability in the stationary "
"distribution constraint.")
if count_model.count_matrix.sum() == 0.0:
raise ValueError(
"The set of states with positive stationary probabilities is not visited by the "
"trajectories. A MarkovStateModel reversible with respect to the given stationary"
" vector can not be estimated")
count_matrix = count_model.count_matrix
# continue sparse or dense?
if not self.sparse and issparse(count_matrix):
# converting count matrices to arrays. As a result the
# transition matrix and all subsequent properties will be
# computed using dense arrays and dense matrix algebra.
count_matrix = count_matrix.toarray()
# restrict stationary distribution to active set
if self.stationary_distribution_constraint is None:
statdist = None
else:
statdist = self.stationary_distribution_constraint[
count_model.state_symbols]
statdist /= statdist.sum() # renormalize
# Estimate transition matrix
if self.allow_disconnected:
P = tmat.estimate_P(count_matrix,
reversible=self.reversible,
fixed_statdist=statdist,
maxiter=self.maxiter,
maxerr=self.maxerr)
else:
opt_args = {}
# TODO: non-rev estimate of msmtools does not comply with its own api...
if statdist is None and self.reversible:
opt_args['return_statdist'] = True
P = msmest.transition_matrix(count_matrix,
reversible=self.reversible,
mu=statdist,
maxiter=self.maxiter,
maxerr=self.maxerr,
**opt_args)
# msmtools returns a tuple for statdist_active=None.
if isinstance(P, tuple):
P, statdist = P
if statdist is None and self.allow_disconnected:
statdist = tmat.stationary_distribution(P, C=count_matrix)
return (P, statdist, counts)
def cktest_resource():
"""Reseed the rng to enforce 'deterministic' behavior"""
rnd_state = np.random.mtrand.get_state()
np.random.mtrand.seed(42)
"""Meta-stable birth-death chain"""
b = 2
q = np.zeros(7)
p = np.zeros(7)
q[1:] = 0.5
p[0:-1] = 0.5
q[2] = 1.0 - 10**(-b)
q[4] = 10**(-b)
p[2] = 10**(-b)
p[4] = 1.0 - 10**(-b)
bdc = BirthDeathChain(q, p)
dtraj = bdc.msm.simulate(10000, start=0)
tau = 1
"""Estimate MSM"""
MSM = estimate_markov_model(dtraj, tau)
P_MSM = MSM.transition_matrix
mu_MSM = MSM.stationary_distribution
"""Meta-stable sets"""
A = [0, 1, 2]
B = [4, 5, 6]
w_MSM = np.zeros((2, mu_MSM.shape[0]))
w_MSM[0, A] = mu_MSM[A] / mu_MSM[A].sum()
w_MSM[1, B] = mu_MSM[B] / mu_MSM[B].sum()
K = 10
P_MSM_dense = P_MSM
p_MSM = np.zeros((K, 2))
w_MSM_k = 1.0 * w_MSM
for k in range(1, K):
w_MSM_k = np.dot(w_MSM_k, P_MSM_dense)
p_MSM[k, 0] = w_MSM_k[0, A].sum()
p_MSM[k, 1] = w_MSM_k[1, B].sum()
"""Assume that sets are equal, A(\tau)=A(k \tau) for all k"""
w_MD = 1.0 * w_MSM
p_MD = np.zeros((K, 2))
eps_MD = np.zeros((K, 2))
p_MSM[0, :] = 1.0
p_MD[0, :] = 1.0
eps_MD[0, :] = 0.0
for k in range(1, K):
"""Build MSM at lagtime k*tau"""
C_MD = count_matrix(dtraj, k * tau, sliding=True) / (k * tau)
lcc_MD = largest_connected_set(C_MD)
Ccc_MD = largest_connected_submatrix(C_MD, lcc=lcc_MD)
c_MD = Ccc_MD.sum(axis=1)
P_MD = transition_matrix(Ccc_MD).toarray()
w_MD_k = np.dot(w_MD, P_MD)
"""Set A"""
prob_MD = w_MD_k[0, A].sum()
c = c_MD[A].sum()
p_MD[k, 0] = prob_MD
eps_MD[k, 0] = np.sqrt(k * (prob_MD - prob_MD**2) / c)
"""Set B"""
prob_MD = w_MD_k[1, B].sum()
c = c_MD[B].sum()
p_MD[k, 1] = prob_MD
eps_MD[k, 1] = np.sqrt(k * (prob_MD - prob_MD**2) / c)
"""Input"""
yield MSM, p_MSM, p_MD
np.random.mtrand.set_state(rnd_state)
def _estimate(self, dtrajs):
""" Estimates the MSM """
# get trajectory counts. This sets _C_full and _nstates_full
dtrajstats = self._get_dtraj_stats(dtrajs)
self._C_full = dtrajstats.count_matrix() # full count matrix
self._nstates_full = self._C_full.shape[0] # number of states
# check for consistency between statdist constraints and core set
if self.core_set is not None and self.statdist_constraint is not None:
if len(self.core_set) != len(self.statdist_constraint):
raise ValueError(
'Number of core sets and stationary distribution '
'constraints do not match.')
# rewrite statdist constraints to full set for compatibility reasons
#TODO: find a more consistent way of dealing with this
import copy
_stdist_constr_coreset = copy.deepcopy(self.statdist_constraint)
self.statdist_constraint = _np.zeros(self._nstates_full)
self.statdist_constraint[self.core_set] = _stdist_constr_coreset
# set active set. This is at the same time a mapping from active to full
if self.connectivity == 'largest':
if self.statdist_constraint is None:
# statdist not given - full connectivity on all states
self.active_set = dtrajstats.largest_connected_set
else:
active_set = self._prepare_input_revpi(
self._C_full, self.statdist_constraint)
self.active_set = active_set
else:
# for 'None' and 'all' all visited states are active
self.active_set = dtrajstats.visited_set
# FIXME: setting is_estimated before so that we can start using the parameters just set, but this is not clean!
# is estimated
self._is_estimated = True
# if active set is empty, we can't do anything.
if _np.size(self.active_set) == 0:
raise RuntimeError('Active set is empty. Cannot estimate MSM.')
# active count matrix and number of states
self._C_active = dtrajstats.count_matrix(subset=self.active_set)
# continue sparse or dense?
if not self.sparse:
# converting count matrices to arrays. As a result the
# transition matrix and all subsequent properties will be
# computed using dense arrays and dense matrix algebra.
self._C_full = self._C_full.toarray()
self._C_active = self._C_active.toarray()
self._nstates = self._C_active.shape[0]
# computed derived quantities
# back-mapping from full to lcs
self._full2active = -1 * _np.ones(dtrajstats.nstates, dtype=int)
self._full2active[self.active_set] = _np.arange(len(self.active_set))
# restrict stationary distribution to active set
if self.statdist_constraint is None:
statdist_active = None
else:
statdist_active = self.statdist_constraint[self.active_set]
statdist_active /= statdist_active.sum() # renormalize
opt_args = {}
# TODO: non-rev estimate of msmtools does not comply with its own api...
if statdist_active is None and self.reversible:
opt_args['return_statdist'] = True
# Estimate transition matrix
if self.connectivity == 'largest':
P = transition_matrix(self._C_active,
reversible=self.reversible,
mu=statdist_active,
maxiter=self.maxiter,
maxerr=self.maxerr,
**opt_args)
elif self.connectivity == 'none':
# reversible mode only possible if active set is connected
# - in this case all visited states are connected and thus
# this mode is identical to 'largest'
if self.reversible and not is_connected(self._C_active):
raise ValueError(
'Reversible MSM estimation is not possible with connectivity mode "none", '
'because the set of all visited states is not reversibly connected'
)
P = transition_matrix(self._C_active,
reversible=self.reversible,
mu=statdist_active,
maxiter=self.maxiter,
maxerr=self.maxerr,
**opt_args)
else:
raise NotImplementedError(
'MSM estimation with connectivity=%s is currently not implemented.'
% self.connectivity)
# msmtools returns a tuple for statdist_active = None.
if isinstance(P, tuple):
P, statdist_active = P
# Done. We set our own model parameters, so this estimator is
# equal to the estimated model.
self._connected_sets = dtrajstats.connected_sets
self.set_model_params(P=P,
pi=statdist_active,
reversible=self.reversible,
dt_model=self.timestep_traj.get_scaled(self.lag))
return self
