Exemplo n.º 1
0
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
Exemplo n.º 2
0
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
Exemplo n.º 3
0
 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()
Exemplo n.º 4
0
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()
Exemplo n.º 5
0
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
Exemplo n.º 6
0
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
Exemplo n.º 7
0
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)
Exemplo n.º 8
0
# 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)
Exemplo n.º 10
0
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)
Exemplo n.º 11
0
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
Exemplo n.º 12
0
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)