def _get_cmap_kwargs(run_diags, variable, **kwargs):
input_data = []
for run in run_diags.runs:
input_data.append(
run_diags.get_variable(run, variable).assign_coords(run=run))
input_data = xr.concat(input_data, dim="run")
return fv3viz.infer_cmap_params(input_data.values, **kwargs)
def plot_spatial_2panel_with_diff(emu: xr.Dataset, base: xr.Dataset, varname: str):
fig = plt.figure()
ax = fig.add_subplot(131, projection=ccrs.Robinson())
ax2 = fig.add_subplot(132, projection=ccrs.Robinson())
ax3 = fig.add_subplot(133, projection=ccrs.Robinson())
fig.set_size_inches(15, 4)
fig.set_dpi(80)
vmin, vmax, cmap = infer_cmap_params(emu[varname], robust=True)
plot_kwargs = dict(vmin=vmin, vmax=vmax, cmap=cmap)
# TODO: can't pass explicit vmin, vmax so cbars arent equivalent
plot_cube(emu, varname, ax=ax, cmap_percentiles_lim=False, **plot_kwargs)
plot_cube(base, varname, ax=ax2, cmap_percentiles_lim=False, **plot_kwargs)
diff = emu.copy()
# Insert diff to keep undiffed grid coordinates
diff[varname] = emu[varname] - base[varname]
plot_cube(diff, varname, ax=ax3)
ax.set_title(f"Emulation: {varname}")
ax2.set_title("Baseline")
ax3.set_title("Diff: Emu - Baseline")
return fig
def plot_meridional(ds: xr.Dataset, varname: str, title="", ax=None, yincrease=False):
meridional = meridional_transect(ds)
if ax is None:
fig, ax = plt.subplots()
fig.set_dpi(120)
if "z_soil" in meridional[varname].dims:
y = "z_soil"
else:
y = "z"
vmin, vmax, cmap = infer_cmap_params(ds[varname], robust=True)
if cmap == "viridis":
cmap = "Blues"
meridional[varname].plot.pcolormesh(
x="lat", y=y, ax=ax, cmap=cmap, vmin=vmin, vmax=vmax, yincrease=yincrease,
)
ax.set_title(title, size=14)
ax.set_ylabel("vertical level", size=12)
ax.set_xlabel("latitude", size=12)
def plot_global_avg_by_height_panel(
da1: xr.DataArray, da2: xr.DataArray, x="time", dpi=80
):
fig, ax = plt.subplots(1, 3)
fig.set_size_inches(12, 4)
fig.set_dpi(dpi)
vmin, vmax, cmap = infer_cmap_params(da2, robust=True)
vkw = dict(vmin=vmin, vmax=vmax, cmap=cmap)
da1.plot.pcolormesh(x=x, y="z", ax=ax[0], yincrease=False, **vkw)
da2.plot.pcolormesh(x=x, y="z", ax=ax[1], yincrease=False, **vkw)
(da1 - da2).plot.pcolormesh(x=x, y="z", ax=ax[2], yincrease=False)
ax[0].set_title("Emulation")
ax[1].set_title("Baseline")
ax[2].set_title("Diff")
for sub_ax in ax:
sub_ax.tick_params(axis="x", labelrotation=15)
plt.tight_layout()
return fig
def test_infer_cmap_params(data, args, expected_result):
result = fv3viz.infer_cmap_params(data, **args)
assert result == expected_result