Update code/modbaseModelAdd.py

code/modbaseModelAdd.py

@@ -10,7 +10,6 @@ def addModbaseModels(dataframe, path_to_input_files, path_to_output_files):
         # GET MODBASE MODELS
         # Get IDs from data to retrieve only their models from MODBASE
         dataframe.reset_index(inplace=True, drop=True)
-
         existing_modbase_models = list(Path(path_to_output_files / 'modbase_structures').glob("*"))
         existing_modbase_models = [str(i) for i in existing_modbase_models]
         existing_modbase_models = [i.split('/')[-1].split('.')[0] for i in existing_modbase_models]
@@ -32,11 +31,15 @@ def addModbaseModels(dataframe, path_to_input_files, path_to_output_files):
         existing_free_sasa = list(Path(path_to_output_files / 'freesasa_files').glob("*"))
         existing_free_sasa = [str(i) for i in existing_free_sasa]
         existing_free_sasa = [i.split('/')[-1].split('.')[0] for i in existing_free_sasa]
+        keep_cols = dataframe.columns
         for i in dataframe.index:
             coordDict = {}
             protein = dataframe.at[i, 'uniprotID']
             varPos = int(dataframe.at[i, 'pos'])
             wt =  dataframe.at[i, 'wt']
+            mut = dataframe.at[i, 'mut']
+            datapoint = dataframe.at[i, 'datapoint']
+            
             if protein not in existing_modbase_models:
                 print('Downloading Modbase models for ', protein)
                 url = 'https://salilab.org/modbase/retrieve/modbase/?databaseID=' + protein
@@ -104,7 +107,7 @@ def addModbaseModels(dataframe, path_to_input_files, path_to_output_files):
                                                'model_id': model_id, 'coordinates': coordDict,
                                                'AAonPDB': AAonPDB, 'coordVAR': coordVAR}
                                     modbase_reduced = modbase_reduced.append(new_row, ignore_index=True)
-                                    modbase_reduced = modbase_reduced[['uniprotID', 'quality_score', 'model_id', 'coordinates', 'AAonPDB', 'coordVAR']]
+                                    modbase_reduced = modbase_reduced[['uniprotID', 'quality_score', 'model_id', 'coordinates', 'AAonPDB', 'coordVAR']]           
                                     modbase = dataframe.merge(modbase_reduced, on='uniprotID', how='left')
                                     modbase.quality_score = modbase.quality_score.astype(float)
                                     modbase = modbase.sort_values(by=['datapoint', 'quality_score'], ascending=False)
@@ -119,13 +122,20 @@ def addModbaseModels(dataframe, path_to_input_files, path_to_output_files):
                                                      '': np.NaN}, inplace=True)
                                 except NameError:
                                     print('This file doesnt have Quality Score. Replacer: -999', model_id)
+                            else:
+                                new_row = {'uniprotID': uniprot_id, 'wt': wt,
+                                               'pos': varPos, 'mut': mut, 'datapoint': datapoint }
+                                no_modbase = no_modbase.append(new_row, ignore_index=True)
+                                
                 else:
-                    k = pd.Series(
-                        dataframe.loc[i])
-                    no_modbase = no_modbase.append(k, ignore_index=True)
+                    new_row = {'uniprotID': uniprot_id, 'wt': wt,
+                                               'pos': varPos, 'mut': mut, 'datapoint': datapoint }
+                    no_modbase = no_modbase.append(new_row, ignore_index=True)
+                    
+
 
     no_modbase_no_Coord = modbase[pd.isna(modbase['coordVAR'])]
     no_modbase = pd.concat([no_modbase, no_modbase_no_Coord])
     modbase = modbase[~pd.isna(modbase['coordVAR'])]
-
+    no_modbase = no_modbase[keep_cols]
     return modbase, no_modbase