diff --git a/mira/metamodel/ops.py b/mira/metamodel/ops.py index ec059f2c..3bae168d 100644 --- a/mira/metamodel/ops.py +++ b/mira/metamodel/ops.py @@ -229,14 +229,12 @@ def stratify( # We now apply the stratum assigned to each controller in this particular # tuple to the controller for controller, c_stratum in zip(stratified_controllers, c_strata_tuple): + stratified_template.name += f"_{c_stratum}" controller.with_context(do_rename=modify_names, inplace=True, **{key: c_stratum}) template_strata.append(c_stratum if param_renaming_uses_strata_names else stratum_index_map[c_stratum]) - tname = stratified_template.name if stratified_template.name \ - else 't' + '_'.join(template_strata) - stratified_template.name = tname # Wew can now rewrite the rate law for this stratified template, # then append the new template param_mappings = rewrite_rate_law(template_model=template_model, diff --git a/tests/test_ops.py b/tests/test_ops.py index fe5bf8f5..79ee8e86 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -42,6 +42,7 @@ def test_stratify_full(self): ) expected_0 = ControlledConversion( + name="t_unvaccinated_unvaccinated", subject=susceptible.with_context(vaccination_status="unvaccinated", do_rename=True), outcome=infected.with_context(vaccination_status="unvaccinated", @@ -54,6 +55,7 @@ def test_stratify_full(self): ) ) expected_1 = ControlledConversion( + name="t_unvaccinated_vaccinated", subject=susceptible.with_context(vaccination_status="unvaccinated", do_rename=True), outcome=infected.with_context(vaccination_status="unvaccinated", @@ -66,6 +68,7 @@ def test_stratify_full(self): ) ) expected_2 = ControlledConversion( + name="t_vaccinated_unvaccinated", subject=susceptible.with_context(vaccination_status="vaccinated", do_rename=True), outcome=infected.with_context(vaccination_status="vaccinated", @@ -78,6 +81,7 @@ def test_stratify_full(self): ) ) expected_3 = ControlledConversion( + name="t_vaccinated_vaccinated", subject=susceptible.with_context(vaccination_status="vaccinated", do_rename=True), outcome=infected.with_context(vaccination_status="vaccinated",