Add 'transforms' argument to plot_cap

[tomicapretto]

This PR adds an optional argument called transforms to plot_cap(). This is useful when the model is fitted on transformed data and we want the visualization in the original scale. See the example

import bambi as bmb
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from bambi.plots import plot_cap

rng = np.random.default_rng(1234)
x = np.abs(rng.normal(size=100, scale=5)) + 5
y = np.exp(2 + 0.3 * x + rng.normal(size=100, scale=0.5))

data = pd.DataFrame({"x": x, "y": y})
data["log_x"] = np.log(x)
data["log_y"] = np.log(y)
model = bmb.Model("log_y ~ 1 + log_x", data)
idata = model.fit()

# Plot CAP on transformed scale. Result is a straight line
fig, ax = plt.subplots()
ax.scatter(data["log_x"], data["log_y"], color="C4", alpha=0.6);
plot_cap(model, idata, "log_x", ax=ax);

# Plot CAP on the original scale, not necessarily a straight line
transforms = {"log_y": np.exp, "log_x": np.exp}
fig, ax = plt.subplots()
ax.scatter(data["x"], data["y"], color="C4", alpha=0.6);
plot_cap(model, idata, "log_x", transforms=transforms, ax=ax);

Update We can also use models that contain inline transformations. Notice the plot is created in the untransformed space. If we want the transformed space, we need to ask for it. While it may look counter-intuitive in the beginning, this is the most sensible approach I think because it works well when the predictor on the horizontal scale is included in more than a single term.

fig, ax = plt.subplots()
ax.scatter(np.log(data["x"]), np.log(data["y"]), color="C4", alpha=0.6);
plot_cap(model, idata, "x", transforms={"x": np.log}, ax=ax);

transforms = {"log(y)": np.exp}
fig, ax = plt.subplots()
ax.scatter(data["x"], data["y"], color="C4", alpha=0.6);
plot_cap(model, idata, "x", transforms=transforms, ax=ax);

Closes #588 and #590

[codecov-commenter]

Codecov Report

Merging #594 (763316a) into main (e61f0d5) will decrease coverage by 0.34%.
The diff coverage is 0.00%.

@@            Coverage Diff             @@
##             main     #594      +/-   ##
==========================================
- Coverage   85.74%   85.39%   -0.35%     
==========================================
  Files          38       38              
  Lines        2932     2944      +12     
==========================================
  Hits         2514     2514              
- Misses        418      430      +12     

Impacted Files	Coverage Δ
bambi/plots/plot_cap.py	0.00% <0.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[tomicapretto]

I think this can be merged. Since it's not that trivial I would like a second pair of eyes @aloctavodia. If you're OK with it, let's merge.

[canyon289]

If I understand correctly this PR has been replaced by #596?

[tomicapretto]

If I understand correctly this PR has been replaced by #596?

This is myself still not knowing git very well haha! Being on this branch, I did git checkout -b the_other_branch because I wanted to have these changes applied as well. Let's say this PR is Part 1, and the other PR is meant to be Part 2. But because of how I used git, the other branch is Part 1 + Part 2. Do you know how one can do what I wanted to do in the beginning?

[canyon289]

it looks like you did this correctly. when you merge this first PR then only the part 2 changes remain

…

On Sunday, November 27, 2022, Tomás Capretto ***@***.***> wrote: If I understand correctly this PR has been replaced by #596 <#596>? This is myself still not knowing git very well haha! Being on this branch, I did git checkout -b the_other_branch because I wanted to have these changes applied as well. Let's say this PR is Part 1, and the other PR is meant to be Part 2. But because of how I used git, the other branch is Part 1 + Part 2. Do you know how one can do what I wanted to do in the beginning? — Reply to this email directly, view it on GitHub <#594 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABXBFYISTVU4XCV2OIPPXGDWKNRVRANCNFSM6AAAAAASLMX224> . You are receiving this because you commented.Message ID: ***@***.***>