Correct bug in representative_ensembles

23be00ea · Javier González-Delgado · 101b0023 · 23be00ea
Commit 23be00ea authored 10 months ago by Javier González-Delgado
--- a/wario/utils.ipynb
+++ b/wario/utils.ipynb
@@ -324,7 +324,7 @@
    "        # Save pdb folder\n",
    "        for j in selected_conf:\n",
    "            traj = md.load_pdb(\"/\".join([\"/\".join([var_dict[\"ensemble_path\"],var_dict[\"folders\"][0]]),conf_list[j]]))\n",
-    "            traj.save_pdb(\"/\".join([clus_k_path, \"\".join([ensemble_name,'_',conf_list[j],'.pdb'])]))\n",
+    "            traj.save_pdb(\"/\".join([save_files, \"\".join([ensemble_name,'_',conf_list[j],'.pdb'])]))\n",
    "\n",
    "    print(\"\\nFiles saved.\\n\")    "
   ]

 %% 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
 ```

 %% 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'])]))

    print("\nFiles saved.\n")
 ```

 %% Cell type:code id:ff9c4941 tags:

 ``` python
 def plot_2umap(embedding_2d, labels_umap, ensemble_name, results_path):

    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)
        plt.savefig("/".join([results_path, "".join(["clusters_2d", ensemble_name, '.png'])]), dpi = 199)
        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(wcont_data, labels_umap, ensemble_name, results_path, subsequence = None):

    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)

    L = int(0.5*(1+np.sqrt(1+8*wcont_data.shape[1]))) # 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):

        # Cluster-specific w-contact matrix
        prop_cluster = pd.Series(labels_umap).value_counts().sort_index()[cluster]/np.shape(wcont_data)[0]
        cont_matrix = pd.DataFrame(np.concatenate([list_pos,np.asarray([wcont_data.loc[labels_umap == cluster,].mean()]).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()
        res = sns.heatmap(cont_matrix.T, cmap='Reds',square=True,  cbar_kws={"shrink": .5,'label':"Contact weight average"})
        plt.suptitle(" ".join([ensemble_name,'contact-based clustering']), fontsize=10)
        plt.title("".join(['Cluster #',str(cluster),' with ',str(round(100*prop_cluster,2)),'% of occupation']), fontsize = 8)

        plt.xlabel('Sequence position')
        plt.ylabel('Sequence position')
        plt.xticks(rotation=0)
        res.set_xticklabels(res.get_xmajorticklabels(), fontsize = 6)
        res.set_yticklabels(res.get_ymajorticklabels(), fontsize = 6)
        plt.savefig("/".join([maps_path,"".join([ensemble_name,'_',str(cluster),'.png'])]),  dpi=199) # Save figure in

    np.save("/".join([results_path,"".join([ensemble_name,'_contdata'])]),cont_data)
 ```

 %% 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([clus_k_path, "".join([ensemble_name,'_',conf_list[j],'.pdb'])]))
+            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):

    # 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"]:

        conf_list = os.listdir(var_dict["folders"][0])
        md_file = md.load_pdb("/".join([var_dict["folders"][0],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

    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=(10, 1.7))
        res = sns.heatmap(prop_dssp[:,:,cluster], cmap='Blues', square = True,  cbar_kws={"shrink": .7,'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 = 8)
        plt.yticks(rotation=0)
        res.set_xticklabels(np.asarray(xlabels).astype(int) + 1, fontsize = 7)
        res.set_yticklabels(['L' if x==' ' else x for x in dssp_types], fontsize = 7)
        plt.savefig("/".join([save_files,"".join([ensemble_name,'_',str(cluster),'_DSSP.png'])]),  dpi=199, bbox_inches='tight')


    print("\nPlots saved.\n")
 ```

 %% Cell type:code id:cc23916d-784e-4204-ab55-2ee8179cc317 tags:

 ``` python
 def compare_families(path_a, path_b, ensemble_name_a, ensemble_name_b):

    contmatrix_a = np.load('/'.join([path_a, "_".join([ensemble_name_a,'contdata.npy'])]))
    contmatrix_b = np.load('/'.join([path_b, "_".join([ensemble_name_b,'contdata.npy'])]))

    a = np.delete(contmatrix_a, -1, axis = 1)
    ma = contmatrix_a[:,-1]/np.sum(contmatrix_a[:,-1])
    b = np.delete(contmatrix_b, -1, axis = 1)
    mb = contmatrix_b[:,-1]/np.sum(contmatrix_b[:,-1])

    if np.shape(a)[1] != np.shape(b)[1]:
        quit('Both sequences must have the same number of residues.')

    M = ot.dist(a, b, metric = 'sqeuclidean') # Cost matrix
    clean = ot.utils.clean_zeros(ma, mb, M)
    w = np.sqrt(ot.emd2(clean[0], clean[1], clean[2])) # 2-Wasserstein distance
    return w
 ```