plot_predictions for a subset of groups

[andins]

I'm fitting a hierarchical model specified as:

model = bmb.Model(
    "y ~ (1|group_id) + x1 + x2",
    data,
    family="t"
)

so I have two fixed regressors and the intercept varying by group.
Now I'd like to visualize model predictions (together with the data) as a function of x1 for a subset of groups (in different panels).
Is it possible to do this using interpret.plot_predictions?
I achieved to plot the predictions for all of the groups but this is a large number (>300) so I'd rather look at a specific subset.
Please let me know if more details are needed.
Thank you in advance.

[GStechschulte]

Hey @andins thanks for raising the discussion. Yes, this should be possible. You would want to pass a specific subset of group_id values to conditional. Optionally, you could also pass x1 and x2. Otherwise, if you don't pass terms to conditional that were in the model, then we compute default values.

I imagine something like the following

bmb.interpret.plot_predictions(
model=...
idata=...
conditional={ 
   "group_id: [50, 51, 52, 53, 54, 55, 150, 151, 152, 153, 154, 155],
   "x1": np.linspace(x1.min(), x1.max(), 50),
  "x2": np.mean(x2)
),
subplot_kwargs={"main": "x1", "group": "group_id", "panel": "group_id"}
)

which would plot a panel for each group_id.

[GStechschulte]

Many thanks for providing a reproducible example. This is interesting. I can't necessarily say anything off the top of my head, but I will have a look at it within the coming days. Thanks! :)

[GStechschulte]

Hey @andins apologies for the delayed response.

If the groups that I pass to the conditional parameter in plot_predictions are not sorted, they get sorted internally and the returned axes will correspond to the sorted groups. Is this the intended behavior?

When plotting a group and panel covariate, Bambi computes the unique values for these covariates before plotting. The default behavior of np.unique is to return a sorted array.

If I plot your data not using your function, it appears everything is working fine.

bmb.interpret.plot_predictions(
    model_synth,
    results_synth,
    conditional={
        "group_id": data.group_id.unique(),
        "x1": np.linspace(data.x1.min(), data.x1.max(), 50),
        "x2": np.mean(data.x2),
    },
    subplot_kwargs={"main": "x1", "group": "group_id", "panel": "group_id"},
    fig_kwargs=dict(figsize=(11, 8)),
    legend=False
)
plt.tight_layout();

This discrepancy is due to the loop for i, p in enumerate(conditional["group_id"]): in your function which iterates over an unsorted array. If you sort and then iterate, your problem will be resolved.

conditional = {
        "group_id": data.group_id.unique(),
        "x1": np.linspace(data.x1.min(), data.x1.max(), 50),
        "x2": np.mean(data.x2),
    }

fig, ax = bmb.interpret.plot_predictions(
    model_synth,
    results_synth,
    conditional=conditional,
    subplot_kwargs={"main": "x1", "group": "group_id", "panel": "group_id"},
    fig_kwargs=dict(figsize=(11, 8)),
    legend=False
)

axes_flat = ax.ravel()
for i, p in enumerate(sorted(conditional["group_id"])):
    ax = axes_flat[i]
    zeroth_idx = np.where(data["group_id"] == p)[0]

    x_values = data.iloc[zeroth_idx]["x1"].values
    y_values = data.iloc[zeroth_idx]["y"].values
    ax.scatter(x_values, y_values)

plt.tight_layout()

[andins]

Thank you very much for your answer!
It makes sense and it's what I suspected: I noticed that I could just sort the groups and get the right pairing of regression plots and data scatter.
As a follow up, is there a way to use enforce a different sorting of covariates in plot_predictions? Not in the toy example above but in my real use case I'm trying to plot groups in order of increasing intercept. Maybe the only way is to change the plot_predictions code...

[GStechschulte]

Ahh cool, good to hear.

is there a way to use enforce a different sorting of covariates in plot_predictions?

Not at the moment. If more users start to request this feature then I will look into this. An alternative way is to: (1) use bmb.interpret.predictions which returns a dataframe, then plot this dataframe how you wish, or (2) follow the advanced interpret usage to access lower level interpret functions.

[andins]

It makes sense. Thank you very much for the suggestion!