Source code for veros.diagnostics.energy

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, )