def plot_tica_landscape(tica_trajs,
ax=None,
figsize=(7, 5),
obs=(0, 1),
cmap='magma'):
"""
Plots a 2D landscape of the tica trajectories
:param tica_trajs: dictionary of tica trajectories
:param ax: mpl.axies
:param figsize: tuple
:param obs: tuple
:param cmap: str
:return: f, ax (figure and axis)
"""
if ax is None:
f, ax = pp.subplots(figsize=figsize)
txx = numpy.concatenate(list(tica_trajs.values()))
ax = msme.plot_free_energy(txx,
ax=ax,
obs=obs,
alpha=1,
n_levels=6,
xlabel='tIC 1',
ylabel='tIC 2',
labelsize=14,
cmap=cmap,
vmin=1e-25)
return f, ax
def plot_overlayed_types(ttrajs,
meta,
obs=(0, 1),
ax=None,
stride=100,
xlabel=None,
ylabel=None,
plot_free_energy_kwargs=None,
plot_kwargs=None):
"""
Overlay each type of system inside the meta object onto the overall tICA
free energy landscape.
:param ttrajs: dict, with keys as integers and values of np.arrays (tICA converted trajectories)
:param meta: an msmbuilder metadata object
:param obs: tuple (optional), the dimensions to plot
:param ax: matplotlib axis to plot in (optional)
:param stride: int (optional, default=100), scatter every `stride` frames to avoid a cluttered plot
:param xlabel: str (optional), the x label
:param ylabel: str (optional), the y label
:param plot_free_energy_kwargs: dict (optional), extra parameters to pass to msme.plot_free_energy
:param plot_kwargs: dict (optional), extra parameters to pass to pyplot.plot
:return ax: drawn matplotlib axis
"""
if ax is None:
ax = pp.gca()
if plot_free_energy_kwargs is None:
plot_free_energy_kwargs = {}
if plot_kwargs is None:
plot_kwargs = {}
txx = np.concatenate(list(ttrajs.values()))
ttrajs_subtypes = split_trajs_by_type(ttrajs, meta)
msme.plot_free_energy(txx, obs=obs, ax=ax, **plot_free_energy_kwargs)
for traj_id, traj_dict in ttrajs_subtypes.items():
system_txx = np.concatenate(list(traj_dict.values()))
ax.scatter(system_txx[::stride, obs[0]],
system_txx[::stride, obs[1]],
label=traj_id,
**plot_kwargs)
pp.legend(loc='best')
if xlabel is not None:
ax.set_xlabel(xlabel)
if ylabel is not None:
ax.set_ylabel(ylabel)
return ax
def draw_MacroCluster_FreeEnergy(self, txx, clusterer, pcca, assignments):
cm = plt.cm.get_cmap('RdYlBu')
msme.plot_free_energy(txx, obs=(0, 1), n_samples=10000, vmin=-0.001,
pi=pcca.populations_[assignments],
xlabel='tIC 1', ylabel='tIC 2')
sc =plt.scatter(clusterer.cluster_centers_[pcca.state_labels_, 0],
clusterer.cluster_centers_[pcca.state_labels_, 1],
s=300,
c=self.microstate_mapping_,
zorder=3,
cmap=cm
)
cbar = plt.colorbar(sc)
pop_percent = self.macro_pop/self.macro_pop.sum()*100
cluster_index = np.array(list(range(self.macro_num)))
cbar.set_ticks(cluster_index)
cbar.set_ticklabels(['{0} {1: .2f}%'.format(i[0]+1, i[1]) for i in zip(cluster_index, pop_percent)])
plt.tight_layout()
def draw_MicroCluster_FreeEnergy(txx, msm, clusterer, assignments):
ax=msme.plot_free_energy(txx, obs=(0, 1), n_samples=10000,
pi=msm.populations_[assignments],
cmap='bone', alpha=0.5,vmin=-.001,
xlabel='tIC 1', ylabel='tIC 2')
plt.scatter(clusterer.cluster_centers_[msm.state_labels_, 0],
clusterer.cluster_centers_[msm.state_labels_, 1],
s=1e4 * msm.populations_, # size by population
c=msm.left_eigenvectors_[:, 1], # color by eigenvector
cmap="coolwarm",
zorder=3)
plt.colorbar(label='First dynamical eigenvector')
plt.tight_layout()
def plot_com_matrix(com_matrix):
"""
Plots the projections of X, Y and Z of a center of mass position matrix
Parameters
----------
com_matrix: np.array of shape (frames, 3)
Returns
-------
fig: plt.figure
"""
fig, axes = plt.subplots(3, 3, figsize=(15, 15))
correspondance = {0: 'x', 1: 'y', 2: 'z'}
# axes is a np.array of shape (3, 3)
for i in range(3):
for j in range(3):
ax = axes[i][j]
if i >= j:
# we'll plot only the lower half
if i == j:
sns.kdeplot(com_matrix[:, i], ax=ax)
ax.set(xlabel='%s (nm)' % correspondance[i],
ylabel='Density')
else:
ax.set(xlabel='%s (nm)' % correspondance[j],
ylabel='%s (nm)' % correspondance[i])
msme.plot_free_energy(com_matrix,
obs=(j, i),
ax=ax,
shade=True,
n_levels=5,
vmin=-1e-12,
cmap='viridis')
else:
ax.set_axis_off()
return fig
def plot_cluster_centers(clusterer,
centers,
txx,
surface=True,
ax=None,
obs=(0, 1),
from_clusterer=True,
msm=None,
add_bit=20,
scatter_kwargs=None):
"""
Plots cluster centers as scatter points on top of a tICA landscape.
Center IDs can be either in MSM labeling
or directly from the clustering labelling.
Parameters
----------
clusterer: a fit clusterer object, with .cluster_centers_ attribute
centers: list of ints, the IDs of the cluster centers to plot
txx: np.array of concatenated tIC trajs, shape = (n_frames, n_features)
surface: bool, Control if energy surface is plotted or not
ax: a matplotlib axes object (optional)
obs: tuple of ints, dimensions to plot (optional)
from_clusterer: bool, are the centers id in clusterer
indexing or msm?
default True means IDs are from clusterer
msm: a MarkovStateModel object which has been fit
add_bit: int, control size of scatter plots
scatter_kwargs: dict, Extra parameters to pass to scatter plot
Returns
-------
ax: a matplotlib axes object
Raises
------
ValueError: if we select from_clusterer=False it means that
the center IDs are in MSM-internal labeling,
so we need an MSM object to map those back to the clusterer naming
"""
if ax is None:
ax = pp.gca()
if scatter_kwargs is None:
scatter_kwargs = {}
if surface:
ax = msme.plot_free_energy(txx,
obs=obs,
n_samples=5000,
gridsize=100,
vmax=5.,
n_levels=8,
cut=5,
xlabel='tIC1',
ylabel='tIC2')
prune = clusterer.cluster_centers_[:, obs]
if from_clusterer:
chosen_centers = prune[centers]
ax.scatter(chosen_centers[:, 0], chosen_centers[:, 1])
else:
if msm is None:
raise ValueError(
'if from_clusterer is False please provide a fit MSM in the msm parameter'
)
else:
# Retrieve cluster centers from clusterer objects that have been used in this MSM
centers_in_clusterer = []
for k, v in msm.mapping_.items():
if v in centers:
centers_in_clusterer.append(k)
scale = 100 / np.max(msm.populations_)
chosen_centers = prune[centers_in_clusterer]
ax.scatter(chosen_centers[:, 0],
chosen_centers[:, 1],
s=add_bit + (scale * msm.populations_),
**scatter_kwargs)
return ax
import msmexplorer as msme
rs = np.random.RandomState(42)
# Load Fs Peptide Data
trajs = MullerPotential().get().trajectories
# Perform Dimensionality Reduction
tica_model = tICA(lag_time=2, n_components=2)
tica_trajs = tica_model.fit_transform(trajs)
# Perform Clustering
clusterer = MiniBatchKMeans(n_clusters=12, random_state=rs)
clustered_trajs = clusterer.fit_transform(tica_trajs)
# Construct MSM
msm = MarkovStateModel(lag_time=2)
assignments = msm.fit_transform(clustered_trajs)
# Plot Free Energy
data = np.concatenate(trajs, axis=0)
pi_0 = msm.populations_[np.concatenate(assignments, axis=0)]
ax = msme.plot_free_energy(data,
obs=(0, 1),
n_samples=100000,
pi=pi_0,
random_state=rs)
# Plot tICA Projection
msme.plot_decomp_grid(tica_model, ax=ax, alpha=0.6)
# Load Fs Peptide Data
trajs = FsPeptide().get().trajectories
# Extract Backbone Dihedrals
featurizer = DihedralFeaturizer(types=['phi', 'psi'])
diheds = featurizer.fit_transform(trajs)
# Perform Dimensionality Reduction
tica_model = tICA(lag_time=2, n_components=2)
tica_trajs = tica_model.fit_transform(diheds)
# Perform Clustering
clusterer = MiniBatchKMeans(n_clusters=12, random_state=rs)
clustered_trajs = clusterer.fit_transform(tica_trajs)
# Construct MSM
msm = MarkovStateModel(lag_time=2)
assignments = msm.fit_transform(clustered_trajs)
# Plot Free Energy
data = np.concatenate(tica_trajs, axis=0)
pi_0 = msm.populations_[np.concatenate(assignments, axis=0)]
msme.plot_free_energy(data, obs=(0, 1), n_samples=100000, pi=pi_0,
random_state=rs,
shade=True,
clabel=True,
clabel_kwargs={'fmt': '%.1f'},
cbar=True,
cbar_kwargs={'format': '%.1f', 'label': 'Free energy (kcal/mol)'}
)
import numpy as np
import msmexplorer as msme
rs = np.random.RandomState(42)
# Load Fs Peptide Data
trajs = FsPeptide().get().trajectories
# Extract Backbone Dihedrals
featurizer = DihedralFeaturizer(types=['phi', 'psi'])
diheds = featurizer.fit_transform(trajs)
# Perform Dimensionality Reduction
tica_model = tICA(lag_time=2, n_components=2)
tica_trajs = tica_model.fit_transform(diheds)
# Perform Clustering
clusterer = MiniBatchKMeans(n_clusters=12, random_state=rs)
clustered_trajs = clusterer.fit_transform(tica_trajs)
# Construct MSM
msm = MarkovStateModel(lag_time=2)
assignments = msm.fit_transform(clustered_trajs)
# Plot Free Energy
data = np.concatenate(tica_trajs, axis=0)
pi_0 = msm.populations_[np.concatenate(assignments, axis=0)]
msme.plot_free_energy(data, n_samples=100000, pi=pi_0,
gridsize=100, random_state=rs)
rs = np.random.RandomState(42)
# Load Fs Peptide Data
trajs = FsPeptide().get().trajectories
# Extract Backbone Dihedrals
featurizer = DihedralFeaturizer(types=['phi', 'psi'])
diheds = featurizer.fit_transform(trajs)
# Perform Dimensionality Reduction
tica_model = tICA(lag_time=2, n_components=2)
tica_trajs = tica_model.fit_transform(diheds)
# Perform Clustering
clusterer = MiniBatchKMeans(n_clusters=12, random_state=rs)
clustered_trajs = clusterer.fit_transform(tica_trajs)
# Construct MSM
msm = MarkovStateModel(lag_time=2)
assignments = msm.fit_transform(clustered_trajs)
# Plot Free Energy
data = np.concatenate(tica_trajs, axis=0)
pi_0 = msm.populations_[np.concatenate(assignments, axis=0)]
msme.plot_free_energy(data,
n_samples=100000,
pi=pi_0,
gridsize=100,
random_state=rs)
def plot_tic_array(tica_trajs,
nrows,
ncols,
ts=0.2,
subplot_kwargs={},
trace_kwargs={},
free_energy_kwargs={}):
'''
Plot a nrows x cols array of the tIC projections, starting from the 1st one.
Parameters
----------
tica_trajs: list, or np.ndarray
The tica transformed trajectory(ies)
nrows: int
Number of rows
ncols: int
Number of cols
ts: float or int (default: 0.2)
Timestep (in microseconds) between each frame in the trajectory
subplot_kwargs: dict, optional
Arguments to pass to plt.subplots
trace_kwargs: dict, optional
Arguments to pass to msme.plot_trace
free_energy_kwargs: dict, optional
Arguments to pass to msme.plot_free_energy
Returns
-------
ax : matplotlib axis
matplotlib figure axis
'''
def convert_to_mus(x, pos):
'function for formatting the x axis of time trace plot'
return x * ts
if isinstance(tica_trajs, list):
tica_trajs = np.concatenate(tica_trajs)
elif not isinstance(tica_trajs, np.ndarray):
raise ValuError('tica_trajs must be of type list or np.ndarray')
fig, ax_list = plt.subplots(nrows=nrows, ncols=ncols, **subplot_kwargs)
for i in range(nrows):
for j in range(ncols):
if (nrows == 1) and (ncols == 1):
# User just passed a 1x1 array...
ax = ax_list
elif nrows == 1:
ax = ax_list[j]
elif ncols == 1:
ax = ax_list[i]
else:
ax = ax_list[i][j]
ax.grid(False)
if i == j:
msme.plot_trace(tica_trajs[:, i], ax=ax, **trace_kwargs)
formatter = FuncFormatter(convert_to_mus)
ax.xaxis.set_major_formatter(formatter)
ax.grid(False, axis='y')
ax.set_xlabel('Time ($\mu$s)')
else:
msme.plot_free_energy(tica_trajs,
obs=(j, i),
ax=ax,
**free_energy_kwargs)
# Bottom row
if i == (nrows - 1):
ax.set_xlabel('tIC{}'.format(j + 1))
# First column
if j == 0:
ax.set_ylabel('tIC{}'.format(i + 1))
fig.tight_layout()
return fig
import msmexplorer as msme
from msmexplorer.example_datasets import FsPeptide
rs = np.random.RandomState(42)
# Load Fs Peptide Data
trajs = FsPeptide().get().trajectories
# Extract Backbone Dihedrals
featurizer = DihedralFeaturizer(types=['phi', 'psi'])
diheds = featurizer.fit_transform(trajs)
# Perform Dimensionality Reduction
tica_model = tICA(lag_time=2, n_components=2)
tica_trajs = tica_model.fit_transform(diheds)
# Perform Clustering
clusterer = MiniBatchKMeans(n_clusters=12, random_state=rs)
clustered_trajs = clusterer.fit_transform(tica_trajs)
# Construct MSM
msm = MarkovStateModel(lag_time=2)
assignments = msm.fit_transform(clustered_trajs)
# Plot Free Energy
data = np.concatenate(tica_trajs, axis=0)
pi_0 = msm.populations_[np.concatenate(assignments, axis=0)]
msme.plot_free_energy(data, obs=(0, 1), n_samples=100000, pi=pi_0,
random_state=rs)