replaced hinv03

hinv03_SPR_components.py

@@ -0,0 +1,87 @@
+import numpy as np
+import pynanigans as pn
+import xarray as xr
+from matplotlib import pyplot as plt
+from cmocean import cm
+from aux01_physfuncs import calculate_filtered_PV
+from aux02_plotting import BuRd, letterize
+plt.rcParams["figure.constrained_layout.use"] = True
+plt.rcParams["font.size"] = 9
+
+modifier = ""
+slice_name = "tafields"
+Fr_h = 0.2
+
+#+++ Read and reindex dataset
+snaps = xr.open_dataset(f"data_post/{slice_name}_snaps{modifier}.nc").chunk(Fr_h=1, Ro_h=1)
+snaps = snaps.reindex(Ro_h = list(reversed(snaps.Ro_h)))
+snaps = snaps.sel(xC = slice(-snaps.headland_intrusion_size_max/3, np.inf),
+                  yC = slice(-snaps.L, np.inf))
+
+try:
+    snaps = snaps.reset_coords(("zC", "zF"))
+except ValueError:
+    pass
+#---
+
+#+++ Options
+cbar_kwargs = dict(location="right", shrink=0.5, fraction=0.012, pad=0.02, aspect=30)
+figsize = (9, 6.5)
+
+#plot_kwargs = dict(vmin=-0.005, vmax=0.005, cmap=plt.cm.RdBu_r, rasterized=True)
+plot_kwargs = dict(vmin=-8e-10, vmax=8e-10, cmap=cm.balance, rasterized=True)
+#---
+
+#+++ Create ageostrophic variables and pick subset of simulations
+snaps["Π"] = snaps.SPR.sum("j")
+snaps["Πₕ"] = snaps.SPR.sel(j=[1,2]).sum("j")
+
+variables = ["Π", "Πₕ"]
+labels = [r"$\Pi$",
+          r"$\Pi_h$"]
+
+snaps = snaps.sel(Fr_h=Fr_h)
+snaps.xC.attrs = dict(long_name="$x$", units="m")
+snaps.yC.attrs = dict(long_name="$y$", units="m")
+#---
+
+#+++ Plotting loop
+fig, axes = plt.subplots(ncols=len(snaps.Ro_h), nrows=2, sharex=True, sharey=True, figsize=figsize)
+snaps = snaps.reindex(Ro_h = list(reversed(snaps.Ro_h)))
+for j_Ro, Ro_h in enumerate(snaps.Ro_h.values):
+    print(f"Plotting Roₕ = {Ro_h}")
+
+    for i, variable in enumerate(variables):
+        ax = axes[i, j_Ro]
+        ct = snaps["q̄"].sel(Ro_h=Ro_h).pncontour(ax=ax, x="x", add_colorbar=False, levels=[0], zorder=10, linestyles="--", colors="purple")
+        im = snaps[variable].sel(Ro_h=Ro_h).pnplot(ax=ax, x="x", add_colorbar=False, **plot_kwargs)
+        if i==0:
+            ax.set_title(f"$Ro_h=$ {Ro_h}")
+            ax.set_xlabel("")
+        else:
+            ax.set_title("")
+
+
+    if j_Ro>0:
+        for ax in axes[:, j_Ro]:
+            ax.set_ylabel("")
+
+    if j_Ro == (len(snaps.Ro_h)-1):
+        for ax, label in zip(axes[:, j_Ro], labels):
+            ax2 = ax.twinx()
+            ax2.set_ylabel(label, fontsize=11)
+            ax2.tick_params(left=False, right=False, bottom=False, labelleft=False, labelright=False, labelbottom=False)
+            ax2.spines['top'].set_visible(False)
+#---
+
+#+++ Prettify and save
+opts_land = dict(cmap="Set2", vmin=0, vmax=1, alpha=1.0, zorder=10,)
+for ax in axes.flatten():
+    ax.pcolormesh(snaps.xC, snaps.yC, snaps.land_mask.where(snaps.land_mask==1), rasterized=True, **opts_land)
+
+fig.colorbar(im, ax=axes.ravel().tolist(), label="Shear production rates", **cbar_kwargs)
+fig.suptitle("")
+fig.get_layout_engine().set(w_pad=0.02, h_pad=0, hspace=0, wspace=0)
+letterize(axes.flatten(), x=0.05, y=0.9)
+fig.savefig(f"figures/SPR_components_comparison_Frh={snaps.Fr_h.item()}_{slice_name}{modifier}.pdf", dpi=200)
+#---

hvid00_facetgrid.py

@@ -79,7 +79,7 @@
     modifiers = ["",]
 
     varnames = ["q̃_norm", "Ro"]
-    varnames = ["q̄"]
+    varnames = ["q̄_norm"]
     contour_variable_name = None #"water_mask_buffered"
     contour_kwargs = dict(colors="y", linewidths=0.8, linestyles="--", levels=[0])
     #---