Setting alias breaks if multiple target formulae are given

[webdrone]

Bug description

Calling set.alias({...}) on a model with multiple target formulae results in assertion error.

For instance, if one wishes to estimate a different noise scale per condition, one would define multiple formulae (targets) in the model. Then trying to set an alias for a parameter fails.

Reproduction

The below two work:

• with no condition-specific sigma

import bambi as bmb

model = bmb.Model(
    formula=bmb.Formula(
        "value ~ T(Condition, 'WT')",
    ),
    data=df_data_fit,
    categorical=["Condition"],
    family="gaussian",
    priors=partial_pooling_priors,
    noncentered=True
    )
model.set_alias({"T(Condition, 'WT')": "Condition"})

• with no aliases

import bambi as bmb

model = bmb.Model(
    formula=bmb.Formula(
        "value ~ T(Condition, 'WT')",
        "sigma ~ Condition"
    ),
    data=df_data_fit,
    categorical=["Condition"],
    family="gaussian",
    priors=partial_pooling_priors,
    noncentered=True
    )

The below fails:

• with condition-specific sigma and aliases

import bambi as bmb

model = bmb.Model(
    formula=bmb.Formula(
        "value ~ T(Condition, 'WT')",
        "sigma ~ Condition"
    ),
    data=df_data_fit,
    categorical=["Condition"],
    family="gaussian",
    priors=partial_pooling_priors,
    noncentered=True
    )
model.set_alias({"T(Condition, 'WT')": "Condition"})

Output

File /opt/homebrew/Caskroom/miniconda/base/envs/dsci/lib/python3.12/site-packages/bambi/models.py:576, in Model.set_alias(self, aliases)
    574     self.response_component.term.alias = component_aliases
    575 else:
--> 576     assert isinstance(component_aliases, dict)
    577     assert component_name in self.distributional_components
    578     component = self.distributional_components[component_name]

AssertionError: 

[tomicapretto]

@webdrone I confirm I can reproduce your issue, using a dataset available in Bambi

import bambi as bmb

df = bmb.load_data("adults")

model = bmb.Model(
    formula=bmb.Formula(
        "hs_week ~ T(sex, 'Female')",
        "sigma ~ age"
    ),
    data=df,
)

model.set_alias({"T(sex, 'Female')": "NewName"})

---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Cell In[7], line 14
      3 df = bmb.load_data("adults")
      5 model = bmb.Model(
      6     formula=bmb.Formula(
      7         "hs_week ~ T(sex, 'Female')",
   (...)
     11     family="gaussian",
     12 )
---> 14 model.set_alias({"T(sex, 'Female')": "NewName"})

File ~/oss/bambinos/bambi/bambi/models.py:576, in Model.set_alias(self, aliases)
    574     self.response_component.term.alias = component_aliases
    575 else:
--> 576     assert isinstance(component_aliases, dict)
    577     assert component_name in self.distributional_components
    578     component = self.distributional_components[component_name]

AssertionError:

However, it's because we need to use .set_allias() in a different way. When you have a distributional model (i.e. more than one response parameter depend on predictors), you need to pass a nested dictionary, where each key is the name of the parameter for which you set the alias.

See this example

import bambi as bmb

df = bmb.load_data("adults")

model = bmb.Model(
    formula=bmb.Formula(
        "hs_week ~ T(sex, 'Female')",
        "sigma ~ age"
    ),
    data=df,
)
model.set_alias({"mu": {"T(sex, 'Female')": "NewName"}})
model.build()
model.graph()

[webdrone]

Wonderful, I was hoping there would be a simple resolution like this. Maybe a less cryptic error message would help others?

[tomicapretto]

@webdrone I think that would help!