import os
from veros import veros_kernel, KernelOutput, runtime_settings
from veros.core.operators import numpy as npx, update_multiply, at
from veros.diagnostics.base import VerosDiagnostic
from veros.variables import Variable
from veros.distributed import global_sum
ENERGY_VARIABLES = dict(
nitts=Variable("nitts", None, write_to_restart=True),
# mean energy content
k_m=Variable("Mean kinetic energy", None, "J", "Mean kinetic energy", write_to_restart=True),
Hd_m=Variable("Mean dynamic enthalpy", None, "J", "Mean dynamic enthalpy", write_to_restart=True),
eke_m=Variable("Meso-scale eddy energy", None, "J", "Meso-scale eddy energy", write_to_restart=True),
iw_m=Variable("Internal wave energy", None, "J", "Internal wave energy", write_to_restart=True),
tke_m=Variable("Turbulent kinetic energy", None, "J", "Turbulent kinetic energy", write_to_restart=True),
# energy changes
dE_tot_m=Variable("Change of total energy", None, "W", "Change of total energy", write_to_restart=True),
dk_m=Variable("Change of KE", None, "W", "Change of kinetic energy", write_to_restart=True),
dHd_m=Variable("Change of Hd", None, "W", "Change of dynamic enthalpy", write_to_restart=True),
deke_m=Variable("Change of EKE", None, "W", "Change of meso-scale eddy energy", write_to_restart=True),
diw_m=Variable("Change of E_iw", None, "W", "Change of internal wave energy", write_to_restart=True),
dtke_m=Variable("Change of TKE", None, "W", "Change of tubulent kinetic energy", write_to_restart=True),
# dissipation
ke_diss_m=Variable("Dissipation of KE", None, "W", "Dissipation of kinetic energy", write_to_restart=True),
Hd_diss_m=Variable("Dissipation of Hd", None, "W", "Dissipation of dynamic enthalpy", write_to_restart=True),
eke_diss_m=Variable(
"Dissipation of EKE", None, "W", "Dissipation of meso-scale eddy energy", write_to_restart=True
),
iw_diss_m=Variable("Dissipation of E_iw", None, "W", "Dissipation of internal wave energy", write_to_restart=True),
tke_diss_m=Variable(
"Dissipation of TKE", None, "W", "Dissipation of turbulent kinetic energy", write_to_restart=True
),
adv_diss_m=Variable("Dissipation by advection", None, "W", "Dissipation by advection", write_to_restart=True),
# external forcing
wind_m=Variable("Wind work", None, "W", "Wind work", write_to_restart=True),
dHd_sources_m=Variable(
"Hd production by ext. sources",
None,
"W",
"Dynamic enthalpy production through external sources",
write_to_restart=True,
),
iw_forc_m=Variable(
"External forcing of E_iw", None, "W", "External forcing of internal wave energy", write_to_restart=True
),
tke_forc_m=Variable(
"External forcing of TKE", None, "W", "External forcing of turbulent kinetic energy", write_to_restart=True
),
# exchange
ke_hd_m=Variable(
"Exchange KE -> Hd", None, "W", "Exchange between kinetic energy and dynamic enthalpy", write_to_restart=True
),
ke_tke_m=Variable(
"Exchange KE -> TKE by vert. friction",
None,
"W",
"Exchange between kinetic energy and turbulent kinetic energy by vertical friction",
write_to_restart=True,
),
ke_iw_m=Variable(
"Exchange KE -> IW by bottom friction",
None,
"W",
"Exchange between kinetic energy and internal wave energy by bottom friction",
write_to_restart=True,
),
tke_hd_m=Variable(
"Exchange TKE -> Hd by vertical mixing",
None,
"W",
"Exchange between turbulent kinetic energy and dynamic enthalpy by vertical mixing",
write_to_restart=True,
),
ke_eke_m=Variable(
"Exchange KE -> EKE by lateral friction",
None,
"W",
"Exchange between kinetic energy and eddy kinetic energy by lateral friction",
write_to_restart=True,
),
hd_eke_m=Variable(
"Exchange Hd -> EKE by GM and lateral mixing",
None,
"W",
"Exchange between dynamic enthalpy and eddy kinetic energy by GM and lateral mixing",
write_to_restart=True,
),
eke_tke_m=Variable(
"Exchange EKE -> TKE", None, "W", "Exchange between eddy and turbulent kinetic energy", write_to_restart=True
),
eke_iw_m=Variable(
"Exchange EKE -> IW",
None,
"W",
"Exchange between eddy kinetic energy and internal wave energy",
write_to_restart=True,
),
# cabbeling
cabb_m=Variable("Cabbeling by vertical mixing", None, "W", "Cabbeling by vertical mixing", write_to_restart=True),
cabb_iso_m=Variable(
"Cabbeling by isopycnal mixing", None, "W", "Cabbeling by isopycnal mixing", write_to_restart=True
),
)
DEFAULT_OUTPUT_VARS = [var for var in ENERGY_VARIABLES.keys() if var not in ("nitts",)]
[docs]
class Energy(VerosDiagnostic):
"""Diagnose globally averaged energy cycle. Also averages energy in time."""
name = "energy" #:
output_path = "{identifier}.energy.nc" #: File to write to. May contain format strings that are replaced with Veros attributes.
output_frequency = None #: Frequency (in seconds) in which output is written.
sampling_frequency = None #: Frequency (in seconds) in which variables are accumulated.
var_meta = ENERGY_VARIABLES
def __init__(self, state):
self.output_variables = DEFAULT_OUTPUT_VARS.copy()
def initialize(self, state):
self.initialize_variables(state)
self.initialize_output(state)
def diagnose(self, state):
energies = diagnose_kernel(state)
# store results
for energy, val in energies._asdict().items():
total_val = self.variables.get(energy)
setattr(self.variables, energy, total_val + val)
self.variables.nitts = self.variables.nitts + 1
def output(self, state):
if not os.path.isfile(self.get_output_file_name(state)):
self.initialize_output(state)
energy_vs = self.variables
nitts = float(energy_vs.nitts or 1)
for key in self.output_variables:
val = getattr(energy_vs, key)
setattr(energy_vs, key, val * state.settings.rho_0 / nitts)
self.write_output(state)
for key in self.output_variables:
setattr(energy_vs, key, 0.0)
energy_vs.nitts = 0
@veros_kernel
def diagnose_kernel(state):
vs = state.variables
settings = state.settings
# changes of dynamic enthalpy
vol_t = vs.area_t[2:-2, 2:-2, npx.newaxis] * vs.dzt[npx.newaxis, npx.newaxis, :] * vs.maskT[2:-2, 2:-2, :]
dP_iso = global_sum(
npx.sum(
vol_t
* settings.grav
/ settings.rho_0
* (
-vs.int_drhodT[2:-2, 2:-2, :, vs.tau] * vs.dtemp_iso[2:-2, 2:-2, :]
- vs.int_drhodS[2:-2, 2:-2, :, vs.tau] * vs.dsalt_iso[2:-2, 2:-2, :]
)
)
)
dP_hmix = global_sum(
npx.sum(
vol_t
* settings.grav
/ settings.rho_0
* (
-vs.int_drhodT[2:-2, 2:-2, :, vs.tau] * vs.dtemp_hmix[2:-2, 2:-2, :]
- vs.int_drhodS[2:-2, 2:-2, :, vs.tau] * vs.dsalt_hmix[2:-2, 2:-2, :]
)
)
)
dP_vmix = global_sum(
npx.sum(
vol_t
* settings.grav
/ settings.rho_0
* (
-vs.int_drhodT[2:-2, 2:-2, :, vs.tau] * vs.dtemp_vmix[2:-2, 2:-2, :]
- vs.int_drhodS[2:-2, 2:-2, :, vs.tau] * vs.dsalt_vmix[2:-2, 2:-2, :]
)
)
)
dP_m = global_sum(
npx.sum(
vol_t
* settings.grav
/ settings.rho_0
* (
-vs.int_drhodT[2:-2, 2:-2, :, vs.tau] * vs.dtemp[2:-2, 2:-2, :, vs.tau]
- vs.int_drhodS[2:-2, 2:-2, :, vs.tau] * vs.dsalt[2:-2, 2:-2, :, vs.tau]
)
)
)
dP_m_all = dP_m + dP_vmix + dP_hmix + dP_iso
# changes of kinetic energy
vol_u = vs.area_u[2:-2, 2:-2, npx.newaxis] * vs.dzt[npx.newaxis, npx.newaxis, :]
vol_v = vs.area_v[2:-2, 2:-2, npx.newaxis] * vs.dzt[npx.newaxis, npx.newaxis, :]
k_m = global_sum(
npx.sum(
vol_t
* 0.5
* (
0.5 * (vs.u[2:-2, 2:-2, :, vs.tau] ** 2 + vs.u[1:-3, 2:-2, :, vs.tau] ** 2)
+ 0.5 * (vs.v[2:-2, 2:-2, :, vs.tau] ** 2)
+ vs.v[2:-2, 1:-3, :, vs.tau] ** 2
)
)
)
p_m = global_sum(npx.sum(vol_t * vs.Hd[2:-2, 2:-2, :, vs.tau]))
dk_m = global_sum(
npx.sum(
vs.u[2:-2, 2:-2, :, vs.tau] * vs.du[2:-2, 2:-2, :, vs.tau] * vol_u
+ vs.v[2:-2, 2:-2, :, vs.tau] * vs.dv[2:-2, 2:-2, :, vs.tau] * vol_v
+ vs.u[2:-2, 2:-2, :, vs.tau] * vs.du_mix[2:-2, 2:-2, :] * vol_u
+ vs.v[2:-2, 2:-2, :, vs.tau] * vs.dv_mix[2:-2, 2:-2, :] * vol_v
)
)
# K*Nsqr and KE and dyn. enthalpy dissipation
vol_w = vs.area_t[2:-2, 2:-2, npx.newaxis] * vs.dzw[npx.newaxis, npx.newaxis, :] * vs.maskW[2:-2, 2:-2, :]
vol_w = update_multiply(vol_w, at[:, :, -1], 0.5)
def mean_w(var):
return global_sum(npx.sum(var[2:-2, 2:-2, :] * vol_w))
mdiss_vmix = mean_w(vs.P_diss_v)
mdiss_nonlin = mean_w(vs.P_diss_nonlin)
mdiss_adv = mean_w(vs.P_diss_adv)
mdiss_hmix = mean_w(vs.P_diss_hmix)
mdiss_iso = mean_w(vs.P_diss_iso)
mdiss_skew = mean_w(vs.P_diss_skew)
mdiss_sources = mean_w(vs.P_diss_sources)
mdiss_h = mean_w(vs.K_diss_h)
mdiss_v = mean_w(vs.K_diss_v)
mdiss_gm = mean_w(vs.K_diss_gm)
mdiss_bot = mean_w(vs.K_diss_bot)
wrhom = global_sum(
npx.sum(
-vs.area_t[2:-2, 2:-2, npx.newaxis]
* vs.maskW[2:-2, 2:-2, :-1]
* (vs.p_hydro[2:-2, 2:-2, 1:] - vs.p_hydro[2:-2, 2:-2, :-1])
* vs.w[2:-2, 2:-2, :-1, vs.tau]
)
)
# wind work
if runtime_settings.pyom_compatibility_mode:
# surface_tau* has different units in PyOM
wind = global_sum(
npx.sum(
vs.u[2:-2, 2:-2, -1, vs.tau]
* vs.surface_taux[2:-2, 2:-2]
* vs.maskU[2:-2, 2:-2, -1]
* vs.area_u[2:-2, 2:-2]
+ vs.v[2:-2, 2:-2, -1, vs.tau]
* vs.surface_tauy[2:-2, 2:-2]
* vs.maskV[2:-2, 2:-2, -1]
* vs.area_v[2:-2, 2:-2]
)
)
else:
wind = global_sum(
npx.sum(
vs.u[2:-2, 2:-2, -1, vs.tau]
* vs.surface_taux[2:-2, 2:-2]
/ settings.rho_0
* vs.maskU[2:-2, 2:-2, -1]
* vs.area_u[2:-2, 2:-2]
+ vs.v[2:-2, 2:-2, -1, vs.tau]
* vs.surface_tauy[2:-2, 2:-2]
/ settings.rho_0
* vs.maskV[2:-2, 2:-2, -1]
* vs.area_v[2:-2, 2:-2]
)
)
# meso-scale energy
if settings.enable_eke:
eke_m = mean_w(vs.eke[..., vs.tau])
deke_m = global_sum(
npx.sum(vol_w * (vs.eke[2:-2, 2:-2, :, vs.taup1] - vs.eke[2:-2, 2:-2, :, vs.tau]) / settings.dt_tracer)
)
eke_diss = mean_w(vs.eke_diss_iw)
eke_diss_tke = mean_w(vs.eke_diss_tke)
else:
eke_m = deke_m = eke_diss_tke = 0.0
eke_diss = mdiss_gm + mdiss_h + mdiss_skew
if not settings.enable_store_cabbeling_heat:
eke_diss += -mdiss_hmix - mdiss_iso
# small-scale energy
if settings.enable_tke:
dt_tke = settings.dt_mom
tke_m = mean_w(vs.tke[..., vs.tau])
dtke_m = mean_w((vs.tke[..., vs.taup1] - vs.tke[..., vs.tau]) / dt_tke)
tke_diss = mean_w(vs.tke_diss)
tke_forc = global_sum(
npx.sum(
vs.area_t[2:-2, 2:-2]
* vs.maskW[2:-2, 2:-2, -1]
* (vs.forc_tke_surface[2:-2, 2:-2] + vs.tke_surf_corr[2:-2, 2:-2])
)
)
else:
tke_m = dtke_m = tke_diss = tke_forc = 0.0
# internal wave energy
if settings.enable_idemix:
iw_m = mean_w(vs.E_iw[..., vs.tau])
diw_m = global_sum(
npx.sum(vol_w * (vs.E_iw[2:-2, 2:-2, :, vs.taup1] - vs.E_iw[2:-2, 2:-2, :, vs.tau]) / vs.dt_tracer)
)
iw_diss = mean_w(vs.iw_diss)
k = npx.maximum(1, vs.kbot[2:-2, 2:-2]) - 1
mask = k[:, :, npx.newaxis] == npx.arange(settings.nz)[npx.newaxis, npx.newaxis, :]
iwforc = global_sum(
npx.sum(
vs.area_t[2:-2, 2:-2]
* (
vs.forc_iw_surface[2:-2, 2:-2] * vs.maskW[2:-2, 2:-2, -1]
+ npx.sum(mask * vs.forc_iw_bottom[2:-2, 2:-2, npx.newaxis] * vs.maskW[2:-2, 2:-2, :], axis=2)
)
)
)
else:
iw_m = diw_m = iwforc = 0.0
iw_diss = eke_diss
if settings.enable_store_bottom_friction_tke:
ke_tke_m = mdiss_v + mdiss_bot
ke_iw_m = 0.0
else:
ke_tke_m = mdiss_v
ke_iw_m = mdiss_bot
hd_eke_m = -mdiss_skew
tke_hd_m = -mdiss_vmix - mdiss_adv
if not settings.enable_store_cabbeling_heat:
hd_eke_m = hd_eke_m - mdiss_hmix - mdiss_iso
tke_hd_m = tke_hd_m - mdiss_nonlin
return KernelOutput(
k_m=k_m,
Hd_m=p_m,
eke_m=eke_m,
iw_m=iw_m,
tke_m=tke_m,
dk_m=dk_m,
dHd_m=dP_m_all + mdiss_sources,
deke_m=deke_m,
diw_m=diw_m,
dtke_m=dtke_m,
dE_tot_m=dk_m + dP_m_all + mdiss_sources + deke_m + diw_m + dtke_m,
wind_m=wind,
dHd_sources_m=mdiss_sources,
iw_forc_m=iwforc,
tke_forc_m=tke_forc,
ke_diss_m=mdiss_h + mdiss_v + mdiss_gm + mdiss_bot,
Hd_diss_m=mdiss_vmix + mdiss_nonlin + mdiss_hmix + mdiss_adv + mdiss_iso + mdiss_skew,
eke_diss_m=eke_diss + eke_diss_tke,
iw_diss_m=iw_diss,
tke_diss_m=tke_diss,
adv_diss_m=mdiss_adv,
ke_hd_m=wrhom,
ke_eke_m=mdiss_h + mdiss_gm,
hd_eke_m=-mdiss_skew,
ke_tke_m=ke_tke_m,
ke_iw_m=ke_iw_m,
tke_hd_m=tke_hd_m,
eke_tke_m=eke_diss_tke,
eke_iw_m=eke_diss,
cabb_m=mdiss_nonlin,
cabb_iso_m=mdiss_hmix + mdiss_iso,
)