Generalize modifier-isoform collapse beyond the hardcoded ubiquitin list

src/logic_network_generator.py

@@ -9,7 +9,7 @@
 from src.neo4j_connector import get_graph
 from src.reaction_generator import (
     _complex_contains_entity_set,
-    _UBIQUITIN_ENTITY_SET_IDS,
+    modifier_isoform_set_ids,
     get_terminal_components,
     MAX_VARIANTS,
 )
@@ -555,7 +555,7 @@ def _complex_variant_leafsets(complex_id: str) -> List[frozenset]:
             per_component_choices.append(_complex_variant_leafsets(member_id))
         elif (
             any(lbl in labels for lbl in ("EntitySet", "DefinedSet", "CandidateSet"))
-            and member_id not in _UBIQUITIN_ENTITY_SET_IDS
+            and member_id not in modifier_isoform_set_ids()
         ):
             choices = [frozenset(get_terminal_components(sm)) for sm in get_set_members(member_id)]
             per_component_choices.append(choices or [frozenset(get_terminal_components(member_id))])
@@ -622,7 +622,7 @@ def _matching_leaves(entity_id: str) -> frozenset:
             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:
+        if entity_id in modifier_isoform_set_ids():
             out = {str(entity_id)}
         else:
             for m in get_set_members(entity_id):
@@ -687,7 +687,7 @@ def _map_annotated_entity_to_nodes(entity_id: str, member_set: Set[str]) -> Set[
         return {f"{entity_id}::variant::{'_'.join(sorted(chosen))}"}
 
     if any(lbl in labels for lbl in ("EntitySet", "DefinedSet", "CandidateSet")):
-        if entity_id in _UBIQUITIN_ENTITY_SET_IDS:
+        if entity_id in modifier_isoform_set_ids():
             return {str(entity_id)}
         # Expand the set to the member SPECIES this VR resolved to. Intersect
         # against the set's members at *matching* granularity (_matching_leaves),
@@ -836,7 +836,7 @@ def _decompose_regulator_entity(
         return result if result else [(entity_id, 1)]
 
     if "EntitySet" in labels or "DefinedSet" in labels or "CandidateSet" in labels:
-        if entity_id in _UBIQUITIN_ENTITY_SET_IDS:
+        if entity_id in modifier_isoform_set_ids():
             return [(entity_id, 1)]
         members = get_set_members(entity_id)
         result = []
@@ -889,7 +889,7 @@ def _expand_complex_variants(complex_id: str) -> List[tuple]:
             per_component_choices.append([vid for vid, _ in sub_variants])
         elif (
             ("EntitySet" in labels or "DefinedSet" in labels or "CandidateSet" in labels)
-            and member_id not in _UBIQUITIN_ENTITY_SET_IDS
+            and member_id not in modifier_isoform_set_ids()
         ):
             alts: List[str] = []
             for set_member in get_set_members(member_id):

src/neo4j_connector.py

@@ -586,6 +586,53 @@ def get_reaction_io_stoichiometry(reaction_id: str, input_or_output: str) -> Dic
     return result
 
 
+# Ubiquitin-like modifier protein gene families (+ RPS27A/UBA52 aliases). An
+# EntitySet whose members are ALL products of these genes is a modifier-isoform
+# set — functionally interchangeable copies of one modifier (ubiquitin's
+# UBB/UBC/RPS27A/UBA52, SUMO1/2/3, ATG8 homologues, …). E3 ligase / enzyme sets
+# are excluded because their members are not these genes.
+_MODIFIER_GENE_NAMES = frozenset({
+    "UBB", "UBC", "RPS27A", "UBA52", "UBA80", "UBCEP1", "UBCEP2",  # ubiquitin
+    "SUMO1", "SUMO2", "SUMO3", "SUMO4",                            # SUMO
+    "NEDD8", "ISG15", "UFM1", "UBD",                               # NEDD8/ISG15/UFM1/FAT10
+    "MAP1LC3A", "MAP1LC3B", "MAP1LC3B2", "MAP1LC3C",
+    "GABARAP", "GABARAPL1", "GABARAPL2",                           # ATG8/LC3 family
+})
+
+_modifier_isoform_set_cache: Optional[Set[str]] = None
+
+
+def get_modifier_isoform_entity_set_ids() -> Set[str]:
+    """stIds of human EntitySets whose members are ALL modifier proteins.
+
+    A modifier-isoform set (ubiquitin's UBB/UBC/RPS27A/UBA52, SUMO1/2/3, ATG8
+    homologues, …) holds functionally-identical copies of one modifier, so
+    decomposing it adds no biology and explodes the variant count. Enzyme sets
+    (e.g. E3 ligases) are excluded — their members aren't modifier genes.
+    Cached for the process. Raises if Neo4j is unreachable; callers wanting an
+    offline fallback should catch and use a hardcoded seed.
+    """
+    global _modifier_isoform_set_cache
+    if _modifier_isoform_set_cache is not None:
+        return _modifier_isoform_set_cache
+    query = """
+        MATCH (s)-[:hasMember|hasCandidate]->(m)
+        WHERE (s:DefinedSet OR s:CandidateSet) AND s.speciesName = 'Homo sapiens'
+        OPTIONAL MATCH (m)-[:referenceEntity]->(re)
+        WITH s, m, size([x IN coalesce(re.geneName, []) WHERE x IN $genes]) AS hit
+        WITH s, count(m) AS total, sum(CASE WHEN hit > 0 THEN 1 ELSE 0 END) AS mods
+        WHERE total >= 2 AND total = mods
+        RETURN collect(s.stId) AS stids
+    """
+    try:
+        data = get_graph().run(query, genes=list(_MODIFIER_GENE_NAMES)).data()
+    except Exception:
+        logger.error("Error in get_modifier_isoform_entity_set_ids", exc_info=True)
+        raise
+    _modifier_isoform_set_cache = set(data[0]["stids"]) if data else set()
+    return _modifier_isoform_set_cache
+
+
 def get_reference_entity_id(entity_id: str) -> Union[str, None]:
     if entity_id in _reference_entity_cache:
         return _reference_entity_cache[entity_id]

src/reaction_generator.py

@@ -168,6 +168,34 @@ def dataframe_for_uids(self, uids: Set[str]) -> pd.DataFrame:
     "R-HSA-9834963",  # Ub [mitochondrial outer membrane]
 }
 
+_modifier_set_cache: Optional[Set[str]] = None
+
+
+def modifier_isoform_set_ids() -> Set[str]:
+    """EntitySet stIds to treat atomically (never decompose into members).
+
+    Superset of _UBIQUITIN_ENTITY_SET_IDS: the full set of modifier-isoform sets
+    (all ubiquitin forms incl. linkage-specific chains, plus SUMO/NEDD8/ATG8/…),
+    discovered from Neo4j by member-gene identity. Falls back to the hardcoded
+    ubiquitin seed when Neo4j is unavailable (offline / tests). Cached; safe to
+    call in hot loops.
+    """
+    global _modifier_set_cache
+    if _modifier_set_cache is not None:
+        return _modifier_set_cache
+    try:
+        from src.neo4j_connector import get_modifier_isoform_entity_set_ids
+        _modifier_set_cache = (
+            get_modifier_isoform_entity_set_ids() | _UBIQUITIN_ENTITY_SET_IDS
+        )
+    except Exception:
+        logger.warning(
+            "Modifier-isoform set discovery failed; falling back to ubiquitin seed",
+            exc_info=True,
+        )
+        _modifier_set_cache = set(_UBIQUITIN_ENTITY_SET_IDS)
+    return _modifier_set_cache
+
 
 def get_component_id_or_reference_entity_id(reactome_id: str) -> str:
     """Get the reference entity ID for a Reactome stable ID, with caching.
@@ -466,7 +494,7 @@ def get_terminal_components(entity_id: str) -> Set[str]:
     if any(s in labels for s in ("EntitySet", "DefinedSet", "CandidateSet")):
         # Ubiquitin sets are atomic at the boundary too: decomposing them
         # would explode every boundary into ~14 indistinguishable copies.
-        if entity_id in _UBIQUITIN_ENTITY_SET_IDS:
+        if entity_id in modifier_isoform_set_ids():
             return {str(entity_id)}
         members = get_set_members(entity_id)
         leaves = set()
@@ -510,7 +538,7 @@ def break_apart_entity(entity_id: str, source_entity_id: Optional[str] = None) -
             return leaves
 
     if "EntitySet" in labels:
-        if entity_id in _UBIQUITIN_ENTITY_SET_IDS:
+        if entity_id in modifier_isoform_set_ids():
             return {str(entity_id)}
 
         member_ids = get_set_members(entity_id)