plt_angles.py

import numpy as np
import mdtraj as md
import pdb
import matplotlib.pyplot as plt
import seaborn as sns
from mdtraj.geometry.order import _compute_director
from matplotlib.colors import normalize

def cat_angle(A, B, C):
    a = (B[:,:,1]-A[:,:,1])*(C[:,:,2]-A[:,:,2])-(C[:,:,1]-A[:,:,1])*(B[:,:,2]-A[:,:,2])
    b = (B[:,:,2]-A[:,:,2])*(C[:,:,0]-A[:,:,0])-(C[:,:,2]-A[:,:,2])*(B[:,:,0]-A[:,:,0])
    c = (B[:,:,0]-A[:,:,0])*(C[:,:,1]-A[:,:,1])-(C[:,:,0]-A[:,:,0])*(B[:,:,1]-A[:,:,1])
    alpha = np.arccos(c/np.sqrt(a**2+b**2+c**2))*180/np.pi
    return alpha

def an_angle(v_ion):
    v_surface = np.zeros_like(v_ion)
    v_surface[:,2] = 1
    alpha = [np.arcsin(np.dot(v, (0, 0, 1))/np.linalg.norm(v))*180/np.pi for v in v_ion]
    if np.isnan(alpha).any():
        pdb.set_trace()
    return alpha
    #return np.arcsin(c/np.sqrt(a**2+b**2+c**2))*180/np.pi

def plt_angles(x, y, bins_, filename):
    H, yedges, xedges, img = plt.hist2d(y, x, bins = [bins_[1], bins_[0]])
    extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]
    fig, ax = plt.subplots(figsize=(10,4))
    #ax = fig.add_subplot(1, 1, 1)
    im = ax.imshow(H, cmap=plt.cm.get_cmap('cubehelix_r'), extent=extent)
    ax.set_aspect('auto')
    fig.colorbar(im, ax=ax)
    ax.set_xlim((0, 3))
    ax.set_xticks([0, 1, 2, 3])
    ax.set_yticks([0, 45, 90, 135, 180])
    plt.ylabel('Angle (deg)')
    plt.xlabel('Distance from surface (nm)')
    plt.savefig(filename, bbox_inches='tight')

gro_file = '/Volumes/Macintosh HD 4/rah-data/science/bi-emim-bf4/check-updated-ff/neutral/trial0/equil.gro'
trj_file = '/Volumes/Macintosh HD 4/rah-data/science/bi-emim-bf4/check-updated-ff/neutral/trial0/equil.xtc'
x_resolution = 0.01
angle_resolution = 1

sns.set_style("whitegrid")

first_frame = md.load_frame(trj_file, index=0, top=gro_file)

ions = dict()
ions['EMIM'] = first_frame.topology.select('resname EMIM')
ions['ACN'] = first_frame.topology.select('resname ACN')
atoms_per_ion = dict()
atoms_per_ion['EMIM'] = 19
atoms_per_ion['ACN'] = 6

for ion in ['EMIM']:
    first_frame = md.load_frame(trj_file, top=gro_file, index=0, atom_indices=ions[ion])
    traj = md.load(trj_file, top=gro_file, atom_indices=ions[ion]) #, atom_indices=atom_ids)
    indices = [[at.index for at in compound.atoms] for compound in list(first_frame.topology.residues)]
    com = list()
    angles = list()
    for i, ids in enumerate(indices):
        sub_traj = traj.atom_slice(ids)
        com_xyz = md.compute_center_of_mass(sub_traj)
        if np.mean(com_xyz[:,2]) < 60: #excessively large cutoff
            A = sub_traj.xyz[:,sub_traj.topology.select('name C5'),:]
            B = sub_traj.xyz[:,sub_traj.topology.select('name C6'),:]
            C = sub_traj.xyz[:,sub_traj.topology.select('name C3'),:]
            com.append(com_xyz[:,2])
            angles.append(cat_angle(A, B, C))
    com = np.array(com)
    angles = np.array(angles)

x = com[:,:].reshape(-1)
y = angles[:,:].reshape(-1)
np.savetxt('data/cat-angles.dat', np.vstack((x, y)))