plot predictions shape error with multinomial model

[GStechschulte]

@AlexAndorra raised a question/issue in the LB stats Slack. Posting here.

---

I'm doing something like this:

bmb.interpret.predictions(
    multinomial_model_bambi,
    idata_multinomial_bambi,
    conditional={
        "cat1": ["A", "B"],
        "name": ["name1", ..., "name10"],
    },
)

which works great, but then when trying to get the plot:

_, ax = bmb.interpret.plot_predictions(
    multinomial_model_bambi,
    idata_multinomial_bambi,
    conditional={
        "cat1": ["A", "B"],
        "name": ["name1", ..., "name10"],
    },
    subplot_kwargs={"main": "cat1", "group": "name", "panel": "name"},
);

that raises a weird error: ValueError: x and y must be the same size

[GStechschulte]

This was my response to Alex.

---

I think you are getting the error because you are predicting multiple outcomes, and not passing the column name estimate_dim to subplot_kwargs. In your first code snippet where you print the summary dataframe, you will see that column.

The following should work:

_, ax = bmb.interpret.plot_predictions(
    multinomial_model_bambi,
    idata_multinomial_bambi,
    conditional={
        "cat1": ["A", "B"],
        "name": ["name1", ..., "name10"],
    },
    subplot_kwargs={"main": "cat1", "group": "name", "panel": "estimate_dim"},
);

---

What I noticed is that it is not obvious that you can pass the column name estimate_dim for model's predicting multiple outcomes, e.g multinomial or categorical. You will most likely want to plot each outcome as a hue or panel (facet), but if you do not know a priori that Bambi adds the column estimate_dim, then you will not pass it as a subplot kwarg.

[tomicapretto]

I'm not sure I understand if this is a bug or something we should document better :sweat_smile:

[GStechschulte]

It's definitely something I would like to document better. From an interface point of view, I think the first plot in my comment above should render because: (1) the interpret.predictions successfully builds the data frame, and (2) none of ValueErrors for .interpret.plot_predictions are raised. So I think this is unexpected behaviour for the user point of view.

I will work on this :)

[tomicapretto]

@GStechschulte amazing, just ping me if you want review or help :)