Assorted Python function for plotting molecular properties

mdpretty.py

# Copyright 2019 João Pedro Rodrigues
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Imports required for the code to work
import collections

import matplotlib.pyplot as plt
import numpy as np
import mdtraj as md

# Assorted functions with proper documentation.

#
# == HBond Viewpoint Explorer ==
# Plots 'barcode' plots of all residues hydrogen bonding
# with a specific viewpoint residue.
#
def hbond_viewpoint(trajectory, residue, intramolecular=False):
    """Produces a 2D numpy array with hydrogen bonding data.
    
    trajectory: mdtraj Trajectory object
    residue: mdtraj Residue object
    
    Returns a ResxFrame 2D numpy binary array with info on
    when two residues are hydrogen bonding over a trajectory.
    
    Also returns a dictionary with row idx (1st Dim) and mdtraj
    Residue object, for labeling purposes.
    """
    t = trajectory
    # 1. Calculate contacts of residue (ca only)
    ca_d_max = 12.0
    atomsel = t.top.select("name CA")
    residsel = {a.index for a in residue.atoms} & set(atomsel)
    t_nl = md.compute_neighbors(trajectory, ca_d_max, 
                                haystack_indices=atomsel, query_indices=residsel)
    
    # 2a. Obtain unique set of atom ids in contact with the viewpoint
    atom_nl = {a for f in t_nl for a in f}
    
    # 2b. Expand selection to residues
    res_nl = [t.top.atom(a).residue for a in atom_nl]
    
    # 2c. Filter intramolecular if requested (default)
    if not intramolecular:
        vp_chain = residue.chain.index
        res_nl = [r for r in res_nl if r.chain.index != vp_chain]
    
    # 2d. Expand selection to all atoms of selected residues
    nl = {a for r in res_nl for a in r.atoms}

    # 2e. Generate reduced trajectory with only these atoms.
    universe_idx = {a.index for a in nl} | {a.index for a in residue.atoms}
    universe = t.atom_slice(list(universe_idx), inplace=False)

    # Create hbond matrix object
    # Make mapping between residue index and hbond row index
    mapping = {r.index: idx for idx, r in enumerate(universe.top.residues)}
    mapping_r = {idx: r for idx, r in enumerate(universe.top.residues)}
    # Shape: ResxFrame
    hbond_mtx = np.zeros((len(mapping), t.n_frames), dtype=np.int32)

    # Atom to Res dictionary
    atomres_dict = {a.index: a.residue for a in universe.top.atoms}
    
    # 3. Get hydrogen bonds between residue and these residues
    vp_idx = residue.index
    for f_idx, f in enumerate(universe):
        hb_list = md.baker_hubbard(f, freq=0.0)
        for hb in hb_list:
            d, _, a = hb
            d_res, a_res = atomres_dict.get(d), atomres_dict.get(a)
            if d_res.index == vp_idx:
                other = a_res
            elif a_res.index == vp_idx:
                other = d_res
            else:
                continue
            o_idx = mapping[other.index]
            hbond_mtx[o_idx, f_idx] += 1
    
    # Remove empty rows in the matrix
    zero_rows = {idx for idx, c in enumerate(hbond_mtx.sum(axis=1)) if c == 0}
    # Update indexes to create labeling dict (idx to residue obj)
    mapping_r = {idx: r for idx, r in mapping_r.items() if idx not in zero_rows}
    sorted_keys = sorted(mapping_r.keys())
    mapping_r = {idx: mapping_r[key] for idx, key in enumerate(sorted_keys)}
    hbond_mtx = np.delete(hbond_mtx, sorted(zero_rows), 0)

    return (mapping_r, hbond_mtx)

def plot_barcode(data, labels):
    """Creates a barcode plot from a 2D array
    
    data: 2D array (preferrably binary values)
    labels: dictionary with row idx to residue object
    """
    n_strips = len(labels)
    fig, axes = plt.subplots(nrows=n_strips, ncols=1)

    barprops = dict(aspect='auto', cmap='Reds', interpolation='nearest')
    for idx, ax in zip(labels, axes):
        data_strip = data[idx, :].reshape(1, -1)
        ax.imshow(data_strip, **barprops)

        ax.yaxis.set_major_locator(plt.NullLocator())
        ax.xaxis.set_major_locator(plt.NullLocator())   
        ax.xaxis.set_minor_locator(plt.NullLocator())   

        # Y labels
        ax.set_yticklabels([])
        r = labels[idx]
        labels_str = f'{r.chain.index}:{r.code}{r.resSeq}'
        ax.set_ylabel(labels_str, rotation="horizontal", ha='right', va='center_baseline')

        # Design
        ax.grid(False)

    # Set parameters for last ax x-axis
    ax.xaxis.set_major_locator(plt.AutoLocator())
    ax.set_xlabel('Simulation Frame')

    fig.align_ylabels()
    fig.tight_layout()