DrugBank column fix

swansonk14 · Dec 12, 2023 · 55010a6 · 55010a6
1 parent be7b1fd
commit 55010a6
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 10 deletions.
diff --git a/admet_ai/constants.py b/admet_ai/constants.py
@@ -9,5 +9,10 @@
     DEFAULT_MODELS_DIR = resources_dir / "models"
 
 # DrugBank columns
+DRUGBANK_ID_COLUMN = "id"
+DRUGBANK_NAME_COLUMN = "name"
+DRUGBANK_SMILES_COLUMN = "smiles"
+DRUGBANK_ATC_PREFIX = "atc"
 DRUGBANK_ATC_NAME_PREFIX = "atc_name"
+DRUGBANK_ATC_CODE_COLUMN = DRUGBANK_ATC_PREFIX
 DRUGBANK_DELIMITER = ";"
diff --git a/admet_ai/web/app/drugbank.py b/admet_ai/web/app/drugbank.py
@@ -5,7 +5,14 @@
 import matplotlib
 import pandas as pd
 
-from admet_ai.constants import DRUGBANK_ATC_NAME_PREFIX, DRUGBANK_DELIMITER
+from admet_ai.constants import (
+    DRUGBANK_ATC_NAME_PREFIX,
+    DRUGBANK_ATC_PREFIX,
+    DRUGBANK_DELIMITER,
+    DRUGBANK_ID_COLUMN,
+    DRUGBANK_NAME_COLUMN,
+    DRUGBANK_SMILES_COLUMN,
+)
 from admet_ai.web.app import app
 from admet_ai.web.app.admet_info import get_admet_id_to_name
 
@@ -95,8 +102,14 @@ def get_drugbank_tasks_ids() -> list[str]:
 
     :return: A list of tasks (properties) predicted in the DrugBank reference set.
     """
-    non_task_columns = ["name", "smiles"] + [
-        column for column in DRUGBANK_DF.columns if column.startswith("atc_")
+    non_task_columns = [
+        DRUGBANK_ID_COLUMN,
+        DRUGBANK_NAME_COLUMN,
+        DRUGBANK_SMILES_COLUMN,
+    ] + [
+        column
+        for column in DRUGBANK_DF.columns
+        if column.startswith(DRUGBANK_ATC_PREFIX)
     ]
     task_columns = set(DRUGBANK_DF.columns) - set(non_task_columns)
     drugbank_task_ids = sorted(task_columns)

diff --git a/scripts/get_drugbank_approved.py b/scripts/get_drugbank_approved.py
@@ -6,9 +6,10 @@
 from rdkit import Chem
 from tqdm import tqdm
 
-from admet_ai.constants import DRUGBANK_ATC_NAME_PREFIX, DRUGBANK_DELIMITER
-from tdc_constants import (
+from admet_ai.constants import (
     DRUGBANK_ATC_CODE_COLUMN,
+    DRUGBANK_ATC_NAME_PREFIX,
+    DRUGBANK_DELIMITER,
     DRUGBANK_ID_COLUMN,
     DRUGBANK_NAME_COLUMN,
     DRUGBANK_SMILES_COLUMN,

diff --git a/scripts/tdc_constants.py b/scripts/tdc_constants.py
@@ -103,8 +103,3 @@
 ADMET_ALL_SMILES_COLUMN = "smiles"
 ADMET_GROUP_SMILES_COLUMN = "Drug"
 ADMET_GROUP_TARGET_COLUMN = "Y"
-
-DRUGBANK_ID_COLUMN = "id"
-DRUGBANK_NAME_COLUMN = "name"
-DRUGBANK_SMILES_COLUMN = "smiles"
-DRUGBANK_ATC_CODE_COLUMN = "atc"