Search code examples
pythoncolorspython-polarsplotnine

Supply hex color values to heat map Plotnine


I am working with some data and would like to make a heat map with a specified, pre-made color palette. The data frame is shown below (I am using polars for all the data frame ops in this analysis).

┌──────────────┬─────────────┬────────────────┬──────────────────────┬─────────┐
│ team1        ┆ team2       ┆ type           ┆ value                ┆ color   │
│ ---          ┆ ---         ┆ ---            ┆ ---                  ┆ ---     │
│ str          ┆ str         ┆ i64            ┆ f64                  ┆ str     │
╞══════════════╪═════════════╪════════════════╪══════════════════════╪═════════╡
│ team1_1      ┆ team2_1     ┆ 18             ┆ 115.850278           ┆ #443b84 │
│ team1_1      ┆ team2_2     ┆ 26             ┆ 24.241389            ┆ #470e61 │
│ team1_2      ┆ team2_2     ┆ 13             ┆ 3.278333             ┆ #440256 │
│ team1_1      ┆ team2_3     ┆ 30             ┆ 94.118333            ┆ #46327e │
│ team1_3      ┆ team2_3     ┆ 13             ┆ 35.186111            ┆ #481467 │
│ …            ┆ …           ┆ …              ┆ …                    ┆ …       │
│ team1_2      ┆ team2_1     ┆ 18             ┆ 67.937778            ┆ #482576 │
│ team1_2      ┆ team2_4     ┆ 22             ┆ 0.0                  ┆ #440154 │
│ team1_4      ┆ team2_2     ┆ 30             ┆ 1.199444             ┆ #440154 │
│ team1_1      ┆ team2_5     ┆ 15             ┆ 0.0                  ┆ #440154 │
│ team1_1      ┆ team2_3     ┆ 11             ┆ 8.345278             ┆ #450559 │
└──────────────┴─────────────┴────────────────┴──────────────────────┴─────────┘

In this dataset, I am trying to plot a single heat map figure for each of the team2 entries, but I want to retain a single palette, so the heat maps are comparable between the different plots (i.e., dark blue in the plot for team2_1 means the same thing as dark blue in the plot for team2_5). I produced the color column using the mizani library:

from mizani.palettes import cmap_pal

palette = cmap_pal("viridis")
palette_values = (df["value"] - df["value"].min()) / (
            df["value"].max() - df["value"].min()
        )
df = df.with_columns(color=pl.Series(palette(palette_values)))

I then loop over all the unique values in the team2 column and produce a heat map using plotnine.

for t2 in df["team2"].unique():
  df1 = df.filter(pl.col("team2") == t2)
  color_dict = {
    key: value
    for key, value in zip(df1["value"], df1["color"])
  }
  plt = (
    pn.ggplot(
      data=df1.to_pandas(),
      mapping=pn.aes(
        x="team1",
        y="type",
        label="value",
        fill="value",
      ),
    )
    + pn.geom_tile(show_legend=False)
    + pn.scale_fill_manual(color_dict)
    + pn.geom_text(show_legend=False, size=9, format_string="{:.2f}")
  )

I am getting an error from plotnine: TypeError: Continuous value supplied to discrete scale. What is the best way to go about what I want here?


Solution

  • You need to map the color column to the fill aesthetic and then use scale_fill_identity.

    plt = (
        pn.ggplot(
          data=df1.to_pandas(),
          mapping=pn.aes(
            x="team1",
            y="type",
            label="value",
            fill="color",
          ),
        )
        + pn.geom_tile(show_legend=False)
        + pn.scale_fill_identity()
        + pn.geom_text(show_legend=False, size=9, format_string="{:.2f}")
      )