[BUG] Unexpected Stratify behaviour when applied more than once

[liunelson]

This scenario came up during the Epi Eval Scenario 2 and the user found Stratify's handling of parameters undesirable and inconsistent with their expectation.

1. I start with a simple compartmental SCRHD model (I renamed to C to avoid imaginary i issue with sympy).

2. I change the rate law between C and H to be the product b * vaxCtoH * ageCtoH * C based on the assumption that the vax status and age are independent demographic features.

for t in model.templates:
    if t.name == "C_to_H":
        t.set_rate_law("C * b * vaxCtoH * ageCtoH")

3. I stratify it once by vax status, on only the S, C states and a, vaxCtoH parameters

model_vax = stratify(
    model, 
    key = "vaxstatus", 
    strata = ["u", "v"], 
    cartesian_control = True, 
    structure = [],
    concepts_to_stratify = ["S", "C"],
    # concepts_to_preserve = ["H", "R", "D"],
    params_to_stratify = ["a", "vaxCtoH"]
)
model_vax.annotations.name = "SCRHD Vax"

4. I optionally add a vaccination process

   model_vax = model_vax.add_template(
   NaturalConversion(
       subject = model_vax.get_concepts_name_map()["S_u"],
       outcome = model_vax.get_concepts_name_map()["S_v"],
       name = "vaccination"
   ).with_mass_action_rate_law("f")
   )
   model_vax.parameters["f"] = Parameter(name = "f", value = 1.0)

5. I stratified the model again, this time by age on all states and parameters, except parameters b, vaxCtoH_*

model_vax_age = stratify(
    model_vax, 
    key = "age",
    strata = ["y", "o"],
    cartesian_control = True, 
    structure = [], 
    params_to_preserve = ["b", "vaxCtoH_0", "vaxCtoH_1"],
    concepts_to_stratify = None
)
model_vax_age.annotations.name = "SCRHD Vax Age"

6. I find that (a) the parameter vaxCtoH correctly depends only on vax status only and (b) the parameter ageCtoH incorrectly depends on both age and vax status. The user expected that ageCtoH to be stratified by age only.

   
   vaxCtoH -> vaxCtoH_0, vaxCtoH_1 -> vaxCtoH_0, vaxCtoH_1

actual but incorrect

ageCtoH -> ageCtoH -> ageCtoH_0, ageCtoH_1, ageCtoH_2, ageCtoH_3

expectation

ageCtoH -> ageCtoH -> ageCtoH_0, ageCtoH_1

Relevant AMR files:
* base model, [scenario2.json](https://github.com/gyorilab/mira/files/15082185/scenario2.json)
* model stratified by vax status, [scenario2_vax.json](https://github.com/gyorilab/mira/files/15082184/scenario2_vax.json)
* model stratified twice, by vax status and age, [scenario2_vax_age.json](https://github.com/gyorilab/mira/files/15082183/scenario2_vax_age.json)

[bgyori]

Thanks for the detailed explanation! This currently happens any time a parameter is used in multiple rate laws. So e.g., if a parameter gamma is used in two templates and is stratified into two parameters, it will result in 4 new parameters being generated. I agree it would be good to assign parameters in a more sophisticated way to have numbering only depend on the stratified concepts - I will come up with a way to achieve this.

[liunelson]

I appreciate the consideration: the logic is probably tricky to implement!

I think the behaviour that the user wants us to support is more or less:

• a single rate parameter r is decomposed into n parameters, r = a * b * c * ...
• the model is stratified n times by N strata each
• the expected outcome is n parameter matrices of size N x N, a -> a_0, a_1, ..., a_{N^2 - 1}
• if the user doesn't do that, then they get a big ugly r matrix of size N^(n-1) x N^(n-1) (relevant in cases where the dynamics is not separable 🤷🏻 )

[bgyori]

@liunelson finally figured this one out. Parameter renaming now depends on the specific strata that were applied to the template that it is involved in. In addition, I added an additional argument to stratify called param_renaming_uses_strata_names which, when set to True, uses strata names in parameters.

Example:

In [1]: from mira.metamodel import *

In [2]: from mira.examples.sir import *

In [3]: for t in sir_parameterized.templates:
   ...:     print(t.get_concept_names(), t.rate_law)
   ...: 
{'susceptible_population', 'infected_population'} beta*infected_population*susceptible_population
{'immune_population', 'infected_population'} gamma*infected_population

In [4]: tm = stratify(sir_parameterized, key='age', strata=['young', 'old'], cartesian_control=True,
                      param_renaming_uses_strata_names=True)

In [5]: for t in tm.templates:
   ...:     print(t.get_concept_names(), t.rate_law)
   ...: 
{'susceptible_population_young', 'infected_population_young'} beta_young_young*infected_population_young*susceptible_population_young
{'susceptible_population_young', 'infected_population_old', 'infected_population_young'} beta_young_old*infected_population_old*susceptible_population_young
{'susceptible_population_old', 'infected_population_old', 'infected_population_young'} beta_old_young*infected_population_young*susceptible_population_old
{'susceptible_population_old', 'infected_population_old'} beta_old_old*infected_population_old*susceptible_population_old
{'infected_population_young', 'immune_population_young'} gamma_young*infected_population_young
{'infected_population_old', 'immune_population_old'} gamma_old*infected_population_old
{'susceptible_population_young', 'susceptible_population_old'} p_young_old*susceptible_population_young
{'susceptible_population_old', 'susceptible_population_young'} p_old_young*susceptible_population_old
{'infected_population_old', 'infected_population_young'} infected_population_young*p_young_old
{'infected_population_old', 'infected_population_young'} infected_population_old*p_old_young
{'immune_population_old', 'immune_population_young'} immune_population_young*p_young_old
{'immune_population_young', 'immune_population_old'} immune_population_old*p_old_young