feat(associations): flexible association-type section matching

backend/src/monarch_py/api/entity.py

@@ -8,7 +8,6 @@
 from monarch_py.api.config import solr
 from monarch_py.api.additional_models import OutputFormat
 from monarch_py.datamodels.model import AssociationTableResults, Node
-from monarch_py.datamodels.category_enums import AssociationCategory
 from monarch_py.utils.format_utils import to_json, to_tsv
 
 router = APIRouter(tags=["entity"], responses={404: {"description": "Not Found"}})
@@ -55,8 +54,9 @@ def _association_table(
         title="ID of the entity to retrieve association table data for",
         examples=["MONDO:0019391"],
     ),
-    category: AssociationCategory = Path(
-        title="Type of association to retrieve association table data for",
+    category: str = Path(
+        title="Association-type section key (an AssociationTypeMapping key). For plain "
+        "single-category sections this is the biolink association category.",
         examples=["biolink:DiseaseToPhenotypicFeatureAssociation"],
     ),
     query: str = Query(default=None, title="query string to limit results to a subset", examples=["thumb"]),

backend/src/monarch_py/datamodels/model.py

@@ -278,14 +278,18 @@ class CompactAssociation(ConfiguredBaseModel):
 
 class AssociationTypeMapping(ConfiguredBaseModel):
     """
-    A data class to hold the necessary information to produce association type counts for given  entities with appropriate directional labels
+    A data class to hold the necessary information to produce association type counts for given entities with appropriate directional labels. Each match criterion (category, predicate, subject_category, object_category, primary_knowledge_source, provided_by) is an optional list; values within a single criterion are OR'd, and criteria are AND'd together. An omitted criterion places no constraint on that field.
     """
+    key: Optional[str] = Field(default=None, description="""A stable identifier for this association-type section, used to accumulate counts and as the table/section key. Defaults to the (single) category when not set.""")
     subject_label: Optional[str] = Field(default=None, description="""A label to describe the subjects of the association type as a whole for use in the UI""")
     object_label: Optional[str] = Field(default=None, description="""A label to describe the objects of the association type as a whole for use in the UI""")
     symmetric: bool = Field(default=False, description="""Whether the association type is symmetric, meaning that the subject and object labels should be interchangeable""")
-    category: str = Field(default=..., description="""The biolink category to use in queries for this association type""")
-    subject_category: Optional[str] = Field(default=None, description="""The biolink category of entities in the subject position of this association type""")
-    object_category: Optional[str] = Field(default=None, description="""The biolink category of entities in the object position of this association type""")
+    category: Optional[list[str]] = Field(default=None, description="""The biolink association category/categories to match (OR'd)""")
+    predicate: Optional[list[str]] = Field(default=None, description="""The predicate(s) to match (OR'd)""")
+    subject_category: Optional[list[str]] = Field(default=None, description="""The biolink category/categories of entities in the subject position of this association type (OR'd)""")
+    object_category: Optional[list[str]] = Field(default=None, description="""The biolink category/categories of entities in the object position of this association type (OR'd)""")
+    primary_knowledge_source: Optional[list[str]] = Field(default=None, description="""The primary knowledge source(s) to match (OR'd)""")
+    provided_by: Optional[list[str]] = Field(default=None, description="""The provided_by ingest source(s) to match (OR'd)""")
 
 
 class ExpandedAssociation(Association):
@@ -465,6 +469,7 @@ class FacetValue(ConfiguredBaseModel):
 
 
 class AssociationCount(FacetValue):
+    key: Optional[str] = Field(default=None, description="""Stable section identifier used as the frontend section key and the association-table key. Equals the category for plain single-category sections.""")
     category: Optional[str] = Field(default=None)
     count_direct: Optional[int] = Field(default=None, description="""Count of direct associations (no closure/descendants)""")
     count_with_orthologs: Optional[int] = Field(default=None, description="""Count including associations from orthologous genes""")

backend/src/monarch_py/datamodels/model.yaml

@@ -24,10 +24,14 @@ classes:
   AssociationCount:
     is_a: FacetValue
     slots:
+    - key
     - category
     - count_direct
     - count_with_orthologs
     slot_usage:
+      key:
+        description: Stable section identifier used as the frontend section key and the
+          association-table key. Equals the category for plain single-category sections.
       category:
         multivalued: false
   AssociationCountList:
@@ -75,15 +79,27 @@ classes:
         range: DirectionalAssociation
   AssociationTypeMapping:
     description: A data class to hold the necessary information to produce association
-      type counts for given  entities with appropriate directional labels
+      type counts for given entities with appropriate directional labels. Each match
+      criterion (category, predicate, subject_category, object_category,
+      primary_knowledge_source, provided_by) is an optional list; values within a
+      single criterion are OR'd, and criteria are AND'd together. An omitted criterion
+      places no constraint on that field.
     slots:
+    - key
     - subject_label
     - object_label
     - symmetric
     - category
+    - predicate
     - subject_category
     - object_category
+    - primary_knowledge_source
+    - provided_by
     slot_usage:
+      key:
+        description: A stable identifier for this association-type section, used to
+          accumulate counts and as the table/section key. Defaults to the (single)
+          category when not set.
       subject_label:
         description: A label to describe the subjects of the association type as a
           whole for use in the UI
@@ -96,15 +112,27 @@ classes:
         ifabsent: false
         required: true
       category:
-        description: The biolink category to use in queries for this association type
-        required: true
-        multivalued: false
+        description: The biolink association category/categories to match (OR'd)
+        required: false
+        multivalued: true
+      predicate:
+        description: The predicate(s) to match (OR'd)
+        required: false
+        multivalued: true
       subject_category:
-        description: The biolink category of entities in the subject position of this
-          association type
+        description: The biolink category/categories of entities in the subject position
+          of this association type (OR'd)
+        multivalued: true
       object_category:
-        description: The biolink category of entities in the object position of this
-          association type
+        description: The biolink category/categories of entities in the object position
+          of this association type (OR'd)
+        multivalued: true
+      primary_knowledge_source:
+        description: The primary knowledge source(s) to match (OR'd)
+        multivalued: true
+      provided_by:
+        description: The provided_by ingest source(s) to match (OR'd)
+        multivalued: true
   CategoryGroupedAssociationResults:
     is_a: Results
     slots:
@@ -555,6 +583,11 @@ slots:
     required: true
   category:
     multivalued: false
+  key:
+    description: A stable identifier for an association-type section, used to accumulate
+      counts and as the table/section key. Defaults to the category when not otherwise set.
+    range: string
+    multivalued: false
   causal_gene:
     description: A list of genes that are known to be causally associated with a disease
     range: Entity

backend/src/monarch_py/implementations/solr/solr_implementation.py

@@ -30,7 +30,7 @@
     EntityCategory,
     MappingPredicate,
 )
-from monarch_py.utils.association_type_utils import AssociationTypeMappings
+from monarch_py.utils.association_type_utils import AssociationTypeMappings, get_solr_criteria_filters
 from monarch_py.implementations.solr.solr_parsers import (
     convert_facet_fields,
     convert_facet_queries,
@@ -776,7 +776,7 @@ def get_association_facets(
     def get_association_table(
         self,
         entity: str,
-        category: AssociationCategory,
+        category: Union[AssociationCategory, str],
         traverse_orthologs: bool = False,
         direct: bool = False,
         q: Optional[str] = None,
@@ -787,6 +787,20 @@ def get_association_table(
         offset: int = 0,
         limit: int = 5,
     ) -> AssociationTableResults:
+        # `category` is really a section key: it resolves to an AssociationTypeMapping
+        # (whose criteria may combine several categories, a predicate, subject/object
+        # categories, etc.). An unknown key is treated as a literal category, which keeps
+        # backward compatibility for existing single-category sections and direct API use.
+        key = category.value if hasattr(category, "value") else category
+        mapping = AssociationTypeMappings.get_mapping_by_key(key)
+        if mapping and mapping.category:
+            categories = mapping.category
+            criteria_filters = get_solr_criteria_filters(mapping)
+        else:
+            categories = [key]
+            criteria_filters = []
+        table_filter_queries = (filter_queries or []) + criteria_filters
+
         entities = [entity]
         if traverse_orthologs:
             ortholog_associations = self.get_associations(
@@ -798,12 +812,12 @@ def get_association_table(
             entities.extend([ent.id for ent in orthologous_entities])
         query = build_association_table_query(
             entity=entities,
-            category=category.value,
+            category=categories,
             direct=direct,
             q=q,
             facet_fields=facet_fields,
             facet_queries=facet_queries,
-            filter_queries=filter_queries,
+            filter_queries=table_filter_queries,
             sort=sort,
             offset=offset,
             limit=limit,
@@ -1077,12 +1091,13 @@ def get_generic_entity_grid(
         mapping = AssociationTypeMappings.get_mapping(first_col_cat)
 
         if mapping and mapping.subject_category and mapping.object_category:
-            # Use YAML metadata to determine direction
-            if context_category == mapping.subject_category:
+            # Use YAML metadata to determine direction. subject_category and
+            # object_category are multivalued (list[str]), so test membership.
+            if context_category in mapping.subject_category:
                 context_field = "subject"
                 column_field = "object"
                 context_closure_field = "subject_closure"
-            elif context_category == mapping.object_category:
+            elif context_category in mapping.object_category:
                 context_field = "object"
                 column_field = "subject"
                 context_closure_field = "object_closure"

backend/src/monarch_py/implementations/solr/solr_parsers.py

@@ -26,7 +26,7 @@
 from monarch_py.datamodels.solr import HistoPhenoKeys, SolrQueryResult
 from monarch_py.service.curie_service import converter
 from monarch_py.implementations.solr.solr_query_utils import build_association_count_suffixes
-from monarch_py.utils.association_type_utils import get_association_type_mapping_by_query_string
+from monarch_py.utils.association_type_utils import AssociationTypeMappings, get_solr_query_fragment
 from monarch_py.utils.utils import get_links_for_field, get_provided_by_link
 
 
@@ -187,55 +187,57 @@ def parse_association_counts(query_result: SolrQueryResult, entities: List[str])
         "orthologs_object": ("orthologs", False),
     }
 
-    # Collect counts into a dict keyed by (label, category)
-    # Each entry accumulates direct, closure, and ortholog counts
-    count_data: Dict[str, Dict] = {}
-
-    def _ensure_entry(label: str, category: str):
-        if label not in count_data:
-            count_data[label] = {"category": category, "direct": 0, "closure": 0, "orthologs": 0}
-
-    for k, v in query_result.facet_counts.facet_queries.items():
-        if v == 0:
-            continue
-
-        # Determine which suffix matches and which count level it represents
-        count_level = None
-        is_subject_direction = None
-        original_query = None
-
+    # Rebuild the exact facet query string for each (mapping, suffix) so each result
+    # can be attributed to its mapping directly, without re-parsing the Solr logic.
+    # This is robust when several mappings share a category (e.g. biolink:Association)
+    # and are distinguished only by predicate / subject / object category.
+    lookup = {}
+    for mapping in AssociationTypeMappings.get_mappings():
+        fragment = get_solr_query_fragment(mapping)
         for suffix_key, suffix_str in suffixes.all_suffixes.items():
-            if k.endswith(suffix_str):
-                original_query = k.replace(f" {suffix_str}", "").lstrip("(").rstrip(")")
-                count_level, is_subject_direction = suffix_metadata[suffix_key]
-                break
-
-        if count_level is None:
-            raise ValueError(f"Unexpected facet query when building association counts: {k}")
+            lookup[f"({fragment}) {suffix_str}"] = (mapping, *suffix_metadata[suffix_key])
 
-        agm = get_association_type_mapping_by_query_string(original_query)
-        label = agm.subject_label if is_subject_direction else agm.object_label
+    # Accumulate counts keyed by the section key
+    count_data: Dict[str, Dict] = {}
 
-        _ensure_entry(label, agm.category)
+    for query_string, value in query_result.facet_counts.facet_queries.items():
+        if value == 0:
+            continue
+        if query_string not in lookup:
+            raise ValueError(f"Unexpected facet query when building association counts: {query_string}")
+
+        mapping, count_level, is_subject_direction = lookup[query_string]
+        key = mapping.key
+        if key not in count_data:
+            count_data[key] = {
+                "key": key,
+                "category": mapping.category[0] if mapping.category else None,
+                "label": None,
+                "direct": 0,
+                "closure": 0,
+                "orthologs": 0,
+            }
+        entry = count_data[key]
+        entry["label"] = mapping.subject_label if is_subject_direction else mapping.object_label
 
         # For symmetric associations, sum both directions; otherwise set the value
-        if agm.symmetric and count_data[label][count_level] > 0:
-            count_data[label][count_level] += v
+        if mapping.symmetric and entry[count_level] > 0:
+            entry[count_level] += value
         else:
-            count_data[label][count_level] = v
+            entry[count_level] = value
 
     # Build AssociationCount objects
-    items = []
-    for label, data in count_data.items():
-        items.append(
-            AssociationCount(
-                label=label,
-                count=data["closure"],
-                count_direct=data["direct"],
-                count_with_orthologs=data["orthologs"] if has_orthologs else None,
-                category=data["category"],
-            )
+    items = [
+        AssociationCount(
+            key=data["key"],
+            label=data["label"],
+            count=data["closure"],
+            count_direct=data["direct"],
+            count_with_orthologs=data["orthologs"] if has_orthologs else None,
+            category=data["category"],
         )
+        for data in count_data.values()
+    ]
 
     return AssociationCountList(items=items)
 

backend/src/monarch_py/implementations/solr/solr_query_utils.py

@@ -147,7 +147,7 @@ def build_association_query(
 
 def build_association_table_query(
     entity: List[str],
-    category: str,
+    category: List[str],
     direct: bool = False,
     q: Optional[str] = None,
     facet_fields: List[str] = None,
@@ -169,7 +169,7 @@ def build_association_table_query(
 
     query = build_association_query(
         entity=entity,
-        category=[category],
+        category=category,
         q=q,
         sort=sort,
         offset=offset,