PyMPDATA_examples.Shipway_and_Hill_2012.plot

View Source

  1import matplotlib
  2import numpy as np
  3from matplotlib import pylab, pyplot
  4
  5from .formulae import convert_to, si
  6
  7
  8def plot(
  9    var,
 10    mult,
 11    label,
 12    output,
 13    rng=None,
 14    threshold=None,
 15    cmap="copper",
 16    rasterized=False,
 17    figsize=None,
 18):
 19    lines = {3: ":", 6: "--", 9: "-", 12: "-."}
 20    colors = {3: "crimson", 6: "orange", 9: "navy", 12: "green"}
 21    fctr = 50  # rebin by fctr in time dimension (https://gist.github.com/zonca/1348792)
 22
 23    dt = (output["t"][1] - output["t"][0]) * fctr
 24    dz = output["z"][1] - output["z"][0]
 25    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"][0::fctr] + dt / 2))
 26    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
 27    convert_to(zgrid, si.km)
 28
 29    fig = pyplot.figure(constrained_layout=True, figsize=figsize)
 30    gs = fig.add_gridspec(25, 5)
 31    ax1 = fig.add_subplot(gs[:-1, 0:3])
 32
 33    data = output[var] * mult
 34    lngth = data.shape[1] - 1
 35    assert (lngth // fctr) * fctr == lngth
 36
 37    tmp = np.empty((data.shape[0], (data.shape[1] - 1) // fctr + 1))
 38    tmp[:, 0] = data[:, 0]
 39    M, N = data[:, 1:].shape
 40    m, n = M, N // fctr
 41    tmp[:, 1:] = data[:, 1:].reshape((m, M // m, n, N // n)).mean(3).mean(1)
 42    data = tmp
 43
 44    if threshold is not None:
 45        data[data < threshold] = np.nan
 46    mesh = ax1.pcolormesh(
 47        tgrid / si.minutes,
 48        zgrid,
 49        data,
 50        cmap=cmap,
 51        rasterized=rasterized,
 52        vmin=None if rng is None else rng[0],
 53        vmax=None if rng is None else rng[1],
 54    )
 55
 56    ax1.set_xlabel("time [min]")
 57    ax1.set_xticks(list(lines.keys()))
 58    ax1.set_ylabel("z [km]")
 59    ax1.grid()
 60
 61    cbar = fig.colorbar(mesh, fraction=0.05, location="top")
 62    cbar.set_label(label)
 63
 64    ax2 = fig.add_subplot(gs[:-1, 3:], sharey=ax1)
 65    ax2.set_xlabel(label)
 66    ax2.grid()
 67    if rng is not None:
 68        ax2.set_xlim(rng)
 69
 70    last_t = -1
 71    for i, t in enumerate(output["t"]):
 72        x, z = output[var][:, i] * mult, output["z"].copy()
 73        convert_to(z, si.km)
 74        params = {"color": "black"}
 75        for line_t, line_s in lines.items():
 76            if last_t < line_t * si.minutes <= t:
 77                params["ls"] = line_s
 78                params["color"] = colors[line_t]
 79                ax2.step(x, z - (dz / si.km) / 2, where="pre", **params)
 80                ax1.axvline(t / si.minutes, **params)
 81        last_t = t
 82
 83
 84def plot_3d(psi, settings, options, r_min, r_max, max_height):
 85    max_psi = np.amax(psi / (1 / si.mg / si.um))
 86    if max_psi > max_height:
 87        raise ValueError(f"max(psi)={max_psi} > {max_height}")
 88    pylab.figure(figsize=(10, 5))
 89    ax = pylab.subplot(projection="3d")
 90    ax.view_init(75, 85)
 91
 92    cmap = matplotlib.colormaps["gray_r"]
 93    min_height = 0
 94
 95    dz = np.rot90(psi, 2).flatten()
 96    rgba = [cmap((k - min_height) / (1 + max_height)) for k in dz]
 97    dz[dz < 0.05 * max_height] = np.nan
 98    factor = 0.9
 99    ax.bar3d(
100        *[
101            arr.flatten()
102            for arr in np.meshgrid(
103                (settings.bin_boundaries[-2::-1] + settings.dr * (1 - factor) / 2)
104                / si.um,
105                (np.arange(settings.nz - 1, -1, -1) + (1 - factor) / 2)
106                * settings.dz
107                / si.km,
108            )
109        ],
110        z=0,
111        dx=factor * settings.dr / si.um,
112        dy=factor * settings.dz / si.km,
113        dz=dz / (1 / si.mg / si.um),
114        shade=True,
115        color=rgba,
116        lightsource=matplotlib.colors.LightSource(azdeg=-64, altdeg=15),
117        zsort="max",
118    )
119    ax.set_title(f"MPDATA iterations: {options.n_iters}")
120    ax.set_xlabel("droplet radius [μm]")
121    ax.set_ylabel("z [km]")
122    ax.set_zlabel("psi [1/μm 1/mg 1/m]")
123    ax.set_zlim(0, max_height / (1 / si.mg / si.um))
124    ax.set_zticks([0, max_height / (1 / si.mg / si.um)])
125    ax.set_ylim(settings.z_max / si.km, 0)
126    ax.set_xlim(r_max / si.um, r_min / si.um)
127    nticks = 6
128    ax.set_xticks((r_min + np.arange(nticks + 1) * (r_max - r_min) / nticks) / si.um)

def plot( var, mult, label, output, rng=None, threshold=None, cmap='copper', rasterized=False, figsize=None): View Source

 9def plot(
10    var,
11    mult,
12    label,
13    output,
14    rng=None,
15    threshold=None,
16    cmap="copper",
17    rasterized=False,
18    figsize=None,
19):
20    lines = {3: ":", 6: "--", 9: "-", 12: "-."}
21    colors = {3: "crimson", 6: "orange", 9: "navy", 12: "green"}
22    fctr = 50  # rebin by fctr in time dimension (https://gist.github.com/zonca/1348792)
23
24    dt = (output["t"][1] - output["t"][0]) * fctr
25    dz = output["z"][1] - output["z"][0]
26    tgrid = np.concatenate(((output["t"][0] - dt / 2,), output["t"][0::fctr] + dt / 2))
27    zgrid = np.concatenate(((output["z"][0] - dz / 2,), output["z"] + dz / 2))
28    convert_to(zgrid, si.km)
29
30    fig = pyplot.figure(constrained_layout=True, figsize=figsize)
31    gs = fig.add_gridspec(25, 5)
32    ax1 = fig.add_subplot(gs[:-1, 0:3])
33
34    data = output[var] * mult
35    lngth = data.shape[1] - 1
36    assert (lngth // fctr) * fctr == lngth
37
38    tmp = np.empty((data.shape[0], (data.shape[1] - 1) // fctr + 1))
39    tmp[:, 0] = data[:, 0]
40    M, N = data[:, 1:].shape
41    m, n = M, N // fctr
42    tmp[:, 1:] = data[:, 1:].reshape((m, M // m, n, N // n)).mean(3).mean(1)
43    data = tmp
44
45    if threshold is not None:
46        data[data < threshold] = np.nan
47    mesh = ax1.pcolormesh(
48        tgrid / si.minutes,
49        zgrid,
50        data,
51        cmap=cmap,
52        rasterized=rasterized,
53        vmin=None if rng is None else rng[0],
54        vmax=None if rng is None else rng[1],
55    )
56
57    ax1.set_xlabel("time [min]")
58    ax1.set_xticks(list(lines.keys()))
59    ax1.set_ylabel("z [km]")
60    ax1.grid()
61
62    cbar = fig.colorbar(mesh, fraction=0.05, location="top")
63    cbar.set_label(label)
64
65    ax2 = fig.add_subplot(gs[:-1, 3:], sharey=ax1)
66    ax2.set_xlabel(label)
67    ax2.grid()
68    if rng is not None:
69        ax2.set_xlim(rng)
70
71    last_t = -1
72    for i, t in enumerate(output["t"]):
73        x, z = output[var][:, i] * mult, output["z"].copy()
74        convert_to(z, si.km)
75        params = {"color": "black"}
76        for line_t, line_s in lines.items():
77            if last_t < line_t * si.minutes <= t:
78                params["ls"] = line_s
79                params["color"] = colors[line_t]
80                ax2.step(x, z - (dz / si.km) / 2, where="pre", **params)
81                ax1.axvline(t / si.minutes, **params)
82        last_t = t

def plot_3d(psi, settings, options, r_min, r_max, max_height): View Source

 85def plot_3d(psi, settings, options, r_min, r_max, max_height):
 86    max_psi = np.amax(psi / (1 / si.mg / si.um))
 87    if max_psi > max_height:
 88        raise ValueError(f"max(psi)={max_psi} > {max_height}")
 89    pylab.figure(figsize=(10, 5))
 90    ax = pylab.subplot(projection="3d")
 91    ax.view_init(75, 85)
 92
 93    cmap = matplotlib.colormaps["gray_r"]
 94    min_height = 0
 95
 96    dz = np.rot90(psi, 2).flatten()
 97    rgba = [cmap((k - min_height) / (1 + max_height)) for k in dz]
 98    dz[dz < 0.05 * max_height] = np.nan
 99    factor = 0.9
100    ax.bar3d(
101        *[
102            arr.flatten()
103            for arr in np.meshgrid(
104                (settings.bin_boundaries[-2::-1] + settings.dr * (1 - factor) / 2)
105                / si.um,
106                (np.arange(settings.nz - 1, -1, -1) + (1 - factor) / 2)
107                * settings.dz
108                / si.km,
109            )
110        ],
111        z=0,
112        dx=factor * settings.dr / si.um,
113        dy=factor * settings.dz / si.km,
114        dz=dz / (1 / si.mg / si.um),
115        shade=True,
116        color=rgba,
117        lightsource=matplotlib.colors.LightSource(azdeg=-64, altdeg=15),
118        zsort="max",
119    )
120    ax.set_title(f"MPDATA iterations: {options.n_iters}")
121    ax.set_xlabel("droplet radius [μm]")
122    ax.set_ylabel("z [km]")
123    ax.set_zlabel("psi [1/μm 1/mg 1/m]")
124    ax.set_zlim(0, max_height / (1 / si.mg / si.um))
125    ax.set_zticks([0, max_height / (1 / si.mg / si.um)])
126    ax.set_ylim(settings.z_max / si.km, 0)
127    ax.set_xlim(r_max / si.um, r_min / si.um)
128    nticks = 6
129    ax.set_xticks((r_min + np.arange(nticks + 1) * (r_max - r_min) / nticks) / si.um)