Merge pull request #35 from bgyori/indra_rels

Add more metadata to INDRA edges
gyorilab · Aug 30, 2021 · 5e5d20a · 5e5d20a
2 parents 042479e + f7a23a5
commit 5e5d20a
Showing 1 changed file with 20 additions and 3 deletions.
diff --git a/src/indra_cogex/sources/indra_db/__init__.py b/src/indra_cogex/sources/indra_db/__init__.py
@@ -35,7 +35,11 @@ class DbProcessor(Processor):
     def __init__(self, path: Union[None, str, Path] = None):
         """Initialize the INDRA database processor.
 
-        :param path: The path to the INDRA database SIF dump pickle. If none given, will look in the default location.
+        Parameters
+        ----------
+        path :
+            The path to the INDRA database SIF dump pickle. If none given,
+            will look in the default location.
         """
         if path is None:
             path = pystow.join("indra", "db", name="sif.pkl")
@@ -85,13 +89,16 @@ def get_nodes(self):  # noqa:D102
             yield Node(db_ns, db_id, ["BioEntity"], dict(name=name))
 
     def get_relations(self):  # noqa:D102
+        rel_type = "indra_rel"
         columns = [
             "agA_ns",
             "agA_id",
             "agB_ns",
             "agB_id",
             "stmt_type",
             "source_counts",
+            "evidence_count",
+            "belief",
             "stmt_hash",
         ]
         for (
@@ -101,17 +108,25 @@ def get_relations(self):  # noqa:D102
             target_id,
             stmt_type,
             source_counts,
+            evidence_count,
+            belief,
             stmt_hash,
         ) in (
             self.df[columns].drop_duplicates().values
         ):
-            data = {"stmt_hash:long": stmt_hash, "source_counts:string": source_counts}
+            data = {
+                "stmt_hash:long": stmt_hash,
+                "source_counts:string": source_counts,
+                "evidence_count:int": evidence_count,
+                "stmt_type:string": stmt_type,
+                "belief:float": belief,
+            }
             yield Relation(
                 source_ns,
                 source_id,
                 target_ns,
                 target_id,
-                stmt_type,
+                rel_type,
                 data,
             )
 
@@ -127,5 +142,7 @@ def fix_id(db_ns: str, db_id: str) -> Tuple[str, str]:
         db_ns = "UPLOC"
     if db_ns == "UP" and "-" in db_id and not db_id.startswith("SL-"):
         db_id = db_id.split("-")[0]
+    if db_ns == "FPLX" and db_id == "TCF-LEF":
+        db_id = "TCF_LEF"
     db_id = ensure_prefix_if_needed(db_ns, db_id)
     return db_ns, db_id