Add torus support to GridGeometry #953

msimberg · 2025-11-19T13:44:51Z

This PR continues the series of PRs to add torus grid support. This is extracted from #853.

This PR:

Adds two subclasses, IcosahedronGridGeometry and TorusGridGeometry, of GridGeometry that handle initializing geometry fields for the two grid types.
Adds a factory method to GridGeometry, with_geometry_type (bikeshedding on name and placement welcome) which initializes GridGeometry based on geometry_type. With a closed set of geometry types, this could possibly also be nicely solved with branching inside a single GridGeometry class, but I think the current setup isn't bad either.
Adds tests for the geometry fields.

msimberg · 2025-11-19T13:47:12Z

model/common/src/icon4py/model/common/grid/geometry_stencils.py

+    )
+    x = edge_orientation * xdiff
+    y = edge_orientation * ydiff
+    z = 0.0 * x  # TODO(msimberg): zeros


@egparedes @havogt is there a nicer way to express this? Essentially I'm looking for np.zeros_like(field) in gt4py, but as far as I could tell that doesn't exist.

We could consider leaving the all-zeros Z field out, but I suspect that it would cause too much complexity for downstream field computations that would have to deal with the missing field.

I think this is the best you can do for now...

The comment should probably be updated to something like:

Suggested change

z = 0.0 * x # TODO(msimberg): zeros

# TODO(msimberg): This should use something like numpy.zeros_like if and when that becomes available in gt4py.

z = 0.0 * x

msimberg · 2025-11-19T13:47:27Z

model/common/src/icon4py/model/common/grid/geometry_stencils.py

+        edge_normal_y: y coordinate of the normal
+        edge_normal_z: y coordinate of the normal
+    """
+    z = 0.0 * edge_tangent_x  # TODO(msimberg): zeros


To do: same as in https://github.com/C2SM/icon4py/pull/953/files#r2542084148.

msimberg · 2025-11-19T13:49:49Z

model/common/src/icon4py/model/common/grid/grid_manager.py

+# TODO(msimberg): x, y, z added temporarily for torus grids, but should they be in a separate dict?
+CoordinateDict: TypeAlias = dict[
+    gtx.Dimension, dict[Literal["lat", "lon", "x", "y", "z"], gtx.Field]
+]


I haven't tried it yet, but I think the places where CoordinateDict is used could be changed to a GeometryDict. We possibly lose a bit of type safety (any GeometryName could be used instead of the closed set of naems currently in CoordinateDict).

@halungge do you have thoughts on what you'd prefer to see here (I see the TODO below, so you may already have thought about this)?

model/common/src/icon4py/model/common/grid/geometry.py

github-actions · 2025-11-19T13:51:34Z

Mandatory Tests

Please make sure you run these tests via comment before you merge!

cscs-ci run default

Optional Tests

To run benchmarks you can use:

cscs-ci run benchmark-bencher

To run tests and benchmarks with the DaCe backend you can use:

cscs-ci run dace

To run test levels ignored by the default test suite (mostly simple datatest for static fields computations) you can use:

cscs-ci run extra

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

msimberg · 2025-11-19T13:57:34Z

model/testing/src/icon4py/model/testing/fixtures/datatest.py

+    params=[
+        definitions.Experiments.MCH_CH_R04B09,
+        definitions.Experiments.EXCLAIM_APE,
+        definitions.Experiments.GAUSS3D,
+    ],


Note that I've added GAUSS3D here which means that the experiment will be used for all tests that use the fixture, including tests that haven't yet been updated to handle torus grids correctly.

My preference would be to keep this here so that we really test with the torus grid everywhere it makes sense. The downside is that we'd have to mark tests that don't handle toruses yet with XFAIL or SKIP.

One alternative is introducing GAUSS3D locally only in test_geometry.py for now.

Preferences?

jcanton · 2025-11-20T13:32:28Z

model/common/src/icon4py/model/common/grid/geometry.py

+        coordinates: gm.CoordinateDict,
+        extra_fields: gm.GeometryDict,
+        metadata: dict[str, model.FieldMetaData],
+    ) -> "GridGeometry":


Suggested change

) -> "GridGeometry":

) -> GridGeometry | None:

Hmm, is this for the type checker, for the assert_never case?

jcanton · 2025-11-20T13:52:35Z

model/common/src/icon4py/model/common/grid/geometry.py


-    def _register_computed_fields(self) -> None:
+    def _register_computed_fields(
+        self, input_fields_provider: factory.PrecomputedFieldProvider


Suggested change

self, input_fields_provider: factory.PrecomputedFieldProvider

self

this input is unused?

Yeah, I think so. It was only a requirement for the torus grid.

jcanton · 2025-11-20T14:00:16Z

model/common/src/icon4py/model/common/grid/geometry.py

Suggested change

These fields are computed and saved in coordinates_ below but are never used anywhere in the code. can we clean them and the create_... method up?
@halungge am I missing something?

Diff is unclear, but are you referring to latitude_of_edge_cell_neighbor_0 and friends? To me they also look unused, but @halungge can hopefully confirm if there was another planned use for them.

jcanton · 2025-11-20T14:01:03Z

model/common/src/icon4py/model/common/grid/geometry.py

Suggested change

see above

jcanton · 2025-11-20T14:20:49Z

model/common/src/icon4py/model/common/grid/geometry.py

+        self.register_provider(input_fields_provider)
+        self._register_computed_fields(input_fields_provider)
+
+    def _register_computed_fields(


_register_computed_fields does not exist in GridGeometry anymore, but it has a lot of code in common between IcosahedronGridGeometry and TorusGridGeometry.
Isn't it better to move back to GridGeometry the common stuff and have here only the differences?

Yeah, this was mostly for a simple first pass, I don't think I've separated things particularly well at this point. As hinted in the PR description, we could even drop the subclasses and do something like what I've done in the interpolation PR: https://github.com/msimberg/icon4py/blob/bc7e2bcd4ffa5d786d4eafe9902eb8cb3a04aeb9/model/common/src/icon4py/model/common/interpolation/interpolation_factory.py#L188-L240. I.e. just dispatch to different implementations field-by-field, and only where needed.

What do you think looks more manageable?

I like the interpolation approach more, dispatching only where needed, and we keep just one GridGeometry class with 2-3 differences in the distance and vector orientation parts

Dropping the subclasses also has the benefit of not requiring going through the special factory method to construct the right instances. It's also more consistent with the rest.

I guess it's very, very, very unlikely that there will be any geometry type other than sphere and torus in the near to medium future?

msimberg · 2025-11-21T09:12:16Z

#954 extracts the interpolation field changes.

@msimberg

Based upon #953 #954 and #957 by @msimberg Adds torus support to geometry, RBF and interpolation. Also adds torus grids to all datatests --------- Co-authored-by: Mikael Simberg <mikael.simberg@iki.fi>

Add torus support to geometry

f21b246

msimberg commented Nov 19, 2025

View reviewed changes

model/common/src/icon4py/model/common/grid/geometry.py Outdated Show resolved Hide resolved

Apply suggestion from @msimberg

5b14541

msimberg commented Nov 19, 2025

View reviewed changes

This was referenced Nov 19, 2025

Initialize geometry fields for torus grids #853

Closed

Add torus support for interpolation fields #954

Closed

jcanton reviewed Nov 20, 2025

View reviewed changes

model/common/src/icon4py/model/common/grid/geometry.py

Comment on lines 130 to 133

Copy link

Contributor

jcanton Nov 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change

see above

jcanton reviewed Nov 20, 2025

View reviewed changes

jcanton closed this Dec 9, 2025

jcanton mentioned this pull request Dec 9, 2025

Better torus support #965

Merged

	z = 0.0 * x # TODO(msimberg): zeros
	# TODO(msimberg): This should use something like numpy.zeros_like if and when that becomes available in gt4py.
	z = 0.0 * x

	self, input_fields_provider: factory.PrecomputedFieldProvider
	self

Add torus support to GridGeometry #953

Add torus support to GridGeometry #953

Uh oh!

Conversation

msimberg commented Nov 19, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions bot commented Nov 19, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

msimberg commented Nov 21, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants