feat[next]: staggered fields

docs/development/ADRs/next/0024-Staggered_Dimensions.md

@@ -0,0 +1,101 @@
+---
+tags: []
+---
+
+# Staggered Dimensions
+
+- **Status**: valid
+- **Authors**: Till Ehrengruber (@tehrengruber)
+- **Created**: 2026-07-08
+- **Updated**: 2026-07-09
+
+A **staggered dimension** is a dimension sitting at the **half-integer**
+positions of a base dimension. For example, in a cell-centered 2D Cartesian grid
+with cells located at `I`, `J`, the edges are located at `I − ½`, `J` and
+`I`, `J − ½`. Differentiating between staggered and non-staggered dimensions
+avoids accidental errors when fields defined on different entities are combined,
+and improves readability by staying close to the mathematical formulation.
+
+## Context
+
+Without a distinction between an entity and its staggered counterpart, fields
+living on, e.g., cells and on edges share the same dimension and can be combined
+silently even though they are geometrically different. We want the type system to
+keep them apart, while still letting users move between the two with an intuitive,
+mathematically-flavoured syntax.
+
+## Indexing convention
+
+A shift by a **half-integer** offset moves between a dimension and its staggered
+counterpart; an integer offset shifts within the (non-staggered) dimension as
+usual.
+
+Let `i_field` be a 1D field of cell values defined on `I`, and `IHalf` the
+corresponding staggered dimension of the edges between those cells. Then:
+
+- `i_field(IHalf + 0.5)` → maps edges to the cell value above: `(i−½)+½ = i`
+- `i_field(IHalf − 0.5)` → maps edges to the cell value below: `(i−½)−½ = i−1`
+
+```
+           ┊      ●      ┊       ●      ┊      ●      ┊
+ I               -1              0             1
+ IHalf    -1½           -½             +½            +1½
+```
+
+Symmetrically, let `ihalf_field` be defined on `I − ½` (i.e. on `IHalf`). Then:
+
+- `ihalf_field(I + 0.5)` → maps cells to the edge above: `i+½`
+- `ihalf_field(I − 0.5)` → maps cells to the edge below: `i−½`
+
+### Storage
+
+`ihalf_field.ndarray[idx]` holds the value for logical index
+`i = domain.start + idx`, which is *interpreted* as sitting at `i − ½`. In the
+example above, `IHalf(0)` therefore sits at `0 − ½ = −½`, i.e. it is the index of
+the edge just below the cell `I(0)`. The memory layout is identical to any normal
+field; the "half" is purely how the index is interpreted geometrically. This
+index arithmetic is encoded in `common.connectivity_for_cartesian_shift`.
+
+## Encoding
+
+A staggered dimension is encoded as its base dimension's name with the internal
+`_Staggered` prefix (`common._STAGGERED_PREFIX`), rather than as a new attribute
+on `Dimension`. The helpers `is_staggered`, `flip_staggered` and
+`as_non_staggered` operate purely on that prefix.
+
+Dimensions are identified by their **name** and appear throughout the toolchain
+in more than one form: as a `common.Dimension` instance, but also as an
+`itir.AxisLiteral` (and as plain strings in generated backend code). Adding a
+"staggered" flag to `Dimension` would have required threading it through every one
+of these representations and every place that constructs, compares, or lowers a
+dimension — a large, invasive change touching the frontend, the IR, and all
+backends. Carrying the marker in the name instead means the frontend and IR pass
+it through unchanged, without any of those formats having to learn about
+staggering. The leading underscore marks the prefix as internal — user code goes
+through the `is_staggered` / `flip_staggered` / `as_non_staggered` helpers
+(exported from `gt4py.next`) rather than manipulating the name directly.
+
+## Dimension Constraints
+
+**A dimension and its staggered counterpart must not appear together** in the
+same field or domain. This follows directly from the geometric meaning of
+staggered dimension and is enforced by `check_dims` (`common.py`), which
+raises otherwise. Other functions may rely on this, e.g `order_dims` just orders 
+by the `as_non_staggered` name.
+
+## gtfn backend
+
+In gtfn a staggered dimension is emitted as a C++ `using` alias of its base
+dimension's tag (`_add_staggered_aliases` in `itir_to_gtfn_ir.py`): e.g.
+`_StaggeredIDim_t` becomes an alias of `IDim_t`, and `visit_AxisLiteral`
+correspondingly emits the base dimension name.
+
+The reason is that gtfn lowers a shift to an integer offset along a SID axis, and
+a half-integer shift produces exactly such an offset (see
+`connectivity_for_cartesian_shift`). But a staggered shift also *relocates* the
+field onto the staggered dimension, a different axis identity. If the staggered
+dimension had its own tag, reproducing that would require not just shifting the
+SID but also *renaming* its axis from the base tag to the staggered tag — which
+the gtfn shift primitive cannot do. Aliasing the two tags makes the base and its
+staggered counterpart a single axis, so the relocation collapses into the plain
+offset shift and no axis has to be renamed.

docs/development/ADRs/next/README.md

@@ -18,6 +18,7 @@ Writing a new ADR is simple:
 
 - [0005 - Extending Iterator IR](0005-Extending_Iterator_IR.md)
 - [0023 - Fingerprinting](0023-Fingerprinting.md)
+- [0024 - Staggered Dimensions](0024-Staggered_Dimensions.md)
 
 ### Frontend and Parsing #frontend
 

src/gt4py/next/__init__.py

@@ -31,7 +31,10 @@
     Field,
     GridType,
     UnitRange,
+    as_non_staggered,
     domain,
+    flip_staggered,
+    is_staggered,
     unit_range,
 )
 from .constructors import FieldConstructor, as_connectivity, as_field, empty, full, ones, zeros
@@ -122,6 +125,9 @@
     "Domain",
     "unit_range",
     "UnitRange",
+    "is_staggered",
+    "flip_staggered",
+    "as_non_staggered",
     # from constructors
     "FieldConstructor",
     "empty",

src/gt4py/next/common.py

@@ -86,10 +86,10 @@ def __str__(self) -> str:
     def __call__(self, val: int) -> NamedIndex:
         return NamedIndex(self, val)
 
-    def __add__(self, offset: int) -> Connectivity:
-        return CartesianConnectivity(self, offset)
+    def __add__(self, offset: int | float) -> Connectivity:
+        return connectivity_for_cartesian_shift(self, offset)
 
-    def __sub__(self, offset: int) -> Connectivity:
+    def __sub__(self, offset: int | float) -> Connectivity:
         return self + (-offset)
 
     def __gt__(self, value: core_defs.IntegralScalar) -> Domain:
@@ -1336,11 +1336,29 @@ def order_dimensions(dims: Iterable[Dimension]) -> list[Dimension]:
     """Find the canonical ordering of the dimensions in `dims`."""
     if sum(1 for dim in dims if dim.kind == DimensionKind.LOCAL) > 1:
         raise ValueError("There are more than one dimension with DimensionKind 'LOCAL'.")
-    return sorted(dims, key=lambda dim: (_DIM_KIND_ORDER[dim.kind], dim.value))
+    return sorted(
+        dims,
+        key=lambda dim: (
+            _DIM_KIND_ORDER[dim.kind],
+            as_non_staggered(dim).value,
+        ),
+    )
 
 
 def check_dims(dims: Sequence[Dimension]) -> None:
-    if dims != order_dimensions(dims):
+    # A dimension and its staggered counterpart (i.e. sharing the same non-staggered base dimension)
+    # denote different grid locations and must not appear together in the same field/domain: mixing
+    # them is ambiguous (it makes `order_dimensions` non-total and produces duplicate backend tags).
+    seen: dict[Dimension, Dimension] = {}
+    for dim in dims:
+        base = as_non_staggered(dim)
+        if base in seen:
+            raise ValueError(
+                f"Dimensions '{seen[base]}' and '{dim}' cannot be combined: a dimension and its "
+                f"staggered counterpart must not appear together in the same field or domain."
+            )
+        seen[base] = dim
+    if list(dims) != order_dimensions(dims):
         raise ValueError(
             f"Dimensions '{', '.join(map(str, dims))}' are not ordered correctly, expected '{', '.join(map(str, order_dimensions(dims)))}'."
         )
@@ -1424,3 +1442,48 @@ def __gt_builtin_func__(cls, /, func: fbuiltins.BuiltInFunction[_R, _P]) -> Call
 #: Equivalent to the `_FillValue` attribute in the UGRID Conventions
 #: (see: http://ugrid-conventions.github.io/ugrid-conventions/).
 _DEFAULT_SKIP_VALUE: Final[int] = -1
+_STAGGERED_PREFIX = "_Staggered"
+
+
+def is_staggered(dim: Dimension) -> bool:
+    """Return whether `dim` is a staggered dimension."""
+    return dim.value.startswith(_STAGGERED_PREFIX)
+
+
+def flip_staggered(dim: Dimension) -> Dimension:
+    """Return the staggered counterpart of `dim`."""
+    if is_staggered(dim):
+        return Dimension(dim.value[len(_STAGGERED_PREFIX) :], dim.kind)
+    else:
+        return Dimension(f"{_STAGGERED_PREFIX}{dim.value}", dim.kind)
+
+
+def as_non_staggered(dim: Dimension) -> Dimension:
+    """Return the non-staggered base dimension of `dim` (`dim` itself if already non-staggered)."""
+    if is_staggered(dim):
+        return flip_staggered(dim)
+    return dim
+
+
+def connectivity_for_cartesian_shift(dim: Dimension, offset: int | float) -> CartesianConnectivity:
+    """
+    Build the connectivity that shifts `dim` by `offset`.
+
+    An integer `offset` shifts within `dim` (the codomain stays `dim`). A half-integer `offset`
+    (fractional part `0.5`) shifts to the staggered counterpart of `dim` (the codomain becomes
+    `flip_staggered(dim)`).
+
+    The half-integer case encodes the convention that a staggered index sits half a cell *below*
+    its base index (see ADR 0024): `IHalf(0)` is the edge below `I(0)`. Because of this asymmetry,
+    shifting out of a non-staggered dimension needs a `+1` index correction that shifting out of a
+    staggered dimension does not, e.g. `I + 0.5` maps `I(i)` to `IHalf(i+1)` (position `i+½`) while
+    `IHalf + 0.5` maps `IHalf(i)` to `I(i)`.
+    """
+    integral_offset, staggered_offset = divmod(offset, 1)
+    if staggered_offset == 0.5:
+        if not is_staggered(dim):
+            integral_offset += 1
+        return CartesianConnectivity(dim, int(integral_offset), codomain=flip_staggered(dim))
+    else:
+        assert staggered_offset == 0
+        return CartesianConnectivity(dim, int(integral_offset), codomain=dim)

src/gt4py/next/ffront/foast_passes/type_deduction.py

@@ -12,7 +12,7 @@
 import gt4py.next.ffront.field_operator_ast as foast
 from gt4py import eve
 from gt4py.eve import NodeTranslator, NodeVisitor, traits
-from gt4py.next import errors
+from gt4py.next import common, errors
 from gt4py.next.common import Dimension, DimensionKind, promote_dims
 from gt4py.next.ffront import (
     dialect_ast_enums,
@@ -607,13 +607,6 @@ def _deduce_compare_type(
     def _deduce_binop_type(
         self, node: foast.BinOp, *, left: foast.Expr, right: foast.Expr, **kwargs: Any
     ) -> Optional[ts.TypeSpec]:
-        # e.g. `IDim+1`
-        if (
-            isinstance(left.type, ts.DimensionType)
-            and isinstance(right.type, ts.ScalarType)
-            and type_info.is_integral(right.type)
-        ):
-            return ts.OffsetType(source=left.type.dim, target=(left.type.dim,))
         if isinstance(left.type, ts.OffsetType):
             raise errors.DSLError(
                 node.location, f"Type '{left.type}' can not be used in operator '{node.op}'."
@@ -679,6 +672,34 @@ def _deduce_binop_type(
                     f"must be one of {', '.join((str(op) for op in logical_ops))}.",
                 )
             return ts.DomainType(dims=promote_dims(left.type.dims, right.type.dims))
+        elif (
+            node.op in (dialect_ast_enums.BinaryOperator.ADD, dialect_ast_enums.BinaryOperator.SUB)
+            and isinstance(left.type, ts.DimensionType)
+            and isinstance(right.type, ts.ScalarType)
+            and type_info.is_arithmetic(right.type)
+        ):
+            # e.g. `IDim+1` or `IDim+0.5`
+            if not isinstance(right, foast.Constant):
+                raise errors.DSLError(
+                    right.location,
+                    "Cartesian offsets are only supported with a literal right-hand side, "
+                    "e.g. 'IDim + 1', but not 'IDim + expr'.",
+                )
+            offset_index = right.value
+            if node.op == dialect_ast_enums.BinaryOperator.SUB:
+                offset_index *= -1
+            if not isinstance(offset_index, (int, float)) or offset_index % 1 not in (0, 0.5):
+                raise errors.DSLError(
+                    right.location,
+                    f"Invalid offset '{right.value}' for a Cartesian shift of dimension "
+                    f"'{left.type.dim.value}'.",
+                    hints=[
+                        "Use an integer offset to shift within the dimension, or a half-integer "
+                        "offset such as '0.5' or '-1.5' to shift to the staggered dimension."
+                    ],
+                )
+            conn = common.connectivity_for_cartesian_shift(left.type.dim, offset_index)
+            return ts.OffsetType(source=conn.codomain, target=(conn.domain_dim,))
         else:
             raise errors.DSLError(node.location, err_msg)
 

src/gt4py/next/ffront/foast_to_gtir.py

@@ -12,7 +12,7 @@
 
 from gt4py import eve
 from gt4py.eve.extended_typing import Never, cast
-from gt4py.next import utils
+from gt4py.next import common, utils
 from gt4py.next.ffront import (
     dialect_ast_enums,
     experimental as experimental_builtins,
@@ -304,19 +304,22 @@ def _visit_shift(self, node: foast.Call, **kwargs: Any) -> itir.Expr:
                     current_expr = im.as_fieldop(
                         im.lambda_("__it")(im.deref(im.shift(offset_name.id, new_index)("__it")))
                     )(current_expr)
-                # `field(Dim + idx)`
+                # `field(Dim + idx)` (where `idx` is integer or half integer)
                 case foast.BinOp(
                     op=dialect_ast_enums.BinaryOperator.ADD | dialect_ast_enums.BinaryOperator.SUB,
-                    left=foast.Name() as dim_name,
+                    left=foast.LocatedNode(type=ts.DimensionType(dim=common.Dimension() as dim)),
                     right=foast.Constant(value=offset_index),
                 ):
                     if arg.op == dialect_ast_enums.BinaryOperator.SUB:
                         offset_index *= -1
-                    assert isinstance(dim_name.type, ts.DimensionType)
-                    dim = dim_name.type.dim
+                    conn = common.connectivity_for_cartesian_shift(dim, offset_index)
                     current_expr = im.as_fieldop(
                         im.lambda_("__it")(
-                            im.deref(im.shift(im.cartesian_offset(dim), offset_index)("__it"))
+                            im.deref(
+                                im.shift(
+                                    im.cartesian_offset(conn.domain_dim, conn.codomain), conn.offset
+                                )("__it")
+                            )
                         )
                     )(current_expr)
                 # `field(Off)`

src/gt4py/next/iterator/embedded.py

@@ -587,16 +587,11 @@ def execute_shift(
         # the assertions above confirm pos is incomplete casting here to avoid duplicating work in a type guard
         return cast(IncompletePosition, pos) | {tag: new_entry}
 
-    # a `CartesianConnectivity` tag is a self-describing cartesian shift (e.g. from a
-    # `CartesianOffset` IR node); named offsets are resolved through the offset provider
     if isinstance(tag, common.CartesianConnectivity):
-        assert tag.domain_dim == tag.codomain  # relocation (staggering) is not supported here
         new_pos = copy.copy(pos)
-        key = tag.domain_dim.value
-        if common.is_int_index(value := new_pos[key]):
-            new_pos[key] = value + index + tag.offset
-        else:
-            raise AssertionError()
+        value = new_pos.pop(tag.domain_dim.value)
+        assert common.is_int_index(value)
+        new_pos[tag.codomain.value] = value + index + tag.offset
         return new_pos
     offset_implementation = common.get_offset(offset_provider, tag)
     if common.is_neighbor_table(offset_implementation):

src/gt4py/next/iterator/ir.py

@@ -88,6 +88,7 @@ def __str__(self):
 InfinityLiteral.POSITIVE = InfinityLiteral(name="POSITIVE")
 
 
+# TODO(tehrengruber): allow int only and create OffsetRef for str instead
 class OffsetLiteral(Expr):
     value: Union[int, str]