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.

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

Allow FacetGrids use in heatmaps #219

Merged
merged 10 commits into from
Jul 10, 2024
Merged

Allow FacetGrids use in heatmaps #219

merged 10 commits into from
Jul 10, 2024

Conversation

coxipi
Copy link
Contributor

@coxipi coxipi commented Jun 10, 2024

Pull Request Checklist:

  • This PR addresses an already opened issue (for bug fixes / features)
    • This PR fixes #xyz
  • (If applicable) Documentation has been added / updated (for bug fixes / features).
  • (If applicable) Tests have been added.
  • CHANGELOG.rst has been updated (with summary of main changes).
    • Link to issue (:issue:number) and pull request (:pull:number) has been added.

What kind of change does this PR introduce?

  • fg.heatmap can now accept plot_kw = {"row":"dim1", "col":"dim2"} like other functions in Figanos!

Does this PR introduce a breaking change?

No

Other information:

Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@coxipi coxipi mentioned this pull request Jun 10, 2024
1 task
Copy link
Collaborator

@sarahclaude sarahclaude left a comment

Choose a reason for hiding this comment

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

Looks great, thanks Eric!

Are "row" and "col" the only arguments we can pass to the Seaborn facetgrids in plot_kw?
I tried using height and aspect in plot_kw and it didn't work (https://seaborn.pydata.org/generated/seaborn.FacetGrid.html).

However, it worked great with fig_kw to change the figure size in the docs.

Comment on lines 1359 to 1376
if ax is None and ("row" not in plot_kw.keys() and "col" not in plot_kw.keys()):
fig, ax = plt.subplots(**fig_kw)
elif ax is not None and ("col" in plot_kw or "row" in plot_kw):
raise ValueError("Cannot use 'ax' and 'col'/'row' at the same time.")
elif ax is None:
if any([k != "figsize" for k in fig_kw.keys()]):
warnings.warn(
"Only figsize arguments can be passed to fig_kw when using facetgrid."
)
plot_kw.setdefault("col", None)
plot_kw.setdefault("row", None)
heatmap_dims = list(
set(da.dims)
- {d for d in [plot_kw["col"], plot_kw["row"]] if d is not None}
)
if da.name is None:
da = da.to_dataset(name="data").data
da_name = da.name
Copy link
Contributor Author

@coxipi coxipi Jun 11, 2024

Choose a reason for hiding this comment

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

This part is a bit more rudimentary than in gridmap for instance, where figure.add_subplot() is also used in fig_kw

"Only figsize and figure.add_subplot() arguments can be passed to fig_kw when using facetgrid."

I don't think I've seen an example on how to properly use figure.add_subplot() in this context, so I prefer not to advertise any of this. If this is a crucial feature I can try to incorporate it too.

Copy link
Collaborator

Choose a reason for hiding this comment

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

No i don't think it is crucial to have it. I was just wondering how much access we have to seaborn facetgrid options via the function heatmap

@coxipi
Copy link
Contributor Author

coxipi commented Jun 12, 2024

Are "row" and "col" the only arguments we can pass to the Seaborn facetgrids in plot_kw?
I tried using height and aspect in plot_kw and it didn't work (https://seaborn.pydata.org/generated/seaborn.FacetGrid.html).

With how I programmed it, yes. I may have misunderstood the general strucutre in Figanos, right now plot_kw is used both in sns.FacetGrid and sns.heatmap, so some options don't work in one function or the other. I will check how other functions handle this.

EDIT: Ok seeing gridmap, I have for instance: im = plot_data.plot.pcolormesh(**plot_kw)
and plot_kw can also contain "col/row" keywords in this context.

Simply put, xr.plot.pcolormesh(da, col="time") works, but not sns.heatmap(da, col="time").

@sarahclaude
Copy link
Collaborator

Usually plot_kw allows you to access the function argument of the xarray plot which also allows the xarray facetgrid arguments.

However, it doesn't always work well and this is why I put some warnings with fig_kw and facetgrids that not all options are available. I think if you add a warning saying only col and row are accepted for the seaborn faccetgrid options in plot_kw it will work great.

@coxipi
Copy link
Contributor Author

coxipi commented Jun 12, 2024

Here is my hacky solution, just look at allowed keywords for each function

        kws_hm = set(inspect.signature(sns.heatmap).parameters)
        kws_fg = set(inspect.signature(sns.FacetGrid).parameters)
        plot_kw_keys = list(plot_kw.keys())
        plot_kw_fg = {k: v for k, v in plot_kw.items() if k in kws_fg.intersection(plot_kw_keys)}
        plot_kw_hm = {k: v for k, v in plot_kw.items() if k in kws_hm.intersection(plot_kw_keys)}
# then use plot_kw_fg for sns.FacetGrid, and plot_kw_hm for sns.heatmap
# height is only sns.FacetGrid, cmap is only sns.heatmap
fg.heatmap(da, plot_kw = {"col": "time", "height":6, "cmap":"Spectral_r"})

image

Comment on lines +1459 to +1474
# When using xarray's FacetGrid, `plot_kw` can be used in the FacetGrid and in the plotting function
# With Seaborn, we need to be more careful and separate keywords.
plot_kw_hm = {
k: v for k, v in plot_kw.items() if k in signature(sns.heatmap).parameters
}
plot_kw_fg = {
k: v for k, v in plot_kw.items() if k in signature(sns.FacetGrid).parameters
}
unused_keys = (
set(plot_kw.keys()) - set(plot_kw_fg.keys()) - set(plot_kw_hm.keys())
)
if unused_keys != set():
raise ValueError(
f"`heatmap` got unexpected keywords in `plot_kw`: {unused_keys}. Keywords in `plot_kw` should be keywords "
"allowed in `sns.heatmap` or `sns.FacetGrid`. "
)
Copy link
Contributor Author

@coxipi coxipi Jun 13, 2024

Choose a reason for hiding this comment

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

Just want to highlight this: This ensures we can use both keywords from heatmap and FacetGrid. With this approach, we would silently ignore wrong keywords, so I have to find if there are unused_keys and if so, add an extra step to throw an error.

A solution to avoid this extra step would be to stuff in illegal keywords in plot_kw_fg for instance, then sns.FacetGrid would handle the error management. But since plot_kw allows for both keywords on FacetGrid and Heatmap, I thought it was better to have our own explicit error message, since we have this peculiar input.

Comment on lines +1440 to +1443
# Any sorting should be performed before sending a DataArray in `fg.heatmap`
else data.pivot_table(
index=args[1], columns=args[0], values=args[2], sort=False
)
Copy link
Contributor Author

Choose a reason for hiding this comment

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

Columns/Rows of individual heatmap are not sorted by pandas in this way. This should be managed by sorting directly the input dataset, I think it's reasonable. Otherwise, it's difficult to manage custom order of coords values in the heatmap.

Copy link
Collaborator

@sarahclaude sarahclaude left a comment

Choose a reason for hiding this comment

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

Looks great, thanks Eric!

Comment on lines +1481 to +1485
cax = g.fig.add_axes([0.95, 0.05, 0.02, 0.9])
g.map_dataframe(
draw_heatmap, *heatmap_dims, da_name, **plot_kw_hm, cbar=True, cbar_ax=cax
)
g.fig.subplots_adjust(right=0.9)
Copy link
Contributor Author

Choose a reason for hiding this comment

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

I tried giving sane defaults to the cax position. I think FacetGrids in xarray handle this better, no need to specify an absolute position. In any case, it can be changed afterwards it this is needed.

@coxipi coxipi merged commit b8811a6 into main Jul 10, 2024
12 checks passed
@coxipi coxipi deleted the heatmap_array branch July 10, 2024 14:40
@coxipi
Copy link
Contributor Author

coxipi commented Jul 11, 2024

Oops, I realized I had not pushed a final commit (one line of code) in my github branch @Zeitsperre anything I could do that doesn't involve simply a new PR? I guess not? Sorry about that

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants