From 5bd994454d5c1fbffdb486b40885b2c645918f04 Mon Sep 17 00:00:00 2001 From: Ben Gyori Date: Tue, 23 Apr 2024 20:41:17 -0400 Subject: [PATCH 1/2] Fix multiple stratification of templates --- mira/metamodel/ops.py | 51 ++++++++++++++++++++----------------------- 1 file changed, 24 insertions(+), 27 deletions(-) diff --git a/mira/metamodel/ops.py b/mira/metamodel/ops.py index 5a6531ed2..513bfb85a 100644 --- a/mira/metamodel/ops.py +++ b/mira/metamodel/ops.py @@ -156,35 +156,32 @@ def stratify( keep_unstratified_parameters = set() for template in template_model.templates: + # If the template doesn't have any concepts that need to be stratified + # then we can just keep it as is and skip the rest of the loop + if not set(template.get_concept_names()) - exclude_concepts: + original_params = template.get_parameter_names() + for param in original_params: + keep_unstratified_parameters.add(param) + templates.append(deepcopy(template)) + continue + # Generate a derived template for each strata for stratum in strata: - new_template = None - if set(template.get_concept_names()) - exclude_concepts: - new_template = template.with_context( - do_rename=modify_names, exclude_concepts=exclude_concepts, - curie_to_name_map=strata_curie_to_name, - **{key: stratum}, - ) - rewrite_rate_law(template_model=template_model, - old_template=template, - new_template=new_template, - params_count=params_count, - params_to_stratify=params_to_stratify, - params_to_preserve=params_to_preserve) - # parameters = list(template_model.get_parameters_from_rate_law(template.rate_law)) - # if len(parameters) == 1: - # new_template.set_mass_action_rate_law(parameters[0]) - templates.append(new_template) - - # This means that the template had no concepts to stratify. This implies - # that the template has no controllers that are stratified either - # so we can skip the rest of this loop - if not new_template: - original_params = template.get_parameter_names() - for param in original_params: - keep_unstratified_parameters.add(param) - templates.append(deepcopy(template)) - continue + new_template = template.with_context( + do_rename=modify_names, exclude_concepts=exclude_concepts, + curie_to_name_map=strata_curie_to_name, + **{key: stratum}, + ) + rewrite_rate_law(template_model=template_model, + old_template=template, + new_template=new_template, + params_count=params_count, + params_to_stratify=params_to_stratify, + params_to_preserve=params_to_preserve) + # parameters = list(template_model.get_parameters_from_rate_law(template.rate_law)) + # if len(parameters) == 1: + # new_template.set_mass_action_rate_law(parameters[0]) + templates.append(new_template) # assume all controllers have to get stratified together # and mixing of strata doesn't occur during control From 72003d130eda911184e9ae2805bc4201f9dfc0a0 Mon Sep 17 00:00:00 2001 From: Ben Gyori Date: Tue, 23 Apr 2024 21:48:28 -0400 Subject: [PATCH 2/2] Add regression test for fix --- tests/test_ops.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_ops.py b/tests/test_ops.py index 5b0c3b12c..d93bb3cc3 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -548,3 +548,16 @@ def condition(template): assert len(tm.templates) == 2 assert tm.templates[1].rate_law.args[0] == sympy.core.numbers.Zero() + + +def test_stratify_excluded_species(): + from mira.examples.sir import sir_parameterized + + tm = stratify(sir_parameterized, + key='vax', + strata=['vax', 'unvax'], + structure=[], + cartesian_control=True, + concepts_to_stratify=['susceptible_population']) + + assert len(tm.templates) == 5, templates