Physics kernels¶

These functions implement the equations and nearby-cell calculations used by advance. Most accept (vs, conf, phys): a State with halos, a fixed static Configuration and a dynamic PhysicalConstants PyTree. State holds the model arrays. Halos are extra cells that hold boundary or neighboring values. The functions return arrays or tuples. The caller assigns results with dataclasses.replace, exchanges halo values when required and calls the stages in the required order.

Single physics timestep¶

veris.advance moves the model forward by one complete timestep and returns State and Diagnostics, which holds step results. Select coupled for ice motion, transport, freezing and melting; dynamics for motion and transport; or thermodynamics for freezing and melting. Keep configuration and the physics choice static when applying JAX transformations; physical constants are dynamic. Supplied forcing, or external inputs, can change and is refreshed before the numerical functions use it. veris.step repeats a supplied step function over multiple timesteps. See Scan rollouts and checkpointing for examples of combining these functions and computing derivatives, which describe how results change with inputs.

Advance ONE physics timestep and return State and coupling Diagnostics.

physics selects coupled, dynamics or thermodynamics. Coupled timesteps must agree. Prescribed forcing accepts scalar or global physical (nx, ny) arrays, in VARIABLES units. Refresh supplied halos before physics and derive wind magnitude when components are supplied without an explicit wSpeed. Missing fields retain their State value; in particular, absent Qnet/Qsw keep Growth’s recursive coupling fluxes. Explicit heat is reapplied on every call.

This function is pure and supports JIT, grad, JVP and VJP with respect to State, independent physical parameters and forcing arrays. Bind configuration and physics as static arguments; pass constants as a dynamic PyTree. Dynamic value validation belongs to host Experiment and initialization; this path checks names, shapes and real dtypes while tracing. For sharding retain the caller’s active x/y mesh. veris.step remains the separate multi-step scan rollout and can consume this transition.

The dynamics_transport() stage is available to callers building specialized step functions.

Dynamics¶

Ice concentration and ice/snow mass at cell centers and velocity points.

veris.physics.dynamics.area_mass.AreaWS(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[Array, Array][source]¶: calculate sea ice cover fraction centered around velocity points

veris.physics.dynamics.area_mass.SeaIceMass(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[Array, Array, Array][source]¶: calculate mass of the ice-snow system centered around c-, u-, and v-points

Wind forcing and dispatch of the free-drift or EVP momentum solver.

veris.numerics.momentum.tauXY(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[Array, Array][source]¶: calculate surface stress from wind and ice velocities

veris.numerics.momentum.WindForcingXY(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[Array, Array][source]¶: calculate surface forcing due to wind and ocean surface tilt

veris.numerics.momentum.IceVelocities(vs: State, conf: Configuration, phys: PhysicalConstants, *, axis_names: tuple[str, ...] = ()) → tuple[Array, Array, Array, Array, Array][source]¶: Calculate ice velocities, reducing EVP diagnostics over supplied mesh axes.

Elastic-viscous-plastic subcycling of ice velocity and stress tensors.

veris.numerics.evp.evp_solver(vs: State, conf: Configuration, phys: PhysicalConstants, *, axis_names: tuple[str, ...] = ()) → tuple[Array, Array, Array, Array, Array][source]¶

solve the momentum equation and calculate u^n, sigma^n from u^(n-1), sigma^(n-1) using subcycling iterations of evp_solver_body.

Inside shard_map, pass its mesh axis names to combine residual diagnostics over all devices/processes. The default reports the serial interior norm.

Solve the local ice momentum balance without internal stress.

veris.numerics.freedrift.freedrift_solver(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[Array, Array][source]¶: calculate ice velocities without taking into account internal ice stress

Ice strength, drag, strain, viscosity and stress on the staggered grid.

veris.physics.dynamics.rheology.SeaIceStrength(vs: State, conf: Configuration, phys: PhysicalConstants) → Array[source]¶: calculate ice strength (= maximum compressive stress) from ice thickness and ice cover fraction

veris.physics.dynamics.rheology.ocean_drag_coeffs(vs: State, conf: Configuration, phys: PhysicalConstants, uIce: ArrayInput, vIce: ArrayInput) → Array[source]¶: calculate linear ice-water drag coefficient from ice and ocean velocities (this coefficient creates a linear relationship between ice-ocean stress difference and ice-ocean velocity difference)

veris.physics.dynamics.rheology.basal_drag_coeffs(vs: State, conf: Configuration, phys: PhysicalConstants, uIce: ArrayInput, vIce: ArrayInput) → Array[source]¶: calculate basal drag coefficient to account for the formation of landfast ice in shallow waters due to the formation of ice keels (Lemieux et al., 2015)

veris.physics.dynamics.rheology.side_drag(vs: State, conf: Configuration, phys: PhysicalConstants, uIce: ArrayInput, vIce: ArrayInput) → tuple[Array, Array][source]¶: calculate the lateral drag coefficient to simulate landfast ice (Liu et al. 2022, A new parameterization of coastal drag to simulate landfast ice in deep marginal seas in the Arctic)

veris.physics.dynamics.rheology.strainrates(vs: State, conf: Configuration, phys: PhysicalConstants, uIce: ArrayInput, vIce: ArrayInput) → tuple[Array, Array, Array][source]¶

Return center normal and corner shear strain rates in inverse seconds.

Velocities use west/south faces of halo-inclusive (x, y) arrays. Spherical normal strain has an exactly conservative trace in the supplied cell areas; its corner shear annihilates solid-body rotation. Planar stencils are kept.

veris.physics.dynamics.rheology.viscosities(vs: State, conf: Configuration, phys: PhysicalConstants, e11: ArrayInput, e22: ArrayInput, e12: ArrayInput) → tuple[Array, Array, Array][source]¶: calculate bulk viscosity zeta, shear viscosity eta, and ice pressure from strain rate tensor components and ice strength. if pressReplFac = 1, a replacement pressure is used to avoid stresses without velocities (Hibler and Ib, 1995). with tensileStrFac != 1, a resistance to tensile stresses can be included (König Beatty and Holland, 2010).

veris.physics.dynamics.rheology.stress(vs: State, conf: Configuration, phys: PhysicalConstants, e11: ArrayInput, e22: ArrayInput, e12: ArrayInput, zeta: ArrayInput, eta: ArrayInput, press: ArrayInput) → tuple[Array, Array, Array][source]¶: calculate stress tensor components

veris.physics.dynamics.rheology.stressdiv(vs: State, conf: Configuration, phys: PhysicalConstants, sig11: ArrayInput, sig22: ArrayInput, sig12: ArrayInput) → tuple[Array, Array][source]¶

Return west/south face forces (N/m2) from integrated stresses (N/m).

Normal stresses are centered and shear stress is at corners. The spherical operator is the negative area-weighted transpose of strainrates, with a factor two for tensor shear work. The original planar arithmetic is kept.

Compute ice-ocean stress with drag and a water turning angle.

veris.physics.dynamics.ocean_stress.OceanStressUV(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[Array, Array][source]¶: calculate stresses on ocean surface from ocean and ice velocities

Transport¶

Flux-limited directional transport of horizontal sea-ice fields.

veris.physics.transport.advection.Advection(vs: State, conf: Configuration) → tuple[Array, Array, Array][source]¶: retrieve changes in sea ice fields

veris.physics.transport.advection.calc_Advection(vs: State, conf: Configuration, field: ArrayInput) → Array[source]¶: calculate change in sea ice field due to advection

veris.physics.transport.advection.calc_flux(vs: State, conf: Configuration, field: ArrayInput, transport: ArrayInput, *, axis: int) → Array[source]¶

Apply the shared second-order limiter stencil along one C-grid axis.

Axis zero is zonal and axis one meridional. Keep storage orientation unchanged for rectangular and sharded grids. The meridional normal-face policy closes global walls; zonal flux extends edges.

Transport takes State and Configuration; it uses no physical-constant inputs. calc_flux selects the zonal stencil with axis=0 and the meridional stencil with axis=1.

veris.physics.transport.advection.limiter(Cr: ArrayInput | float) → Array[source]¶: Apply the Superbee slope limiter to a scalar or horizontal slope ratio.

Interpolate cell-centered JAX fields to staggered grid corners.

veris.grid.averaging.c_point_to_z_point(vs: State, conf: Configuration, phys: PhysicalConstants, Cfield: Array | ndarray | bool | number | bool | int | float | complex) → Array[source]¶: calculates value at z-point by averaging c-point values

Redistribute thin ice and bound transported concentration and thickness.

veris.physics.transport.cleanup.clean_up_advection(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[Array, Array, Array, Array, Array, Array][source]¶: clean up overshoots and other pathological cases after advection

veris.physics.transport.cleanup.ridging(vs: State, conf: Configuration, phys: PhysicalConstants) → Array[source]¶: cut off ice cover fraction at 1 after advection to account for ridging

Thermodynamics¶

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.

veris.physics.thermodynamics.growth.Growth(vs: State, conf: Configuration, phys: PhysicalConstants) → tuple[GrowthResult, Array][source]¶

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.

values, status = Growth(state, configuration, constants) returns the fifteen physical arrays and the formation status separately. Check that status is zero before accepting the updated column. advance carries this status in Diagnostics and Simulation checks it before committing a timestep.

See Ocean boundary budgets for mass, salt, enthalpy and status conventions.

Actual ocean mass and salt exchange for the native zero-layer ice column.

For solid volume change dV per horizontal area and external atmospheric/land mass input B, the ocean input is B - (rho_i*dV_i + rho_s*dV_s)/dt. Uniform bulk ice salinity S_i in g/kg retains rho_i*dV_i*S_i/1000 kg salt/m2 in ice; snow and external water are fresh. Both outputs are downward-positive grid arrays in kg/m2/s. These are physical transfers, before a fixed-volume ocean’s separate virtual-salinity conversion. Transport redistribution is excluded.

class veris.physics.thermodynamics.budgets.OceanExchangeFluxes(mass_input: Array, salt_mass_input: Array)[source]¶

Bases: object

Grid-shaped downward mass and actual salt input in kg/m2/s.

veris.physics.thermodynamics.budgets.ocean_exchange_fluxes(*, ice_change: Array, snow_change: Array, external_mass_input: Array, cleanup_ice: Array, cleanup_snow: Array, configuration: Configuration, constants: PhysicalConstants) → OceanExchangeFluxes[source]¶

Account for thermodynamics and signed raw-minus-prepared cleanup.

Changes are final minus prepared pre-thermodynamic thickness in metres. Cleanup is raw post-advection minus prepared thickness, also in metres; subtracting it includes both clipped negative stock and removed thin stock. The caller supplies finite, zeroed inactive cells and boundary-event mass rates independently of final inventories. Numerical timestep is static; physical densities and salinity remain differentiable leaves.

Solve the original Veris conductive and atmospheric ice/snow heat balance.

Horizontal thickness, temperature, and forcing arrays produce surface temperature, two net heat fluxes, penetrating shortwave radiation, and sublimation. Six Newton iterations retain the reference melting-temperature cap and albedo switches; gradients follow their selected smooth branches.

veris.numerics.surface_temperature.solve4temp(vs: State, conf: Configuration, phys: PhysicalConstants, hIceActual: ArrayInput, hSnowActual: ArrayInput, TSurfIn: ArrayInput, TempFrz: ArrayInput) → SurfaceFluxResult[source]¶: calculate heat fluxes through the ice and ice surface temperature

Atmospheric fluxes¶

Standalone JAX bulk heat-flux kernels, retaining the original CESM equations.

Array inputs preserve the original shapes and units documented per function. Configuration and physical constants are supplied as separate frozen dataclasses. The return casts describe JIT’s array outputs for formulas whose eager NumPy or Python inputs would otherwise infer NumPy arrays or scalars. They perform no conversion and leave the original equations unchanged.

The saturation humidity of air (kg/m^3)

Argument:: tk (ndarray): temperature (K)

Saturated specific humidity (kg/kg)

Parameters:

ps (ndarray) – atm sfc pressure (Pa)
tk (ndarray) – atm temperature (K)

Returns:

ndarray

Reference:: Barnier B., L. Siefridt, P. Marchesiello, (1995): Thermal forcing for a global ocean circulation model using a three-year climatology of ECMWF analyses, Journal of Marine Systems, 6, p. 363-380.

veris.physics.thermodynamics.heat_flux_cesm.get_press_levs(sp: ArrayInput, hya: ArrayInput, hyb: ArrayInput) → Array[source]¶

Compute pressure levels

Parameters:

sp (ndarray) – Atmospheric surface pressure
hya (ndarray) – Hybrid sigma level A coefficient for vertical grid
hyb (ndarray) – Hybrid sigma level B coefficient for vertical grid

Returns:

ndarray

veris.physics.thermodynamics.heat_flux_cesm.compute_z_level(phys: PhysicalConstants, t: ArrayInput, q: ArrayInput, ph: ArrayInput) → Array[source]¶

Computes the altitudes at ECMWF Integrated Forecasting System (ECMWF-IFS) model half- and full-levels (for 137 levels model reanalysis: L137)

Parameters:

t (ndarray) – Atmospheric temperture [K]
q (ndarray) – Atmospheric specific humidity [kg/kg]
ph (ndarray) – Pressure at half model levels

Note

The top level of the atmosphere is excluded

Reference:

Returns:: Altitude of the atmospheric near surface layer (second IFS level)
Return type:: ndarray

veris.physics.thermodynamics.heat_flux_cesm.dqnetdt(conf: Configuration, phys: PhysicalConstants, mask: MaskInput, ps: ArrayInput, rbot: ArrayInput, sst: ArrayInput, ubot: ArrayInput, vbot: ArrayInput, us: ArrayInput, vs: ArrayInput) → tuple[Array, Array, Array][source]¶

Calculates correction term of net ocean heat flux (W/m^2)

Parameters:

mask (ndarray) – ocean mask (0-1)
ps (ndarray) – surface pressure (Pa)
rbot (ndarray) – atm density at full model level (kg/m^3)
sst (ndarray) – surface temperature (K)
vmag (ndarray) – atm wind speed at full model level (m/s)

Returns:

tuple(ndarray, ndarray, ndarray)

Reference:: Barnier B., L. Siefridt, P. Marchesiello, (1995): Thermal forcing for a global ocean circulation model using a three-year climatology of ECMWF analyses, Journal of Marine Systems, 6, p. 363-380.

veris.physics.thermodynamics.heat_flux_cesm.net_lw_ocn(conf: Configuration, phys: PhysicalConstants, mask: MaskInput, lat: ArrayInput, qbot: ArrayInput, sst: ArrayInput, tbot: ArrayInput, tcc: ArrayInput) → Array[source]¶

Compute net downward LW radiation at the ocean surface (W/m^2)

Parameters:

mask (ndarray) – ocn domain mask 0 <=> out of domain
lat (ndarray) – latitude coordinates (deg)
qbot (ndarray) – atm specific humidity (kg/kg)
sst (ndarray) – sea surface temperature (K)
tbot (ndarray) – atm T (K)
tcc (ndarray) – total cloud cover (0-1)

Returns:

ndarray

Reference:: Clark, N.E., L.Eber, R.M.Laurs, J.A.Renner, and J.F.T.Saur, (1974): Heat exchange between ocean and atmosphere in the eastern North Pacific for 1961-71, NOAA Technical report No. NMFS SSRF-682.

Neutral drag coeff at 10m

Argument:: umps (ndarray): wind speed (m/s)

Unstable part of psimh

Argument:: xd (ndarray): model level height devided by Obukhov length

Unstable part of psimx

Argument:: xd (ndarray): model level height devided by Obukhov length

veris.physics.thermodynamics.heat_flux_cesm.flux_atmOcn(conf: Configuration, phys: PhysicalConstants, mask: MaskInput, rbot: ArrayInput, zbot: ArrayInput, ubot: ArrayInput, vbot: ArrayInput, qbot: ArrayInput, tbot: ArrayInput, thbot: ArrayInput, us: ArrayInput, vs: ArrayInput, ts: ArrayInput) → tuple[Array, Array, Array, Array, Array, Array, Array, Array, Array, Array, Array, Array][source]¶

atm/ocn fluxes calculation

Parameters:

mask (ndarray) – ocn domain mask 0 <=> out of domain
rbot (ndarray) – atm density (kg/m^3)
zbot (ndarray) – atm level height (m)
ubot (ndarray) – atm u wind (m/s)
vbot (ndarray) – atm v wind (m/s)
qbot (ndarray) – atm specific humidity (kg/kg)
tbot (ndarray) – atm T (K)
thbot (ndarray) – atm potential T (K)
us (ndarray) – ocn u-velocity (m/s)
vs (ndarray) – ocn v-velocity (m/s)
ts (ndarray) – ocn temperature (K)

Returns:

heat flux: sensible (W/m^2) lat (ndarray): heat flux: latent (W/m^2) lwup (ndarray): heat flux: lw upward (W/m^2) evap (ndarray): water flux: evap ((kg/s)/m^2) taux (ndarray): surface stress, zonal (N) tauy (ndarray): surface stress, maridional (N)

tref (ndarray): diag: 2m ref height T (K) qref (ndarray): diag: 2m ref humidity (kg/kg) duu10n(ndarray): diag: 10m wind speed squared (m/s)^2

ustar_sv(ndarray): diag: ustar re_sv (ndarray): diag: sqrt of exchange coefficient (water) ssq_sv (ndarray): diag: sea surface humidity (kg/kg)

Return type:

sen (ndarray)

Reference:

Large, W. G., & Pond, S. (1981). Open Ocean Momentum Flux Measurements in Moderate to Strong Winds, Journal of Physical Oceanography, 11(3), pp. 324-336
Large, W. G., & Pond, S. (1982). Sensible and Latent Heat Flux Measurements over the Ocean, Journal of Physical Oceanography, 12(5), 464-482.
https://svn-ccsm-release.cgd.ucar.edu/model_versions/cesm1_0_5/models/csm_share/shr/shr_flux_mod.F90

veris.physics.thermodynamics.heat_flux_cesm.flux_atmOcn_simple(conf: Configuration, phys: PhysicalConstants, mask: MaskInput, ps: ArrayInput, qbot: ArrayInput, rbot: ArrayInput, ubot: ArrayInput, vbot: ArrayInput, tbot: ArrayInput, us: ArrayInput, vs: ArrayInput, ts: ArrayInput) → tuple[Array, Array, Array][source]¶

Calculates bulk net heat flux

Parameters:

mask (ndarray) – ocn domain mask 0 <=> out of domain
ps (ndarray) – surface pressure (Pa)
qbot (ndarray) – atm specific humidity (kg/kg)
rbot (ndarray) – atm density at full model level (kg/m^3)
tbot (ndarray) – temperature at full model level (K)
ubot (ndarray) – atm u wind (m/s)
vbot (ndarray) – atm v wind (m/s)
qbot – atm specific humidity (kg/kg)
us (ndarray) – ocn u-velocity (m/s)
vs (ndarray) – ocn v-velocity (m/s)
ts (ndarray) – surface temperature (K)

Returns:

tuple(ndarray, ndarray, ndarray)

Reference:: Barnier B., L. Siefridt, P. Marchesiello, (1995): Thermal forcing for a global ocean circulation model using a three-year climatology of ECMWF analyses, Journal of Marine Systems, 6, p. 363-380.

Standalone JAX bulk heat-flux kernels, retaining the original MITgcm LANL equations.

Array inputs preserve the original shapes and units documented per function. Configuration and physical constants are supplied as separate frozen dataclasses.

Calculate bulk formula fluxes over open ocean.

The bulk formulae are:

wind stress = (ust,vst) = rhoA * Cd * Ws * (del.u,del.v)
Sensib Heat flux = fsha = rhoA * Ch * Ws * del.T * CpAir
Latent Heat flux = flha = rhoA * Ce * Ws * del.Q * Lvap
                = -Evap * Lvap
with Ws = wind speed = sqrt(del.u^2 +del.v^2) ;
    del.T = Tair - Tsurf ; del.Q = Qair - Qsurf
    Cd,Ch,Ce = transfer coefficient for momentum, sensible
            & latent heat flux [no units]

Parameters:

uw (ndarray) – zonal wind speed (at grid center) [m/s]
vw (ndarray) – meridional wind speed (at grid center) [m/s]
ta (ndarray) – air temperature [K] at height ht
qa (ndarray) – specific humidity [kg/kg] at heigth ht
tsf (ndarray) – sea surface temperature [K]
ocn_mask (ndarray) – 0=land, nonzero=ocean (not an area weight). Land diagnostics and input sensitivities are zero, including when land inputs contain missing values. MITgcm bulkf_forcing.F gates calls to the scalar LANL formula on nonzero maskC.

Returns:

upward long wave radiation (>0 upward) [W/m2] flha (ndarray): latent heat flux (>0 downward) [W/m2] fsha (ndarray): sensible heat flux (>0 downward) [W/m2] df0dT (ndarray): derivative of heat flux with respect to Tsf [W/m2/K] ust (ndarray): zonal wind stress (at grid center) [N/m2] vst (ndarray): meridional wind stress (at grid center)[N/m2] evp (ndarray): evaporation rate (over open water) [kg/m2/s] ssq (ndarray): surface specific humidity [kg/kg] dEvdT (ndarray): derivative of evap. with respect to tsf [kg/m2/s/K]

Return type:

flwupa (ndarray)