Update pipeline.py

Author: cthoyt

src/semra/api.py

@@ -14,11 +14,19 @@
 import networkx as nx
 import pandas as pd
 import ssslm
+from pydantic import BaseModel, Field
 from ssslm import LiteralMapping
 from tqdm.auto import tqdm
 
 from semra.io.graph import _from_digraph_edge, to_digraph
-from semra.rules import EXACT_MATCH, FLIP, INVERSION_MAPPING, SubsetConfiguration
+from semra.rules import (
+    DB_XREF,
+    EXACT_MATCH,
+    FLIP,
+    INVERSION_MAPPING,
+    KNOWLEDGE_MAPPING,
+    SubsetConfiguration,
+)
 from semra.struct import (
     Evidence,
     Mapping,
@@ -35,6 +43,7 @@
     "IdentifierIndex",
     "Index",
     "M2MIndex",
+    "Mutation",
     "PrefixIdentifierDict",
     "PrefixIdentifierDict",
     "PrefixPairCounter",
@@ -60,6 +69,7 @@
     "get_terms",
     "get_test_evidence",
     "get_test_reference",
+    "handle_mutations",
     "hydrate_subsets",
     "keep_object_prefixes",
     "keep_prefixes",
@@ -1190,3 +1200,57 @@ def get_asymmetric_counter(
             for (left_prefix, right_prefix), identifiers in index.items()
         }
     )
+
+
+class Mutation(BaseModel):
+    """Represents a mutation operation on a mapping set."""
+
+    source: str = Field(..., description="The source type")
+    target: str | list[str] | None = Field(None, description="limit mutation to these")
+    confidence: float = 1.0
+    old: Reference = Field(default=DB_XREF)
+    new: Reference = Field(default=EXACT_MATCH)
+
+    def should_apply_to(self, mapping: Mapping) -> bool:
+        """Check if the mutation should be applied."""
+        if mapping.subject.prefix != self.source:
+            return False
+        if mapping.predicate != self.old:
+            return False
+        if self.target is None:
+            return True
+        elif isinstance(self.target, str):
+            return self.target == mapping.object.prefix
+        elif isinstance(self.target, list):
+            return any(t == mapping.object.prefix for t in self.target)
+        raise NotImplementedError
+
+
+def handle_mutations(
+    mappings: Iterable[Mapping], mutations: Iterable[Mutation], *, progress: bool = True
+) -> Iterable[Mapping]:
+    """Apply mutations."""
+    mutation_index = {}
+    for mutation__ in mutations:
+        if mutation__.source in mutation_index:
+            raise KeyError(f"got multiple configured mutations for source: {mutation__.source}")
+        mutation_index[mutation__.source] = mutation__
+    for mapping in tqdm(mappings, disable=not progress):
+        mutation = mutation_index.get(mapping.subject.prefix)
+        if not mutation:
+            yield mapping
+        elif not mutation.should_apply_to(mapping):
+            yield mapping
+        else:
+            yield Mapping(
+                subject=mapping.subject,
+                predicate=mutation.new,
+                object=mapping.object,
+                evidence=[
+                    ReasonedEvidence(
+                        justification=KNOWLEDGE_MAPPING,
+                        mappings=[mapping],
+                        confidence_factor=mutation.confidence,
+                    )
+                ],
+            )

src/semra/pipeline.py

@@ -18,11 +18,13 @@
 from typing_extensions import Self
 
 from semra.api import (
+    Mutation,
     assemble_evidences,
     filter_mappings,
     filter_prefixes,
     filter_self_matches,
     filter_subsets,
+    handle_mutations,
     hydrate_subsets,
     keep_prefixes,
     prioritize,
@@ -40,7 +42,7 @@
     write_neo4j,
     write_sssom,
 )
-from semra.rules import DB_XREF, EXACT_MATCH, IMPRECISE, SubsetConfiguration
+from semra.rules import IMPRECISE, SubsetConfiguration
 from semra.sources import SOURCE_RESOLVER
 from semra.sources.biopragmatics import (
     from_biomappings_negative,
@@ -100,16 +102,6 @@ class Input(BaseModel):
     extras: dict[str, Any] = Field(default_factory=dict)
 
 
-class Mutation(BaseModel):
-    """Represents a mutation operation on a mapping set."""
-
-    source: str = Field(..., description="The source type")
-    target: str | list[str] | None = Field(None, description="limit mutation to these")
-    confidence: float = 1.0
-    old: Reference = Field(default=DB_XREF)
-    new: Reference = Field(default=EXACT_MATCH)
-
-
 class Configuration(BaseModel):
     """Represents the steps taken during mapping assembly."""
 
@@ -974,7 +966,11 @@ def process(
     # _log_diff(before, mappings, verb="Filtered source internal", elapsed=time.time() - start)
 
     if mutations:
-        raise NotImplementedError
+        logger.info("Applying mutations")
+        before = len(mappings)
+        start = time.time()
+        mappings = list(handle_mutations(mappings, mutations, progress=progress))
+        _log_diff(before, mappings, verb="Applied mutations", elapsed=time.time() - start)
 
     if upgrade_prefixes and len(upgrade_prefixes) > 1:
         logger.info("Inferring mapping upgrades")