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