<jupyter_start><jupyter_text>Column integrated nudging tendencies<jupyter_code>import matplotlib.pyplot as plt
import matplotlib as mpl

import xarray as xr
import fsspec
import cftime
import intake
import vcm
import numpy as np

from vcm import  mass_integrate
from vcm.calc.thermo import column_integrated_heating_from_isochoric_transition, minus_column_integrated_moistening, latent_heat_vaporization
from vcm.catalog import catalog
from fv3viz import plot_cube
from cartopy import crs as ccrs
mpl.rcParams.update({'font.size': 8})
N2F_TENDENCIES_PATH = "gs://vcm-ml-experiments/2021-04-13-n2f-c3072/3-hrly-ave-rad-precip-setting-30-min-rad-timestep-shifted-start-tke-edmf/nudging_tendencies.zarr"
MAPPABLE_VAR_KWARGS = {
    "coord_x_center": "x",
    "coord_y_center": "y",
    "coord_x_outer": "x_interface",
    "coord_y_outer": "y_interface",
    "coord_vars": {
        "lonb": ["y_interface", "x_interface", "tile"],
        "latb": ["y_interface", "x_interface", "tile"],
        "lon": ["y", "x", "tile"],
        "lat": ["y", "x", "tile"],
    },
}

def zonal_mean(
    ds: xr.Dataset, latitude: xr.DataArray, bins=np.arange(-90, 91, 2)
) -> xr.Dataset:
    zm = ds.groupby_bins(latitude, bins=bins).mean().rename(lat_bins="latitude")
    latitude_midpoints = [x.item().mid for x in zm["latitude"]]
    return zm.assign_coords(latitude=latitude_midpoints)


n2f_tendencies = intake.open_zarr(N2F_TENDENCIES_PATH).to_dask()<jupyter_output><empty_output><jupyter_text>moist static energy nudging tendency $ \langle nQ_{MSE} \rangle =    \langle  nQ1 \rangle + L_{v}  \langle nQ2 \rangle$$$ heating is for constant volume<jupyter_code>Cp =  1004.
Rd = 287.05
Cv = Cp - Rd

def moist_static_energy_tendency(T: xr.DataArray, dQ1: xr.DataArray, dQ2: xr.DataArray): 
    return Cv * dQ1 + latent_heat_vaporization(T) * dQ2<jupyter_output><empty_output><jupyter_text>nudged to fine wind tendencies are on the interfaces, need to center and rotate<jupyter_code>wind_rotation_matrix = catalog["wind_rotation/c48"].read()
grid=catalog["grid/c48"].read()
n2f_tendencies["dQu"], n2f_tendencies["dQv"] = vcm.cubedsphere.center_and_rotate_xy_winds(
    wind_rotation_matrix,
    n2f_tendencies["x_wind_tendency_due_to_nudging"],
    n2f_tendencies["y_wind_tendency_due_to_nudging"],
)<jupyter_output><empty_output><jupyter_text>Need the pressure level thicknesses<jupyter_code>n2f_state = intake.open_zarr("gs://vcm-ml-experiments/2021-04-13-n2f-c3072/3-hrly-ave-rad-precip-setting-30-min-rad-timestep-shifted-start-tke-edmf/state_after_timestep.zarr").to_dask().sel(time=n2f_tendencies.time)
n2f_tendencies["pressure_thickness_of_atmospheric_layer"] = n2f_state["pressure_thickness_of_atmospheric_layer"]
n2f_tendencies["air_temperature"] = n2f_state["air_temperature"]
n2f_tendencies["column_integrated_dQv"] = mass_integrate(
    n2f_tendencies["dQv"],
    n2f_tendencies["pressure_thickness_of_atmospheric_layer"],
    dim="z"
).mean("time").load()


n2f_tendencies["column_integrated_dQu"] = mass_integrate(
    n2f_tendencies["dQu"],
    n2f_tendencies["pressure_thickness_of_atmospheric_layer"],
    dim="z"
).mean("time").load()
 
n2f_tendencies["column_integrated_dQv"] = mass_integrate(
    n2f_tendencies["dQv"],
    n2f_tendencies["pressure_thickness_of_atmospheric_layer"],
    dim="z"
).mean("time").load()

n2f_tendencies["column_integrated_heating_from_isochoric_transition"] = column_integrated_heating_from_isochoric_transition(
    n2f_tendencies["air_temperature_tendency_due_to_nudging"],
    n2f_tendencies["pressure_thickness_of_atmospheric_layer"],
).mean("time").load()

n2f_tendencies["column_integrated_drying"] = minus_column_integrated_moistening(
    n2f_tendencies["specific_humidity_tendency_due_to_nudging"],
    n2f_tendencies["pressure_thickness_of_atmospheric_layer"],
).mean("time").load()



n2f_tendencies["column_integrated_mse_tendency"] = mass_integrate(
    moist_static_energy_tendency(n2f_tendencies["air_temperature"], n2f_tendencies["air_temperature_tendency_due_to_nudging"], n2f_tendencies["specific_humidity_tendency_due_to_nudging"]),
    n2f_tendencies["pressure_thickness_of_atmospheric_layer"],
    dim="z"
).mean("time").load()<jupyter_output><empty_output><jupyter_text>dQu pressure level zonal mean time average<jupyter_code>pressure_level_northward_wind = vcm.interpolate_to_pressure_levels(
         n2f_tendencies["dQv"], delp=n2f_tendencies["pressure_thickness_of_atmospheric_layer"], dim="z"
)

dQv_time_mean_zonal_pressure_level = zonal_mean(pressure_level_northward_wind, grid.lat).mean("time").load()

dQv_time_mean_zonal_pressure_level = dQv_time_mean_zonal_pressure_level.assign_coords(
    {'pressure': 
     (dQv_time_mean_zonal_pressure_level.pressure/100).assign_attrs({'units': 'hPa'})
    }
)<jupyter_output><empty_output><jupyter_text>Composite figure<jupyter_code>fig, axes = plt.subplots(2, 2, subplot_kw={"projection":ccrs.Robinson()}, figsize=(7.6, 4))

# a) heating
plot_cube(
    xr.merge([n2f_tendencies["column_integrated_heating_from_isochoric_transition"], grid]),
    "column_integrated_heating_from_isochoric_transition",
    cbar_label="W/m$^2$",
    ax=axes[0][0],
    rasterized=True,
)
axes[0][0].set_title(r"a) $\langle \Delta Q_1 \rangle$")

# b) MSE
plot_cube(
    xr.merge([n2f_tendencies["column_integrated_mse_tendency"], grid]),
    "column_integrated_mse_tendency",
    cbar_label="W/m$^2$",
    ax=axes[0][1],
    rasterized=True
)
axes[0][1].set_title(r"b) $\langle \Delta Q_m \rangle$")

# c) eastward wind
plot_cube(
    xr.merge([n2f_tendencies["column_integrated_dQu"], grid]),
    "column_integrated_dQu",
    cbar_label="Pa",
    ax=axes[1][0],
    rasterized=True
    
)
axes[1][0].set_title(r"c) $\langle \Delta Q_u \rangle$")

# d) northward wind
axes[1][1].set_axis_off()

left = 0.58
bottom = 0.125
width = 0.375
height = 0.28
ax = fig.add_axes([left, bottom, width, height])

f = (
        dQv_time_mean_zonal_pressure_level * (3 * 3600)
    ).plot(robust=True, yincrease=False, ax=ax, rasterized=True)
ax.set_title("d) $\Delta Q_v$")
ax.set_ylabel("pressure [hPa]")
f.colorbar.set_label("m/s/3-hr")


fig.subplots_adjust(hspace=0.7, wspace=0.7)
plt.tight_layout()

plt.savefig("figures/Figure_8_composite_nudging_tendencies_col_int.pdf", bbox_inches = "tight")<jupyter_output><empty_output>