PyMPDATA_examples.Jarecka_et_al_2015.plot_output

  1import numpy as np
  2from matplotlib import pylab
  3from PyMPDATA_examples.Jarecka_et_al_2015 import formulae
  4
  5
  6def plot_output(times, output, settings, return_data=True):
  7    lambdas_analytic = formulae.d2_el_lamb_lamb_t_evol(
  8        times=[0] + list(times), lamb_x0=settings.lx0, lamb_y0=settings.ly0
  9    )[1:]
 10
 11    data = {}
 12    cuts = ("x", "y")
 13    _, axs = pylab.subplots(
 14        nrows=len(times), ncols=len(cuts), figsize=(9, 9), constrained_layout=True
 15    )
 16    momenta = {"x": "uh", "y": "vh"}
 17    for i_cut, cut in enumerate(cuts):
 18        if cut == "x":
 19            n = settings.nx
 20            idx = (slice(None, None), slice(n // 2, n // 2 + 1))
 21            coord = settings.dx * (np.linspace(-n // 2, n // 2, n) + 0.5)
 22            x = coord
 23            y = 0
 24        else:
 25            n = settings.ny
 26            idx = (slice(n // 2, n // 2 + 1), slice(None, None))
 27            coord = settings.dy * (np.linspace(-n // 2, n // 2, n) + 0.5)
 28            x = 0
 29            y = coord
 30        for i_t, t in enumerate(times):
 31            key = f"cut={cut} t={t}"
 32            data[key] = {"coord": coord, "x": x, "y": y}
 33            datum = data[key]
 34            datum["h_analytic"] = formulae.amplitude(
 35                x, y, *lambdas_analytic[i_t, (0, 2)]
 36            )
 37            where = "mid"
 38            axs[i_t, i_cut].set_xlim(-8.5, 8.5)
 39            axs[i_t, i_cut].set_xticks((-5, 0, 5))
 40            axs[i_t, i_cut].step(
 41                datum["coord"],
 42                output[0]["h"][idx].squeeze(),
 43                color="black",
 44                label="h @ t=0",
 45                where=where,
 46            )
 47            axs[i_t, i_cut].step(
 48                datum["coord"],
 49                output[t]["h"][idx].squeeze(),
 50                color="red",
 51                label="MPDATA: h",
 52                where=where,
 53            )
 54            axs[i_t, i_cut].plot(
 55                datum["coord"],
 56                datum["h_analytic"],
 57                color="blue",
 58                linestyle=":",
 59                label="analytic: h",
 60            )
 61            twin = axs[i_t, i_cut].twinx()
 62            datum["h_numeric"] = output[t]["h"][idx].squeeze()
 63            mask = datum["h_numeric"] > settings.eps
 64            datum["q_h_numeric"] = np.where(mask, np.nan, 0)
 65            np.divide(
 66                output[t][momenta[cut]][idx].squeeze(),
 67                datum["h_numeric"],
 68                where=mask,
 69                out=datum["q_h_numeric"],
 70            )
 71            twin.step(
 72                datum["coord"],
 73                datum["q_h_numeric"],
 74                color="red",
 75                linestyle="--",
 76                label="MPDATA: q/h",
 77                where=where,
 78            )
 79            datum["q_h_analytic"] = np.where(
 80                datum["h_analytic"] > settings.eps,
 81                datum["coord"]
 82                / lambdas_analytic[i_t, 0 + 2 * i_cut]
 83                * lambdas_analytic[i_t, 1 + 2 * i_cut],
 84                0,
 85            )
 86            twin.plot(
 87                datum["coord"],
 88                datum["q_h_analytic"],
 89                linestyle="-.",
 90                label="analytic: q/h",
 91                color="blue",
 92            )
 93            twin.set_ylabel("u" if cut == "x" else "v")
 94            twin.set_ylim(-1.2, 1.2)
 95            twin.set_yticks((-1, 0, 1))
 96            axs[i_t, i_cut].set_xlabel(cut)
 97            axs[i_t, i_cut].set_ylim(-0.6, 0.6)
 98            axs[i_t, i_cut].set_yticks((-0.5, 0, 0.5))
 99            axs[i_t, i_cut].set_ylabel("h")
100            axs[i_t, i_cut].set_title(f"t={t}", y=1.0, pad=-14, x=0.075)
101            axs[i_t, i_cut].grid()
102
103    legend_props = {"frameon": False, "fontsize": "small"}
104    axs[-1, -1].legend(loc="lower right", **legend_props)
105    twin.legend(loc="lower center", **legend_props)
106    if return_data:
107        return data
108    return None
def plot_output(times, output, settings, return_data=True):
  7def plot_output(times, output, settings, return_data=True):
  8    lambdas_analytic = formulae.d2_el_lamb_lamb_t_evol(
  9        times=[0] + list(times), lamb_x0=settings.lx0, lamb_y0=settings.ly0
 10    )[1:]
 11
 12    data = {}
 13    cuts = ("x", "y")
 14    _, axs = pylab.subplots(
 15        nrows=len(times), ncols=len(cuts), figsize=(9, 9), constrained_layout=True
 16    )
 17    momenta = {"x": "uh", "y": "vh"}
 18    for i_cut, cut in enumerate(cuts):
 19        if cut == "x":
 20            n = settings.nx
 21            idx = (slice(None, None), slice(n // 2, n // 2 + 1))
 22            coord = settings.dx * (np.linspace(-n // 2, n // 2, n) + 0.5)
 23            x = coord
 24            y = 0
 25        else:
 26            n = settings.ny
 27            idx = (slice(n // 2, n // 2 + 1), slice(None, None))
 28            coord = settings.dy * (np.linspace(-n // 2, n // 2, n) + 0.5)
 29            x = 0
 30            y = coord
 31        for i_t, t in enumerate(times):
 32            key = f"cut={cut} t={t}"
 33            data[key] = {"coord": coord, "x": x, "y": y}
 34            datum = data[key]
 35            datum["h_analytic"] = formulae.amplitude(
 36                x, y, *lambdas_analytic[i_t, (0, 2)]
 37            )
 38            where = "mid"
 39            axs[i_t, i_cut].set_xlim(-8.5, 8.5)
 40            axs[i_t, i_cut].set_xticks((-5, 0, 5))
 41            axs[i_t, i_cut].step(
 42                datum["coord"],
 43                output[0]["h"][idx].squeeze(),
 44                color="black",
 45                label="h @ t=0",
 46                where=where,
 47            )
 48            axs[i_t, i_cut].step(
 49                datum["coord"],
 50                output[t]["h"][idx].squeeze(),
 51                color="red",
 52                label="MPDATA: h",
 53                where=where,
 54            )
 55            axs[i_t, i_cut].plot(
 56                datum["coord"],
 57                datum["h_analytic"],
 58                color="blue",
 59                linestyle=":",
 60                label="analytic: h",
 61            )
 62            twin = axs[i_t, i_cut].twinx()
 63            datum["h_numeric"] = output[t]["h"][idx].squeeze()
 64            mask = datum["h_numeric"] > settings.eps
 65            datum["q_h_numeric"] = np.where(mask, np.nan, 0)
 66            np.divide(
 67                output[t][momenta[cut]][idx].squeeze(),
 68                datum["h_numeric"],
 69                where=mask,
 70                out=datum["q_h_numeric"],
 71            )
 72            twin.step(
 73                datum["coord"],
 74                datum["q_h_numeric"],
 75                color="red",
 76                linestyle="--",
 77                label="MPDATA: q/h",
 78                where=where,
 79            )
 80            datum["q_h_analytic"] = np.where(
 81                datum["h_analytic"] > settings.eps,
 82                datum["coord"]
 83                / lambdas_analytic[i_t, 0 + 2 * i_cut]
 84                * lambdas_analytic[i_t, 1 + 2 * i_cut],
 85                0,
 86            )
 87            twin.plot(
 88                datum["coord"],
 89                datum["q_h_analytic"],
 90                linestyle="-.",
 91                label="analytic: q/h",
 92                color="blue",
 93            )
 94            twin.set_ylabel("u" if cut == "x" else "v")
 95            twin.set_ylim(-1.2, 1.2)
 96            twin.set_yticks((-1, 0, 1))
 97            axs[i_t, i_cut].set_xlabel(cut)
 98            axs[i_t, i_cut].set_ylim(-0.6, 0.6)
 99            axs[i_t, i_cut].set_yticks((-0.5, 0, 0.5))
100            axs[i_t, i_cut].set_ylabel("h")
101            axs[i_t, i_cut].set_title(f"t={t}", y=1.0, pad=-14, x=0.075)
102            axs[i_t, i_cut].grid()
103
104    legend_props = {"frameon": False, "fontsize": "small"}
105    axs[-1, -1].legend(loc="lower right", **legend_props)
106    twin.legend(loc="lower center", **legend_props)
107    if return_data:
108        return data
109    return None