"""Update ice/snow thickness, concentration, and ocean thermodynamic fluxes.
The original Veris Growth equations average surface fluxes over nITC thickness
categories and combine ice/snow melt, mixed-layer exchange, precipitation, and
salt/freshwater coupling. Inputs, fifteen physical outputs and formation status
are horizontal arrays; the internal category axis is last. Surface temperatures
come from solve4temp.
"""
from __future__ import annotations
from dataclasses import replace
from functools import partial
import jax.numpy as jnp
from jax import Array
from veris.model.configuration import Configuration
from veris.model.physical_constants import PhysicalConstants
from veris.model.types import GrowthResult, State, jit
from veris.numerics.surface_properties import (
_category_factors,
_prepare_thickness_geometry,
)
from veris.numerics.surface_temperature import solve4temp
from veris.physics.thermodynamics.budgets import ocean_exchange_fluxes
[docs]
@partial(jit, static_argnames=["conf"])
def Growth(
vs: State, conf: Configuration, phys: PhysicalConstants
) -> tuple[GrowthResult, Array]:
"""Return fifteen column/boundary arrays and a same-precision status array.
Physical results contain updated ice/snow thickness, concentration, surface
temperature, freshwater and salt forcing, shortwave and net ocean heat,
ice/snow load, penetrating shortwave, regularized reciprocal thickness,
ocean mass and salt mass inputs, and external mass and heat inputs.
Boundary inputs are downward positive; native Qnet is upward positive.
Status is zero for valid formation and two where new ice would contain
more salt than the source ocean. Callers must check status before accepting
the physical results. Bulk ice salinity is fixed for a trajectory.
"""
# Dry operands are neutral before arithmetic, including NaN prescribed data.
wet = vs.iceMask != 0
vs = replace(
vs,
**{
name: jnp.where(
wet,
getattr(vs, name),
jnp.where(jnp.isfinite(vs.TSurf), vs.TSurf, phys.celsius2K)
if name == "TSurf"
else 0,
)
for name in (
"hIceMean",
"hSnowMean",
"Area",
"TSurf",
"LWdown",
"SWdown",
"ATemp",
"aqh",
"wSpeed",
"fCori",
"ocSalt",
"theta",
"snowfall",
"precip",
"runoff",
"evap",
"Qnet",
"Qsw",
"os_hIceMean",
"os_hSnowMean",
)
},
)
initial_ice, initial_snow = vs.hIceMean, vs.hSnowMean
##### initializations #####
# F_ia_net: net heat flux divergence at the sea ice/snow surface [W/m2]
# = 0: surface heat loss is balanced by upward conductive fluxes
# < 0: net heat flux convergence at the ice/ snow surface
# F_io_net: net upward conductive heat flux through sea ice and snow [W/m2]
# JAX updates return new arrays, so category work fields can share zeros.
zeros3d = jnp.zeros((*vs.iceMask.shape, conf.nITC), dtype=vs.iceMask.dtype)
hIceActual_mult = zeros3d
hSnowActual_mult = zeros3d
F_io_net_mult = zeros3d
F_ia_net_mult = zeros3d
IcePenetSW_mult = zeros3d
FWsublim_mult = zeros3d
# constants for converting heat fluxes into growth rates
qi = 1 / (phys.rhoIce * phys.lhFusion)
qs = 1 / (phys.rhoSnow * phys.lhFusion)
geometry = _prepare_thickness_geometry(vs, phys)
hIceMeanpreTH = geometry.ice_mean
hSnowMeanpreTH = geometry.snow_mean
AreapreTH = geometry.area
recip_regArea = geometry.reciprocal_area
recip_hIceActual = geometry.reciprocal_ice_thickness
hIceActual = geometry.ice_thickness
hSnowActual = geometry.snow_thickness
##### calculate heat fluxes through the ice #####
# set ice and snow thickness categories to account for thicknes variations in one grid cell
TIce_mult = zeros3d
category_factors = _category_factors(conf)
for l in range(conf.nITC):
# the ice categories all have the same initial temperature
TIce_mult = TIce_mult.at[:, :, l].set(vs.TSurf)
# set relative thickness of ice and snow categories
pFac = category_factors[l]
# actual snow and ice thickness within each category
hIceActual_mult = hIceActual_mult.at[:, :, l].set(hIceActual * pFac)
hSnowActual_mult = hSnowActual_mult.at[:, :, l].set(hSnowActual * pFac)
# freezing temperature
TempFrz = phys.tempFrz + phys.dtempFrz_dS * vs.ocSalt + phys.celsius2K
# calculate heat fluxes
for l in range(conf.nITC):
output = solve4temp(
vs,
conf,
phys,
hIceActual_mult[:, :, l],
hSnowActual_mult[:, :, l],
TIce_mult[:, :, l],
TempFrz,
)
TIce_mult = TIce_mult.at[:, :, l].set(output[0])
F_io_net_mult = F_io_net_mult.at[:, :, l].set(output[1])
F_ia_net_mult = F_ia_net_mult.at[:, :, l].set(output[2])
IcePenetSW_mult = IcePenetSW_mult.at[:, :, l].set(output[3])
FWsublim_mult = FWsublim_mult.at[:, :, l].set(output[4])
# update surface temperature and fluxes
TSurf = jnp.sum(TIce_mult, axis=2) * conf.recip_nITC
# multplying the fluxes with the area changes them from mean fluxes
# for the ice part of the cell to mean fluxes for the whole cell
F_io_net = jnp.sum(F_io_net_mult, axis=2) * conf.recip_nITC * AreapreTH
F_ia_net = jnp.sum(F_ia_net_mult, axis=2) * conf.recip_nITC * AreapreTH
IcePenetSW = jnp.sum(IcePenetSW_mult, axis=2) * conf.recip_nITC * AreapreTH
sublimation = jnp.sum(FWsublim_mult, axis=2) * conf.recip_nITC * AreapreTH
# Diagnose the actual boundary event before phase change. Sublimation first
# consumes existing snow, then ice; deposition is fresh frost. The fluxes
# and geometric coefficients use the original column for this split step.
available = phys.rhoIce * hIceMeanpreTH + phys.rhoSnow * hSnowMeanpreTH
sublimation = jnp.minimum(sublimation, available * conf.recip_deltatTherm)
removed = jnp.maximum(sublimation, 0) * conf.deltatTherm
snow_removed = jnp.minimum(removed, phys.rhoSnow * hSnowMeanpreTH)
hSnowMeanpreTH = (
hSnowMeanpreTH
+ (-snow_removed - jnp.minimum(sublimation, 0) * conf.deltatTherm)
/ phys.rhoSnow
)
hIceMeanpreTH = jnp.maximum(
hIceMeanpreTH - (removed - snow_removed) / phys.rhoIce, 0
)
# Independent atmospheric/land boundary inputs, before phase changes.
# Snow enters with enthalpy -L*m relative to liquid at freezing. The +L*m
# vapor term converts surface sublimation heat to that same reference.
external_mass = (
phys.rhoFresh
* (vs.precip + vs.snowfall + vs.runoff + (1 - AreapreTH) * vs.evap)
- sublimation
)
external_heat = -(
F_ia_net + F_io_net + IcePenetSW + (1 - AreapreTH) * vs.Qnet
) + phys.lhFusion * (sublimation - phys.rhoFresh * vs.snowfall)
##### evaluate precipitation as snow or rain #####
# if there is ice and the temperature is below the freezing point,
# the precipitation falls and accumulates as snow
tmp = (AreapreTH > 0) & (jnp.mean(TIce_mult, axis=2) < phys.celsius2K)
# snow accumulation rate over ice [m/s]
# the snowfall is given in water equivalent, therefore it also needs to be muliplied with rhoFresh2rhoSnow
SnowAccRateOverIce = vs.snowfall
SnowAccRateOverIce = (
jnp.where(tmp, SnowAccRateOverIce + vs.precip, SnowAccRateOverIce)
* phys.rhoFresh2rhoSnow
)
##### calculate growth rates of ice and snow #####
# the ice growth rate beneath ice is given by the upward conductive
# flux F_io_net and qi:
IceGrowthRateUnderExistingIce = F_io_net * qi
IceGrowthRateUnderExistingIce = jnp.where(
AreapreTH == 0, 0, IceGrowthRateUnderExistingIce
)
# the potential snow melt rate if all snow surface heat flux
# convergence (F_ia_net < 0) goes to melting snow [m/s]
PotSnowMeltRateFromSurf = -F_ia_net * qs
# the thickness of snow that can be melted in one time step:
PotSnowMeltFromSurf = PotSnowMeltRateFromSurf * conf.deltatTherm
# if the heat flux convergence could melt more snow than is actually
# there, the excess is used to melt ice
# case 1: snow will remain after melting, i.e. all of the heat flux
# convergence will be used up to melt snow
# case 2: all snow will be melted if the potential snow melt
# height is larger or equal to the actual snow height. if there is
# an excess of heat flux convergence after snow melting, it will
# be used to melt ice
# (use hSnowActual for comparison with the MITgcm)
allSnowMelted = PotSnowMeltFromSurf >= hSnowMeanpreTH
# the actual thickness of snow to be melted by snow surface
# heat flux convergence [m]
SnowMeltFromSurface = jnp.where(allSnowMelted, hSnowMeanpreTH, PotSnowMeltFromSurf)
# the actual snow melt rate due to snow surface heat flux convergence [m/s]
SnowMeltRateFromSurface = jnp.where(
allSnowMelted,
SnowMeltFromSurface * conf.recip_deltatTherm,
PotSnowMeltRateFromSurf,
)
# the actual surface heat flux convergence used to melt snow [W/m2]
SurfHeatFluxConvergToSnowMelt = jnp.where(
allSnowMelted, -hSnowMeanpreTH * conf.recip_deltatTherm / qs, F_ia_net
)
# the surface heat flux convergence is reduced by the amount that
# is used for melting snow:
F_ia_net = F_ia_net - SurfHeatFluxConvergToSnowMelt
# the remaining heat flux convergence is used to melt ice:
IceGrowthRateFromSurface = F_ia_net * qi
# the total ice growth rate is then:
# (remove * recip_regArea for comparison with the MITgcm)
NetExistingIceGrowthRate = (
IceGrowthRateUnderExistingIce + IceGrowthRateFromSurface
) * recip_regArea
##### calculate the ice melting due to mixed layer temperature #####
tmpscal0 = phys.McPheeTaperArea
tmpscal1 = phys.McPheeTaperSteepness / tmpscal0
tmpscal2 = phys.stantonNr * phys.uStarBase * phys.rhoSea * phys.cpWater
# the ocean temperature cannot be lower than the freezing temperature
surf_theta = jnp.maximum(vs.theta, TempFrz)
# mltf = mixed layer turbulence factor (determines how much of the temperature
# difference is used for heat flux)
mltf = 1 + (phys.McPheeTaperFac - 1) / (
1 + jnp.exp((AreapreTH - tmpscal0) * tmpscal1)
)
# heat flux from ocean to the ice (+ = upward)
F_oi = -tmpscal2 * (surf_theta - TempFrz) * mltf
IceGrowthRateMixedLayer = F_oi * qi
##### calculate the ice growth in open water due to heat loss to atmosphere #####
# in this configuration, all of the incoming shortwave radiation is absorbed in the
# uppermost ocean cell. If this is not the case, the shortwave radiation that passes
# the ocean surface cell is to be distracted from Qnet as it is not part of the heat
# flux convergence
# define swFracPass as the fraction of the shortwave radiation that passes the ocean
# surface cell
# Qsw_in_first_layer = vs.Qsw * ( 1 - swFracPass )
# IceGrowthRateOpenWater = qi * ( vs.Qnet - vs.Qsw + Qsw_in_first_layer )
# using swFracPass = 0 yields the same result as just using
# IceGrowthRateOpenWater = qi * vs.Qnet
# Open-water snow melts immediately, consuming L*P before phase change.
# At freezing that loss reforms ice; warm mixed-layer heat may melt it.
IceGrowthRateOpenWater = qi * (
vs.Qnet + phys.lhFusion * phys.rhoFresh * vs.snowfall
)
##### calculate change in ice, snow thicknesses and area #####
# calculate thickness derivatives of ice and snow
dhIceMean_dt = (
NetExistingIceGrowthRate * AreapreTH
+ IceGrowthRateOpenWater * (1 - AreapreTH)
+ IceGrowthRateMixedLayer
)
# Accumulation is per ice area; the melt flux is already per grid-cell area.
dhSnowMean_dt = SnowAccRateOverIce * AreapreTH - SnowMeltRateFromSurface
tmpscal0 = phys.lateralMeltAreaFactor * recip_hIceActual
# ice growth open water (due to ocean-atmosphere fluxes)
# reduce ice cover if the open water growth rate is negative
dArea_oaFlux = tmpscal0 * IceGrowthRateOpenWater * (1 - AreapreTH)
# increase ice cover if the open water growth rate is positive
tmp = (IceGrowthRateOpenWater > 0) & ((AreapreTH > 0) | (dhIceMean_dt > 0))
dArea_oaFlux = jnp.where(
tmp, IceGrowthRateOpenWater * (1 - AreapreTH), dArea_oaFlux
)
# multiply with lead closing factor
dArea_oaFlux = jnp.where(
(tmp & (vs.fCori < 0)),
dArea_oaFlux * phys.recip_h0_south,
dArea_oaFlux * phys.recip_h0,
)
# ice growth mixed layer (due to ocean-ice fluxes)
# (if the ocean is warmer than the ice: IceGrowthRateMixedLayer > 0.
# the supercooled state of the ocean is ignored/ does not lead to ice
# growth. ice growth is only due to fluxes calculated by solve4temp)
dArea_oiFlux = jnp.where(
IceGrowthRateMixedLayer <= 0, tmpscal0 * IceGrowthRateMixedLayer, 0
)
# ice growth over ice (due to ice-atmosphere fluxes)
# (NetExistingIceGrowthRate leads to vertical and lateral melting but
# only to vertical growing. lateral growing is covered by
# IceGrowthRateOpenWater)
dArea_iaFlux = jnp.where(
(NetExistingIceGrowthRate <= 0) & (hIceMeanpreTH > 0),
tmpscal0 * NetExistingIceGrowthRate * AreapreTH,
0,
)
# total change in area
dArea_dt = dArea_oaFlux + dArea_oiFlux + dArea_iaFlux
###### update ice, snow thickness and area #####
Area = AreapreTH + dArea_dt * vs.iceMask * conf.deltatTherm
hIceMean = hIceMeanpreTH + dhIceMean_dt * vs.iceMask * conf.deltatTherm
hSnowMean = hSnowMeanpreTH + dhSnowMean_dt * vs.iceMask * conf.deltatTherm
# set boundaries:
Area = jnp.clip(Area, 0, 1)
hIceMean = jnp.maximum(hIceMean, 0)
hSnowMean = jnp.maximum(hSnowMean, 0)
noIce = (hIceMean == 0) | (Area == 0)
Area *= ~noIce
hSnowMean *= ~noIce
hIceMean *= ~noIce
# change of ice thickness due to conversion of snow to ice if snow
# is submerged with water
h_sub = (hSnowMean * phys.rhoSnow + hIceMean * phys.rhoIce) * phys.recip_rhoSea
d_hIceMeanByFlood = jnp.maximum(0, h_sub - hIceMean)
hIceMean = hIceMean + d_hIceMeanByFlood
hSnowMean = hSnowMean - d_hIceMeanByFlood * phys.rhoIce2rhoSnow
recip_hIceMean = 1 / jnp.sqrt(hIceMean**2 + phys.hIce_reg)
##### calculate output to ocean #####
# effective shortwave heating rate
# IcePenetSW was converted to a grid-cell mean after category averaging.
Qsw = IcePenetSW + vs.Qsw * (1 - AreapreTH)
ice_change = hIceMean - initial_ice
snow_change = hSnowMean - initial_snow
fluxes = ocean_exchange_fluxes(
ice_change=ice_change,
snow_change=snow_change,
external_mass_input=external_mass,
cleanup_ice=-vs.os_hIceMean,
cleanup_snow=-vs.os_hSnowMean,
configuration=conf,
constants=phys,
)
EmPmR = -fluxes.mass_input
forc_salt_surface = (
EmPmR * vs.ocSalt + 1000 * fluxes.salt_mass_input
) / phys.rhoFresh
# The residual ocean heat follows the zero-layer first law. External heat
# above is diagnosed from boundary events, independently of final stock.
solid_mass_change = phys.rhoIce * (ice_change + vs.os_hIceMean) + phys.rhoSnow * (
snow_change + vs.os_hSnowMean
)
Qnet = -external_heat - phys.lhFusion * solid_mass_change * conf.recip_deltatTherm
gross_formation = (
jnp.maximum(IceGrowthRateUnderExistingIce, 0)
+ jnp.maximum(IceGrowthRateFromSurface, 0)
+ jnp.maximum(IceGrowthRateOpenWater, 0) * (1 - AreapreTH)
+ d_hIceMeanByFlood * conf.recip_deltatTherm
)
status = jnp.where(
wet & (gross_formation > 0) & (vs.ocSalt < phys.saltIce_ref), 2, 0
).astype(vs.Area.dtype)
# sea ice + snow load on the sea surface
SeaIceLoad = hIceMean * phys.rhoIce + hSnowMean * phys.rhoSnow
return (
(
hIceMean,
hSnowMean,
Area,
TSurf,
EmPmR,
forc_salt_surface,
Qsw,
Qnet,
SeaIceLoad,
IcePenetSW,
recip_hIceMean,
fluxes.mass_input,
fluxes.salt_mass_input,
external_mass,
external_heat,
),
status,
)