Add the tmx scalar diffusion component by havogt · Pull Request #1490 · C2SM/icon4py

havogt · 2026-09-24T00:27:30Z

Part of the tmx port. This PR adds the scalar diffusion of qv, qc and qi and of the energy/temperature, with its stage interface.

Changes

tmx/stencils/scalar_diffusion.py: four programs.

assemble_scalar_diffusion_matrix: Add the tmx vertical diffusion operators and a common tridiagonal solver #1488's _assemble_vertical_diffusion_matrix_on_cells with 1 / air_mass, for the tracers. It is a separate program because the matrix is reused for qv, qc and qi.
diffuse_tracer: the diffusion core _diffuse_scalar for one tracer, then the update. _diffuse_scalar is shared with the energy and does three things:
- puts the surface flux on the bottom row of the right-hand side;
- runs the implicit vertical solve (Add the tmx vertical diffusion operators and a common tridiagonal solver #1488's _solve_implicit_vertical_diffusion_on_cells);
- adds the explicit horizontal diffusion flux and its divergence.
compute_energy_from_temperature: _compute_internal_energy_from_temperature.
diffuse_energy_and_update_temperature: does three things in one program:
- assembles the energy matrix (prefactor zfactor);
- diffuses the energy with the same core, the surface flux being _compute_surface_internal_energy_flux;
- converts it back with _compute_temperature_from_internal_energy, using the new qv, qc and qi.

The energy is always the internal energy. ICON's dry static energy option (energy_type=1) is not ported: it has not been used since internal energy became ICON's default (icon-mpim MR !326, 2024).

tmx/scalar_diffusion.py: the ScalarDiffusion component.

It follows Diagnostics and has run_hydrometeor_diffusion and run_temperature_diffusion.
The halo exchange of the input tracers starts before the tracer matrix assembly and finishes before its first reader, so qv, qc and qi need one exchange instead of three.

tmx/config.py:

SolverType has only IMPLICIT and EnergyType only INTERNAL. ICON's explicit solver was only used for early testing.
TmxConfig rejects solver_type=1 and energy_type=1 with a clear error.
The Fortran-namelist converter (scripts/python/fortran_config_converter.py) reads both as plain ints, so the same error reaches it instead of a bare enum conversion failure.

Tendency outputs. The solve writes the tendency instead of accumulating onto it, so there is no per-step zero fill. Rows outside the computed domain (cells NUDGING..LOCAL, full column) are not written and keep whatever the caller's buffer holds.

States and test infrastructure:

tmx_states.py:
- TmxInputState gains the tracers and air_mass.
- New TmxNewState and TmxTendencyState.
- TmxSurfaceFluxState has only the two fluxes this PR reads.
model/testing/.../serialbox.py: accessors for the tmx entry tracers and mair, the hydrometeor and temperature exit savepoints, and the surface-flux savepoint (this PR's fields only). Additive only; nothing existing changes.
model/testing/.../definitions.py: ExperimentDescription gains dates, set for EXCLAIM_APE_AES. The tmx datatests use dates[1:] (the first step is the initialization call), and the muphys datatest uses dates.
The tmx test files drop the tmx_ prefix, so the integration and unit test names match the stencil tests (test_diagnostics.py, test_scalar_diffusion.py, test_namelist_config.py, test_config.py).

The energy/temperature conversions reuse what is already in model/common. There are no changes to common or dycore code.

Tridiagonal matrix assembly for full-level cell, half-level cell and full-level edge fields, the implicit solve (Thomas algorithm as two vertical scans) on the same three grids, and the explicit update on full-level cells. Half-level fields are typed on KHalfDim. Co-authored-by: Jacopo Canton <jacopo.canton@gmail.com>

model/common/math/tridiagonal.py holds the Thomas algorithm as a forward sweep and a back substitution scan, in three variants: full levels (wp), half levels (wp) and half levels in mixed precision. The forward sweep carries q = -c' in all of them, as ICON's z_q does. The mixed-precision pair is dycore's former w solver, unchanged; the dycore stencils call it from common. tmx vertical diffusion uses the two wp pairs instead of its own scans. Co-authored-by: Jacopo Canton <jacopo.canton@gmail.com>

Its only caller is the scalar diffusion's explicit solver.

_solve_tridiagonal_matrix_on_{cells,cell_half_levels,edges} run the forward sweep and the back substitution in one call. Field operators are not generic over dimensions, so there is one per field type. The tmx implicit vertical diffusion uses them.

Matrix assembly on V's operators, one program per tracer for the implicit vertical and the horizontal diffusion, and the energy path of the temperature diffusion. Co-authored-by: Jacopo Canton <jacopo.canton@gmail.com>

Co-authored-by: Jacopo Canton <jacopo.canton@gmail.com>

…iffusion datatests

…the field

…ion adds them

# Conflicts: # model/atmosphere/subgrid_scale_physics/tmx/src/icon4py/model/atmosphere/subgrid_scale_physics/tmx/stencils/vertical_diffusion.py # model/common/src/icon4py/model/common/math/tridiagonal.py

nfarabullini

I just skimmed through this PR and made a few comments. More thorough review to follow

nfarabullini

more comments

…cit solver The energy type decides three operations (energy from temperature, surface energy flux, temperature from energy); each is now a named pair of field operators selected by the static use_internal_energy, and the energy goes through the same diffusion core as the tracers. The energy matrix is assembled inside the energy program. SolverType has only IMPLICIT; TmxConfig and namelist parsing reject the explicit solver, which was only used for early testing in ICON.

… type

The explicit solver is removed from SolverType in the scalar-diffusion PR (#1490).

nfarabullini

just a couple of small things

jcanton · 2026-09-28T09:00:22Z

wait for me on this one please, I have a couple of comments queued since Friday and will finish today

jcanton

more comments than I though. the ones about the ser_data and its version are unfortunately blocking

jcanton · 2026-09-28T13:05:10Z

+    fields = (
+        (setup.tendency_state.tend_qv, exit_savepoint.tend_qv(), "tend_qv", 5.0e-20),
+        (setup.tendency_state.tend_qc, exit_savepoint.tend_qc(), "tend_qc", 5.0e-21),
+        (setup.tendency_state.tend_qi, exit_savepoint.tend_qi(), "tend_qi", 3.0e-22),
+        (setup.new_state.qv, exit_savepoint.qv_new(), "qv_new", 2.0e-17),
+        (setup.new_state.qc, exit_savepoint.qc_new(), "qc_new", 2.0e-18),
+        (setup.new_state.qi, exit_savepoint.qi_new(), "qi_new", 7.0e-20),


should we have a more strict/precise per-backend dict of atol+rtol so crappier backends don't hide potential regressions in better backends? (again, since we're making this a template for future ports and we can measure atol/rtol at port-time with ICON4PY_DALLCLOSE_PRINT_INSTEAD_OF_FAIL: true

Should we do it in the end after all components are in, because there are already the tests from the first PR which didn't do that?

OngChia · 2026-09-28T14:24:07Z

+        )
+
+        log.debug("communication of energy (cells): start")
+        self._exchange.exchange(dims.CellDim, self.energy)


Just wondering whether we can remove this exchange by putting energy calculation into the subsequent stencil and gt4py does the rest.

Let's maybe do it later in an optimization round, because currently we don't have mpi tests.

# Conflicts: # model/atmosphere/subgrid_scale_physics/tmx/src/icon4py/model/atmosphere/subgrid_scale_physics/tmx/config.py # model/atmosphere/subgrid_scale_physics/tmx/tests/tmx/fixtures.py # model/atmosphere/subgrid_scale_physics/tmx/tests/tmx/unit_tests/test_tmx_config.py

@jcanton

Catches up with #1466, which moved the Fortran-to-Python mapping out of the config classes and into `scripts/python/fortran_config_converter.py`, and makes the tmx datatests take their configuration from the experiment's `config.yml`: - drop `TmxConfig.from_fortran_dict` and the `icon_equivalent` annotations; the positional `aes_vdf_nml` pins live in the converter's `TMX` mapping. The converter reads `solver_type` as a plain int, so that `TmxConfig.__post_init__` reports the unported explicit solver rather than the bare enum conversion failing first. The `from_fortran_dict` unit tests go: the converter tests already cover the member count and `use_tmx` checks, and gain the explicit-solver rejection. - drop the `tmx_config` fixture: the datatests read `experiment.config.tmx`. - drop the `tmx_dtime` fixture, which read `dt_vdf` as the time step of the solver. ICON runs tmx with the model time step (`init_tmx(p_patch(jg), dt_loc)` with `dt_loc = get_model_timestep_sec(...)` in `mo_atmo_nonhydrostatic.f90`); `dt_vdf` only sets how often the process fires. The datatests now use `experiment.config.driver.dtime`. The two coincide in the archive (300 s), which is why the fixture passed. 🤖 Written by an agent on behalf of @jcanton

jcanton · 2026-09-29T16:03:08Z

I took the liberty of pushing ad84851 updating with the changes merged in #1466 and fixing lots of my own review comments (which I'll now mark as resolved)

there is still the non-blocking question of whether we drop the dry static energy entirely for which I wanted an opinion from @OngChia as well, but as I wrote to @havogt we can remove it already, it's easy to bring it back if needed.

Remove the dry-static-energy path (EnergyType keeps only INTERNAL, with a clear error for 1), take the experiment dates from EXCLAIM_APE_AES.dates, add TODOs for the energy-function naming and the energy exchange, and drop the tmx_ prefix from the tmx integration and unit test file names.

…-scalar-diffusion # Conflicts: # model/atmosphere/subgrid_scale_physics/tmx/src/icon4py/model/atmosphere/subgrid_scale_physics/tmx/config.py # model/atmosphere/subgrid_scale_physics/tmx/tests/tmx/fixtures.py # model/atmosphere/subgrid_scale_physics/tmx/tests/tmx/unit_tests/test_tmx_config.py # scripts/tests/python/test_fortran_config_converter.py

github-actions · 2026-09-29T18:02:54Z

When developing, you can test your changes on CSCS CI before merge with the default pipeline: cscs-ci run default. This will run a default subset of tests.

You can pass options to override pipeline variables, for example:

cscs-ci run default;BACKENDS=gtfn_cpu;LEVELS=unit
cscs-ci run default;MODEL_SUBPACKAGES=common:driver;SESSIONS=model

Avoid running the pipeline for all tests when you are developing.

Available options are:

SESSIONS: model, model_mpi, or tools (correspond to nox sessions)
MODEL_SUBSETS: datatest, basic, or stencils (correspond to nox session selections)
MODEL_SUBPACKAGES: subpackages for non-MPI tests (last component, e.g. diffusion, driver)
MODEL_MPI_SUBPACKAGES: subpackages for MPI tests (as above)
BACKENDS: backends
GRIDS: grids for stencil tests (simple, icon_regional, or icon_global)
LEVELS: testing level for non-stencil tests (unit or integration)

For each option, all can be used as a shorthand for all possible values of that variable, e.g. LEVELS=all.

Multiple values can be given to each option with : used as the separator (; separates options and , separates pipelines).

See scripts/python/generate_ci_pipeline.py and noxfile.py for available values for each option.

The all pipeline can be run with cscs-ci run all. This will run all icon4py tests in CSCS CI which can be expensive. This pipeline runs on a schedule on main, and can be run when extensive validation is needed (e.g. before releases).

Merging

Once your PR is approved and ready for merging, add it to the merge queue. The merge CSCS CI pipeline will run automatically on the merge-queue branch and must pass before the PR is merged. A dummy merge check will be triggered on the PR itself since it's required to add a PR to the merge queue.

Optional Tests

To run benchmarks you can use:

cscs-ci run benchmark-bencher

For more detailed information please look at CI in the EXCLAIM universe.

havogt · 2026-09-29T18:07:17Z

cscs-ci run default;SESSIONS=model;MODEL_SUBSETS=stencils:datatest;MODEL_SUBPACKAGES=tmx:muphys;BACKENDS=gtfn_gpu:dace_gpu;LEVELS=unit:integration

jcanton

lgtm now

jcanton · 2026-09-30T06:36:51Z

        grid=Grids.R02B04_GLOBAL,
-        version=8,
+        version=11,
+        dates=("2008-09-01T00:00:00.000", "2008-09-01T00:05:00.000", "2008-09-01T00:10:00.000"),


ok for now, but should definitely be extracted from the ser_data metadata, not hardcoded here
will unify all tests in one separate PR

Take the experiment dates from EXCLAIM_APE_AES.dates (the same change as in #1490), align the TmxNewState/TmxTendencyState docstrings with #1490, use single backticks for inline code, and drop the tmx_ prefix from the datatest file name.

havogt and others added 15 commits September 23, 2026 13:20

Test the tmx vertical-diffusion operators against numpy

a8d135c

Leave the explicit vertical diffusion to the scalar diffusion PR

126704c

Its only caller is the scalar diffusion's explicit solver.

Merge remote-tracking branch 'upstream/main' into tmx-vertical-diffusion

955842c

Add the tmx scalar-diffusion operators

3233b73

Matrix assembly on V's operators, one program per tracer for the implicit vertical and the horizontal diffusion, and the energy path of the temperature diffusion. Co-authored-by: Jacopo Canton <jacopo.canton@gmail.com>

Add the tmx scalar-diffusion component and its states

45f54ce

Co-authored-by: Jacopo Canton <jacopo.canton@gmail.com>

Test the tmx scalar diffusion against numpy and ICON savepoints

f1ddbda

Co-authored-by: Jacopo Canton <jacopo.canton@gmail.com>

Read the tmx time step from the experiment's namelist in the scalar-d…

8715a4b

…iffusion datatests

Reject the explicit solver in the tmx scalar diffusion

4f96556

Test the tmx scalar-diffusion programs on a horizontal domain inside …

29f8f8e

…the field

Keep only the surface fluxes the tmx scalar diffusion reads

d9c5dc3

Keep the surface-flux savepoint and constructor where the wind diffus…

9d55484

…ion adds them

Merge remote-tracking branch 'upstream/main' into tmx-scalar-diffusion

8899113

# Conflicts: # model/atmosphere/subgrid_scale_physics/tmx/src/icon4py/model/atmosphere/subgrid_scale_physics/tmx/stencils/vertical_diffusion.py # model/common/src/icon4py/model/common/math/tridiagonal.py

havogt mentioned this pull request Sep 24, 2026

Add the tmx wind diffusion component #1491

Open

nfarabullini reviewed Sep 24, 2026

View reviewed changes

Comment thread model/atmosphere/subgrid_scale_physics/tmx/tests/tmx/stencil_tests/test_scalar_diffusion.py Outdated

havogt commented Sep 24, 2026

View reviewed changes

nfarabullini reviewed Sep 24, 2026

View reviewed changes

OngChia reviewed Sep 24, 2026

View reviewed changes

Comment thread ...scale_physics/tmx/src/icon4py/model/atmosphere/subgrid_scale_physics/tmx/scalar_diffusion.py Outdated

Comment thread ...scale_physics/tmx/src/icon4py/model/atmosphere/subgrid_scale_physics/tmx/scalar_diffusion.py

havogt added 2 commits September 25, 2026 16:35

Parametrize the scalar-diffusion energy stencil tests over the energy…

44f879c

… type

havogt added a commit that referenced this pull request Sep 25, 2026

Drop the wind diffusion's solver-type guard and use unary minus

3be86a2

The explicit solver is removed from SolverType in the scalar-diffusion PR (#1490).

havogt marked this pull request as ready for review September 25, 2026 17:07

havogt requested review from OngChia, jcanton and nfarabullini September 28, 2026 07:07

nfarabullini approved these changes Sep 28, 2026

View reviewed changes

docs: act on the second #1490 review round

5c856d5

jcanton requested changes Sep 28, 2026

View reviewed changes

docs: single backticks for inline code, as in the dycore

7d018ec

OngChia reviewed Sep 28, 2026

View reviewed changes

havogt and others added 2 commits September 29, 2026 13:50

havogt added 2 commits September 29, 2026 18:07

havogt requested a review from jcanton September 29, 2026 18:33

jcanton approved these changes Sep 30, 2026

View reviewed changes

havogt added this pull request to the merge queue Sep 30, 2026

Merged via the queue into main with commit 08297fa Sep 30, 2026
48 checks passed

havogt mentioned this pull request Sep 30, 2026

Prescribed surface fluxes for tmx #1458

Closed

Conversation

havogt commented Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Uh oh!

nfarabullini left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nfarabullini left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nfarabullini left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jcanton commented Sep 28, 2026

Uh oh!

jcanton left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jcanton Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

OngChia Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

havogt Sep 29, 2026

Choose a reason for hiding this comment

Uh oh!

jcanton commented Sep 29, 2026

Uh oh!

github-actions Bot commented Sep 29, 2026

Uh oh!

havogt commented Sep 29, 2026

Uh oh!

jcanton left a comment

Choose a reason for hiding this comment

Uh oh!

jcanton Sep 30, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

havogt commented Sep 24, 2026 •

edited

Loading

nfarabullini left a comment •

edited

Loading