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
- 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.)
- Replace the bodies of
_extrapolate_quadratically_to_surface (cells) and
_extrapolate_at_top (edges) with calls to the common operators.
- 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.
- 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
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
wgtfacqfield whoseabsolute
KDimindex is the full level the coefficient multiplies (KDim ∈ [nlev-3, nlev)for the surface weights,
[0, 3)for the top weights). That is whatmetrics_factoryemits 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) andthe slicing helper (
_coefficient_fields) are gone, and the common operators now takethe 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
vpfloatvariants of the four operators incommon/math/vertical_operations.py, alongside the existingwpfloatones —same split as
interpolate_to_cell_center_vp/interpolate_to_cell_center_wp.(Or keep one
wpfloatoperator andastypeat the dycore call sites; thevp/wppair matches the surrounding convention better.)_extrapolate_quadratically_to_surface(cells) and_extrapolate_at_top(edges) with calls to the common operators.@gtx.programwrappers in those two modules are stillneeded, 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.
_extrapolate_at_topis misnamed. It is launched atKDim: (vertical_end - 1, vertical_end)(
compute_diagnostics_from_normal_wind.py:213), i.e. the surface half level, andits 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.pyrepeats the samethree-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 defaultacrossbackends because it touches
solve_nonhydro/velocity_advection.🤖 Written by an agent on behalf of @jcanton