python-tutorial/tutorial/pandas_helpers.py at 7dee72bad34f46fc2a64d81e5da7a530a08a15da · empa-scientific-it/python-tutorial · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
"""Miscellaneous visual aids"""

import matplotlib.pyplot as plt
import pandas as pd


def low_med_high_bins_viz(data, column, ylabel, title, figsize=(15, 3)):
    """Visualize the low, medium, and high equal-width bins."""
    ax = data.plot(y=column, figsize=figsize, color="black", title=title)

    xlims = ax.get_xlim()

    for bin_name, hatch, bounds in zip(
        ["low", "med", "high"],
        ["///", "", "\\\\\\"],
        pd.cut(data[column], bins=3).unique().categories.values,
        strict=False,
    ):
        plt.axhspan(
            bounds.left,
            bounds.right,
            alpha=0.2,
            label=bin_name,
            hatch=hatch,
            color="black",
        )
        plt.annotate(
            f"  {bin_name}",
            xy=(xlims[0], (bounds.left + bounds.right) / 2.1),
            ha="left",
        )

    ax.set(xlabel="", ylabel=ylabel)
    plt.legend(bbox_to_anchor=(1, 0.75), frameon=False)

    return ax


def quartile_bins_viz(data, column, ylabel, title, figsize=(15, 8)):
    """Visualize quartile bins."""
    ax = data.plot(y=column, figsize=figsize, color="black", title=title)

    xlims = ax.get_xlim()

    for bin_name, hatch, bounds in zip(
        [r"$Q_1$", r"$Q_2$", r"$Q_3$", r"$Q_4$"],
        ["\\\\\\", "", "///", "||||"],
        pd.qcut(data.volume, q=4).unique().categories.values,
        strict=False,
    ):
        plt.axhspan(
            bounds.left,
            bounds.right,
            alpha=0.2,
            label=bin_name,
            hatch=hatch,
            color="black",
        )
        plt.annotate(
            f"  {bin_name}",
            xy=(xlims[0], (bounds.left + bounds.right) / 2.1),
            fontsize=11,
        )

    ax.set(xlabel="", ylabel=ylabel)
    plt.legend(bbox_to_anchor=(1, 0.67), frameon=False, fontsize=14)

    return ax