Support grid of heatmaps

[coxipi]

Addressing a Problem?

I was trying to use fg.heatmap in the same way as fg.gridmap, supplying plot_kw = dict(col="time"), but it seems in this case we need a two-dimensional dataset, already prepared for a sns.heatmap.

I took the example ds_space from the tutorial, but I take a small square subset of a map, and identify lat -> model, lon -> prop, just to imitate a heatmap with this part of the tutorial:

ds_space = opened[['tx_max_p50']].isel(time=[0, 1, 2]).sel(lat=slice(40,65), lon=slice(-90,-55))
# subset and select variable
sl = slice(100,100+5)
da = ds_space.isel(lat=sl, lon=sl).drop("horizon").tx_max_p50
da = da.rename({"lat":"model", "lon":"prop"})
da = da.assign_coords(model=[f"s{n}" for n  in np.arange(da.model.size)], prop=[f"p{m}" for m in np.arange(da.prop.size)])

If I select a single time point, everything is good for a heatmap, fg.heatmap(da.isel(time=0).drop("time"))

But I can't naively ask for fg.heatmap(da, plot_kw={"col":"time"}) as mentioned above, I get

ValueError: DataArray must have exactly two dimensions

Hoping fg.gridmap may handle my use case, I try: fg.gridmap(da, plot_kw={"col":"time"}), but the strings I put in my coordinates are posing a problem:

UFuncTypeError: ufunc 'subtract' did not contain a loop with signature matching types (dtype('<U2'), dtype('<U2')) -> None

so this is not intended for this use. I can still replace my coordinates fg.gridmap(da.assign_coords(model=np.arange(da.model.size), prop=np.arange(da.prop.size)), plot_kw={"col":"time"}) , I get: the coordinates are dropped altogether

Potential Solution

I'm not sure if this use case is intended to be use gridmap for this purpose? Currently I'm just using plt.imshow directly da.plot.imshow(col="time"), which works well enough for me :

Additional context

Have other people used grids of heatmaps with figanos? Is there something obvious I'm not seeing?

Contribution

• [X] I would be willing/able to open a Pull Request to contribute this feature.

[juliettelavoie]

The multiple subplots mechanics comes from the xarray plotting library (https://docs.xarray.dev/en/latest/user-guide/plotting.html#faceting). The fg.heatmap function uses the seaborn library. Hence, subplots are not available in figanos (yet?). Your type of data is really meant to be plotted with heatmap, not gridmap.

Your solution with imshow, through xarray, seems to work well for you, but doesn't translate directly to figanos as heatmap uses sns.heatmap, not imshow.

It would be cool (and consistent) to be able to declare a subplot in the same way for heatmap, than gridmap and timeseries. @sarahclaude did you look at using seaborn factegrids when working on the subplots?

[sarahclaude]

Yes, it would indeed be nice if all our seaborn based functions could also be used for subplots.

I looked a bit into seaborn facetgrid at first since I was confused between their facetgrid and xarray facetgrids. Altough the two are similar, I believe xarray inspired themselves from seaborn to make their own, they are a bit different. The main one being seaborn works with pd.Dataframe. If we want to keep using seaborn and add subplots we would have to add conversion from xarray.Da to pd.Dataframe.

So we could either continue with seaborn and transform xarray.Da components into dataframes or switch to .imshow()

[juliettelavoie]

I personnaly like seaborn heatmap more. It has useful features like annot that imshow doesn't have.

[coxipi]

It has useful features like annot

I tried playing with ax.annotate, but having a visually appealing, with the text not too big, etc etc, is messy. It's probably better to let seaborn/matplotlib handle these lower level operations and focus on xarray -> pandas conversion if needed. I find that the integration with xarray.Datarray / facetgrids is really neat, it's too bad there's not a fully featured "heatmap" function in xarray / matplotlib.

[coxipi]

generate figure below with new figanos

import matplotlib.pyplot as plt
import xarray as xr
import numpy as np
import figanos.matplotlib as fg
# create xarray object from a NetCDF
url = 'https://pavics.ouranos.ca//twitcher/ows/proxy/thredds/dodsC/birdhouse/disk2/cccs_portal/indices/Final/BCCAQv2_CMIP6/tx_max/YS/ssp585/ensemble_percentiles/tx_max_ann_BCCAQ2v2+ANUSPLIN300_historical+ssp585_1950-2100_30ymean_percentiles.nc'
opened = xr.open_dataset(url, decode_timedelta=False)
ds_space = opened[['tx_max_p50']].isel(time=[0, 1, 2]).sel(lat=slice(40,65), lon=slice(-90,-55))
# subset and select variable
sl = slice(100,100+5)
da = ds_space.isel(lat=sl, lon=sl).drop("horizon").tx_max_p50
da = da.rename({"lat":"model", "lon":"prop"})
da = da.assign_coords(model=[f"s{n}" for n  in np.arange(da.model.size)], prop=[f"p{m}" for m in np.arange(da.prop.size)])
fg.heatmap(da, plot_kw = {"col": "time", "annot":True}, fig_kw={"figsize":(14,4)})

I'm getting this with a relatively simple modification of fg.heatmap ... I have to play with figsize to ensure annot has enough place to appear well, is that expected or something that should be handled automatically?

the heart of the change in figanos

    # plot
    if ax is not None: 
        sns.heatmap(df, ax=ax, **plot_kw)
        # format
        plt.xticks(rotation=45, ha="right", rotation_mode="anchor")
        ax.tick_params(axis="both", direction="out")

        set_plot_attrs(
            use_attrs,
            da,
            ax,
            title_loc="center",
            wrap_kw={"min_line_len": 35, "max_line_len": 44},
        )

        return ax
    else: 
        def draw_heatmap(*args, **kwargs):
            data = kwargs.pop('data')
            d = data.pivot(index=args[1], columns=args[0], values=args[2])
            sns.heatmap(d, **kwargs)
        plt.figure(**fig_kw)
        g = sns.FacetGrid(df, col=plot_kw["col"], row=plot_kw["row"])
        plot_kw.pop("col")
        plot_kw.pop("row")
        cax = g.fig.add_axes([.92, .12, .02, .8])
        ax = g.map_dataframe(draw_heatmap, *heatmap_dims, da_name, **plot_kw, cbar=True, cbar_ax=cax)
        g.fig.subplots_adjust(right=.9)
        if "figsize" in fig_kw.keys():
            g.fig.set_size_inches(*fig_kw["figsize"])
        plt.xticks(rotation=45, ha="right", rotation_mode="anchor")
        return g

[juliettelavoie]

nice ! Does this mean that heatmap will return a FacetGrid even if it is 1 subplot ? I suggest changing the else for elif "row" in plot_kw or "col" in plot_kw.

I think it's okay that we have to play with figsize or fmt to make the annotation fit.

[coxipi]

this mean that heatmap will return a FacetGrid even if it is 1 subplot

If you specify some col/row, yes, otherwise it's a normal ax:

out = fg.heatmap(da.isel(time=0), plot_kw={"col":"time"})
print(type(out))
>>> <class 'seaborn.axisgrid.FacetGrid'>
out = fg.heatmap(da.isel(time=0))
print(type(out))
>>><class 'matplotlib.axes._axes.Axes'>

It's similar to fg.gridmap. However, I tried doing the same with fg.gridmap, and it won't work if I try a FacetGrid with only one member in the FacetGrid, maybe it's a difference between xr-FacetGrids and sns-FacetGrids?

I suggest changing the else for elif "row" in plot_kw or "col" in plot_kw

Ok, got it, I changed this in my PR (#219 ). There was already a check before, this condition is True is ax is None, so in this context it was equivalent, but maybe it's not the best practice, especially if the function changes, maybe there could be a new formulation where this is not the case.