Skip to content

Deduplicate the quadratic boundary extrapolation between dycore and common #1449

Description

@jcanton

Context

The quadratic extrapolation of a full-level field onto the top / surface half level is
implemented twice in the codebase with identical arithmetic:

  • model/atmosphere/dycore/.../stencils/extrapolate_quadratically_to_surface.py
    (_extrapolate_quadratically_to_surface, cells, vpfloat)
  • model/atmosphere/dycore/.../stencils/extrapolate_at_top.py
    (_extrapolate_at_top, edges, vpfloat)
  • model/common/src/icon4py/model/common/math/vertical_operations.py
    (extrapolate_quadratically_to_{top,surface}_on_{cells,edges}, wpfloat)

Both now use the same coefficient representation: a three-row wgtfacq field whose
absolute KDim index is the full level the coefficient multiplies (KDim ∈ [nlev-3, nlev)
for the surface weights, [0, 3) for the top weights). That is what
metrics_factory emits and what the driver hands dycore
(driver_utils.py:325), so nothing needs converting at either call site.

The divergence that made a shared operator impossible was on the tmx side and is
resolved in #1359: tmx used to flip the coefficient rows back into Fortran order and
slice them into three KDim-less 2D fields. That flip (_reverse_coefficient_rows) and
the slicing helper (_coefficient_fields) are gone, and the common operators now take
the aligned weight field directly.

What is left to do

Have the dycore stencils call the common field operators instead of repeating the
three-term expression.

Scope

  1. Add vpfloat variants of the four operators in
    common/math/vertical_operations.py, alongside the existing wpfloat ones —
    same split as interpolate_to_cell_center_vp / interpolate_to_cell_center_wp.
    (Or keep one wpfloat operator and astype at the dycore call sites; the
    vp/wp pair matches the surrounding convention better.)
  2. Replace the bodies of _extrapolate_quadratically_to_surface (cells) and
    _extrapolate_at_top (edges) with calls to the common operators.
  3. Decide whether the thin @gtx.program wrappers in those two modules are still
    needed, or whether the dycore callers
    (compute_cell_diagnostics_for_dycore.py, compute_diagnostics_from_normal_wind.py,
    compute_horizontal_velocity_quantities.py,
    vertically_implicit_dycore_solver.py,
    set_theta_v_prime_ic_at_lower_boundary.py,
    compute_contravariant_correction_of_w_for_lower_boundary.py)
    can call the common field operator directly.
  4. Drive-by rename: _extrapolate_at_top is misnamed. It is launched at
    KDim: (vertical_end - 1, vertical_end)
    (compute_diagnostics_from_normal_wind.py:213), i.e. the surface half level, and
    its shifts are KDim - 1 .. - 3. It is the edge counterpart of
    _extrapolate_quadratically_to_surface, not a top extrapolation.

Note compute_contravariant_correction_of_w_for_lower_boundary.py repeats the same
three-term expression on pre-shifted operands rather than on a single interpolant, so it
may not fit the shared operator as-is — worth checking, not worth forcing.

Why this is its own PR

Steps 2 and 3 change dycore's hot path. #1359 is a tmx-scoped refactor and should not
carry a dycore change that needs its own validation run.

Risk

Low: pure deduplication, no semantic change. The representation is already shared, so
there is no data-layout migration involved. Needs the usual cscs-ci run default across
backends because it touches solve_nonhydro / velocity_advection.

🤖 Written by an agent on behalf of @jcanton

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions