Merge pull request #47 from bcgsc/bugfix/KBDEV-1266-slow-matching

Fix redundand disease matching

Author: mathieulemieux

pori_python/ipr/annotate.py

@@ -57,7 +57,7 @@ def get_second_pass_variants(
 
 
 def get_ipr_statements_from_variants(
-    graphkb_conn: GraphKBConnection, matches: List[Variant], disease_name: str
+    graphkb_conn: GraphKBConnection, matches: List[Variant], disease_matches: List[str]
 ) -> List[KbMatch]:
     """IPR upload formatted GraphKB statements from the list of variants.
 
@@ -72,7 +72,7 @@ def get_ipr_statements_from_variants(
     existing_statements = {s["@rid"] for s in statements}
 
     for ipr_row in convert_statements_to_alterations(
-        graphkb_conn, statements, disease_name, convert_to_rid_set(matches)
+        graphkb_conn, statements, disease_matches, convert_to_rid_set(matches)
     ):
         rows.append(ipr_row)
 
@@ -88,7 +88,7 @@ def get_ipr_statements_from_variants(
     for ipr_row in convert_statements_to_alterations(
         graphkb_conn,
         inferred_statements,
-        disease_name,
+        disease_matches,
         convert_to_rid_set(inferred_matches),
     ):
         ipr_row["kbData"]["inferred"] = True
@@ -99,15 +99,17 @@ def get_ipr_statements_from_variants(
 
 def annotate_expression_variants(
     graphkb_conn: GraphKBConnection,
+    disease_matches: List[str],
     variants: List[IprExprVariant],
-    disease_name: str,
     show_progress: bool = False,
 ) -> List[KbMatch]:
     """Annotate expression variants with GraphKB in the IPR alterations format.
 
     Args:
-        graphkb_conn: the graphkb api connection object
-        variants: list of variants
+        graphkb_conn (GraphKBConnection): the graphkb api connection object
+        disease_matches (list.str): GraphKB disease RIDs
+        variants (list.IprExprVariant): list of variants.
+        show_progress (bool): Progressbar displayed for long runs.
 
     Returns:
         list of kbMatches records for IPR
@@ -127,7 +129,7 @@ def annotate_expression_variants(
             continue
         try:
             matches = gkb_match.match_expression_variant(graphkb_conn, gene, variant)
-            for ipr_row in get_ipr_statements_from_variants(graphkb_conn, matches, disease_name):
+            for ipr_row in get_ipr_statements_from_variants(graphkb_conn, matches, disease_matches):
                 ipr_row["variant"] = row["key"]
                 ipr_row["variantType"] = row.get("variantType", "exp")
                 alterations.append(ipr_row)
@@ -150,15 +152,17 @@ def annotate_expression_variants(
 
 def annotate_copy_variants(
     graphkb_conn: GraphKBConnection,
+    disease_matches: List[str],
     variants: List[IprCopyVariant],
-    disease_name: str,
     show_progress: bool = False,
 ) -> List[KbMatch]:
     """Annotate allowed copy variants with GraphKB in the IPR alterations format.
 
     Args:
-        graphkb_conn: the graphkb api connection object
-        variants: list of variants
+        graphkb_conn (GraphKBConnection): the graphkb api connection object
+        disease_matches (list.str): GraphKB disease RIDs
+        variants (list.IprCopyVariant): list of variants.
+        show_progress (bool): Progressbar displayed for long runs.
 
     Returns:
         list of kbMatches records for IPR
@@ -180,7 +184,7 @@ def annotate_copy_variants(
             continue
         try:
             matches = gkb_match.match_copy_variant(graphkb_conn, gene, variant)
-            for ipr_row in get_ipr_statements_from_variants(graphkb_conn, matches, disease_name):
+            for ipr_row in get_ipr_statements_from_variants(graphkb_conn, matches, disease_matches):
                 ipr_row["variant"] = row["key"]
                 ipr_row["variantType"] = row.get("variantType", "cnv")
                 alterations.append(ipr_row)
@@ -206,16 +210,16 @@ def annotate_copy_variants(
 def annotate_positional_variants(
     graphkb_conn: GraphKBConnection,
     variants: Sequence[IprStructuralVariant] | Sequence[Hashabledict],
-    disease_name: str,
+    disease_matches: List[str],
     show_progress: bool = False,
 ) -> List[Hashabledict]:
     """Annotate SNP, INDEL or fusion variant calls with GraphKB and return in IPR match format.
 
     Hashable type is required to turn lists into sets.
     Args:
         graphkb_conn (GraphKBConnection): the graphkb api connection object
-        variants (list.<dict>): list of variants. Defaults to [].
-        disease_name (str): GraphKB disease name for statement matching.  'cancer' is most general
+        variants (Sequence): list of variants.
+        disease_matches (list.str): GraphKB disease RIDs
         show_progress (bool): Progressbar displayed for long runs.
 
     Returns:
@@ -259,7 +263,9 @@ def annotate_positional_variants(
                         raise parse_err
 
                 for ipr_row in get_ipr_statements_from_variants(
-                    graphkb_conn, matches, disease_name
+                    graphkb_conn,
+                    matches,
+                    disease_matches,
                 ):
                     ipr_row["variant"] = row["key"]
                     ipr_row["variantType"] = row.get(
@@ -304,15 +310,16 @@ def annotate_positional_variants(
 
 def annotate_msi(
     graphkb_conn: GraphKBConnection,
-    disease_name: str = "cancer",
+    disease_matches: List[str],
     msi_category: str = "microsatellite instability",
 ) -> List[KbMatch]:
     """Annotate microsatellite instablity from GraphKB in the IPR alterations format.
 
     Match to GraphKb Category variants with similar names
     Args:
-        graphkb_conn: the graphkb api connection object
-        msi_category: such as 'microsatellite instability'
+        graphkb_conn (GraphKBConnection): the graphkb api connection object
+        disease_matches (list.str): GraphKB disease RIDs
+        msi_category (str): such as 'microsatellite instability'
 
     Returns:
         list of kbMatches records for IPR
@@ -335,7 +342,9 @@ def annotate_msi(
     )
     if msi_categories:
         msi_variants = [cast(Variant, var) for var in msi_categories]
-        for ipr_row in get_ipr_statements_from_variants(graphkb_conn, msi_variants, disease_name):
+        for ipr_row in get_ipr_statements_from_variants(
+            graphkb_conn, msi_variants, disease_matches
+        ):
             ipr_row["variant"] = msi_category
             ipr_row["variantType"] = "msi"
             gkb_matches.append(ipr_row)
@@ -344,15 +353,15 @@ def annotate_msi(
 
 def annotate_tmb(
     graphkb_conn: GraphKBConnection,
-    disease_name: str = "cancer",
+    disease_matches: List[str],
     category: str = TMB_HIGH_CATEGORY,
 ) -> List[KbMatch]:
     """Annotate Tumour Mutation Burden (tmb) categories from GraphKB in the IPR alterations format.
 
     Match to GraphKb Category variants with similar names
     Args:
-        graphkb_conn: the graphkb api connection object
-        disease_name: oncotree disease name for graphkb matching.
+        graphkb_conn (GraphKBConnection): the graphkb api connection object
+        disease_matches (list.str): GraphKB disease RIDs
         category: such as 'high mutation burden'
 
     Returns:
@@ -376,7 +385,9 @@ def annotate_tmb(
     )
     if categories:
         cat_variants = [cast(Variant, var) for var in categories]
-        for ipr_row in get_ipr_statements_from_variants(graphkb_conn, cat_variants, disease_name):
+        for ipr_row in get_ipr_statements_from_variants(
+            graphkb_conn, cat_variants, disease_matches
+        ):
             ipr_row["variant"] = category
             ipr_row["variantType"] = "tmb"
             gkb_matches.append(ipr_row)
@@ -385,19 +396,19 @@ def annotate_tmb(
 
 def annotate_signature_variants(
     graphkb_conn: GraphKBConnection,
+    disease_matches: List[str],
     variants: List[IprSignatureVariant] = [],
-    disease_name: str = "cancer",
     show_progress: bool = False,
 ) -> List[KbMatch]:
     """Annotate Signature variants with GraphKB in the IPR alterations format.
 
     Match to corresponding GraphKB Variants, then to linked GraphKB Statements
 
     Args:
-        graphkb_conn: the graphkb api connection object
-        variants: list of signature variants
-        disease_name: oncotree disease name for graphkb matching
-        show_progress: progressbar displayed for long runs; default to False
+        graphkb_conn (GraphKBConnection): the graphkb api connection object
+        disease_matches (list.str): GraphKB disease RIDs
+        variants (list.IprSignatureVariant): list of signature variants
+        show_progress (bool): progressbar displayed for long runs; default to False
 
     Returns:
         list of kbMatches records for IPR
@@ -416,7 +427,7 @@ def annotate_signature_variants(
             )
             # Matching GKB Variants to GKB Statements
             for ipr_row in get_ipr_statements_from_variants(
-                graphkb_conn, matched_variants, disease_name
+                graphkb_conn, matched_variants, disease_matches
             ):
                 ipr_row["variant"] = variant["key"]
                 ipr_row["variantType"] = "sigv"

pori_python/ipr/ipr.py

@@ -116,15 +116,15 @@ def link_refs(refs) -> Tuple[str, str]:
 def convert_statements_to_alterations(
     graphkb_conn: GraphKBConnection,
     statements: List[Statement],
-    disease_name: str,
+    disease_matches: List[str],
     variant_matches: Iterable[str],
 ) -> List[KbMatch]:
     """Convert statements matched from graphkb into IPR equivalent representations.
 
     Args:
         graphkb_conn: the graphkb connection object
         statements: list of statement records from graphkb
-        disease_name: name of the cancer type for the patient being reported on
+        disease_matches: GraphKB disease RIDs
         variant_matches: the list of RIDs the variant matched for these statements
 
     Raises:
@@ -136,14 +136,6 @@ def convert_statements_to_alterations(
     Notes:
         - only report disease matched prognostic markers https://www.bcgsc.ca/jira/browse/GERO-72 and GERO-196
     """
-    disease_matches = {
-        r["@rid"]
-        for r in gkb_vocab.get_term_tree(graphkb_conn, disease_name, ontology_class="Disease")
-    }
-
-    if not disease_matches:
-        raise ValueError(f"failed to match disease ({disease_name}) to graphkb")
-
     rows = []
     ev_map = get_evidencelevel_mapping(graphkb_conn)
     # GERO-318 - add all IPR-A evidence equivalents to the approvedTherapy flag
@@ -632,3 +624,33 @@ def get_kb_matches_sections(
         "kbMatchedStatements": kb_matched_statements,
         "kbStatementMatchedConditions": kb_statement_matched_conditions,
     }
+
+
+def get_kb_disease_matches(
+    graphkb_conn: GraphKBConnection, kb_disease_match: str = None, verbose: bool = True
+) -> list[str]:
+
+    if not kb_disease_match:
+        kb_disease_match = 'cancer'
+        if verbose:
+            logger.warning(f"No disease provided; will use '{kb_disease_match}'")
+
+    if verbose:
+        logger.info(f"Matching disease ({kb_disease_match}) to graphkb")
+
+    disease_matches = {
+        r["@rid"]
+        for r in gkb_vocab.get_term_tree(
+            graphkb_conn,
+            kb_disease_match,
+            ontology_class="Disease",
+        )
+    }
+
+    if not disease_matches:
+        msg = f"failed to match disease ({kb_disease_match}) to graphkb"
+        if verbose:
+            logger.error(msg)
+        raise ValueError(msg)
+
+    return disease_matches

pori_python/ipr/main.py

@@ -48,6 +48,7 @@
     create_key_alterations,
     filter_structural_variants,
     germline_kb_matches,
+    get_kb_disease_matches,
     get_kb_matches_sections,
     select_expression_plots,
 )
@@ -338,8 +339,6 @@ def ipr_report(
         logger.error("Failed schema check - report variants may be corrupted or unmatched.")
         logger.error(f"Failed schema check: {err}")
 
-    kb_disease_match = content["kbDiseaseMatch"]
-
     # validate the input variants
     signatureVariants: List[IprSignatureVariant] = preprocess_signature_variants(
         [
@@ -380,6 +379,14 @@ def ipr_report(
 
     graphkb_conn.login(gkb_user, gkb_pass)
 
+    # DISEASE
+    # Disease term from bioapps; expected OncoTree term
+    kb_disease_match: str = content["kbDiseaseMatch"]
+
+    # Matching disease RIDs from GraphKB using term tree
+    # (Will raise uncatched error if no match)
+    disease_matches: list[str] = get_kb_disease_matches(graphkb_conn, kb_disease_match)
+
     # GKB MATCHING
     gkb_matches: List[Hashabledict] = []
 
@@ -405,7 +412,7 @@ def ipr_report(
                 tmb["kbCategory"] = TMB_HIGH_CATEGORY
 
             # GERO-296 - try matching to graphkb
-            tmb_matches = annotate_tmb(graphkb_conn, kb_disease_match, TMB_HIGH_CATEGORY)
+            tmb_matches = annotate_tmb(graphkb_conn, disease_matches, TMB_HIGH_CATEGORY)
             if tmb_matches:
                 tmb_variant["kbCategory"] = TMB_HIGH_CATEGORY  # type: ignore
                 tmb_variant["variant"] = TMB_HIGH_CATEGORY
@@ -431,7 +438,7 @@ def ipr_report(
             msi_cat = msi.get("kbCategory")
             msi_variant = msi.copy()
         logger.info(f"Matching GKB msi {msi_cat}")
-        msi_matches = annotate_msi(graphkb_conn, kb_disease_match, msi_cat)
+        msi_matches = annotate_msi(graphkb_conn, disease_matches, msi_cat)
         if msi_matches:
             msi_variant["kbCategory"] = msi_cat  # type: ignore
             msi_variant["variant"] = msi_cat
@@ -445,7 +452,7 @@ def ipr_report(
     logger.info(f"annotating {len(signatureVariants)} signatures")
     gkb_matches.extend(
         annotate_signature_variants(
-            graphkb_conn, signatureVariants, kb_disease_match, show_progress=interactive
+            graphkb_conn, disease_matches, signatureVariants, show_progress=interactive
         )
     )
     logger.debug(f"\tgkb_matches: {len(gkb_matches)}")
@@ -454,7 +461,7 @@ def ipr_report(
     logger.info(f"annotating {len(small_mutations)} small mutations")
     gkb_matches.extend(
         annotate_positional_variants(
-            graphkb_conn, small_mutations, kb_disease_match, show_progress=interactive
+            graphkb_conn, small_mutations, disease_matches, show_progress=interactive
         )
     )
     logger.debug(f"\tgkb_matches: {len(gkb_matches)}")
@@ -465,7 +472,7 @@ def ipr_report(
         annotate_positional_variants(
             graphkb_conn,
             structural_variants,
-            kb_disease_match,
+            disease_matches,
             show_progress=interactive,
         )
     )
@@ -477,7 +484,7 @@ def ipr_report(
         [
             Hashabledict(copy_var)
             for copy_var in annotate_copy_variants(
-                graphkb_conn, copy_variants, kb_disease_match, show_progress=interactive
+                graphkb_conn, disease_matches, copy_variants, show_progress=interactive
             )
         ]
     )
@@ -490,8 +497,8 @@ def ipr_report(
             Hashabledict(exp_var)
             for exp_var in annotate_expression_variants(
                 graphkb_conn,
+                disease_matches,
                 expression_variants,
-                kb_disease_match,
                 show_progress=interactive,
             )
         ]
@@ -550,7 +557,7 @@ def ipr_report(
         graphkb_comments = auto_analyst_comments(
             graphkb_conn,
             gkb_matches,
-            disease_name=kb_disease_match,
+            disease_matches=set(disease_matches),
             variants=all_variants,
         )
         comments_list.append(graphkb_comments)