PySDM_examples.deJong_Mackay_et_al_2023.plot_rates 

 1import numpy as np
 2from matplotlib import colors, pyplot
 3
 4from PySDM.physics.constants import convert_to, si
 5
 6
 7def log_kwargs(clog, cmin, cmax):
 8    if clog:
 9        return {"norm": colors.LogNorm(vmin=cmin, vmax=cmax)}
10    return {"vmin": cmin, "vmax": cmax}
11
12
13def plot_ax(
14    ax,
15    var,
16    qlabel,
17    output,
18    *,
19    contour_var1=None,
20    contour_lvl1=None,
21    contour_var2=None,
22    contour_lvl2=None,
23    cmin=None,
24    cmax=None,
25    clog=False,
26):
27    tgrid = output["t"].copy()
28    zgrid = output["z"].copy()
29    convert_to(zgrid, si.km)
30
31    if clog:
32        data = output[var].copy()
33        data[data == 0] = np.nan
34    else:
35        data = output[var]
36
37    mesh = ax.pcolormesh(
38        tgrid,
39        zgrid,
40        data,
41        cmap="BuPu",
42        shading="nearest",
43        **log_kwargs(clog, cmin, cmax),
44    )
45
46    if contour_var1 is not None and contour_lvl1 is not None:
47        ax.contour(
48            tgrid,
49            zgrid,
50            output[contour_var1],
51            contour_lvl1,
52            colors="k",
53            linestyles="--",
54        )
55    if contour_var2 is not None and contour_lvl2 is not None:
56        ax.contour(
57            tgrid,
58            zgrid,
59            output[contour_var2],
60            contour_lvl2,
61            colors="r",
62            linestyles="--",
63        )
64
65    ax.set_xlabel("time [s]")
66    ax.set_ylabel("z [km]")
67    ax.set_ylim(0, None)
68
69    if clog:
70        cbar_levels = np.logspace(np.log10(cmin), np.log10(cmax), 5, endpoint="True")
71    else:
72        cbar_levels = np.linspace(cmin, cmax, 5, endpoint="True")
73    cbar = pyplot.colorbar(mesh, fraction=0.05, location="top", extend="max", ax=ax)
74    cbar.set_ticks(cbar_levels)
75    cbar.set_label(qlabel)
76
77
78def plot_zeros_ax(ax, var, qlabel, output, cmin=None, cmax=None, clog=False):
79    dt = output["t"][1] - output["t"][0]
80    dz = output["z"][1] - output["z"][0]
81    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"] + dt / 2))
82    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
83    convert_to(zgrid, si.km)
84
85    zeros = np.zeros_like(output[var])
86    kwargs = log_kwargs(clog, cmin, cmax)
87    cmap = "BuPu"
88    mesh = ax.pcolormesh(tgrid, zgrid, zeros, cmap=cmap, **kwargs)
89
90    ax.set_xlabel("time [s]")
91    ax.set_ylabel("z [km]")
92    ax.set_ylim(0, None)
93
94    cbar = pyplot.colorbar(mesh, fraction=0.05, location="top", ax=ax)
95    cbar.set_label(qlabel)
def log_kwargs(clog, cmin, cmax): 
 8def log_kwargs(clog, cmin, cmax):
 9    if clog:
10        return {"norm": colors.LogNorm(vmin=cmin, vmax=cmax)}
11    return {"vmin": cmin, "vmax": cmax}
def plot_ax( ax, var, qlabel, output, *, contour_var1=None, contour_lvl1=None, contour_var2=None, contour_lvl2=None, cmin=None, cmax=None, clog=False): 
14def plot_ax(
15    ax,
16    var,
17    qlabel,
18    output,
19    *,
20    contour_var1=None,
21    contour_lvl1=None,
22    contour_var2=None,
23    contour_lvl2=None,
24    cmin=None,
25    cmax=None,
26    clog=False,
27):
28    tgrid = output["t"].copy()
29    zgrid = output["z"].copy()
30    convert_to(zgrid, si.km)
31
32    if clog:
33        data = output[var].copy()
34        data[data == 0] = np.nan
35    else:
36        data = output[var]
37
38    mesh = ax.pcolormesh(
39        tgrid,
40        zgrid,
41        data,
42        cmap="BuPu",
43        shading="nearest",
44        **log_kwargs(clog, cmin, cmax),
45    )
46
47    if contour_var1 is not None and contour_lvl1 is not None:
48        ax.contour(
49            tgrid,
50            zgrid,
51            output[contour_var1],
52            contour_lvl1,
53            colors="k",
54            linestyles="--",
55        )
56    if contour_var2 is not None and contour_lvl2 is not None:
57        ax.contour(
58            tgrid,
59            zgrid,
60            output[contour_var2],
61            contour_lvl2,
62            colors="r",
63            linestyles="--",
64        )
65
66    ax.set_xlabel("time [s]")
67    ax.set_ylabel("z [km]")
68    ax.set_ylim(0, None)
69
70    if clog:
71        cbar_levels = np.logspace(np.log10(cmin), np.log10(cmax), 5, endpoint="True")
72    else:
73        cbar_levels = np.linspace(cmin, cmax, 5, endpoint="True")
74    cbar = pyplot.colorbar(mesh, fraction=0.05, location="top", extend="max", ax=ax)
75    cbar.set_ticks(cbar_levels)
76    cbar.set_label(qlabel)
def plot_zeros_ax(ax, var, qlabel, output, cmin=None, cmax=None, clog=False): 
79def plot_zeros_ax(ax, var, qlabel, output, cmin=None, cmax=None, clog=False):
80    dt = output["t"][1] - output["t"][0]
81    dz = output["z"][1] - output["z"][0]
82    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"] + dt / 2))
83    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
84    convert_to(zgrid, si.km)
85
86    zeros = np.zeros_like(output[var])
87    kwargs = log_kwargs(clog, cmin, cmax)
88    cmap = "BuPu"
89    mesh = ax.pcolormesh(tgrid, zgrid, zeros, cmap=cmap, **kwargs)
90
91    ax.set_xlabel("time [s]")
92    ax.set_ylabel("z [km]")
93    ax.set_ylim(0, None)
94
95    cbar = pyplot.colorbar(mesh, fraction=0.05, location="top", ax=ax)
96    cbar.set_label(qlabel)