generate metric wedge plot; inspired by https://www.nature.com/articles/s41597-021-01079-3/figures/3

#wedge-quadrant-plot

 generate metric wedge plot; inspired by https://www.nature.com/articles/s41597-021-01079-3/figures/3

df_metric_inset.csv

wedge_quadrant.py

"""
Author: Ting Sun (ting.sun@reading.ac.uk)
wedge_quadrant.py (c) 2021
Desc: generate metric plot; inspired by https://www.nature.com/articles/s41597-021-01079-3/figures/3
Created:  2021-12-11T09:31:26.685Z
"""


# %%
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.collections import PatchCollection
from matplotlib.patches import Circle, Polygon, Wedge

# %%
# dict for labelling annotations
dict_lbl_var = {
    "KDOWN": r"$K_\downarrow$",
    "LDOWN": r"$L_\downarrow$",
    "KUP": r"$K_\uparrow$",
    "LUP": r"$L_\uparrow$",
    "QN": r"$Q^*$",
    "QF": r"$Q_F$",
    "QH": r"$Q_H$",
    "QE": r"$Q_E$",
    "Bo": r"$Bo$",
}
dict_lbl_var = {
    var: lbl + r" (W m$^{-2}$)" if var != "Bo" else lbl
    for var, lbl in dict_lbl_var.items()
}

dict_lbl_metric = {
    "MBE": r"(W m$^{-2}$)",
    "MAE": r"(W m$^{-2}$)",
    "HR": r"",
}

df_metric_inset = pd.read_csv(
    "./df_metric_inset.csv",
    index_col=[0, 1],
    header=[0, 1, 2],
)
df_metric_inset


# %%
# helper function to create a quadrant wedge
def plot_wedge(ax, ar_val, cmap_func):
    p = [
        Wedge(
            (0.5, 0.5),
            1,
            -45 + 90 * s,
            45 + 90 * s,
            facecolor=cmap_func.to_rgba(v),
            edgecolor="grey",
        )
        for s, v in enumerate(ar_val)
    ]
    p = PatchCollection(p, match_original=True)
    _ = ax.add_collection(p)
    _ = ax.xaxis.set_ticks([])
    _ = ax.yaxis.set_ticks([])


# list for looping
list_var = df_metric_inset.index.get_level_values("var").unique()
list_cat = df_metric_inset.index.get_level_values("cat").unique()
list_site = df_metric_inset.columns.get_level_values("site").unique()
list_metric = df_metric_inset.columns.get_level_values("metric").unique().drop("N")

# set up figure container
n_row, n_col = len(list_var) * len(list_cat), len(list_metric) * len(list_site)
size_base = 1.0

fig, axes = plt.subplots(
    n_row,
    n_col,
    figsize=(n_col * 1 * size_base, n_row * 0.6 * size_base),
    sharey=True,
    sharex=True,
    constrained_layout=True,
)


dict_cmap = {
    "MAE": "Reds",
    "MBE": "RdBu_r",
    "HR": "Purples",
    # "HR": "Blues",
}
dict_norm = {}
for metric in dict_cmap.keys():
    ar_var = df_metric_inset[metric].values
    dict_norm[metric] = mpl.colors.Normalize(vmin=ar_var.min(), vmax=ar_var.max())

# =============================================================================
# colour patches
for r0, var in enumerate(list_var):
    ar_var = df_metric_inset.loc[var].values
    for r1, cat in enumerate(list_cat):
        r = r0 * len(list_cat) + r1
        for c0, metric in enumerate(list_metric):
            norm = dict_norm[metric]
            cmap = mpl.cm.get_cmap(dict_cmap[metric])
            cmap_func = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
            for c1, site in enumerate(list_site):
                c = c0 * len(list_site) + c1
                ax = axes[r, c]

                # draw wedge into quadrant
                df_plot = df_metric_inset.loc[(var, cat), (metric, site)]
                plot_wedge(ax, df_plot.values, cmap_func)

                # label: site
                if r == 0:
                    _ = ax.set_xlabel(site)
                    _ = ax.xaxis.set_label_position("top")

                # label: category
                if c == 0:
                    _ = ax.set_ylabel(cat)


# draw colorbar
list_cax = []
for c, metric in enumerate(list_metric[-1::-1]):
    norm = dict_norm[metric]
    cmap = mpl.cm.get_cmap(dict_cmap[metric])
    cmap_func = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    cbar = fig.colorbar(
        cmap_func,
        ax=axes,
        orientation="horizontal",
        fraction=0.04,
        panchor=(0.1, 0.0),
        anchor=(1, 1),
        shrink=0.7,
        pad=0.01,
    )

    # colorbar label: hacked version as the label can't be set in the other end
    cax = cbar.ax
    lbl = metric + "\n" + dict_lbl_metric[metric]
    _ = cax.text(
        -0.08,
        0.0,
        lbl,
        size=11,
        va="center",
        ha="center",
        transform=cax.transAxes,
    )
    list_cax.append(cax)

# ============================================================
# note:
# > Constrained Layout will be effectuated each time the figure is drawn.
# > Hence you would need to draw the figure first, then you can get the actual position of the axes inside of it.
# ref: https://stackoverflow.com/a/54585004/920789
fig.canvas.draw()

# ============================================================
# outskirt labels: variables
for r0, var in enumerate(list_var):
    y = []
    for r1, cat in enumerate(list_cat):
        r = r0 * len(list_cat) + r1
        ax_test = axes[r, 0]
        bbox = ax_test.get_position()
        x_ax, y_ax, w, h = bbox.bounds
        # print(r, x_ax, y_ax, w, h, y_ax + h / 2)
        y.append(y_ax + h / 2)
        # y_ax_mid = y_ax + h / 2
    # print(y)
    x = -0.01
    y = np.mean(y)
    lbl = dict_lbl_var[var].split(" ")[0]
    _ = fig.text(
        x,
        y,
        lbl,
        rotation="vertical",
        ha="right",
        va="center",
        fontsize=13,
    )

# outskirt labels: metrics
for c0, metric in enumerate(list_metric):
    x = []
    for c1, cat in enumerate(list_cat):
        c = c0 * len(list_cat) + c1
        ax_test = axes[0, c]
        bbox = ax_test.get_position()
        x_ax, y_ax, w, h = bbox.bounds
        # print(c, x_ax, y_ax, w, h, x_ax + w / 2)
        x.append(x_ax + w / 2)
        # x_ax_mid = x_ax + w / 2
    # print(x)
    x = np.mean(x)
    lbl = metric
    _ = fig.text(
        x,
        1,
        lbl,
        rotation="horizontal",
        ha="center",
        va="bottom",
        fontsize=13,
    )

# ============================================================
# # note:
# # > Constrained Layout will be effectuated each time the figure is drawn.
# # > Hence you would need to draw the figure first, then you can get the actual position of the axes inside of it.
# # ref: https://stackoverflow.com/a/54585004/920789
fig.canvas.draw()


# ============================================================
# draw legend: quartrant wedge

# location of legend: use y of cax
ar_bbox_cax = np.array([cax.get_position().bounds for cax in list_cax])
y0 = np.mean(ar_bbox_cax[:, 1])

# location of legend: use x of axes
x0, _, w, h = axes[-1, 0].get_position().bounds

# bbox size of legend: use bbox of axes
w, h = w * 1.3, h * 1.6

# new axes to hold legend
ax = fig.add_axes([x0 - 0.02, y0 - h / 2, w, h])

# shift distance of month labels
offset = 0.28
x = [0, offset, 0, -offset]
y = [offset, 0, -offset, 0]
p = []
for s, month in enumerate(["Jan", "Apr", "Jul", "Oct"]):
    p += [
        Wedge(
            (0.5, 0.5),
            1,
            -45 + 90 * s,
            45 + 90 * s,
            facecolor="none",
            edgecolor="grey",
        )
    ]
    # month label
    _ = ax.text(x[s] + 0.5, y[s] + 0.5, month, ha="center", va="center")

p = PatchCollection(p, match_original=True)
_ = ax.add_collection(p)
_ = ax.xaxis.set_ticks([])
_ = ax.yaxis.set_ticks([])

fig.savefig("metric_wedge_quadrant.pdf", bbox_inches="tight")