example_plots.py

"""Functions to plot some images using matplotlib.

The generated plots are in the `assets` directory.
"""

from dataclasses import asdict

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np

from catppuccin import PALETTE
from catppuccin.models import Flavor

SEED = 0
POINTS = 50


def plot_palette(palette: Flavor) -> plt.Figure:  # type: ignore [name-defined]
    """Plot a palette with color names and hex values."""
    colors = asdict(palette.colors)

    # Create figure and adjust figure height to number of colormaps
    nrows = len(colors)
    figh = 0.35 + 0.15 + (nrows + (nrows - 1) * 0.1) * 0.22
    fig, axs = plt.subplots(nrows=nrows, figsize=(6.4, figh))
    fig.subplots_adjust(top=1 - 0.35 / figh, bottom=0.15 / figh, left=0.2, right=0.99)

    axs[0].set_title(palette.name, fontsize=14)

    for ax, color_name in zip(axs, colors, strict=True):
        ax.hlines(0, 0, 1, colors=colors[color_name]["hex"], linewidth=15)
        ax.text(
            -0.01,
            0.5,
            f"{color_name} {colors[color_name]['hex']}",
            va="center",
            ha="right",
            fontsize=10,
            transform=ax.transAxes,
        )

    # Turn off *all* ticks & spines, not just the ones with colormaps.
    for ax in axs:
        ax.set_axis_off()

    fig.tight_layout()

    return fig


def example_plot() -> plt.Figure:  # type: ignore [name-defined]
    """Generate a plot with multiple sin functions with phase shifts."""
    x = np.linspace(0.0, 1.0, num=101)
    phases = np.linspace(0.0, -0.8, num=5)
    np.sin(2 * np.pi * x)

    fig = plt.figure()
    for idx, phase in enumerate(phases):
        plt.plot(x, np.sin(2 * np.pi * x + phase), label=f"Color {idx + 1}")
    plt.grid()
    plt.legend()

    return fig


def example_scatter() -> plt.Figure:  # type: ignore [name-defined]
    """Generate a scatter plot with two sets of random data."""
    rng = np.random.default_rng(SEED)
    x = rng.random(POINTS)

    fig = plt.figure()
    plt.scatter(x, rng.random(POINTS))
    plt.scatter(x, rng.random(POINTS))

    return fig


def example_boxplot() -> plt.Figure:  # type: ignore [name-defined]
    """Generate a boxplot with random data."""
    rng = np.random.default_rng(SEED)
    bars = 4
    nominal_values = rng.random(bars)
    distributions = rng.random((bars, POINTS))

    fig = plt.figure()
    plt.boxplot(nominal_values + distributions.T, patch_artist=True)
    return fig


def example_bar() -> plt.Figure:  # type: ignore [name-defined]
    """Generate a bar plot with random data."""
    rng = np.random.default_rng(SEED)
    bars = 10
    x = np.arange(bars).astype(np.float64) + 0.5
    y = rng.random(bars)

    fig = plt.figure()
    plt.bar(x, y)
    return fig


def example_patches() -> plt.Figure:  # type: ignore [name-defined]
    """Generate a plot with two arrows."""
    fig, ax = plt.subplots()
    arrow_1 = mpatches.FancyArrowPatch((0, 1), (1, 0), mutation_scale=100)
    arrow_2 = mpatches.FancyArrowPatch((0, 0), (1, 1), mutation_scale=100)
    ax.set_xlim(-0.1, 1.1)
    ax.set_ylim(-0.1, 1.1)
    ax.add_patch(arrow_1)
    ax.add_patch(arrow_2)

    return fig


def example_imshow() -> plt.Figure:  # type: ignore [name-defined]
    """Generate an image plot with random data."""
    rng = np.random.default_rng(SEED)
    data = rng.random((30, 30))

    fig, ax = plt.subplots()
    im = ax.imshow(data)
    ax.tick_params(
        left=False, right=False, labelleft=False, labelbottom=False, bottom=False
    )

    fig.colorbar(im, ax=ax, ticks=[])

    return fig


if __name__ == "__main__":
    flavor = PALETTE.mocha
    plt.style.use(flavor.matplotlib_style)
    plot_palette(flavor)
    example_plot()
    example_scatter()
    example_boxplot()
    example_bar()
    example_patches()
    example_imshow()
    plt.show()