1
"""Matplotlib rankplot."""
2 4
import matplotlib.pyplot as plt
3 4
import numpy as np
4

5 4
from ....stats.density_utils import histogram
6 4
from ...plot_utils import _scale_fig_size, make_label, compute_ranks
7 4
from . import backend_kwarg_defaults, backend_show, create_axes_grid
8

9

10 4
def plot_rank(
11
    axes,
12
    length_plotters,
13
    rows,
14
    cols,
15
    figsize,
16
    plotters,
17
    bins,
18
    kind,
19
    colors,
20
    ref_line,
21
    labels,
22
    ref_line_kwargs,
23
    bar_kwargs,
24
    vlines_kwargs,
25
    marker_vlines_kwargs,
26
    backend_kwargs,
27
    show,
28
):
29
    """Matplotlib rankplot.."""
30 4
    if ref_line_kwargs is None:
31 4
        ref_line_kwargs = {}
32 4
    ref_line_kwargs.setdefault("linestyle", "--")
33 4
    ref_line_kwargs.setdefault("color", "k")
34

35 4
    if bar_kwargs is None:
36 4
        bar_kwargs = {}
37 4
    bar_kwargs.setdefault("align", "center")
38

39 4
    if vlines_kwargs is None:
40 4
        vlines_kwargs = {}
41 4
    vlines_kwargs.setdefault("lw", 2)
42

43 4
    if marker_vlines_kwargs is None:
44 4
        marker_vlines_kwargs = {}
45 4
        marker_vlines_kwargs.setdefault("marker", "o")
46 4
        marker_vlines_kwargs.setdefault("lw", 0)
47

48 4
    if backend_kwargs is None:
49 4
        backend_kwargs = {}
50

51 4
    backend_kwargs = {
52
        **backend_kwarg_defaults(),
53
        **backend_kwargs,
54
    }
55

56 4
    figsize, ax_labelsize, titlesize, _, _, _ = _scale_fig_size(figsize, None, rows=rows, cols=cols)
57 4
    backend_kwargs.setdefault("figsize", figsize)
58 4
    backend_kwargs.setdefault("squeeze", True)
59 4
    if axes is None:
60 4
        _, axes = create_axes_grid(
61
            length_plotters,
62
            rows,
63
            cols,
64
            backend_kwargs=backend_kwargs,
65
        )
66

67 4
    for ax, (var_name, selection, var_data) in zip(np.ravel(axes), plotters):
68 4
        ranks = compute_ranks(var_data)
69 4
        bin_ary = np.histogram_bin_edges(ranks, bins=bins, range=(0, ranks.size))
70 4
        all_counts = np.empty((len(ranks), len(bin_ary) - 1))
71 4
        for idx, row in enumerate(ranks):
72 4
            _, all_counts[idx], _ = histogram(row, bins=bin_ary)
73 4
        gap = all_counts.max() * 1.05
74 4
        width = bin_ary[1] - bin_ary[0]
75

76 4
        bar_kwargs.setdefault("width", width)
77 4
        bar_kwargs.setdefault("edgecolor", ax.get_facecolor())
78
        # Center the bins
79 4
        bin_ary = (bin_ary[1:] + bin_ary[:-1]) / 2
80

81 4
        y_ticks = []
82 4
        if kind == "bars":
83 4
            for idx, counts in enumerate(all_counts):
84 4
                y_ticks.append(idx * gap)
85 4
                ax.bar(
86
                    bin_ary,
87
                    counts,
88
                    bottom=y_ticks[-1],
89
                    color=colors[idx],
90
                    **bar_kwargs,
91
                )
92 4
                if ref_line:
93 4
                    ax.axhline(y=y_ticks[-1] + counts.mean(), **ref_line_kwargs)
94 4
            if labels:
95 4
                ax.set_ylabel("Chain", fontsize=ax_labelsize)
96 4
        elif kind == "vlines":
97 4
            ymin = all_counts.mean()
98

99 4
            for idx, counts in enumerate(all_counts):
100 4
                ax.plot(bin_ary, counts, color=colors[idx], **marker_vlines_kwargs)
101 4
                ax.vlines(bin_ary, ymin, counts, colors=colors[idx], **vlines_kwargs)
102 4
            ax.set_ylim(0, all_counts.mean() * 2)
103 4
            if ref_line:
104 4
                ax.axhline(y=ymin, **ref_line_kwargs)
105

106 4
        if labels:
107 4
            ax.set_xlabel("Rank (all chains)", fontsize=ax_labelsize)
108 4
            ax.set_yticks(y_ticks)
109 4
            ax.set_yticklabels(np.arange(len(y_ticks)))
110 4
            ax.set_title(make_label(var_name, selection), fontsize=titlesize)
111
        else:
112 0
            ax.set_yticks([])
113

114 4
    if backend_show(show):
115 0
        plt.show()
116

117 4
    return axes

Read our documentation on viewing source code .

Loading