rasterizing the heatmap layer into a tiny (blurry) bitmap when writing PDF #426

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# /// script
# dependencies = ["pandas", "numpy", "plotly", "kaleido"]
# ///


from __future__ import annotations

from pathlib import Path
from importlib.metadata import version


import numpy as np
import pandas as pd
import plotly.graph_objects as go


print(f"plotly={version('plotly')}")
print(f"kaleido={version('kaleido')}")


def load_csv(csv_path: Path) -> pd.DataFrame:
    df = pd.read_csv(csv_path, index_col=0)
    df.index = pd.to_datetime(df.index.astype(str), format="%Y-%m", errors="raise")
    df.index.name = "time"
    return df


def prepare_heatmap_data(df_git: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
    df_binned = df_git.resample("6ME").sum().rename_axis("time_binned")

    heatmap_data = df_binned.T
    heatmap_data.columns = heatmap_data.columns.to_series().dt.strftime("%Y-%m")

    log_data = np.log10(heatmap_data.replace(0, np.nan))
    return heatmap_data, log_data


def main() -> None:
    df_git = load_csv(Path("data.csv"))
    heatmap_data, log_data = prepare_heatmap_data(df_git)

    fig = go.Figure()
    fig.add_heatmap(
        z=log_data.to_numpy(),
        x=log_data.columns.to_list(),
        y=log_data.index,
        customdata=heatmap_data.to_numpy(),
    )
    fig.write_image("repro.pdf")


if __name__ == "__main__":
    main()