Expand EntitySets to member species (no bare set nodes)

src/logic_network_generator.py

@@ -591,6 +591,50 @@ def _complex_variant_leafsets(complex_id: str) -> List[frozenset]:
     return result
 
 
+_matching_leaves_cache: Dict[str, frozenset] = {}
+
+
+def _matching_leaves(entity_id: str) -> frozenset:
+    """The member species of `entity_id` at the MATCHING layer's granularity.
+
+    Mirrors :func:`reaction_generator.break_apart_entity` as a pure read (no
+    store writes): EntitySets expand to their members; a Complex is kept atomic
+    UNLESS it contains an EntitySet (then it decomposes); leaves return
+    themselves. Crucially this preserves modified species (e.g. p-ERK), unlike
+    :func:`get_terminal_components`, which collapses to the base reference
+    protein. Used so emission can line a set participant up with the terminal
+    reactome-ids the VR actually resolved to.
+    """
+    from src.neo4j_connector import get_labels, get_complex_components, get_set_members
+
+    if entity_id in _matching_leaves_cache:
+        return _matching_leaves_cache[entity_id]
+    try:
+        labels = get_labels(entity_id)
+    except IndexError:
+        labels = []
+
+    out: Set[str] = set()
+    if "Complex" in labels:
+        if not _complex_contains_entity_set(entity_id):
+            out = {str(entity_id)}
+        else:
+            for m in get_complex_components(entity_id):
+                out |= _matching_leaves(m)
+    elif any(s in labels for s in ("EntitySet", "DefinedSet", "CandidateSet")):
+        if entity_id in _UBIQUITIN_ENTITY_SET_IDS:
+            out = {str(entity_id)}
+        else:
+            for m in get_set_members(entity_id):
+                out |= _matching_leaves(m)
+    else:
+        out = {str(entity_id)}
+
+    result = frozenset(out or {str(entity_id)})
+    _matching_leaves_cache[entity_id] = result
+    return result
+
+
 def _map_annotated_entity_to_nodes(entity_id: str, member_set: Set[str]) -> Set[str]:
     """Map one reaction-annotated input/output entity to its emission node(s).
 
@@ -637,7 +681,19 @@ def _map_annotated_entity_to_nodes(entity_id: str, member_set: Set[str]) -> Set[
     if any(lbl in labels for lbl in ("EntitySet", "DefinedSet", "CandidateSet")):
         if entity_id in _UBIQUITIN_ENTITY_SET_IDS:
             return {str(entity_id)}
-        present = get_terminal_components(entity_id) & member_set
+        # Expand the set to the member SPECIES this VR resolved to. Intersect
+        # against the set's members at *matching* granularity (_matching_leaves),
+        # NOT get_terminal_components: the latter collapses modified species down
+        # to their base reference protein (e.g. p-ERK R-HSA-109844 -> ERK
+        # R-HSA-109842), while the matching layer keeps the modified form, so the
+        # intersection came up empty and the set fell back to a bare, producer-less
+        # node. Aligning the two granularities lets the set expand to its members
+        # (sets should never survive as nodes — see docs/DESIGN_DECISIONS.md).
+        if os.environ.get("LNG_SET_EXPAND", "1") == "0":
+            # Old behavior (leaf-granularity intersection) for A/B comparison.
+            present = member_set & get_terminal_components(entity_id)
+        else:
+            present = member_set & _matching_leaves(entity_id)
         return present if present else {str(entity_id)}
 
     return {str(entity_id)}  # simple entity (protein / small molecule / …)

tests/test_logic_network_generator.py

@@ -690,3 +690,56 @@ def test_under_cap_still_expands(self, monkeypatch):
             f"3x3 cartesian under cap should yield 9 variants, got {len(result)}"
         )
         assert all(vid.startswith("C::variant::") for vid, _ in result)
+
+
+class TestMatchingLeavesGranularity:
+    """_matching_leaves mirrors the MATCHING layer's decomposition depth so a
+    set participant lines up with the terminal reactome-ids a VR resolved to.
+
+    The bug it fixes: get_terminal_components collapses modified species down to
+    their base reference protein (p-ERK -> ERK), while the matching layer keeps
+    the modified form. Intersecting the two came up empty, so EntitySets fell
+    back to bare, producer-less nodes. _matching_leaves keeps Complexes atomic
+    (unless they contain a set) — the same rule break_apart_entity uses.
+    """
+
+    def _clear(self):
+        import src.logic_network_generator as m
+        m._matching_leaves_cache.clear()
+
+    def test_set_of_atomic_complexes_stops_at_the_complexes(self, monkeypatch):
+        import src.logic_network_generator as m
+        from src import neo4j_connector as nc
+        self._clear()
+        # Set S -> {C1, C2}; C1,C2 are Complexes WITHOUT internal sets (atomic).
+        monkeypatch.setattr(nc, "get_labels", lambda e: {
+            "S": ["EntitySet", "DefinedSet"], "C1": ["Complex"], "C2": ["Complex"],
+        }.get(e, ["EntityWithAccessionedSequence"]))
+        monkeypatch.setattr(nc, "get_set_members",
+                            lambda e: {"C1", "C2"} if e == "S" else set())
+        monkeypatch.setattr(nc, "get_complex_components",
+                            lambda e: {"P1": 1, "P2": 1} if e in ("C1", "C2") else {})
+        monkeypatch.setattr(m, "_complex_contains_entity_set", lambda e: False)
+        result = m._matching_leaves("S")
+        assert result == frozenset({"C1", "C2"}), (
+            "a set of atomic complexes must resolve to the complexes (matching "
+            f"granularity), not their proteins; got {set(result)}"
+        )
+
+    def test_complex_with_internal_set_decomposes(self, monkeypatch):
+        import src.logic_network_generator as m
+        from src import neo4j_connector as nc
+        self._clear()
+        # Complex C contains set S={A,B} plus protein P -> decompose to {A,B,P}.
+        monkeypatch.setattr(nc, "get_labels", lambda e: {
+            "C": ["Complex"], "S": ["EntitySet"],
+        }.get(e, ["EntityWithAccessionedSequence"]))
+        monkeypatch.setattr(nc, "get_complex_components",
+                            lambda e: {"S": 1, "P": 1} if e == "C" else {})
+        monkeypatch.setattr(nc, "get_set_members",
+                            lambda e: {"A", "B"} if e == "S" else set())
+        monkeypatch.setattr(m, "_complex_contains_entity_set", lambda e: e == "C")
+        result = m._matching_leaves("C")
+        assert result == frozenset({"A", "B", "P"}), (
+            f"complex-with-set must decompose to members, got {set(result)}"
+        )