monobiome/monobiome/plotting.py

import numpy as np
import matplotlib.pyplot as plt

from monobiome.constants import (
    h_map,
    L_space,
    L_points,
    accent_h_map,
    monotone_h_map,
    Lspace_Cmax_Hmap,
    Lpoints_Cstar_Hmap,
)


def plot_hue_chroma_bounds() -> None:
    name_h_map = {}
    ax_h_map = {}
    fig, axes = plt.subplots(
        len(monotone_h_map),
        1,
        sharex=True,
        sharey=True,
        figsize=(4, 10)
    )

    for i, h_str in enumerate(Lpoints_Cstar_Hmap):
        _h = h_map[h_str]

        l_space_Cmax = Lspace_Cmax_Hmap[h_str]
        l_points_Cstar = Lpoints_Cstar_Hmap[h_str]
        
        if _h not in ax_h_map:
            ax_h_map[_h] = axes[i]
        ax = ax_h_map[_h]
        
        if _h not in name_h_map:
            name_h_map[_h] = []
        name_h_map[_h].append(h_str)

        # plot Cmax and Cstar
        ax.plot(L_space, l_space_Cmax, c="g", alpha=0.3, label="Cmax")
        
        cstar_label = f"{'accent' if h_str in accent_h_map else 'monotone'} C*"
        ax.plot(L_points, l_points_Cstar, alpha=0.7, label=cstar_label)
        
        ax.title.set_text(f"Hue [${_h}$] - {'|'.join(name_h_map[_h])}")
        
    axes[-1].set_xlabel("Lightness (%)")
    axes[-1].set_xticks([L_points[0], L_points[-1]])

    fig.tight_layout()
    fig.subplots_adjust(top=0.9)

    handles, labels = axes[-1].get_legend_handles_labels()
    unique = dict(zip(labels, handles))
    fig.legend(unique.values(), unique.keys(), loc='lower center', bbox_to_anchor=(0.5, -0.06), ncol=3)

    plt.suptitle("$C^*$ curves for hue groups")
    plt.show()


def plot_hue_chroma_star() -> None:
    fig, ax = plt.subplots(1, 1, figsize=(8, 6))

    # uncomment to preview 5 core term colors
    colors = accent_h_map.keys()
    #colors = set(["red", "orange", "yellow", "green", "blue"])

    for h_str in Lpoints_Cstar_Hmap:
        if h_str not in accent_h_map or h_str not in colors:
            continue
        ax.fill_between(
            L_points,
            Lpoints_Cstar_Hmap[h_str],
            alpha=0.2,
            color='grey',
            label=h_str
        )

        x, y = L_points, Lpoints_Cstar_Hmap[h_str]
        n = int(0.45*len(x))
        ax.text(x[n], y[n]-0.01, h_str, rotation=10, va='center', ha='left')
        
    ax.set_xlabel("Lightness (%)")
    ax.set_xticks([L_points[0], 45, 50, 55, 60, 65, 70, L_points[-1]])
    plt.suptitle("$C^*$ curves (v1.4.0)")
    fig.show()


def palette_image(palette, cell_size=40, keys=None):
    if keys is None:
        names = list(palette.keys())
    else:
        names = keys
        
    row_count = len(names)
    col_counts = [len(palette[n]) for n in names]
    max_cols = max(col_counts)

    h = row_count * cell_size
    w = max_cols * cell_size
    img = np.ones((h, w, 3), float)

    lightness_keys_per_row = []

    for r, name in enumerate(names):
        shades = palette[name]
        keys = sorted(shades.keys())
        lightness_keys_per_row.append(keys)
        for c, k in enumerate(keys):
            col = Color(shades[k]).convert("srgb").fit(method="clip")
            rgb = [col["r"], col["g"], col["b"]]
            r0, r1 = r * cell_size, (r + 1) * cell_size
            c0, c1 = c * cell_size, (c + 1) * cell_size
            img[r0:r1, c0:c1, :] = rgb

    return img, names, lightness_keys_per_row, cell_size, max_cols


def show_palette(palette, cell_size=40, keys=None):
    img, names, keys, cell_size, max_cols = palette_image(palette, cell_size, keys=keys)

    fig_w = img.shape[1] / 100
    fig_h = img.shape[0] / 100
    fig, ax = plt.subplots(figsize=(fig_w, fig_h))

    ax.imshow(img, interpolation="none", origin="upper")
    ax.set_xticks([])

    ytick_pos = [(i + 0.5) * cell_size for i in range(len(names))]
    ax.set_yticks(ytick_pos)
    ax.set_yticklabels(names)

    ax.set_ylim(img.shape[0], 0)   # ensures rows render correctly without half-cells

    plt.show()


if __name__ == "__main__":
    from monobiome.constants import OKLCH_hL_dict

    keys = [
        "alpine",
        "badlands",
        "chaparral",
        "savanna",
        "grassland",
        "reef",
        "tundra",
        "heathland",
        "moorland",
        "orange",
        "yellow",
        "green",
        "cyan",
        "blue",
        "violet",
        "magenta",
        "red",
    ]
    term_keys = [
        "alpine",
        "badlands",
        "chaparral",
        "savanna",
        "grassland",
        "tundra",
        "red",
        "orange",
        "yellow",
        "green",
        "blue",
    ]

    show_palette(OKLCH_hL_dict, cell_size=25, keys=keys)
    # show_palette(OKLCH_hL_dict, cell_size=1, keys=term_keys)