Drug Categorization Code Added

argnorm/drug_categorization.py

@@ -0,0 +1,83 @@
+# These functions are run in the run function of the BaseNormalizer.
+# After ARO numbers are obtained, these functions can be executed - hence should be independent of which db is used.
+
+import pronto
+from typing import List, Tuple
+
+# Load the ArgNorm ontology from the 'aro.obo' file
+ARO = pronto.Ontology.from_obo_library('aro.obo')
+
+def get_immediate_drug_classes(aro_num: str) -> List[Tuple]:
+    '''
+    Description: Gets the drug classes to which a gene confers resistance to.
+    Only lists the drug class column in the CARD db.
+
+    Parameters:
+        aro_num (str): ARO number. Needs to be in the form 'ARO:number'.
+
+    Returns:
+        drug_classes_list (list[tuple]): 
+            A two-dimensional list where each inner list represents a drug class.
+            Each inner list contains the ARO number and name of the drug class in that order. [ARO:number, name].
+    '''
+
+    # Some databases don't provide aro numbers as strings. 
+    # Converting those aro numbers to pronto's desired format.
+    if type(aro_num) == float or type(aro_num) == int:
+        aro_num = 'ARO:' + str(aro_num)
+
+    # If dealing with aro nans, the final drug class categorization will give [].
+    # Hence, immediate drug classes categorization also gives empty list.
+    if aro_num == 'ARO:nan':
+        return []
+
+    # Returning empty list if aro number not in ARO ontology.
+    if aro_num not in ARO.terms():
+        return []
+
+    gene = ARO[aro_num]
+
+    confers_resistance_to_drug_class = any(r.name == 'confers_resistance_to_drug_class' for r in gene.relationships)
+    confers_resistance_to_antibiotic = any(r.name == 'confers_resistance_to_antibiotic' for r in gene.relationships)
+
+    drug_classes = []
+
+    if confers_resistance_to_drug_class:
+        for drug_class in gene.relationships[ARO.get_relationship('confers_resistance_to_drug_class')]:
+            drug_classes.append((drug_class.id, drug_class.name))
+
+    if confers_resistance_to_antibiotic:
+        for drug_class in gene.relationships[ARO.get_relationship('confers_resistance_to_antibiotic')]:
+            drug_classes.append((drug_class.id, drug_class.name))
+
+    return drug_classes
+
+def get_drug_class_category(drug_classes_list: List[Tuple]) -> List[str]:
+    '''
+    Description: Gives a list of categories of drug classes, e.g. cephem and penam are categorized as beta_lactam antibiotics.
+
+    Parameters:
+        drug_classes_list (list[tuple]): 
+            A two-dimensional list where each inner list represents a drug class.
+            Each inner list contains the ARO number and name of the drug class in that order. [ARO:number, name].
+            Designed to use the return value of the function 'get_immediate_drug_classes'.
+
+    Returns:
+        drug_class_categories (list[str]):
+            A list containing the names of the drug class categories of each drug class given as input to the function.
+            Order of the names of the drug class categories corresponds to the order in which the drug classes were given
+            to the function in the drug_classes_list.
+    '''
+    drug_class_categories = []
+
+    for drug_class in drug_classes_list:
+        drug_class_instance = ARO[drug_class[0]]
+        drug_class_instance_superclasses = list(drug_class_instance.superclasses())
+        superclasses_len = len(drug_class_instance_superclasses)
+
+        if superclasses_len >= 3:
+            drug_class_categories.append(drug_class_instance_superclasses[superclasses_len - 3].name)
+        else:
+            drug_class_categories.append(drug_class_instance_superclasses[0].name)
+
+    return drug_class_categories