[Python] How to create a heatmap from a tidy / long polars dataframe

I need to create a heatmap on the basis of a tidy/long pl.DataFrame. Consider the following example, where I used pandas and plotly to create a heatmap.

import plotly.express as px
import polars as pl

tidy_df_pl = pl.DataFrame(
{
"x": [10, 10, 10, 20, 20, 20, 30, 30, 30],
"y": [3, 4, 5, 3, 4, 5, 3, 4, 5],
"value": [5, 8, 2, 4, 10, 14, 10, 8, 9],
}
)

print(tidy_df_pl)

shape: (9, 3)
┌─────┬─────┬───────┐
│ x ┆ y ┆ value │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═══════╡
│ 10 ┆ 3 ┆ 5 │
│ 10 ┆ 4 ┆ 8 │
│ 10 ┆ 5 ┆ 2 │
│ 20 ┆ 3 ┆ 4 │
│ 20 ┆ 4 ┆ 10 │
│ 20 ┆ 5 ┆ 14 │
│ 30 ┆ 3 ┆ 10 │
│ 30 ┆ 4 ┆ 8 │
│ 30 ┆ 5 ┆ 9 │
└─────┴─────┴───────┘


Transforming to a wide pd.DataFrame:

pivot_df_pd = (
tidy_df_pl.pivot(index="x", on="y", values="value").to_pandas().set_index("x")
)
print(pivot_df_pd)

3 4 5
x
10 5 8 2
20 4 10 14
30 10 8 9


Creating the heatmap using plotly.

fig = px.imshow(pivot_df_pd)
fig.show()




This all seems a bit cumbersome. I am looking for polars-only. How can I create this heatmap directly from polars without going through a third library?

Continue reading...