#!/usr/bin/env python3 import sys import argparse import pickle try: import numpy as np except: print("error: please install 'numpy' for python3") sys.exit(1) try: import matplotlib import matplotlib.pyplot as plt except: print("error: please install the 'matplotlib' library for python3") sys.exit(1) def rama2rgb(phi, psi): def pdist2(ref_phi, ref_psi, phi, psi): return np.min([(ref_phi -phi)**2 + (ref_psi-psi)**2 , (ref_phi-360-phi)**2 + (ref_psi-psi)**2 , (ref_phi+360-phi)**2 + (ref_psi-psi)**2 , (ref_phi -phi)**2 + (ref_psi-360-psi)**2 , (ref_phi -phi)**2 + (ref_psi+360-psi)**2 , (ref_phi-360-phi)**2 + (ref_psi-360-psi)**2 , (ref_phi+360-phi)**2 + (ref_psi-360-psi)**2 , (ref_phi-360-phi)**2 + (ref_psi+360-psi)**2 , (ref_phi+360-phi)**2 + (ref_psi+360-psi)**2]) r = pdist2(-120, 120, phi, psi) g = pdist2(-60, -60, phi, psi) b = pdist2( 60, 60, phi, psi) rescale = lambda x: 1-(np.sqrt(x) / np.sqrt(2)/180.0) return [rescale(r), rescale(g), rescale(b)] def plot_overview(export_img=None): img_dat = [] for i in range(-180,180, 1): img_dat.append([]) for j in range(-180,180, 1): img_dat[-1].append(rama2rgb(j, i)) img_dat = np.array(img_dat) ext = [-180.0, 180.0, -180.0, 180.0] plt.imshow(img_dat, extent=ext, origin="lower", interpolation="none") plt.xlabel("$\phi$") plt.ylabel("$\psi$") if export_img: plt.savefig(export_img, bbox_inches='tight') else: plt.show() def plot_state_classification(classification_tuples, export_img=None): img_data = classification_tuples[0][0] state_names = classification_tuples[0][1] n_residues = classification_tuples[0][2] for i in range(1, len(classification_tuples)): img_data_tmp, state_names_tmp, n_residues_tmp = classification_tuples[i] img_data = np.hstack((img_data, img_data_tmp)) state_names.extend(state_names_tmp) x_scaling = 3 n_states = len(state_names) plt.imshow(img_data[::-1], extent=[0, x_scaling*n_states, 0, n_residues], interpolation="none", aspect="auto") # state labels x_ticks_pos = [f+x_scaling*0.5 for f in range(0, x_scaling*n_states, x_scaling)] x_ticks_names = ["%s" % state_names[i] for i in range(0, n_states)] plt.xticks(x_ticks_pos, x_ticks_names) # residue labels #ticks_steps = 5 ticks_steps = 2 y_ticks_pos = [f+0.5 for f in range(0,n_residues, ticks_steps)] y_ticks_names = ["res. %d" % (f+1.5) for f in y_ticks_pos] plt.yticks(y_ticks_pos, y_ticks_names) if export_img: plt.savefig(export_img, bbox_inches='tight') else: plt.show() def plot_state_differences(classification_tuples, export_img=None): img_data = classification_tuples[0][0] state_names = classification_tuples[0][1] n_residues = classification_tuples[0][2] x_scaling = 3 y_scaling = 5 for i in range(1, len(classification_tuples)): img_data_tmp, state_names_tmp, n_residues_tmp = classification_tuples[i] img_data = np.hstack((img_data, img_data_tmp)) state_names.extend(state_names_tmp) n_states = len(state_names) # compute diff. img diff_img = [] for i_res in range(n_residues): diff_img.append([]) i_state1 = 0 for i_state2 in range(i_state1, n_states): if i_state1 != i_state2: diff_img[-1].append(np.abs(img_data[i_res][i_state1] - img_data[i_res][i_state2])) # state labels x_ticks_names = [] i_state1 = 0 for i_state2 in range(i_state1, n_states): if i_state1 != i_state2: x_ticks_names.append("(%s, %s)" % (state_names[i_state1], state_names[i_state2])) n_states = len(x_ticks_names) x_ticks_pos = [f+x_scaling*0.5 for f in range(0, x_scaling*n_states, x_scaling)] # residue labels ticks_steps = 5 y_ticks_pos = [f+0.5 for f in range(0,n_residues, ticks_steps)] y_ticks_names = ["res. %d" % (f+1.5) for f in y_ticks_pos] # plot it plt.imshow(diff_img[::-1], extent=[0, x_scaling*n_states, 0, n_residues], interpolation="none", aspect="auto") plt.xticks(x_ticks_pos, x_ticks_names) plt.yticks(y_ticks_pos, y_ticks_names) if export_img: plt.savefig(export_img, bbox_inches='tight') else: plt.show() def compute_state_classification(filename_dih, filename_states, selected_states=None): if filename_states: states = np.loadtxt(filename_states, dtype=int) else: n_lines = 0 for line in open(filename_dih, 'r').readlines(): n_lines += 1 states = np.ones(n_lines) if not selected_states: selected_states = np.unique(states) dih = open(filename_dih, 'r') n_res = int(len(dih.readline().split()) / 2) # reset file to first line dih.seek(0) img = {} n_state_elems = {} for state in selected_states: n_state_elems[state] = 0 img[state] = [] for i in range(n_res): img[state].append(np.array([0, 0, 0], dtype=float)) # parse dihedrals and accumulate colors for the different states for state in states: dih_frame = np.array([float(f) for f in dih.readline().split()], dtype=float) if state in selected_states: n_state_elems[state] += 1 for i_res in range(n_res): img[state][i_res] += rama2rgb(dih_frame[2*i_res], dih_frame[2*i_res+1]) # normalize color code for state in selected_states: for i_res in range(n_res): img[state][i_res] /= n_state_elems[state] # transform img to numpy array n_states = len(selected_states) img_array = [x[:] for x in [[0]*n_states]*n_res] for i_res in range(n_res): for i_state in range(n_states): img_array[i_res][i_state] = img[selected_states[i_state]][i_res] return img_array, [str(k) for k in selected_states], n_res ### if __name__ == "__main__": parser = argparse.ArgumentParser("ramacolor") parser.add_argument("selection", metavar="SELECTION", default=None, nargs="*", type=str, help="selected states.") parser.add_argument("-d", "--dihedrals", dest="dihedrals", type=str, help="ASCII file with dihedrals (in degrees!).") parser.add_argument("-s", "--states", dest="states", type=str, help="single-column ASCII file with states (i.e. clustered trajectory).") parser.add_argument("--overview", dest="overview", action="store_true", help="plot color palette overview.") parser.add_argument("--store", dest="store", type=str, help="store state classification to file.") parser.add_argument("--load", dest="load", type=str, nargs="+", help="load state classification from file.") parser.add_argument("--difference", metavar="CLASSIFICATION_FILES", dest="difference", nargs="*", help="plot difference of first against all other classifications.") parser.add_argument("--export", dest="export_img", default=None, help="export as image. (e.g. --export plot.png). supported formats: pdf, png.") args = parser.parse_args() if args.overview: plot_overview(export_img=args.export_img) elif args.dihedrals and args.states: if args.selection: selected_states = [int(s) for s in args.selection] else: print("\nerror: please provide one or more selected states!\n") parser.print_help() sys.exit(1) classification_tuple = compute_state_classification(args.dihedrals, args.states, selected_states) if args.store: pickle.dump(classification_tuple, open(args.store, 'wb')) else: plot_state_classification([classification_tuple], export_img=args.export_img) elif args.load: classifications = [] for fname in args.load: classifications.append(pickle.load(open(fname, 'rb'))) plot_state_classification(classifications, export_img=args.export_img) elif args.difference: classifications = [] for fname in args.difference: classifications.append(pickle.load(open(fname, 'rb'))) plot_state_differences(classifications, export_img=args.export_img) else: parser.print_help()