PySDM_examples.Shipway_and_Hill_2012.plot

  1import numpy as np
  2from matplotlib import pyplot
  3from open_atmos_jupyter_utils import show_plot
  4
  5from PySDM.physics import convert_to, si
  6
  7
  8def plot(
  9    var, qlabel, fname, output, vmin=None, vmax=None, cmin=None, cmax=None, line=None
 10):
 11    line = line or {15: ":", 20: "--", 25: "-", 30: "-."}
 12    dt = output["t"][1] - output["t"][0]
 13    dz = output["z"][1] - output["z"][0]
 14    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"] + dt / 2))
 15    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
 16    convert_to(zgrid, si.km)
 17
 18    fig = pyplot.figure(constrained_layout=True)
 19    gs = fig.add_gridspec(25, 5)
 20    ax1 = fig.add_subplot(gs[:-1, 0:4])
 21    if cmin is not None and cmax is not None:
 22        mesh = ax1.pcolormesh(
 23            tgrid, zgrid, output[var], cmap="BuPu", vmin=cmin, vmax=cmax
 24        )
 25    else:
 26        mesh = ax1.pcolormesh(tgrid, zgrid, output[var], cmap="BuPu")
 27
 28    ax1.set_xlabel("time [s]")
 29    ax1.set_ylabel("z [km]")
 30    ax1.set_ylim(0, None)
 31
 32    cbar = fig.colorbar(mesh, fraction=0.05, location="top")
 33    cbar.set_label(qlabel)
 34
 35    ax2 = fig.add_subplot(gs[:-1, 4:], sharey=ax1)
 36    ax2.set_xlabel(qlabel)
 37    # ax2.set_yticks([], minor=False)
 38
 39    last_t = 0
 40    for i, t in enumerate(output["t"]):
 41        x, z = output[var][:, i], output["z"].copy()
 42        convert_to(z, si.km)
 43        params = {"color": "black"}
 44        for line_t, line_s in line.items():
 45            if last_t < line_t * si.min <= t:
 46                params["ls"] = line_s
 47                ax2.plot(x, z, **params)
 48                if vmin is not None and vmax is not None:
 49                    ax1.axvline(t, ymin=vmin, ymax=vmax, **params)
 50                else:
 51                    ax1.axvline(t, **params)
 52        last_t = t
 53
 54    show_plot(filename=fname, inline_format="png")
 55
 56
 57def plot_plusminus(
 58    plus_vars,
 59    minus_vars,
 60    qlabel,
 61    fname,
 62    output,
 63    vmin=None,
 64    vmax=None,
 65    line=None,
 66    cmin=0.0,
 67    cmax=None,
 68):
 69    line = line or {15: ":", 20: "--", 25: "-", 30: "-."}
 70    dt = output["t"][1] - output["t"][0]
 71    dz = output["z"][1] - output["z"][0]
 72    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"] + dt / 2))
 73    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
 74    convert_to(zgrid, si.km)
 75
 76    fig = pyplot.figure(constrained_layout=True)
 77    gs = fig.add_gridspec(25, 5)
 78    ax1 = fig.add_subplot(gs[:-1, 0:4])
 79    plot_var = 0.0 * output[plus_vars[0]]
 80    for plus_var in plus_vars:
 81        plot_var += output[plus_var]
 82    for minus_var in minus_vars:
 83        plot_var -= output[minus_var]
 84
 85    mesh = ax1.pcolormesh(tgrid, zgrid, plot_var, cmap="BuPu", vmin=cmin, vmax=cmax)
 86
 87    ax1.set_xlabel("time [s]")
 88    ax1.set_ylabel("z [km]")
 89    ax1.set_ylim(0, None)
 90
 91    cbar = fig.colorbar(mesh, fraction=0.05, location="top")
 92    cbar.set_label(qlabel)
 93
 94    ax2 = fig.add_subplot(gs[:-1, 4:], sharey=ax1)
 95    ax2.set_xlabel(qlabel)
 96    # ax2.set_yticks([], minor=False)
 97
 98    last_t = 0
 99    for i, t in enumerate(output["t"]):
100        x, z = plot_var[:, i], output["z"].copy()
101        convert_to(z, si.km)
102        params = {"color": "black"}
103        for line_t, line_s in line.items():
104            if last_t < line_t * si.min <= t:
105                params["ls"] = line_s
106                ax2.plot(x, z, **params)
107                if vmin is not None and vmax is not None:
108                    ax1.axvline(t, ymin=vmin, ymax=vmax, **params)
109                else:
110                    ax1.axvline(t, **params)
111        last_t = t
112
113    show_plot(filename=fname, inline_format="png")