Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

plot_cap default args. not working for categorical regression #673

Closed
wants to merge 5 commits into from
Closed

plot_cap default args. not working for categorical regression #673

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f4c7d19
add indexing for N-dim. y_hat_bounds for ax.fill_between()
GStechschulte May 20, 2023
2968285
legend shows class names
GStechschulte May 22, 2023
029263f
legend now shows class names for categorical outcome. Ran Black.
GStechschulte May 22, 2023
90835bd
_plot_cap_categorical show class names for categorical outcome
GStechschulte May 22, 2023
7a8aefd
_plot_cap_categorical show class names for categorical outcome
GStechschulte May 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 101 additions & 25 deletions bambi/plots/plot_cap.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,10 @@ def plot_cap(
fig = axes[0].get_figure()

main = covariates.get("horizontal")

if is_categorical_dtype(model.data[response_name]):
legend = {response_name: (model.data[response_name].cat.categories.tolist())}

if is_numeric_dtype(cap_data[main]):
axes = _plot_cap_numeric(
covariates, cap_data, y_hat_mean, y_hat_bounds, transforms, legend, axes
Expand All @@ -262,12 +266,17 @@ def plot_cap(
def _plot_cap_numeric(covariates, cap_data, y_hat_mean, y_hat_bounds, transforms, legend, axes):
main = covariates.get("horizontal")
transform_main = transforms.get(main, identity)
y_hat_bounds_dim = y_hat_bounds.ndim
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Can we call it y_ndim or y_hat_bounds_ndim?


if len(covariates) == 1:
ax = axes[0]
values_main = transform_main(cap_data[main])
ax.plot(values_main, y_hat_mean, solid_capstyle="butt")
ax.fill_between(values_main, y_hat_bounds[0], y_hat_bounds[1], alpha=0.4)
if y_hat_bounds_dim > 2:
for dim in range(y_hat_bounds_dim):
ax.fill_between(values_main, y_hat_bounds[0][dim], y_hat_bounds[1][dim], alpha=0.4)
else:
ax.fill_between(values_main, y_hat_bounds[0], y_hat_bounds[1], alpha=0.4)
elif "color" in covariates and not "panel" in covariates:
ax = axes[0]
color = covariates.get("color")
Expand All @@ -276,21 +285,40 @@ def _plot_cap_numeric(covariates, cap_data, y_hat_mean, y_hat_bounds, transforms
idx = (cap_data[color] == clr).to_numpy()
values_main = transform_main(cap_data.loc[idx, main])
ax.plot(values_main, y_hat_mean[idx], color=f"C{i}", solid_capstyle="butt")
ax.fill_between(
values_main,
y_hat_bounds[0][idx],
y_hat_bounds[1][idx],
alpha=0.4,
color=f"C{i}",
)
if y_hat_bounds_dim > 2:
for dim in range(y_hat_bounds_dim):
ax.fill_between(
values_main,
y_hat_bounds[0][dim][idx],
y_hat_bounds[1][dim][idx],
alpha=0.4,
color=f"C{i}",
)
else:
ax.fill_between(
values_main,
y_hat_bounds[0][idx],
y_hat_bounds[1][idx],
alpha=0.4,
color=f"C{i}",
)
elif not "color" in covariates and "panel" in covariates:
panel = covariates.get("panel")
panels = get_unique_levels(cap_data[panel])
for ax, pnl in zip(axes.ravel(), panels):
idx = (cap_data[panel] == pnl).to_numpy()
values_main = transform_main(cap_data.loc[idx, main])
ax.plot(values_main, y_hat_mean[idx], solid_capstyle="butt")
ax.fill_between(values_main, y_hat_bounds[0][idx], y_hat_bounds[1][idx], alpha=0.4)
if y_hat_bounds_dim > 2:
for dim in range(y_hat_bounds_dim):
ax.fill_between(
values_main,
y_hat_bounds[0][dim][idx],
y_hat_bounds[1][dim][idx],
alpha=0.4,
)
else:
ax.fill_between(values_main, y_hat_bounds[0][idx], y_hat_bounds[1][idx], alpha=0.4)
ax.set(title=f"{panel} = {pnl}")
elif "color" in covariates and "panel" in covariates:
color = covariates.get("color")
Expand All @@ -302,30 +330,63 @@ def _plot_cap_numeric(covariates, cap_data, y_hat_mean, y_hat_bounds, transforms
idx = (cap_data[panel] == pnl).to_numpy()
values_main = transform_main(cap_data.loc[idx, main])
ax.plot(values_main, y_hat_mean[idx], color=f"C{i}", solid_capstyle="butt")
ax.fill_between(
values_main,
y_hat_bounds[0][idx],
y_hat_bounds[1][idx],
alpha=0.4,
color=f"C{i}",
)
ax.set(title=f"{panel} = {pnl}")
else:
for ax, pnl in zip(axes.ravel(), panels):
for i, clr in enumerate(colors):
idx = ((cap_data[panel] == pnl) & (cap_data[color] == clr)).to_numpy()
values_main = transform_main(cap_data.loc[idx, main])
ax.plot(values_main, y_hat_mean[idx], color=f"C{i}", solid_capstyle="butt")
if y_hat_bounds_dim > 2:
for dim in range(y_hat_bounds_dim):
ax.fill_between(
values_main,
y_hat_bounds[0][dim][idx],
y_hat_bounds[1][dim][idx],
alpha=0.4,
color=f"C{i}",
)
else:
ax.fill_between(
values_main,
y_hat_bounds[0][idx],
y_hat_bounds[1][idx],
alpha=0.4,
color=f"C{i}",
)
ax.set(title=f"{panel} = {pnl}")
else:
for ax, pnl in zip(axes.ravel(), panels):
for i, clr in enumerate(colors):
idx = ((cap_data[panel] == pnl) & (cap_data[color] == clr)).to_numpy()
values_main = transform_main(cap_data.loc[idx, main])
ax.plot(values_main, y_hat_mean[idx], color=f"C{i}", solid_capstyle="butt")
if y_hat_bounds_dim > 2:
for dim in range(y_hat_bounds_dim):
ax.fill_between(
values_main,
y_hat_bounds[0][dim][idx],
y_hat_bounds[1][dim][idx],
alpha=0.4,
color=f"C{i}",
)
else:
ax.fill_between(
values_main,
y_hat_bounds[0][idx],
y_hat_bounds[1][idx],
alpha=0.4,
color=f"C{i}",
)
ax.set(title=f"{panel} = {pnl}")

if "color" in covariates and legend:
if not isinstance(legend, bool):
title = list(legend.keys())[0]
values = list(legend.values())[0]
handles = [
(
Line2D([], [], color=f"C{i}", solid_capstyle="butt"),
Patch(color=f"C{i}", alpha=0.4, lw=1),
)
for i in range(len(values))
]
for ax in axes.ravel():
ax.set_label(values)
ax.legend(handles, values, title=title, handlelength=1.3, handleheight=1, loc="best")
elif "color" in covariates and legend:
handles = [
(
Line2D([], [], color=f"C{i}", solid_capstyle="butt"),
Expand All @@ -337,6 +398,7 @@ def _plot_cap_numeric(covariates, cap_data, y_hat_mean, y_hat_bounds, transforms
ax.legend(
handles, tuple(colors), title=color, handlelength=1.3, handleheight=1, loc="best"
)

return axes


Expand Down Expand Up @@ -391,7 +453,21 @@ def _plot_cap_categoric(covariates, cap_data, y_hat_mean, y_hat_bounds, legend,
ax.vlines(idxs, y_hat_bounds[0][idx], y_hat_bounds[1][idx], color=f"C{i}")
ax.set(title=f"{panel} = {pnl}")

if "color" in covariates and legend:
if not isinstance(legend, bool):
title = list(legend.keys())[0]
values = list(legend.values())[0]
handles = [
(
Line2D([], [], color=f"C{i}", solid_capstyle="butt"),
Patch(color=f"C{i}", alpha=0.4, lw=1),
)
for i in range(len(values))
]
for ax in axes.ravel():
ax.set_label(values)
ax.legend(handles, values, title=title, handlelength=1.3, handleheight=1, loc="best")

elif "color" in covariates and legend:
handles = [
Line2D([], [], c=f"C{i}", marker="o", label=level) for i, level in enumerate(colors)
]
Expand Down