Skip to content
Snippets Groups Projects
Commit 8c9ae834 authored by Javier González-Delgado's avatar Javier González-Delgado
Browse files

Update utils.ipynb

parent 59fdfd3a
No related branches found
No related tags found
No related merge requests found
%% Cell type:code id:f76b47d6 tags:
``` python
import mdtraj as md
import os
import numpy as np
from ipywidgets import widgets
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import itertools
from tqdm import tqdm
import ot
import MDAnalysis
import h5py
```
%% Cell type:code id:1219123d-f44f-4e08-9e71-4dc45cd81ced tags:
``` python
def multiframe_pdb_to_xtc(pdb_file, save_path, prot_name):
u = MDAnalysis.core.universe.Universe(pdb_file)
at = u.atoms
os.chdir(save_path)
# Write the trajectory in .xtc format
at.write(".".join([prot_name,'xtc']), frames='all')
# Write a frame of the trajectory in .pdb format for topology information
at.write(".".join([prot_name,'pdb']))
```
%% Cell type:code id:6f669fe3 tags:
``` python
def get_cluster_files(ensemble_path, ensemble_name, labels_umap):
# Initial parameters
var_dict = {'multiframe' : 'n', 'check_folder' : True, 'do_xtc' : False, 'do_pdb' : False,
'ensemble_name' : ensemble_name, 'ensemble_path' : ensemble_path}
var_dict['xtc_files'] = [file for file in os.listdir(ensemble_path) if file.endswith(".xtc")]
var_dict['pdb_files'] = [file for file in os.listdir(ensemble_path) if file.endswith(".pdb") or file.endswith(".prmtop") or file.endswith(".parm7") or file.endswith(".gro")]
var_dict['folders'] = [file for file in os.listdir(ensemble_path) if (os.path.isdir("/".join([ensemble_path,file])) and not file.startswith('.') and not file.startswith("results"))]
# File processing
if len(var_dict["xtc_files"]) + len(var_dict["folders"]) + len(var_dict["pdb_files"]) == 0:
sys.exit("".join(['Folder for ', var_dict["ensemble_name"], ' ensemble is empty...']))
# .xtc file with a .pdb topology file
if len(var_dict["xtc_files"]) >= len(var_dict["pdb_files"]) and len(var_dict["pdb_files"]) == 1:
print('\nTaking as input:\n')
print("".join([str(var_dict["xtc_files"][0]),' : trajectory of ',var_dict["ensemble_name"],',']))
print("".join([str(var_dict["pdb_files"][0]),' : topology file of ',var_dict["ensemble_name"],'.']))
if len(var_dict["xtc_files"]) > 1:
print("\nMore than one .xtc file were found. Taking the first as the trajectory file.\n")
var_dict["do_xtc"] = True
var_dict["xtc_root_path"] = var_dict["ensemble_path"]
var_dict['check_folder'] = False
# multiframe .pdb files
if var_dict['multiframe'] == 'y' or (len(var_dict["pdb_files"]) >= 1 and len(var_dict["xtc_files"]) == 0):
print('\nTaking as input:\n')
print("".join([str(var_dict["pdb_files"][0]),' : trajectory of ',var_dict["ensemble_name"],'.']))
if len(var_dict["pdb_files"]) > 1:
print("\nMore than one multiframe .pdb file were found. Taking the first as the trajectory file.\n")
print("\nTaking the previously converted files.\n")
var_dict["do_xtc"] = True
var_dict["xtc_root_path"] = "/".join([var_dict["ensemble_path"],'converted_files'])
var_dict["xtc_files"] = [file for file in os.listdir(var_dict["xtc_root_path"]) if file.endswith(".xtc")]
var_dict["pdb_files"] = [file for file in os.listdir(var_dict["xtc_root_path"]) if file.endswith(".pdb")]
var_dict['check_folder'] = False
# folder with .pdb files
if len(var_dict["folders"]) >= 1 and var_dict['check_folder'] == True:
print('\nTaking as input:\n')
print("".join([var_dict["folders"][0],' folder contains: trajectory of ',var_dict["ensemble_name"],"."]))
if len(var_dict["folders"]) > 1:
print("\nMore than one .pdb folder were found. Taking the first as the trajectory folder.\n")
var_dict["do_pdb"] = True
if not var_dict["do_pdb"] and not var_dict["do_xtc"]:
sys.exit("".join(['\n Sorry, I did not understood the input. Please follow the guidelines described in the function documentation to create ',ensemble_name,' folder.\n']))
print("\n----------------------------------------------------------------------------------\n")
print("\nCreating cluster-specific files...\n")
results_path = "/".join([os.path.abspath(ensemble_path),"_".join(['results',ensemble_name])])
save_files = "/".join([results_path, "cluster_files"])
if not os.path.exists(save_files):
os.mkdir(save_files)
if var_dict["do_xtc"]:
traj_file = md.load_xtc("/".join([var_dict["xtc_root_path"],var_dict["xtc_files"][0]]), top = "/".join([var_dict["xtc_root_path"],var_dict["pdb_files"][0]]))
# Save .xtc cluster files
for k in tqdm(range(len(np.unique(labels_umap[labels_umap >= 0])))):
traj_file[np.where(labels_umap == k)].save_xtc("/".join([save_files, "".join([ensemble_name,'_',str(k),'.xtc'])]))
if var_dict["do_pdb"]:
conf_list = os.listdir("/".join([var_dict["ensemble_path"],var_dict["folders"][0]]))
for k in tqdm(range(len(np.unique(labels_umap[labels_umap >= 0])))):
clus_k_path = "/".join([save_files, "_".join(['clus',str(k)])])
if not os.path.exists(clus_k_path):
os.mkdir(clus_k_path)
clus_k = np.where(labels_umap == k)[0]
for j in range(len(clus_k)):
traj = md.load_pdb("/".join(["/".join([var_dict["ensemble_path"],var_dict["folders"][0]]),conf_list[clus_k[j]]]))
traj.save_pdb("/".join([clus_k_path, "".join([ensemble_name,'_',str(clus_k[j]),'.pdb'])]))
traj.save_pdb("/".join([clus_k_path, "".join([ensemble_name,'_',str(conf_list[clus_k[j]]),'.pdb'])]))
print("\nFiles saved.\n")
```
%% Cell type:code id:ff9c4941 tags:
``` python
def plot_2umap(embedding_2d, labels_umap, ensemble_name, results_path, pdf = False, dpi_png = 200):
classified = np.where(labels_umap >= 0)[0]
output1 = widgets.Output()
with output1:
fig, ax = plt.subplots()
ax.scatter(embedding_2d[~classified, 0],
embedding_2d[~classified, 1],
color=(0.5, 0.5, 0.5),
s=0.5,
alpha=0.5)
scatter = ax.scatter(embedding_2d[classified, 0],
embedding_2d[classified, 1],
c=labels_umap[classified],
s=0.5,
alpha = 1,
cmap='Spectral')
plt.xlabel('UMAP coordinate 1')
plt.ylabel('UMAP coordinate 2')
plt.title("".join(['UMAP 2-dimensional projection after contact clustering for ',ensemble_name,' ensemble']), fontsize = 8)
if pdf:
plt.savefig("/".join([results_path, "".join(["clusters_2d", ensemble_name, '.pdf'])]))
else:
plt.savefig("/".join([results_path, "".join(["clusters_2d", ensemble_name, '.png'])]), dpi = dpi_png)
plt.show()
output2 = widgets.Output()
with output2:
repartition = pd.Series(labels_umap).value_counts()
repartition.index = ["Unclassified" if i == -1 else i for i in repartition.index]
display(pd.DataFrame({"Cluster" : np.array(repartition.index), "Occupancy (%)" : 100*np.array(repartition.values)/len(labels_umap)}))
two_columns = widgets.HBox([output1, output2])
display(two_columns)
```
%% Cell type:code id:00dfbe24 tags:
``` python
def get_wmaps(labels_umap, ensemble_name, results_path, subsequence = None, marks = None,
pdf = False, dpi_png = 500, fontsize_title = 10, fontsize_axis = 5,
fontsize_suptitle = 12, shrink_cbar = .5, xticks_angle = 90):
maps_path = "/".join([results_path,"wcont_maps"]) # Path to save files
if not os.path.exists(maps_path): # Create if doesn't exist
os.mkdir(maps_path)
h5f = h5py.File("/".join([results_path, "_".join([ensemble_name,'wcontmatrix.h5'])]),'r')
n, p = np.shape(h5f['data'])
L = int(0.5*(1+np.sqrt(1+8*p))) # Sequence length
list_pos = np.asarray(list(itertools.combinations(range(1,L+1), 2))) # List of position pairs
repartition = pd.Series(labels_umap).value_counts() # Clustering partition
repartition.index = ["Unclassified" if i == -1 else i for i in repartition.index]
repartition = repartition.drop("Unclassified")
cont_data = np.zeros([len(repartition),len(list_pos)+1])
for cluster in tqdm(repartition.index):
# Load cluster-specific data
cluster_data = h5f['data'][labels_umap == cluster,]
# Cluster-specific w-contact matrix
prop_cluster = pd.Series(labels_umap).value_counts().sort_index()[cluster]/n
cont_matrix = pd.DataFrame(np.concatenate([list_pos,np.asarray([cluster_data.mean(axis = 0)]).T], axis = 1), columns=['pos1','pos2','cp'])
cont_matrix.pos1 = cont_matrix.pos1.astype(int)
cont_matrix.pos2 = cont_matrix.pos2.astype(int)
cont_data[np.where(repartition.index == cluster)[0][0],:] = np.append(cont_matrix.cp.to_numpy(), prop_cluster)
cont_matrix = cont_matrix.pivot(index='pos1',columns='pos2',values='cp')
if subsequence is not None:
list_pos = list_pos - 1
cont_matrix.index = np.unique(subsequence[list_pos[:,0]]).astype(int)
cont_matrix.columns = np.unique(subsequence[list_pos[:,1]]).astype(int)
fig = plt.figure(rasterized=True)
res = sns.heatmap(cont_matrix.T, cmap='Reds',square=True, cbar_kws={"shrink": shrink_cbar,'label':"Contact weight average"})
plt.suptitle(" ".join([ensemble_name,'contact-based clustering']), fontsize = fontsize_suptitle)
plt.title("".join(['Cluster #',str(cluster),' with ',str(round(100*prop_cluster,2)),'% of occupation']), fontsize = fontsize_title)
plt.xlabel('Sequence position')
plt.ylabel('Sequence position')
plt.xticks(rotation = xticks_angle)
res.set_xticklabels(res.get_xmajorticklabels(), fontsize = fontsize_axis)
res.set_yticklabels(res.get_ymajorticklabels(), fontsize = fontsize_axis)
if marks is not None:
res.hlines(marks, *res.get_xlim())
res.vlines(marks, *res.get_ylim())
if pdf:
plt.savefig("/".join([maps_path,"".join([ensemble_name,'_',str(cluster),'.pdf'])])) # Save figure in pdf
else:
plt.savefig("/".join([maps_path,"".join([ensemble_name,'_',str(cluster),'.png'])]), dpi = dpi_png) # Save figure in png
np.save("/".join([results_path,"".join([ensemble_name,'_contdata'])]),cont_data) # Save contact map values
```
%% Cell type:code id:a4c98c97 tags:
``` python
def representative_ensemble(size, ensemble_path, ensemble_name, labels_umap):
# Initial parameters
var_dict = {'multiframe' : 'n', 'check_folder' : True, 'do_xtc' : False, 'do_pdb' : False,
'ensemble_name' : ensemble_name, 'ensemble_path' : ensemble_path}
var_dict['xtc_files'] = [file for file in os.listdir(ensemble_path) if file.endswith(".xtc")]
var_dict['pdb_files'] = [file for file in os.listdir(ensemble_path) if file.endswith(".pdb") or file.endswith(".prmtop") or file.endswith(".parm7") or file.endswith(".gro")]
var_dict['folders'] = [file for file in os.listdir(ensemble_path) if (os.path.isdir("/".join([ensemble_path,file])) and not file.startswith('.') and not file.startswith("results"))]
# File processing
if len(var_dict["xtc_files"]) + len(var_dict["folders"]) + len(var_dict["pdb_files"]) == 0:
sys.exit("".join(['Folder for ', var_dict["ensemble_name"], ' ensemble is empty...']))
# .xtc file with a .pdb topology file
if len(var_dict["xtc_files"]) >= len(var_dict["pdb_files"]) and len(var_dict["pdb_files"]) == 1:
print('\nTaking as input:\n')
print("".join([str(var_dict["xtc_files"][0]),' : trajectory of ',var_dict["ensemble_name"],',']))
print("".join([str(var_dict["pdb_files"][0]),' : topology file of ',var_dict["ensemble_name"],'.']))
if len(var_dict["xtc_files"]) > 1:
print("\nMore than one .xtc file were found. Taking the first as the trajectory file.\n")
var_dict["do_xtc"] = True
var_dict["xtc_root_path"] = var_dict["ensemble_path"]
var_dict['check_folder'] = False
# multiframe .pdb files
if var_dict['multiframe'] == 'y' or (len(var_dict["pdb_files"]) >= 1 and len(var_dict["xtc_files"]) == 0):
print('\nTaking as input:\n')
print("".join([str(var_dict["pdb_files"][0]),' : trajectory of ',var_dict["ensemble_name"],'.']))
if len(var_dict["pdb_files"]) > 1:
print("\nMore than one multiframe .pdb file were found. Taking the first as the trajectory file.\n")
print("\nTaking the previously converted files.\n")
var_dict["do_xtc"] = True
var_dict["xtc_root_path"] = "/".join([var_dict["ensemble_path"],'converted_files'])
var_dict["xtc_files"] = [file for file in os.listdir(var_dict["xtc_root_path"]) if file.endswith(".xtc")]
var_dict["pdb_files"] = [file for file in os.listdir(var_dict["xtc_root_path"]) if file.endswith(".pdb")]
var_dict['check_folder'] = False
# folder with .pdb files
if len(var_dict["folders"]) >= 1 and var_dict['check_folder'] == True:
print('\nTaking as input:\n')
print("".join([var_dict["folders"][0],' folder contains: trajectory of ',var_dict["ensemble_name"],"."]))
if len(var_dict["folders"]) > 1:
print("\nMore than one .pdb folder were found. Taking the first as the trajectory folder.\n")
var_dict["do_pdb"] = True
if not var_dict["do_pdb"] and not var_dict["do_xtc"]:
sys.exit("".join(['\n Sorry, I did not understood the input. Please follow the guidelines described in the function documentation to create ',ensemble_name,' folder.\n']))
print("\n----------------------------------------------------------------------------------\n")
print("\nSampling representative family...\n")
repartition = pd.Series(labels_umap).value_counts() # Clustering partition
repartition.index = ["Unclassified" if i == -1 else i for i in repartition.index]
repartition = repartition.drop("Unclassified")
probas = repartition.values/np.sum(repartition.values)
selected_conf = np.zeros(size)
for i in range(size):
choose_cluster = np.random.choice(repartition.index, size = 1, p = probas)[0]
selected_conf[i] = np.random.choice(np.where(labels_umap == choose_cluster)[0], size = 1)[0]
selected_conf = np.ndarray.astype(selected_conf, int)
results_path = "/".join([os.path.abspath(ensemble_path),"_".join(['results',ensemble_name])])
save_files = "/".join([results_path, "representative_family"])
if not os.path.exists(save_files):
os.mkdir(save_files)
if var_dict["do_xtc"]:
traj_file = md.load_xtc("/".join([var_dict["xtc_root_path"],var_dict["xtc_files"][0]]), top = "/".join([var_dict["xtc_root_path"],var_dict["pdb_files"][0]]))
# Save .xtc file
traj_file[selected_conf].save_xtc("/".join([save_files, "".join([ensemble_name,'_repfam.xtc'])]))
if var_dict["do_pdb"]:
conf_list = os.listdir("/".join([var_dict["ensemble_path"],var_dict["folders"][0]]))
# Save pdb folder
for j in selected_conf:
traj = md.load_pdb("/".join(["/".join([var_dict["ensemble_path"],var_dict["folders"][0]]),conf_list[j]]))
traj.save_pdb("/".join([save_files, "".join([ensemble_name,'_',conf_list[j],'.pdb'])]))
print("\nFiles saved.\n")
```
%% Cell type:code id:7069e00f tags:
``` python
def cluster_descriptors(ensemble_path, ensemble_name, labels_umap, subsequence = None,
fig_width = 10, fig_height = 1.7, shrink_cbar = .7, yticks_angle = 0,
fontsize_title = 8, fontsize_axis = 7, pdf = False, dpi_png = 200):
# Initial parameters
var_dict = {'multiframe' : 'n', 'check_folder' : True, 'do_xtc' : False, 'do_pdb' : False,
'ensemble_name' : ensemble_name, 'ensemble_path' : ensemble_path}
var_dict['xtc_files'] = [file for file in os.listdir(ensemble_path) if file.endswith(".xtc")]
var_dict['pdb_files'] = [file for file in os.listdir(ensemble_path) if file.endswith(".pdb") or file.endswith(".prmtop") or file.endswith(".parm7") or file.endswith(".gro")]
var_dict['folders'] = [file for file in os.listdir(ensemble_path) if (os.path.isdir("/".join([ensemble_path,file])) and not file.startswith('.') and not file.startswith("results"))]
# File processing
if len(var_dict["xtc_files"]) + len(var_dict["folders"]) + len(var_dict["pdb_files"]) == 0:
sys.exit("".join(['Folder for ', var_dict["ensemble_name"], ' ensemble is empty...']))
# .xtc file with a .pdb topology file
if len(var_dict["xtc_files"]) >= len(var_dict["pdb_files"]) and len(var_dict["pdb_files"]) == 1:
print('\nTaking as input:\n')
print("".join([str(var_dict["xtc_files"][0]),' : trajectory of ',var_dict["ensemble_name"],',']))
print("".join([str(var_dict["pdb_files"][0]),' : topology file of ',var_dict["ensemble_name"],'.']))
if len(var_dict["xtc_files"]) > 1:
print("\nMore than one .xtc file were found. Taking the first as the trajectory file.\n")
var_dict["do_xtc"] = True
var_dict["xtc_root_path"] = var_dict["ensemble_path"]
var_dict['check_folder'] = False
# multiframe .pdb files
if var_dict['multiframe'] == 'y' or (len(var_dict["pdb_files"]) >= 1 and len(var_dict["xtc_files"]) == 0):
print('\nTaking as input:\n')
print("".join([str(var_dict["pdb_files"][0]),' : trajectory of ',var_dict["ensemble_name"],'.']))
if len(var_dict["pdb_files"]) > 1:
print("\nMore than one multiframe .pdb file were found. Taking the first as the trajectory file.\n")
print("\nTaking the previously converted files.\n")
var_dict["do_xtc"] = True
var_dict["xtc_root_path"] = "/".join([var_dict["ensemble_path"],'converted_files'])
var_dict["xtc_files"] = [file for file in os.listdir(var_dict["xtc_root_path"]) if file.endswith(".xtc")]
var_dict["pdb_files"] = [file for file in os.listdir(var_dict["xtc_root_path"]) if file.endswith(".pdb")]
var_dict['check_folder'] = False
# folder with .pdb files
if len(var_dict["folders"]) >= 1 and var_dict['check_folder'] == True:
print('\nTaking as input:\n')
print("".join([var_dict["folders"][0],' folder contains: trajectory of ',var_dict["ensemble_name"],"."]))
if len(var_dict["folders"]) > 1:
print("\nMore than one .pdb folder were found. Taking the first as the trajectory folder.\n")
var_dict["do_pdb"] = True
if not var_dict["do_pdb"] and not var_dict["do_xtc"]:
sys.exit("".join(['\n Sorry, I did not understood the input. Please follow the guidelines described in the function documentation to create ',ensemble_name,' folder.\n']))
print("\n----------------------------------------------------------------------------------\n")
print("\nComputing cluster-specific descriptors...\n")
results_path = "/".join([os.path.abspath(ensemble_path),"_".join(['results',ensemble_name])])
save_files = "/".join([results_path, "cluster_descriptors"])
if not os.path.exists(save_files):
os.mkdir(save_files)
if var_dict["do_xtc"]:
traj_file = md.load_xtc("/".join([var_dict["xtc_root_path"],var_dict["xtc_files"][0]]), top = "/".join([var_dict["xtc_root_path"],var_dict["pdb_files"][0]]))
L = traj_file.n_residues
Nconf = traj_file.n_frames
dssp_types = ['H','B','E','G','I','T','S',' ']
prop_dssp = np.zeros([len(dssp_types),L,len(labels_umap)-1])
rg = np.zeros([len(labels_umap)-1])
for k in range(len(np.unique(labels_umap[labels_umap >= 0]))):
prop_dssp_k = np.zeros([len(dssp_types),L])
dssp_k = md.compute_dssp(traj_file[np.where(labels_umap == k)], simplified = False)
rg[k] = np.mean(md.compute_rg(traj_file[np.where(labels_umap == k)]))
for dt in range(len(dssp_types)):
prop_dssp_k[dt,:] = (dssp_k == dssp_types[dt]).sum(axis = 0)/len(np.where(labels_umap == k)[0])
prop_dssp[:,:,k] = prop_dssp_k
if var_dict["do_pdb"]:
pdb_folder = "/".join([var_dict["ensemble_path"],var_dict["folders"][0]])
conf_list = os.listdir(pdb_folder)
md_file = md.load_pdb("/".join([pdb_folder,conf_list[0]]))
L = md_file.topology.n_residues
Nconf = len(conf_list)
dssp_types = ['H','B','E','G','I','T','S',' ']
prop_dssp = np.zeros([len(dssp_types),L,len(labels_umap)-1])
rg = np.zeros([len(labels_umap)-1])
for k in range(len(np.unique(labels_umap[labels_umap >= 0]))):
prop_dssp_k = np.zeros([len(dssp_types),L])
clus_k = np.where(labels_umap == k)[0]
dssp_k = np.zeros([len(clus_k),L]).astype(str)
rg_k = np.zeros([len(clus_k)])
for l in range(len(clus_k)):
dssp_k[l,:] = md.compute_dssp(md.load_pdb("/".join([pdb_folder,conf_list[clus_k[l]]])), simplified = False)[0].astype(str)
rg_k[l] = md.compute_rg(md.load_pdb("/".join([pdb_folder,conf_list[clus_k[l]]])))
rg[k] = np.mean(rg_k)
for dt in range(len(dssp_types)):
prop_dssp_k[dt,:] = (dssp_k == dssp_types[dt]).sum(axis = 0)/len(np.where(labels_umap == k)[0])
prop_dssp[:,:,k] = prop_dssp_k
if subsequence is None:
subsequence = np.arange(L)
for cluster in tqdm(range(len(np.unique(labels_umap[labels_umap >= 0])))):
prop_cluster = round(100*len(np.where(labels_umap == cluster)[0])/Nconf,2)
fig = plt.figure(figsize=(fig_width, fig_height))
res = sns.heatmap(prop_dssp[:,subsequence,cluster], cmap='Blues', square = True, cbar_kws={"shrink": shrink_cbar,'label':"Class prop."})
xlabels = [item.get_text() for item in res.get_xmajorticklabels()]
plt.xlabel('Sequence position')
plt.ylabel('DSSP class')
plt.title("".join([ensemble_name, ' - cluster #',str(cluster),' (',str(prop_cluster),'% oc.). Average RG = ', str(round(10*rg[cluster],2)),r'$\AA$.']), fontsize = fontsize_title)
plt.yticks(rotation = yticks_angle)
res.set_xticklabels(np.asarray(xlabels).astype(int) + 1, fontsize = fontsize_axis)
res.set_yticklabels(['L' if x==' ' else x for x in dssp_types], fontsize = fontsize_axis)
if pdf:
plt.savefig("/".join([save_files,"".join([ensemble_name,'_',str(cluster),'_DSSP.pdf'])]), bbox_inches='tight')
else:
plt.savefig("/".join([save_files,"".join([ensemble_name,'_',str(cluster),'_DSSP.png'])]), dpi = dpi_png, bbox_inches='tight')
print("\nPlots saved.\n")
```
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment