Update app.py

app.py

@@ -6,7 +6,7 @@ import textwrap
 from email.mime.multipart import MIMEMultipart
 from email.mime.text import MIMEText
 from email.utils import formatdate, make_msgid
-from functools import cache
+from functools import cache, partial
 from math import pi
 from time import sleep, time
 from uuid import uuid4
@@ -25,6 +25,7 @@ import hydra
 import pandas as pd
 from pandarallel import pandarallel
 import requests
+from rdkit.DataStructs import BulkTanimotoSimilarity
 from requests.adapters import HTTPAdapter, Retry
 from markdown import markdown
 from rdkit import Chem, DataStructs
@@ -291,29 +292,24 @@ def check_expiry():
                 send_email(job)
 
 
-@cache
-def max_tanimoto_similarity(smi, seen_smiles):
+def max_tanimoto_similarity(smi, seen_smiles_with_fp):
     if smi is None:
         return 0
+    if smi in seen_smiles_with_fp['X1'].values:
+        return 1
     mol = Chem.MolFromSmiles(smi)
     if mol is None:
         return 0
     mol_ecfp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048)
-    max_sim = 0
-    for smiles in seen_smiles:
-        mol_seen = Chem.MolFromSmiles(smiles)
-        mol_seen_ecfp = AllChem.GetMorganFingerprintAsBitVect(mol_seen, radius=2, nBits=2048)
-        sim = DataStructs.TanimotoSimilarity(mol_ecfp, mol_seen_ecfp)
-        if sim == 1:
-            return 1
-        max_sim = max(sim, max_sim)
-    return max_sim
+    sims = BulkTanimotoSimilarity(mol_ecfp, seen_smiles_with_fp['FP'])
+    return max(sims)
 
 
-@cache
 def max_sequence_identity(seq, seen_fastas):
     if seq is None:
         return 0
+    if seq in seen_fastas:
+        return 1
     aligner = PairwiseAligner()
     aligner.mode = 'local'
     max_id = 0
@@ -328,16 +324,24 @@ def max_sequence_identity(seq, seen_fastas):
 
 @cache
 def get_seen_smiles(family, task):
+    if family == 'General':
+        family = 'all_families_full'
+    else:
+        family = TARGET_FAMILY_MAP[family.title()]
     seen_smiles = pd.read_csv(
-        f'data/benchmarks/seen_compounds/{TARGET_FAMILY_MAP[family.title()]}_{task.lower()}_random_split.csv')
-    return seen_smiles['X1'].tolist()
+        f'data/benchmarks/seen_compounds/{family}_{task.lower()}_random_split.csv')
+    return seen_smiles
 
 
 @cache
 def get_seen_fastas(family, task):
+    if family == 'General':
+        family = 'all_families_full'
+    else:
+        family = TARGET_FAMILY_MAP[family.title()]
     seen_fastas = pd.read_csv(
-        f'data/benchmarks/seen_targets/{TARGET_FAMILY_MAP[family.title()]}_{task.lower()}_random_split.csv')
-    return seen_fastas['X2'].tolist()
+        f'data/benchmarks/seen_targets/{family}_{task.lower()}_random_split.csv')
+    return seen_fastas
 
 
 @cache
@@ -709,7 +713,6 @@ def submit_predict(predict_filepath, task, preset, target_family, opts, state):
     error = None
     task_file_abbr = {'Compound-Protein Interaction': 'CPI', 'Compound-Protein Binding Affinity': 'CPA'}
     predictions_file = None
-
     df_training = pd.read_csv(f'data/complete_{TASK_MAP[task].lower()}_dataset.csv')
     orig_df = pd.read_csv(predict_filepath)
     alignment_df = get_fasta_family_map()
@@ -737,12 +740,9 @@ def submit_predict(predict_filepath, task, preset, target_family, opts, state):
     if 'Target Family' not in orig_df.columns:
         orig_df['Target Family'] = None
     if orig_df['Target Family'].isna().any():
-        orig_df.loc[
-            orig_df['Target Family'].isna(), 'Target Family'
-        ] = orig_df.loc[
-            orig_df['Target Family'].isna(), 'X2'
-        ].parallel_apply(detect_family)
-
+        orig_df.loc[orig_df['Target Family'].isna(), 'Target Family'] = (
+            orig_df.loc[orig_df['Target Family'].isna(), 'X2'].parallel_apply(detect_family)
+        )
     detect_family.cache_clear()
 
     orig_df = orig_df.merge(df_training[['X1', 'X2', 'Y']], on=['X1', 'X2'], how='left', indicator=False)
@@ -783,76 +783,82 @@ def submit_predict(predict_filepath, task, preset, target_family, opts, state):
             prediction_df = pd.concat([prediction_df, predictions])
 
         else:
-            predictions_file = f'{SERVER_DATA_DIR}/{job_id}_{task_file_abbr[task]}_{preset}_auto_predictions.csv'
+            predictions_file = f'{SERVER_DATA_DIR}/{job_id}_{task_file_abbr[task]}_family-recommended_predictions.csv'
             task_value = TASK_MAP[task]
             score = TASK_METRIC_MAP[task]
             benchmark_df = pd.read_csv(f'data/benchmarks/{task_value}_test_metrics.csv')
             predict_df = pd.read_csv(predict_filepath)
 
             for family, subset in predict_df.groupby('Target Family'):
-                predict_subset_filepath = f'{SERVER_DATA_DIR}/{job_id}_{family}_input.csv'
+                predict_subset_filepath = os.path.join(
+                    os.path.dirname(predict_filepath), f'{job_id}_{family}_input.csv'
+                )
                 subset.to_csv(predict_subset_filepath, index=False, na_rep='')
-                seen_compounds = get_seen_smiles(family, task_value)
 
+                seen_compounds = get_seen_smiles(family, task_value)['X1'].values
                 if subset['X1'].iloc[0] in seen_compounds:
                     scenario = "Seen Compound"
                 else:
                     scenario = "Unseen Compound"
 
                 filtered_df = benchmark_df[(benchmark_df['Family'] == family.title())
-                                           & (benchmark_df['Scenario'] == scenario)]
+                                           & (benchmark_df['Scenario'] == scenario)
+                                           & (benchmark_df['Type'] == 'Family')]
 
-                preset = filtered_df.loc[filtered_df[score].idxmax(), 'Model']
-                preset_value = PRESET_MAP[preset]
+                seen_compounds = get_seen_smiles('General', task_value)['X1'].values
+                if subset['X1'].iloc[0] in seen_compounds:
+                    scenario = "Seen Compound"
+                else:
+                    scenario = "Unseen Compound"
+
+                filtered_df = pd.concat([
+                    filtered_df,
+                    benchmark_df[(benchmark_df['Family'] == family.title())
+                                 & (benchmark_df['Scenario'] == scenario)
+                                 & (benchmark_df['Type'] == 'General')]
+                ])
 
-                target_family = TARGET_FAMILY_MAP[family.title()]
+                row = filtered_df.loc[filtered_df[score].idxmax()]
+                preset_value = PRESET_MAP[row['Model']]
+                target_family = TARGET_FAMILY_MAP[family.title()] if row['Type'] == 'Family' else 'general'
                 cfg = hydra.compose(
                     config_name="webserver_inference",
                     overrides=[f"task={task_value}",
                                f"preset={preset_value}",
-                               # f"ckpt_path=D:/checkpoints/{preset_value}-{task_value}-{target_family}.ckpt",
                                f"ckpt_path=resources/checkpoints/{preset_value}-{task_value}-{target_family}.ckpt",
                                f"data.data_file='{str(predict_subset_filepath)}'"])
 
                 predictions, _ = predict(cfg)
                 predictions = pd.concat([pd.DataFrame(prediction) for prediction in predictions], ignore_index=True)
-                predictions['Source'] = f'Predicted ({preset} {family})'
+                predictions['Source'] = (f'Predicted ({row["Model"]} '
+                                         f'{family.title() if row["Type"] == "Family" else "General"})')
                 prediction_df = pd.concat([prediction_df, predictions])
 
         prediction_df = prediction_df.merge(orig_df, on=['X1', 'X2'], how='left', indicator=False)
         prediction_df = pd.concat([prediction_df, annotated_df], ignore_index=True)
 
-        # prediction_df['Max. Tanimoto Similarity'] = prediction_df.groupby('Target Family')['X1'].apply(
-        #     lambda group: group.parallel_apply(
-        #         max_tanimoto_similarity,
-        #         seen_smiles=tuple(get_seen_smiles(family=group.name, task=task_value))
-        #     )
-        # ).values
-        #
-        # prediction_df['Max. Sequence Identity'] = prediction_df.groupby('Target Family')['X2'].apply(
-        #     lambda group: group.parallel_apply(
-        #         max_sequence_identity,
-        #         seen_fastas=tuple(get_seen_fastas(family=group.name, task=task_value))
-        #     )
-        # ).values
         if "Include Max. Tanimoto Similarity" in opts:
             for family in prediction_df['Target Family'].unique():
-                prediction_df.loc[
-                    prediction_df['Target Family'] == family, 'Max. Tanimoto Similarity'] = prediction_df.loc[
-                    prediction_df['Target Family'] == family, 'X1'].parallel_apply(
-                    max_tanimoto_similarity,
-                    seen_smiles=tuple(get_seen_smiles(family=family, task=task_value))
+                family_smiles_df = get_seen_smiles(family=family, task=task_value)
+                family_smiles_df['FP'] = family_smiles_df['X1'].parallel_apply(
+                    lambda smiles: AllChem.GetMorganFingerprintAsBitVect(
+                        Chem.MolFromSmiles(smiles), radius=2, nBits=2048)
+                )
+                max_sim = cache(partial(max_tanimoto_similarity, seen_smiles_with_fp=family_smiles_df))
+                prediction_df.loc[prediction_df['Target Family'] == family, 'Max. Tanimoto Similarity'] = (
+                    prediction_df.loc[prediction_df['Target Family'] == family, 'X1'].parallel_apply(max_sim)
                 )
-            max_tanimoto_similarity.cache_clear()
+                max_sim.cache_clear()
+
         if "Include Max. Sequence Identity" in opts:
             for family in prediction_df['Target Family'].unique():
-                prediction_df.loc[
-                    prediction_df['Target Family'] == family, 'Max. Sequence Identity'] = prediction_df.loc[
-                    prediction_df['Target Family'] == family, 'X2'].parallel_apply(
-                    max_sequence_identity,
-                    seen_fastas=tuple(get_seen_fastas(family=family, task=task_value))
+                family_fastas_df = get_seen_fastas(family=family, task=task_value)
+                max_id = cache(partial(max_sequence_identity, seen_fastas=family_fastas_df['X2'].values))
+                prediction_df.loc[prediction_df['Target Family'] == family, 'Max. Sequence Identity'] = (
+                    prediction_df.loc[prediction_df['Target Family'] == family, 'X2'].parallel_apply(max_id)
                 )
-            max_sequence_identity.cache_clear()
+                max_id.cache_clear()
+
         prediction_df.drop(['N'], axis=1).to_csv(predictions_file, index=False, na_rep='')
         status = "COMPLETED"
 
@@ -1968,9 +1974,8 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
             return [None, family]
 
         if family == 'General':
-            seen_targets = pd.read_csv(
-                f'data/benchmarks/seen_targets/all_families_full_{task.lower()}_random_split.csv')
-            if process_target_fasta(fasta) in seen_targets['X2'].values:
+            seen_targets = get_seen_fastas('General', task)['X2'].values
+            if process_target_fasta(fasta) in seen_targets:
                 scenario = "Seen Target"
             else:
                 scenario = "Unseen Target"
@@ -1979,16 +1984,14 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
                                        & (benchmark_df['Type'] == 'General')]
 
         else:
-            seen_targets_general = pd.read_csv(
-                f'data/benchmarks/seen_targets/all_families_full_{task.lower()}_random_split.csv')
-            if process_target_fasta(fasta) in seen_targets_general['X2'].values:
+            seen_targets_general = get_seen_fastas('General', task)['X2'].values
+            if process_target_fasta(fasta) in seen_targets_general:
                 scenario_general = "Seen Target"
             else:
                 scenario_general = "Unseen Target"
 
-            seen_targets_family = pd.read_csv(
-                f'data/benchmarks/seen_targets/{TARGET_FAMILY_MAP[family.title()]}_{task.lower()}_random_split.csv')
-            if process_target_fasta(fasta) in seen_targets_family['X2'].values:
+            seen_targets_family = get_seen_fastas(family, task)['X2'].values
+            if process_target_fasta(fasta) in seen_targets_family:
                 scenario_family = "Seen Target"
             else:
                 scenario_family = "Unseen Target"
@@ -2008,10 +2011,9 @@ QALAHAYFAQYHDPDDEPVADPYDQSFESRDLLIDEWKSLTYDEVISFVPPPLDQEEMES
             scenario = "Unseen Target (<0.85 sequence identity)"
 
         return {drug_screen_preset:
-                    gr.Dropdown(value=row['Model'],
-                                info=f"Reason: {row['Scenario']} in training; we recommend the {row['Type']}-trained "
-                                     f"model with the best {score} in the {scenario} scenario "
-                                     f"on {row['Family']}."),
+                gr.Dropdown(value=row['Model'],
+                            info=f"Reason: {row['Scenario']} in training; we recommend the {row['Type']}-trained "
+                                 f"model with the best {score} in the {scenario} scenario on {row['Family']}."),
                 drug_screen_target_family:
                     gr.Dropdown(value='General') if row['Type'] == 'General' else gr.Dropdown(value=family)}
 
@@ -2569,4 +2571,4 @@ if __name__ == "__main__":
     hydra.initialize(version_base="1.3", config_path="configs", job_name="webserver_inference")
     demo.queue(default_concurrency_limit=None, max_size=10).launch(show_api=False)
     scheduler.add_job(check_expiry, 'interval', hours=1)
-    scheduler.start()
+    scheduler.start()