classify_disease_precision

related-sciences · Oct 16, 2023 · 136b3b0 · 136b3b0
1 parent c1f0812
commit 136b3b0
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/nxontology_data/efo/efo.py b/nxontology_data/efo/efo.py
@@ -293,6 +293,22 @@ def create_nxo(self) -> NXOntology[str]:
         )
         return nxo
 
+    def classify_disease_precision(self, nxo: NXOntology[str]) -> pd.DataFrame:
+        """
+        Use nxontology-ml to classify nodes in EFO OTAR Slim based on their disease precision.
+        Modifies nxo node attributes in place. Returns a pd.DataFrame of the predictions and features.
+        """
+        assert nxo.name == "efo_otar_slim"
+        nxo.freeze()
+        logger.info("Beginning nxontology-ml disease precision classification.")
+        precision_df = nxontology_ml_train_predict(nxo=nxo)
+        id_to_precision = {
+            row.identifier: row.precision for row in precision_df.itertuples()
+        }
+        for node, data in nxo.graph.nodes(data=True):
+            data["disease_precision"] = id_to_precision.get(node, "non_disease")
+        return precision_df
+
     def write_outputs(self) -> None:
         output_dir = get_source_output_dir("efo")
         nxo = self.create_nxo()
@@ -314,18 +330,11 @@ def write_outputs(self) -> None:
         if nxo.name == "efo_otar_profile":
             nxo_slim = self.create_slim_nxo(nxo)
             # classify EFO node/disease precision using nxontology-ml
-            nxo_slim.freeze()
-            logger.info("Beginning nxontology-ml disease precision classification.")
-            precision_df = nxontology_ml_train_predict(nxo=nxo_slim)
+            precision_df = self.classify_disease_precision(nxo_slim)
             write_dataframe(
                 precision_df,
                 output_dir.joinpath(f"{self.name}_precision_classifications.json.gz"),
             )
-            # TODO: set value for non-diseases
-            for row in precision_df.itertuples():
-                nxo_slim.graph.nodes[row.identifier][
-                    "disease_precision"
-                ] = row.precision
             write_ontology(nxo_slim, output_dir, compression_threshold_mb=30.0)
 
     @staticmethod