utils.py 3.72 KB
Newer Older
Thomas Purcell's avatar
Thomas Purcell committed
1
# Copyright 2021 Thomas A. R. Purcell
2
#
Thomas Purcell's avatar
Thomas Purcell committed
3
4
5
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
6
#
Thomas Purcell's avatar
Thomas Purcell committed
7
#     http://www.apache.org/licenses/LICENSE-2.0
8
#
Thomas Purcell's avatar
Thomas Purcell committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Utility functions used for generating plots

Functions:

setup_plot_ax: Sets up the fig and ax object
adjust_box_widths: Adjust the widths of a seaborn-generated boxplot.
latexify: Convert a string s into a latex string
"""
from copy import deepcopy
23
import numpy as np
Thomas Purcell's avatar
Thomas Purcell committed
24

25
from sissopp.postprocess.plot import plt, config
26
from matplotlib.patches import PathPatch
Thomas Purcell's avatar
Thomas Purcell committed
27
28


29
def setup_plot_ax(fig_settings, single_axis=True):
Thomas Purcell's avatar
Thomas Purcell committed
30
31
32
33
    """Sets up the fig and ax object

    Args:
        fig_settings (dict): Settings used to augment the plot
34
        single_axis (bool): If True use the single axis subplots adjust
Thomas Purcell's avatar
Thomas Purcell committed
35
    Returns:
36
37
38
39
        tuple: A tuple containing:
            - fig_config (dict): The dictionary of the figure configuration settings
            - fig (mpl.Figure): The figure object of the plot
            - ax (mpl.Axes): The axes object of the plot
Thomas Purcell's avatar
Thomas Purcell committed
40
41
42
43
44
45
46
47
48
49
    """
    fig_config = deepcopy(config)
    if fig_settings:
        fig_config.update(fig_settings)

    fig, ax = plt.subplots(
        nrows=1,
        ncols=1,
        figsize=[fig_config["size"]["width"], fig_config["size"]["height"]],
    )
50
51
52
53
    if single_axis:
        fig.subplots_adjust(**fig_config["subplots_adjust_single_axis"])
    else:
        fig.subplots_adjust(**fig_config["subplots_adjust"])
Thomas Purcell's avatar
Thomas Purcell committed
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

    return fig_config, fig, ax


def adjust_box_widths(ax, fac):
    """Adjust the widths of a seaborn-generated boxplot.

    Args:
        ax (pyplot.Axis): Axis to replot
        fac(float): Rescaling factor
    """
    # iterating through axes artists:
    for c in ax.get_children():

        # searching for PathPatches
        if isinstance(c, PathPatch):
            # getting current width of box:
            p = c.get_path()
            verts = p.vertices
            verts_sub = verts[:-1]
            xmin = np.min(verts_sub[:, 0])
            xmax = np.max(verts_sub[:, 0])
            xmid = 0.5 * (xmin + xmax)
            xhalf = 0.5 * (xmax - xmin)

            # setting new width of box
            xmin_new = xmid - fac * xhalf
            xmax_new = xmid + fac * xhalf
            verts_sub[verts_sub[:, 0] == xmin, 0] = xmin_new
            verts_sub[verts_sub[:, 0] == xmax, 0] = xmax_new

            # setting new width of median line
            for l in ax.lines:
                if np.all(l.get_xdata() == [xmin, xmax]):
                    l.set_xdata([xmin_new, xmax_new])


def latexify(s):
    """Convert a string s into a latex string"""
    power_split = s.split("^")
    if len(power_split) == 1:
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
        temp_s = s
    else:
        power_split[0] += "$"
        for pp in range(1, len(power_split)):
            unit_end = power_split[pp].split(" ")
            unit_end[0] = "{" + unit_end[0] + "}$"
            unit_end[-1] += "$"
            power_split[pp] = " ".join(unit_end)
        temp_s = "^".join(power_split)[:-1]

    subscript_split = temp_s.split("_")
    if len(subscript_split) == 1:
        return temp_s

    subscript_split[0] += "$"
    for pp in range(1, len(subscript_split)):
        unit_end = subscript_split[pp].split(" ")
        unit_end[0] = "\\mathrm{" + unit_end[0] + "}$"
Thomas Purcell's avatar
Thomas Purcell committed
113
        unit_end[-1] += "$"
114
        subscript_split[pp] = " ".join(unit_end)
Thomas Purcell's avatar
Thomas Purcell committed
115

116
    return "_".join(subscript_split)[:-1]