PySDM/examples/PySDM_examples/Shipway_and_Hill_2012/plot.py at 4af2e969e173dcff6ede2d616d9ca1d2bd3ec9d1 · open-atmos/PySDM · 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
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
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import numpy as np
from matplotlib import pyplot
from open_atmos_jupyter_utils import show_plot

from PySDM.physics import convert_to, si


def plot(
    var, qlabel, fname, output, *, vmin=None, vmax=None, cmin=None, cmax=None, line=None, colors=None,
):
    line = line or {15: ":", 20: "--", 25: "-", 30: "-."}
    colors = colors or {15: 'k', 20: 'tab:blue', 25: 'tab:orange', 30: 'tab:green'}
    dt = output["t"][1] - output["t"][0]
    dz = output["z"][1] - output["z"][0]
    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"] + dt / 2))
    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
    convert_to(zgrid, si.km)

    fig = pyplot.figure(constrained_layout=True)
    gs = fig.add_gridspec(25, 5)
    ax1 = fig.add_subplot(gs[:-1, 0:4])
    if cmin is not None and cmax is not None:
        mesh = ax1.pcolormesh(
            tgrid, zgrid, output[var], cmap="BuPu", vmin=cmin, vmax=cmax
        )
    else:
        mesh = ax1.pcolormesh(tgrid, zgrid, output[var], cmap="BuPu")

    ax1.set_xlabel("time [s]")
    ax1.set_ylabel("z [km]")
    ax1.set_ylim(0, None)

    cbar = fig.colorbar(mesh, fraction=0.05, location="top")
    cbar.set_label(qlabel)

    ax2 = fig.add_subplot(gs[:-1, 4:], sharey=ax1)
    ax2.set_xlabel(qlabel)
    # ax2.set_yticks([], minor=False)

    last_t = 0
    for i, t in enumerate(output["t"]):
        x, z = output[var][:, i], output["z"].copy()
        convert_to(z, si.km)
        params = {"color": "black"}
        for line_t, line_s in line.items():
            if last_t < line_t * si.min <= t:
                params["ls"] = line_s
                ax2.plot(x, z, color=colors[line_t])
                if vmin is not None and vmax is not None:
                    ax1.axvline(t, ymin=vmin, ymax=vmax, color=colors[line_t])
                else:
                    ax1.axvline(t, color=colors[line_t])
        last_t = t

    show_plot(filename=fname, inline_format="png")


def plot_plusminus(
    plus_vars,
    minus_vars,
    qlabel,
    fname,
    output,
    *,
    vmin=None,
    vmax=None,
    line=None,
    cmin=0.0,
    cmax=None,
):
    line = line or {15: ":", 20: "--", 25: "-", 30: "-."}
    dt = output["t"][1] - output["t"][0]
    dz = output["z"][1] - output["z"][0]
    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"] + dt / 2))
    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
    convert_to(zgrid, si.km)

    fig = pyplot.figure(constrained_layout=True)
    gs = fig.add_gridspec(25, 5)
    ax1 = fig.add_subplot(gs[:-1, 0:4])
    plot_var = 0.0 * output[plus_vars[0]]
    for plus_var in plus_vars:
        plot_var += output[plus_var]
    for minus_var in minus_vars:
        plot_var -= output[minus_var]

    mesh = ax1.pcolormesh(tgrid, zgrid, plot_var, cmap="BuPu", vmin=cmin, vmax=cmax)

    ax1.set_xlabel("time [s]")
    ax1.set_ylabel("z [km]")
    ax1.set_ylim(0, None)

    cbar = fig.colorbar(mesh, fraction=0.05, location="top")
    cbar.set_label(qlabel)

    ax2 = fig.add_subplot(gs[:-1, 4:], sharey=ax1)
    ax2.set_xlabel(qlabel)
    # ax2.set_yticks([], minor=False)

    last_t = 0
    for i, t in enumerate(output["t"]):
        x, z = plot_var[:, i], output["z"].copy()
        convert_to(z, si.km)
        params = {"color": "black"}
        for line_t, line_s in line.items():
            if last_t < line_t * si.min <= t:
                params["ls"] = line_s
                ax2.plot(x, z, **params)
                if vmin is not None and vmax is not None:
                    ax1.axvline(t, ymin=vmin, ymax=vmax, **params)
                else:
                    ax1.axvline(t, **params)
        last_t = t

    show_plot(filename=fname, inline_format="png")