Further refine concept composition

mira/metamodel/composition.py

@@ -74,10 +74,6 @@ def compose_two_models(tm0, tm1):
                                                       tm0.observables)
                        for key, value in d.items()}
 
-    # Create a copy of templates in case we need to modify template rate laws
-    # when doing time substitutions
-    tm0.templates = deepcopy(tm0.templates)
-    tm1.templates = deepcopy(tm1.templates)
 
     new_annotations = annotation_composition(tm0.annotations,
                                              tm1.annotations)
@@ -89,13 +85,69 @@ def compose_two_models(tm0, tm1):
         # get the union of all template model attributes
         # as the models are 100% distinct
         # prioritize tm0
-        new_templates = tm0.templates + tm1.templates
-        new_parameters = {**tm1.parameters, **tm0.parameters}
-        new_initials = {**tm1.initials, **tm0.initials}
+        new_templates = deepcopy(tm0.templates) + deepcopy(tm1.templates)
+        combined_parameters = {**tm1.parameters, **tm0.parameters}
+        combined_initials = {**tm1.initials, **tm0.initials}
+
+        outer_tm_concepts_name_map = tm0.get_concepts_name_map()
+        inner_tm_concepts_name_map = tm1.get_concepts_name_map()
+
+        shared_concept_names = set(outer_tm_concepts_name_map.keys()) & set(
+            inner_tm_concepts_name_map.keys()
+        )
+
+        for shared_concept_name in shared_concept_names:
+            outer_tm_concept = outer_tm_concepts_name_map[shared_concept_name]
+            inner_tm_concept = inner_tm_concepts_name_map[shared_concept_name]
+
+            # handles 1 case where neither concept has an identifier
+            if (
+                not outer_tm_concept.identifiers
+                and not inner_tm_concept.identifiers
+            ):
+                continue
+
+            # handles 1 case where the inner has an identifier
+            elif (
+                not outer_tm_concept.identifiers
+                and inner_tm_concept.identifiers
+            ):
+                for new_template in new_templates:
+                    if hasattr(new_template, "subject"):
+                        if new_template.subject.name == shared_concept_name:
+                            new_template.subject.identifiers = inner_tm_concept.identifiers
+                    if hasattr(new_template, "outcome"):
+                        if new_template.outcome.name == shared_concept_name:
+                            new_template.outcome.identifiers = inner_tm_concept.identifiers
+                    if hasattr(new_template, "controller"):
+                        if new_template.controller.name == shared_concept_name:
+                            new_template.controller.identifiers = inner_tm_concept.identifiers
+
+            # handles 3 cases: 1. if the outer concept has an identifier and the inner doesn't
+            # 2. if the outer concept and inner concept both have an identifier but
+            # aren't equal then we prioritize the outer concept
+            # 3. if both concepts' identifiers are equal then we do nothing
+            elif outer_tm_concept.identifiers:
+                if (
+                    inner_tm_concept.identifiers
+                    and outer_tm_concept.identifiers
+                    != outer_tm_concept.identifiers
+                ) or not inner_tm_concept.identifiers:
+                    for new_template in new_templates:
+                        if hasattr(new_template, "subject"):
+                            if new_template.subject.name == shared_concept_name:
+                                new_template.subject.identifiers = outer_tm_concept.identifiers
+                        if hasattr(new_template, "outcome"):
+                            if new_template.outcome.name == shared_concept_name:
+                                new_template.outcome.identifiers = outer_tm_concept.identifiers
+                        if hasattr(new_template, "controller"):
+                            if new_template.controller.name == shared_concept_name:
+                                new_template.controller.identifiers = outer_tm_concept.identifiers
+
 
         composed_tm = TemplateModel(templates=new_templates,
-                                    parameters=new_parameters,
-                                    initials=new_initials,
+                                    parameters=combined_parameters,
+                                    initials=combined_initials,
                                     observables=new_observables,
                                     annotations=new_annotations,
                                     time=new_time)
@@ -108,7 +160,7 @@ def compose_two_models(tm0, tm1):
         new_templates = []
         new_parameters = {}
         new_initials = {}
-        concept_map = {}
+        replaced_concept_map = {}
 
         # We wouldn't have an edge from a template to a concept node,
         # so we only need to check if the source edge tuple contains a template
@@ -141,8 +193,7 @@ def compose_two_models(tm0, tm1):
             replaced_concept = compare_graph.concept_nodes[replaced_tm_id][
                 replaced_concept_id]
             new_concept = compare_graph.concept_nodes[new_tm_id][new_concept_id]
-            concept_map.setdefault(replaced_concept.name, set())
-            concept_map[replaced_concept.name].add(new_concept.name)
+            replaced_concept_map[replaced_concept.name] = new_concept.name
 
         # process templates that are present in a relation first
         # we only process the source template because either it's a template
@@ -154,11 +205,10 @@ def compose_two_models(tm0, tm1):
             new_tm, new_template = compare_graph.template_models[
                 new_tm_id], \
                 compare_graph.template_nodes[new_tm_id][new_template_id]
-
             process_template(new_templates, new_template, new_tm,
-                             new_parameters, new_initials)
+                             new_parameters, new_initials, replaced_concept_map)
 
-        update_observable_expressions(new_observables, concept_map)
+        update_observable_expressions(new_observables, replaced_concept_map)
 
         for outer_template_id, outer_template in enumerate(tm0.templates):
             for inner_template_id, inner_template in enumerate(tm1.templates):
@@ -168,16 +218,36 @@ def compose_two_models(tm0, tm1):
 
                 # process inner template first such that outer_template from
                 # tm0 take priority
+                # replace template concept names if applicable
+                new_inner_template = deepcopy(inner_template)
+                new_outer_template = deepcopy(outer_template)
+                if hasattr(inner_template, "subject"):
+                    if inner_template.subject.name in replaced_concept_map:
+                        new_inner_template.subject.name = replaced_concept_map[
+                            inner_template.subject.name
+                        ]
+                if hasattr(inner_template, "outcome"):
+                    if inner_template.outcome.name in replaced_concept_map:
+                        new_inner_template.outcome.name = replaced_concept_map[
+                            inner_template.outcome.name
+                        ]
+                if hasattr(inner_template, "controller"):
+                    if inner_template.controller.name in replaced_concept_map:
+                        new_inner_template.controller.name = replaced_concept_map[
+                            inner_template.controller.name
+                        ]
                 if not check_template_in_inter_edge_dict(inter_template_edges,
                                                          inner_tm_id,
                                                          inner_template_id):
-                    process_template(new_templates, inner_template, tm1,
-                                     new_parameters, new_initials)
+                    process_template(new_templates, new_inner_template, tm1,
+                                         new_parameters, new_initials,
+                                         replaced_concept_map)
                 if not check_template_in_inter_edge_dict(inter_template_edges,
                                                          outer_tm_id,
                                                          outer_template_id):
-                    process_template(new_templates, outer_template, tm0,
-                                     new_parameters, new_initials)
+                    process_template(new_templates, new_outer_template, tm0,
+                                     new_parameters, new_initials,
+                                     replaced_concept_map)
 
     composed_tm = TemplateModel(templates=new_templates,
                                 parameters=new_parameters,
@@ -220,7 +290,7 @@ def check_template_in_inter_edge_dict(inter_edge_dict, tm_id, template_id):
 
 
 def process_template(templates, added_template, added_tm, parameters,
-                     initials):
+                     initials, replaced_concept_map):
     """Helper method that updates the dictionaries that contain the attributes
     to be used for the new composed template model
 
@@ -240,9 +310,15 @@ def process_template(templates, added_template, added_tm, parameters,
     initials :
         The dictionary of initials to update that will be used for the
         composed template model
+    replaced_concept_map:
+        A dictionary mapping replaced concept names to their new name
     """
     if added_template not in templates:
         templates.append(added_template)
+        if added_template.rate_law:
+            for old_concept_name, new_concept_name in replaced_concept_map.items():
+                added_template.rate_law = added_template.rate_law.subs(sympy.Symbol(
+                    old_concept_name), sympy.Symbol(new_concept_name))
         parameters.update({param_name: added_tm.parameters[param_name] for
                            param_name
                            in added_template.get_parameter_names()})
@@ -263,11 +339,10 @@ def update_observable_expressions(observables, concept_map):
         The mapping of old concepts to the list of new concepts
     """
     for observable in observables.values():
-        for old_concept_name, new_concept_list in concept_map.items():
-            new_expression = sum([sympy.Symbol(new_concept_name) for
-                                  new_concept_name in new_concept_list])
-            observable.expression = observable.expression.subs(sympy.Symbol(
-                old_concept_name), new_expression)
+        if observable.expression:
+            for old_concept_name, new_concept_name in concept_map.items():
+                observable.expression = observable.expression.subs(sympy.Symbol(
+                    old_concept_name), sympy.Symbol(new_concept_name))
 
 
 def substitute_time(tm, time_0, time_1):

tests/test_model_composition.py

@@ -63,6 +63,62 @@
     ]
 )
 
+S1 = Concept(name="Susceptible", identifiers={"ido": "0000514"})
+I1 = Concept(name="Infected", identifiers={"ido": "0000511"})
+R1 = Concept(name="Recovery", identifiers={"ido": "0000592"})
+
+S2 = Concept(name="Susceptible", identifiers={"ido": "0000513"})
+I2 = Concept(name="Infected", identifiers={"ido": "0000512"})
+R2 = Concept(name="Recovery", identifiers={"ido": "0000593"})
+
+S3 = Concept(name="S", identifiers={"ido": "0000514"})
+I3 = Concept(name="I", identifiers={"ido": "0000511"})
+R3 = Concept(name="R", identifiers={"ido": "0000592"})
+
+S4 = Concept(name="S")
+I4 = Concept(name="I")
+R4 = Concept(name="R")
+
+model_A1 = TemplateModel(
+    templates=[
+        ControlledConversion(
+            name="Infection", subject=S1, outcome=I1, controller=I1
+        )
+    ]
+)
+
+model_B1 = TemplateModel(
+    templates=[NaturalConversion(name="Recovery", subject=I1, outcome=R1)]
+)
+
+model_B2 = TemplateModel(
+    templates=[NaturalConversion(name="Recovery", subject=I2, outcome=R2)]
+)
+
+model_B3 = TemplateModel(
+    templates=[NaturalConversion(name="Recovery", subject=I3, outcome=R3)]
+)
+
+model_B4 = TemplateModel(
+    templates=[NaturalConversion(name="Recovery", subject=I4, outcome=R4)]
+)
+
+model_A3 = TemplateModel(
+    templates=[
+        ControlledConversion(
+            name="Infection", subject=S3, outcome=I3, controller=I3
+        )
+    ]
+)
+
+model_A4 = TemplateModel(
+    templates=[
+        ControlledConversion(
+            name="Infection", subject=S4, outcome=I4, controller=I4
+        )
+    ]
+)
+
 
 def test_compose_two_models():
     composed_model = compose_two_models(sir_reinfection, sir)
@@ -116,3 +172,70 @@ def test_template_inclusion():
     composed template model"""
     composed_tm = compose_two_models(mini_sir, sir)
     assert len(composed_tm.templates) == 2
+
+
+def test_concept_composition():
+    """We test model composition's consolidation of concepts using test-cases
+    described in this issue here: https://github.com/gyorilab/mira/issues/409
+    """
+    model_ab11 = compose([model_A1, model_B1])
+    model_ab12 = compose([model_A1, model_B2])
+    model_ab13 = compose([model_A1, model_B3])
+    model_ab34 = compose([model_A3, model_B4])
+    model_ab44 = compose([model_A4, model_B4])
+
+    assert len(model_ab11.get_concepts_map()) == 3
+    assert len(model_ab12.get_concepts_map()) == 4
+    assert len(model_ab13.get_concepts_map()) == 3
+    assert len(model_ab34.get_concepts_map()) == 3
+    assert len(model_ab44.get_concepts_name_map()) == 3
+
+    # model_A1 and model_B1's respective infected compartments have same name
+    # and identifiers, they should be composed
+    assert model_ab11.templates[0].subject.name == "Infected"
+    assert model_ab11.templates[0].subject.identifiers == {"ido": "0000511"}
+
+    assert model_ab11.templates[1].outcome.name == "Infected"
+    assert model_ab11.templates[1].controller.name == "Infected"
+    assert model_ab11.templates[1].outcome.identifiers == {"ido": "0000511"}
+    assert model_ab11.templates[1].controller.identifiers == {"ido": "0000511"}
+
+    # model_A1 and model_B2's respective infected compartments have separate identifiers
+    # they should be treated as separate concepts and not composed
+    assert model_ab12.templates[0].outcome.name == "Infected"
+    assert model_ab12.templates[0].controller.name == "Infected"
+    assert model_ab12.templates[0].outcome.identifiers == {"ido": "0000511"}
+    assert model_ab12.templates[0].controller.identifiers == {"ido": "0000511"}
+
+    assert model_ab12.templates[1].subject.name == "Infected"
+    assert model_ab12.templates[1].subject.identifiers == {"ido": "0000512"}
+
+    # model_A1 and model_B3 share an infected compartment with the same identifiers
+    # "Infected" for model_A1 and "I" for model_B3
+    # since model_A1 is the first model passed in, we prioritize model_A1's
+    # infected compartment and test to see if appropriate concept replacement has taken place
+    assert model_ab13.templates[0].subject.name == "Infected"
+    assert model_ab13.templates[1].outcome.name == "Infected"
+    assert model_ab13.templates[1].controller.name == "Infected"
+
+    # model_A3 contains infected component "I" with identifiers
+    # model_B4 contains infected compartment "Infected" with no identifiers
+    # model_A3 is the first model passed in so we expect the composed model's
+    # infected compartment's name to be "I" with identifier present
+    assert model_ab34.templates[0].outcome.name == "I"
+    assert model_ab34.templates[0].controller.name == "I"
+    assert model_ab34.templates[1].subject.name == "I"
+    assert model_ab34.templates[0].outcome.identifiers == {"ido": "0000511"}
+    assert model_ab34.templates[0].controller.identifiers == {"ido": "0000511"}
+    assert model_ab34.templates[1].subject.identifiers == {"ido": "0000511"}
+
+    # model_A4 and model_B4's respective infected compartment share the same name,
+    # "I" and neither has an identifier.
+    # They should be treated as the same compartment.
+
+    # This tests to see if the two infected compartments are composed into one
+    # If treated as separate compartments, then the number of templates will be 3
+    assert len(model_ab44.templates) == 2
+    assert model_ab44.templates[0].subject.name == "I"
+    assert model_ab44.templates[1].outcome.name == "I"
+    assert model_ab44.templates[1].controller.name == "I"