Budgets

Conservation-budget diagnostics. For a tracer \(\theta\) with flux \(\mathbf{F}\) and source \(Q\), the budget residual measures how far a dataset departs from exact conservation:

\[ r = \frac{\partial\theta}{\partial t} + \nabla\cdot\mathbf{F} - Q. \]

ControlVolumeIntegral and BoundaryFlux assemble the volume and surface terms; the tracer-specific residuals specialise \(\theta\) to heat, salt, volume, or kinetic energy.

Building blocks

ControlVolumeIntegral

ControlVolumeIntegral(variable: str, *, volume_metrics: Dataset, region: DataArray | None = None, dims: Sequence[str] = ('z', 'lat', 'lon'), cell_volume_var: str = 'cell_volume')

Bases: Operator

Volume-weighted integral of a field over a control volume.

Integrates variable over dims weighted by the cell volumes in volume_metrics (optionally masked to region). The metrics are explicit, never auto-derived — build them with :func:xrops.grad.grid_metrics_from_coords if the data does not ship them.

Parameters:

Name	Type	Description	Default
variable	str	Name of the field to integrate.	required
volume_metrics	Dataset	Dataset carrying the cell_volume_var weights.	required
region	DataArray | None	Optional boolean mask selecting the control volume.	None
dims	Sequence[str]	Dimensions integrated over.	('z', 'lat', 'lon')
cell_volume_var	str	Name of the cell-volume variable in volume_metrics.	'cell_volume'

Returns:

Type	Description
	The volume integral as a DataArray (reduced over dims).

Source code in packages/xrtoolz-ops/src/xrops/budgets/operators.py

def __init__(
    self,
    variable: str,
    *,
    volume_metrics: xr.Dataset,
    region: xr.DataArray | None = None,
    dims: Sequence[str] = ("z", "lat", "lon"),
    cell_volume_var: str = "cell_volume",
) -> None:
    self.variable = variable
    self.volume_metrics = volume_metrics
    self.region = region
    self.dims = tuple(dims)
    self.cell_volume_var = cell_volume_var

BoundaryFlux

BoundaryFlux(*, variable: str | None, velocity_vars: dict[str, str], face_metrics: Dataset, region: DataArray | None = None)

Bases: Operator

Advective flux through the faces of a control volume.

Integrates the velocity_vars (optionally carrying variable) over the face areas in face_metrics, optionally restricted to region. The metrics are explicit, never auto-derived.

Parameters:

Name	Type	Description	Default
variable	str | None	Tracer carried by the flow, or None for a pure volume (velocity) flux.	required
velocity_vars	dict[str, str]	Mapping of face direction → velocity variable name.	required
face_metrics	Dataset	Dataset carrying the face-area weights.	required
region	DataArray | None	Optional boolean mask selecting the control volume.	None

Returns:

Type	Description
	Dataset of per-boundary fluxes.

Source code in packages/xrtoolz-ops/src/xrops/budgets/operators.py

def __init__(
    self,
    *,
    variable: str | None,
    velocity_vars: dict[str, str],
    face_metrics: xr.Dataset,
    region: xr.DataArray | None = None,
) -> None:
    self.variable = variable
    self.velocity_vars = dict(velocity_vars)
    self.face_metrics = face_metrics
    self.region = region

BudgetResidual

Bases: Operator

Generic conservation residual ∂φ/∂t + ∇·F − source + sink.

Combines a precomputed tendency and flux divergence (with optional source / sink terms) into the budget-closure residual; a perfectly closed budget is zero. Inputs are passed to __call__, not the constructor.

Parameters:

Name	Type	Description	Default
tendency		Local time tendency ∂φ/∂t.	required
flux_divergence		Divergence of the flux ∇·F.	required
source		Optional source term.	required
sink		Optional sink term.	required

Returns:

Type	Description
	The budget-closure residual DataArray.

Tracer-specific residuals

HeatBudgetResidual

HeatBudgetResidual(*, temp_var: str = 'theta', **kw: Any)

Bases: _TracerBudgetOp

Per-cell heat-budget residual ∂θ/∂t + ∇·(u θ) − Q.

Uses the spherical-metric divergence from :mod:xrops.grad; returns the per-cell residual (zero for a closed budget). It does not consume volume/face metrics — multiply by cell_volume and integrate with :class:ControlVolumeIntegral for a control-volume closure.

Parameters:

Name	Type	Description	Default
temp_var	str	Temperature (θ) variable name.	'theta'
u_var		Eastward velocity variable name.	required
v_var		Northward velocity variable name.	required
w_var		Vertical velocity variable name, or None.	required
surface_flux_var		Optional surface heat-flux variable name.	required
time_dim		Time dimension name.	required
lat		Latitude coordinate name.	required
lon		Longitude coordinate name.	required
depth		Vertical coordinate name, or None for a 2-D budget.	required

Returns:

Type	Description
	The per-cell heat-budget residual DataArray.

Source code in packages/xrtoolz-ops/src/xrops/budgets/operators.py

def __init__(self, *, temp_var: str = "theta", **kw: Any) -> None:
    super().__init__(tracer_var=temp_var, **kw)

SaltBudgetResidual

SaltBudgetResidual(*, salt_var: str = 'so', **kw: Any)

Bases: _TracerBudgetOp

Per-cell salt-budget residual ∂S/∂t + ∇·(u S) − F.

Like :class:HeatBudgetResidual but for salinity; returns the per-cell residual (zero for a closed budget).

Parameters:

Name	Type	Description	Default
salt_var	str	Salinity (S) variable name.	'so'
u_var		Eastward velocity variable name.	required
v_var		Northward velocity variable name.	required
w_var		Vertical velocity variable name, or None.	required
surface_flux_var		Optional surface salt-flux variable name.	required
time_dim		Time dimension name.	required
lat		Latitude coordinate name.	required
lon		Longitude coordinate name.	required
depth		Vertical coordinate name, or None for a 2-D budget.	required

Returns:

Type	Description
	The per-cell salt-budget residual DataArray.

Source code in packages/xrtoolz-ops/src/xrops/budgets/operators.py

def __init__(self, *, salt_var: str = "so", **kw: Any) -> None:
    super().__init__(tracer_var=salt_var, **kw)

VolumeBudgetResidual

VolumeBudgetResidual(*, u_var: str = 'u', v_var: str = 'v', w_var: str | None = None, lat: str = 'lat', lon: str = 'lon', depth: str | None = 'depth')

Bases: Operator

Per-cell volume (continuity) residual ∇·u.

The horizontal + vertical velocity divergence; zero for a non-divergent (volume-conserving) flow.

Parameters:

Name	Type	Description	Default
u_var	str	Eastward velocity variable name.	'u'
v_var	str	Northward velocity variable name.	'v'
w_var	str | None	Vertical velocity variable name, or None.	None
lat	str	Latitude coordinate name.	'lat'
lon	str	Longitude coordinate name.	'lon'
depth	str | None	Vertical coordinate name, or None for a 2-D budget.	'depth'

Returns:

Type	Description
	The per-cell continuity residual DataArray.

Source code in packages/xrtoolz-ops/src/xrops/budgets/operators.py

def __init__(
    self,
    *,
    u_var: str = "u",
    v_var: str = "v",
    w_var: str | None = None,
    lat: str = "lat",
    lon: str = "lon",
    depth: str | None = "depth",
) -> None:
    self.u_var = u_var
    self.v_var = v_var
    self.w_var = w_var
    self.lat = lat
    self.lon = lon
    self.depth = depth

KineticEnergyBudgetResidual

KineticEnergyBudgetResidual(*, u_var: str = 'u', v_var: str = 'v', forcing_vars: Sequence[str] | None = None, time_dim: str = 'time', lat: str = 'lat', lon: str = 'lon', depth: str | None = 'depth', w_var: str | None = None)

Bases: Operator

Per-cell kinetic-energy budget residual.

Residual of ∂KE/∂t + ∇·(u KE) − forcing with KE = ½(u² + v²); zero for a closed KE budget.

Parameters:

Name	Type	Description	Default
u_var	str	Eastward velocity variable name.	'u'
v_var	str	Northward velocity variable name.	'v'
forcing_vars	Sequence[str] | None	Optional forcing / dissipation variable names summed as the source term.	None
time_dim	str	Time dimension name.	'time'
lat	str	Latitude coordinate name.	'lat'
lon	str	Longitude coordinate name.	'lon'
depth	str | None	Vertical coordinate name, or None for a 2-D budget.	'depth'
w_var	str | None	Vertical velocity variable name, or None.	None

Returns:

Type	Description
	The per-cell KE-budget residual DataArray.

Source code in packages/xrtoolz-ops/src/xrops/budgets/operators.py

def __init__(
    self,
    *,
    u_var: str = "u",
    v_var: str = "v",
    forcing_vars: Sequence[str] | None = None,
    time_dim: str = "time",
    lat: str = "lat",
    lon: str = "lon",
    depth: str | None = "depth",
    w_var: str | None = None,
) -> None:
    self.u_var = u_var
    self.v_var = v_var
    self.forcing_vars = None if forcing_vars is None else tuple(forcing_vars)
    self.time_dim = time_dim
    self.lat = lat
    self.lon = lon
    self.depth = depth
    self.w_var = w_var

Functional primitives (Layer 0)

These pure functions back the operators above; each takes xr.DataArray/xr.Dataset and a dim: argument.

control_volume_integral

control_volume_integral(ds: Dataset, *, variable: str, volume_metrics: Dataset, region: DataArray | None = None, dims: Sequence[str] = ('z', 'lat', 'lon'), cell_volume_var: str = 'cell_volume') -> DataArray

Volume-weighted integral ∫ var dV over a control volume.

Parameters:

Name	Type	Description	Default
ds	Dataset	Dataset containing variable.	required
variable	str	Variable to integrate.	required
volume_metrics	Dataset	Dataset carrying cell_volume_var (m³) on the same horizontal / vertical grid as ds[variable].	required
region	DataArray | None	Optional boolean mask selecting the integration volume. True cells are kept. None integrates over the full grid.	None
dims	Sequence[str]	Spatial dims to integrate over. Any other dims (e.g. time) are preserved on the result.	('z', 'lat', 'lon')
cell_volume_var	str	Name of the volume variable in volume_metrics. Default "cell_volume".	'cell_volume'

Returns:

Type	Description
DataArray	DataArray with the integrated value, reduced over dims.

Notes

For 2-D budgets without a vertical coordinate, pass volume_metrics whose cell_volume equals cell_area (the V4.4 helper does this automatically when depth=None) and set dims=("lat", "lon").

Source code in packages/xrtoolz-ops/src/xrops/budgets/_src/volume.py

def control_volume_integral(
    ds: xr.Dataset,
    *,
    variable: str,
    volume_metrics: xr.Dataset,
    region: xr.DataArray | None = None,
    dims: Sequence[str] = ("z", "lat", "lon"),
    cell_volume_var: str = "cell_volume",
) -> xr.DataArray:
    """Volume-weighted integral ``∫ var dV`` over a control volume.

    Args:
        ds: Dataset containing ``variable``.
        variable: Variable to integrate.
        volume_metrics: Dataset carrying ``cell_volume_var`` (m³) on
            the same horizontal / vertical grid as ``ds[variable]``.
        region: Optional boolean mask selecting the integration volume.
            ``True`` cells are kept. ``None`` integrates over the full
            grid.
        dims: Spatial dims to integrate over. Any other dims (e.g.
            ``time``) are preserved on the result.
        cell_volume_var: Name of the volume variable in
            ``volume_metrics``. Default ``"cell_volume"``.

    Returns:
        DataArray with the integrated value, reduced over ``dims``.

    Notes:
        For 2-D budgets without a vertical coordinate, pass
        ``volume_metrics`` whose ``cell_volume`` equals ``cell_area``
        (the V4.4 helper does this automatically when ``depth=None``)
        and set ``dims=("lat", "lon")``.
    """
    if cell_volume_var not in volume_metrics:
        raise ValueError(
            f"volume_metrics is missing variable {cell_volume_var!r}; "
            f"got {tuple(volume_metrics.data_vars)}."
        )
    field = ds[variable]
    cv = volume_metrics[cell_volume_var]
    weighted = field * cv
    if region is not None:
        weighted = weighted.where(region)
    int_dims = [d for d in dims if d in weighted.dims]
    out = weighted.sum(dim=int_dims, skipna=True)
    return out.rename(f"{variable}_volume_integral").astype(np.float64)

boundary_flux

boundary_flux(ds: Dataset, *, variable: str | None, velocity_vars: dict[str, str], face_metrics: Dataset, region: DataArray | None = None, time_dim: str = 'time') -> Dataset

Region-integrated advective face flux.

Computes ∑ φ u A_face over each face direction. By Gauss's theorem the same sum equals the closed-boundary flux when φ u is divergence-free over the integration domain; for general flows the result is the area-weighted total of the face-flux field rather than a pure boundary integral. See the module docstring.

Parameters:

Name	Type	Description	Default
ds	Dataset	Dataset containing variable (or None for volume flux only) and the velocity components.	required
variable	str | None	Tracer name or None. If None the flux is the volume flux u A_face (used by the volume budget).	required
velocity_vars	dict[str, str]	Mapping {"u": "u_var", "v": "v_var", "w": "w_var"}; "w" is optional. Names are looked up in ds.	required
face_metrics	Dataset	Dataset carrying area_e (east face area, m²), area_n (north face area, m²), and optionally area_top.	required
region	DataArray | None	Optional boolean mask. True cells contribute; the mask should be shifted to face centres if a boundary-only integral is desired.	None
time_dim	str	Name of the time coordinate to preserve as a non- spatial dim. Default "time".	'time'

Returns:

Type	Description
Dataset	Dataset with one scalar (or time_dim-indexed) variable per
Dataset	active face direction: flux_x, flux_y, optionally
Dataset	flux_z.

Source code in packages/xrtoolz-ops/src/xrops/budgets/_src/flux.py

def boundary_flux(
    ds: xr.Dataset,
    *,
    variable: str | None,
    velocity_vars: dict[str, str],
    face_metrics: xr.Dataset,
    region: xr.DataArray | None = None,
    time_dim: str = "time",
) -> xr.Dataset:
    """Region-integrated advective face flux.

    Computes ``∑ φ u A_face`` over each face direction. By Gauss's
    theorem the same sum equals the closed-boundary flux when ``φ u``
    is divergence-free over the integration domain; for general flows
    the result is the area-weighted total of the face-flux field
    rather than a pure boundary integral. See the module docstring.

    Args:
        ds: Dataset containing ``variable`` (or ``None`` for volume
            flux only) and the velocity components.
        variable: Tracer name or ``None``. If ``None`` the flux is the
            volume flux ``u A_face`` (used by the volume budget).
        velocity_vars: Mapping ``{"u": "u_var", "v": "v_var", "w": "w_var"}``;
            ``"w"`` is optional. Names are looked up in ``ds``.
        face_metrics: Dataset carrying ``area_e`` (east face area, m²),
            ``area_n`` (north face area, m²), and optionally ``area_top``.
        region: Optional boolean mask. ``True`` cells contribute; the
            mask should be shifted to face centres if a boundary-only
            integral is desired.
        time_dim: Name of the time coordinate to preserve as a non-
            spatial dim. Default ``"time"``.

    Returns:
        Dataset with one scalar (or ``time_dim``-indexed) variable per
        active face direction: ``flux_x``, ``flux_y``, optionally
        ``flux_z``.
    """
    out: dict[str, xr.DataArray] = {}
    tracer = ds[variable] if variable is not None else None

    for axis, area_key in (("u", "area_e"), ("v", "area_n"), ("w", "area_top")):
        if axis not in velocity_vars:
            continue
        if area_key not in face_metrics:
            continue
        u = ds[velocity_vars[axis]]
        a = face_metrics[area_key]
        flux = u if tracer is None else tracer * u
        flux = flux * a
        if region is not None:
            flux = flux.where(region)
        flux_dim = {"u": "x", "v": "y", "w": "z"}[axis]
        spatial_dims = [d for d in flux.dims if d != time_dim]
        out[f"flux_{flux_dim}"] = flux.sum(dim=spatial_dims, skipna=True)

    return xr.Dataset(out)

budget_residual

budget_residual(tendency: DataArray, flux_divergence: DataArray, *, source: DataArray | None = None, sink: DataArray | None = None) -> DataArray

Combine the four canonical budget terms into a residual.

Sign convention (matches the V4.3 tracer / KE residuals)::

residual = tendency + flux_divergence - source + sink

where flux_divergence is ∇·(u φ) — the divergence of the advective flux out of each cell. The conservation equation ∂φ/∂t + ∇·(u φ) = source - sink rearranges to residual ≈ 0.

A closed budget has residual ≈ 0 within float tolerance.

Source code in packages/xrtoolz-ops/src/xrops/budgets/_src/residual.py

def budget_residual(
    tendency: xr.DataArray,
    flux_divergence: xr.DataArray,
    *,
    source: xr.DataArray | None = None,
    sink: xr.DataArray | None = None,
) -> xr.DataArray:
    """Combine the four canonical budget terms into a residual.

    Sign convention (matches the V4.3 tracer / KE residuals)::

        residual = tendency + flux_divergence - source + sink

    where ``flux_divergence`` is ``∇·(u φ)`` — the divergence of the
    advective flux *out* of each cell. The conservation equation
    ``∂φ/∂t + ∇·(u φ) = source - sink`` rearranges to ``residual ≈ 0``.

    A closed budget has ``residual ≈ 0`` within float tolerance.
    """
    res = tendency + flux_divergence
    if source is not None:
        res = res - source
    if sink is not None:
        res = res + sink
    return res.rename("budget_residual")

heat_budget_residual

heat_budget_residual(ds: Dataset, *, temp_var: str = 'theta', u_var: str = 'u', v_var: str = 'v', w_var: str | None = None, surface_flux_var: str | None = None, time_dim: str = 'time', lat: str = 'lat', lon: str = 'lon', depth: str | None = 'depth') -> DataArray

Per-cell heat-budget residual.

Computes ∂θ/∂t + ∇·(u θ) - F_surface where F_surface is an optional surface heat-flux source field. ρ c_p factors are omitted: the residual is in units of [θ] / s, not W/m³. For closure tests this is sufficient.

Parameters:

Name	Type	Description	Default
ds	Dataset	Dataset with potential temperature + velocities.	required
temp_var, u_var, v_var, w_var		Variable names. w_var=None silently skips the vertical-advection term — the residual is then a surface-only diagnostic.	required
surface_flux_var	str | None	Optional source field on the surface plane.	None
time_dim	str	Name of the time coordinate.	'time'
lat, lon		Names of horizontal coords (degrees).	required
depth	str | None	Vertical coord name. None skips vertical advection.	'depth'

Returns:

Type	Description
DataArray	Per-cell residual, same dims as ds[temp_var].

Notes

ρ c_p θ would be the proper heat content; we use θ directly so the closure test does not require a density field. Multiply the output by ρ_0 c_p upstream if you need W/m³.

Source code in packages/xrtoolz-ops/src/xrops/budgets/_src/heat.py

def heat_budget_residual(
    ds: xr.Dataset,
    *,
    temp_var: str = "theta",
    u_var: str = "u",
    v_var: str = "v",
    w_var: str | None = None,
    surface_flux_var: str | None = None,
    time_dim: str = "time",
    lat: str = "lat",
    lon: str = "lon",
    depth: str | None = "depth",
) -> xr.DataArray:
    """Per-cell heat-budget residual.

    Computes ``∂θ/∂t + ∇·(u θ) - F_surface`` where ``F_surface`` is an
    optional surface heat-flux source field. ``ρ c_p`` factors are
    omitted: the residual is in units of ``[θ] / s``, not W/m³. For
    closure tests this is sufficient.

    Args:
        ds: Dataset with potential temperature + velocities.
        temp_var, u_var, v_var, w_var: Variable names. ``w_var=None``
            silently skips the vertical-advection term — the residual
            is then a surface-only diagnostic.
        surface_flux_var: Optional source field on the surface plane.
        time_dim: Name of the time coordinate.
        lat, lon: Names of horizontal coords (degrees).
        depth: Vertical coord name. ``None`` skips vertical advection.

    Returns:
        Per-cell residual, same dims as ``ds[temp_var]``.

    Notes:
        ``ρ c_p θ`` would be the proper heat content; we use ``θ``
        directly so the closure test does not require a density field.
        Multiply the output by ``ρ_0 c_p`` upstream if you need W/m³.
    """
    return _tracer_budget_residual(
        ds,
        tracer_var=temp_var,
        u_var=u_var,
        v_var=v_var,
        w_var=w_var,
        surface_flux_var=surface_flux_var,
        time_dim=time_dim,
        lat=lat,
        lon=lon,
        depth=depth,
    )

salt_budget_residual

salt_budget_residual(ds: Dataset, *, salt_var: str = 'so', u_var: str = 'u', v_var: str = 'v', w_var: str | None = None, surface_flux_var: str | None = None, time_dim: str = 'time', lat: str = 'lat', lon: str = 'lon', depth: str | None = 'depth') -> DataArray

Per-cell salt-budget residual ∂S/∂t + ∇·(u S) - F_surface.

Source code in packages/xrtoolz-ops/src/xrops/budgets/_src/salt.py

def salt_budget_residual(
    ds: xr.Dataset,
    *,
    salt_var: str = "so",
    u_var: str = "u",
    v_var: str = "v",
    w_var: str | None = None,
    surface_flux_var: str | None = None,
    time_dim: str = "time",
    lat: str = "lat",
    lon: str = "lon",
    depth: str | None = "depth",
) -> xr.DataArray:
    """Per-cell salt-budget residual ``∂S/∂t + ∇·(u S) - F_surface``."""
    return _tracer_budget_residual(
        ds,
        tracer_var=salt_var,
        u_var=u_var,
        v_var=v_var,
        w_var=w_var,
        surface_flux_var=surface_flux_var,
        time_dim=time_dim,
        lat=lat,
        lon=lon,
        depth=depth,
    )

volume_budget_residual

volume_budget_residual(ds: Dataset, *, u_var: str = 'u', v_var: str = 'v', w_var: str | None = None, lat: str = 'lat', lon: str = 'lon', depth: str | None = 'depth') -> DataArray

∇·u with spherical curvature + ∂w/∂z if w_var is given.

Closes to ≈ 0 for an incompressible flow. w_var=None returns the surface 2-D divergence only. Uses :func:xrops.grad.divergence so the spherical curvature term is included; when the vertical component is present it goes in the same call as a third rectilinear axis rather than a separately added partial.

Source code in packages/xrtoolz-ops/src/xrops/budgets/_src/volume_budget.py

def volume_budget_residual(
    ds: xr.Dataset,
    *,
    u_var: str = "u",
    v_var: str = "v",
    w_var: str | None = None,
    lat: str = "lat",
    lon: str = "lon",
    depth: str | None = "depth",
) -> xr.DataArray:
    """``∇·u`` with spherical curvature ``+ ∂w/∂z`` if ``w_var`` is given.

    Closes to ≈ 0 for an incompressible flow. ``w_var=None`` returns
    the surface 2-D divergence only. Uses :func:`xrops.grad.divergence`
    so the spherical curvature term is included; when the vertical
    component is present it goes in the same call as a third
    ``rectilinear`` axis rather than a separately added partial.
    """
    if w_var is not None and depth is not None and depth in ds[w_var].dims:
        return grad.divergence(
            ds[[u_var, v_var, w_var]],
            (u_var, v_var, w_var),
            dims=(lon, lat, depth),
            geometry=("spherical", "spherical", "rectilinear"),
            lon=lon,
            lat=lat,
        ).rename("volume_budget_residual")
    return grad.divergence(
        ds[[u_var, v_var]],
        (u_var, v_var),
        dims=(lon, lat),
        geometry="spherical",
        lon=lon,
        lat=lat,
    ).rename("volume_budget_residual")

kinetic_energy_budget_residual

kinetic_energy_budget_residual(ds: Dataset, *, u_var: str = 'u', v_var: str = 'v', forcing_vars: Sequence[str] | None = None, time_dim: str = 'time', lat: str = 'lat', lon: str = 'lon', depth: str | None = 'depth', w_var: str | None = None) -> DataArray

Per-cell KE-budget residual.

KE per unit mass E = 0.5 (u² + v²). Residual::

∂E/∂t + ∇·(u E) - Σ forcings

Forcings (e.g. wind-stress work) are supplied as additional source fields summed into the source term.

Source code in packages/xrtoolz-ops/src/xrops/budgets/_src/ke.py

def kinetic_energy_budget_residual(
    ds: xr.Dataset,
    *,
    u_var: str = "u",
    v_var: str = "v",
    forcing_vars: Sequence[str] | None = None,
    time_dim: str = "time",
    lat: str = "lat",
    lon: str = "lon",
    depth: str | None = "depth",
    w_var: str | None = None,
) -> xr.DataArray:
    """Per-cell KE-budget residual.

    KE per unit mass ``E = 0.5 (u² + v²)``. Residual::

        ∂E/∂t + ∇·(u E) - Σ forcings

    Forcings (e.g. wind-stress work) are supplied as additional source
    fields summed into the source term.
    """
    u = ds[u_var]
    v = ds[v_var]
    e = 0.5 * (u**2 + v**2)
    tendency = _time_derivative(e, time_dim)
    w = ds[w_var] if w_var is not None else None
    flux_div = _flux_divergence(e, u=u, v=v, w=w, lat=lat, lon=lon, depth=depth)
    res = tendency + flux_div
    if forcing_vars:
        for name in forcing_vars:
            res = res - ds[name]
    return res.rename("kinetic_energy_budget_residual")