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