From 29c3ac34079af0d0cffa352ed906e0c6225176da Mon Sep 17 00:00:00 2001 From: Hannes Vogt Date: Wed, 5 Aug 2026 16:27:12 +0200 Subject: [PATCH] idea[havogt]: align staggering with the shipped shift convention A shift names the dimension the *result* lives on (the rule unstructured shifts already follow), not the field's own; staggered point i sits at i - 1/2. Flips the Staggered[D] overloads, the runtime index rule, the prototype and its static/runtime checks. Also corrects three claims in Part II: gtfn is not unaffected (it needs _add_staggered_aliases), the I + 0.7 gap is mypy-only, and check_dims (a dim and its counterpart may not co-occur) is not statically expressible. --- .gitignore | 5 + .../dimension-generic-fields.md | 136 ++++++++++++------ .../havogt/dimension-generic-fields/mypy.ini | 2 +- .../dimension-generic-fields/static_checks.py | 42 +++--- .../dimension-generic-fields/static_errors.py | 6 +- .../test_typed_dimensions.py | 44 +++--- .../typed_dimensions.py | 64 ++++++--- 7 files changed, 193 insertions(+), 106 deletions(-) diff --git a/.gitignore b/.gitignore index d2a56e2..c4d0e73 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,11 @@ Thumbs.db *.swo *~ +# Python (prototypes vendored next to proposals are runnable) +__pycache__/ +*.pyc +.pytest_cache/ + # Quartz local preview artifacts (build happens in CI from a cloned Quartz) public/ .quartz-cache/ diff --git a/content/personal/havogt/dimension-generic-fields/dimension-generic-fields.md b/content/personal/havogt/dimension-generic-fields/dimension-generic-fields.md index 8ad41fe..4935f8b 100644 --- a/content/personal/havogt/dimension-generic-fields/dimension-generic-fields.md +++ b/content/personal/havogt/dimension-generic-fields/dimension-generic-fields.md @@ -9,11 +9,11 @@ status: draft > **TL;DR** Redesign `Dimension` so that concrete dimensions are *types* > (`class I(gtx.Dimension)`), usable directly in static type checking without a > mypy plugin. On top of that: statically typed **staggering** -> (`b: Field[Dims[Staggered[I]]] = a(I + 1/2)`) via a `Staggered[D]` type +> (`b: Field[Dims[Staggered[I]]] = a(Staggered[I] + 1/2)`) via a `Staggered[D]` type > constructor, and **dimension variables** (`TypeVar`/`TypeVarTuple` over > dimensions) so field operators can be generic in their dimensions. -> **Prototypes**: the self-contained static-expressibility prototype (mypy 1.19, +> **Prototypes**: the self-contained static-expressibility prototype (mypy 2.3, > no plugin) is vendored alongside this document in > [`dimension-generic-fields/`](dimension-generic-fields/) (see §8.A). The > type-system extension (`ts.DimensionVar`/`ts.DimsVar`, real `src/` changes plus @@ -34,7 +34,7 @@ status: draft ([[personal/havogt/dtype-generic-fields|dtype-generic fields]]) to **dimensions**: (I) redesigning `Dimension` so concrete dimensions are *types* usable in static type checking, (II) statically typed **staggering** - (`b: Field[Dims[Staggered[I]]] = a(I + 1/2)`), and (III) dimension + (`b: Field[Dims[Staggered[I]]] = a(Staggered[I] + 1/2)`), and (III) dimension *variables* (dim-generic operators) in the DSL type system. - **Prototypes**: the critical pieces are implemented and tested — see §8 for what exactly is proven by the vendored @@ -60,8 +60,8 @@ class K(gtx.Dimension, kind=gtx.DimensionKind.VERTICAL): ... a: gtx.Field[gtx.Dims[I], gtx.float64] -b: gtx.Field[gtx.Dims[gtx.Staggered[I]], gtx.float64] = a(I + 1 / 2) # (II) staggering -c: gtx.Field[gtx.Dims[I], gtx.float64] = b(gtx.Staggered[I] + 1 / 2) +b: gtx.Field[gtx.Dims[gtx.Staggered[I]], gtx.float64] = a(gtx.Staggered[I] + 1 / 2) # (II) +c: gtx.Field[gtx.Dims[I], gtx.float64] = b(I + 1 / 2) T = TypeVar("T", gtx.float32, gtx.float64) Ds = TypeVarTuple("Ds") @@ -212,7 +212,7 @@ which we reject anyway). ### 4.1 Why it was "not expressible", and what changed -A shift `a(I + 1/2)` maps `Field[Dims[..., I, ...]]` to +A shift `a(Staggered[I] + 1/2)` maps `Field[Dims[..., I, ...]]` to `Field[Dims[..., Istag, ...]]` — a **type-level substitution inside a variadic tuple**. Python typing has no type-level `Map`/`Replace` over a `TypeVarTuple`, and a tuple type may contain at most one unpacked @@ -251,6 +251,34 @@ unstructured remaps) statically typed, not just staggering. ### 4.2 `Staggered[D]`: the dual grid as a type constructor +**Convention: a shift is precomposition.** A field is a map `f: I → value` and a +connectivity is a map `φ: D → I`; applying one is `f ∘ φ`, a field on `D`. The +connectivity's domain is therefore the *result's* dimension and its codomain the +field's — forced by composition, not chosen. Concretely, the index expression in +`f(...)` is written in the result's index space and must *evaluate* to an index +of the field's own dimension: for `a: Field[Dims[I]]`, `a(Staggered[I] + 1/2)` +is `p ↦ a[p + ½]` with `p` on `Staggered[I]`, so it lands on `Staggered[I]`. + +This is the rule gt4py already follows everywhere else; the unstructured +connectivities encode it in their very names. `V2E` means "for each Vertex, its +Edges" — domain `Vertex`, codomain `Edge` — so it consumes an edge field and +produces a vertex field. Staggering is merely the first place the convention +becomes *observable*, because integer Cartesian shifts have domain == codomain. +Correspondingly, staggered point `i` sits at position `i - 1/2` (ADR 0026): +`Staggered[I](0)` is the point just *below* `I(0)`. + +**Discarded: source-naming.** The alternative — name the dimension the field is +*on*, `a(I + 1/2)` for `a` on `I` — reads more naturally under the "roll the +data along its own axis" intuition, and it would remove the `dual(dim)` call +from every dual-generic body below (the operator has its own dimension in hand; +under precomposition it must name the dual of it). It was rejected because it +gives no rule for the unstructured case: `a(E2C[0])` cannot name the field's +dimension, so unstructured shifts would have to become the exception. The +apparent third argument — that `u_{i+½}` in the literature *is* source-naming — +does not survive: the subscript denotes a value at a location, and reading the +same formula as a field definition, `b_j = a_{j+½}`, makes the free index `j` +the result's. Only that reading composes. + ```python class Staggered(Dimension, Generic[D]): base: ClassVar[type[Dimension]] @@ -262,8 +290,8 @@ class Staggered(Dimension, Generic[D]): `Ihalf: TypeAlias = Staggered[I]`. - **Staggering is an involution**, encoded in overload pairs on `DimensionMeta.__add__` (the `Staggered[D]` overload *before* the generic - `D` overload, so `Staggered[I] + 1/2` yields `Connectivity[I, Staggered[I]]` - and `Staggered[Staggered[I]]` never arises; constructing it explicitly + `D` overload, so a shift naming a staggered dimension consumes an unstaggered + field and `Staggered[Staggered[I]]` never arises; constructing it explicitly raises): ```python @@ -272,9 +300,9 @@ def __add__(cls: type[Staggered[D]], offset: int) -> Connectivity[Staggered[D], @overload def __add__(cls: type[D], offset: int) -> Connectivity[D, D]: ... @overload -def __add__(cls: type[Staggered[D]], offset: float) -> Connectivity[D, Staggered[D]]: ... +def __add__(cls: type[Staggered[D]], offset: float) -> Connectivity[Staggered[D], D]: ... @overload -def __add__(cls: type[D], offset: float) -> Connectivity[Staggered[D], D]: ... +def __add__(cls: type[D], offset: float) -> Connectivity[D, Staggered[D]]: ... ``` - **Doubly staggered types are unrepresentable.** The hierarchy has two @@ -294,10 +322,13 @@ def __add__(cls: type[D], offset: float) -> Connectivity[Staggered[D], D]: ... nominal types. Making the type unrepresentable sidesteps the equation entirely: the only fixed points are `D` and `Staggered[D]`, matching the physics (there are exactly two grids per dimension). -- The notation is exactly the desired `a(I + 1/2)`: `1/2` is a `float`, and +- The notation is the desired `a(Staggered[I] + 1/2)`: `1/2` is a `float`, and `int` vs `float` is statically distinguishable. **Caveat**: float *values* - are not (no `Literal` for floats), so `I + 0.7` statically claims to be a - staggering shift and is only rejected at runtime (`ValueError`). If this is + are not (no `Literal` for floats), so `Staggered[I] + 0.7` statically claims + to be a staggering shift. This is a **mypy-only** gap, not a runtime-only + rejection: the DSL already refuses it at FOAST time with a proper `DSLError` + and a hint (`type_deduction.py`, `offset_index % 1 not in (0, 0.5)`); only the + static layer and direct embedded use fall back to a runtime `ValueError`. If this is considered too weak, a dedicated `HALF` singleton type (`I + HALF`, `I - HALF + 1`) restores full static soundness at the cost of the literal notation; both can coexist. @@ -319,10 +350,16 @@ def __add__(cls: type[D], offset: float) -> Connectivity[Staggered[D], D]: ... @overload def avg(f: Field[Dims[D], float]) -> Field[Dims[Staggered[D]], float]: ... def avg(f): - (dim,) = f.dims - return f(dim - 1 / 2) + f(dim + 1 / 2) + to = dual(f.dims[0]) # a shift names the result grid: the dual of the field's + return f(to - 1 / 2) + f(to + 1 / 2) ``` + This is where precomposition costs something: the body cannot shift on its + own dimension and needs `dual(dim)` — the typed counterpart of gt4py's + value-level `flip_staggered`. The prototype gives `dual` the same overload + pair, so it stays precise under a dimension variable (`to_staggered` in + `static_checks.py`). + The overload pair does **not** have to be user-written, though: a `Dual[X]` marker type plus a decorator whose argument type *recognizes* dual-generic signatures and whose return type is a library-provided protocol carrying @@ -364,20 +401,23 @@ def __add__(cls: type[D], offset: float) -> Connectivity[Staggered[D], D]: ... ### 4.3 Semantics (value level) -Convention: staggered point `i` of `Staggered[I]` sits at position `i + 1/2` -of `I`. `b = a(conn)` means `b[p] = a[p + conn.offset]` in *position* space; -in index space this is a shift by `ceil(offset)` reading from an unstaggered -and `floor(offset)` reading from a staggered dimension. Useful identities -(all covered by runtime tests in the prototype): - -- `a(I + 1/2)(Staggered[I] + 1/2) == a(I + 1)` — two half shifts = one full. -- `a(I + 1/2)(Staggered[I] - 1/2) == a` — round trip. -- `a(I + 1/2) - a(I - 1/2)` — C-grid finite difference, lives on - `Staggered[I]`. - -`a(I + 1/2)` is *pure relabeling + translation* (no interpolation); +Staggered point `i` sits at position `i - 1/2`, so the point of `a(conn)` at +index `p` reads `a` at position `pos(p) + frac(offset)`. In index space that is +a shift by `ceil(offset)` when the shift names an **unstaggered** dimension and +`floor(offset)` when it names a **staggered** one — keyed on the dimension the +result lives on. This is what `connectivity_for_cartesian_shift` already +implements (`divmod`, plus a `+1` correction exactly when the named dimension is +unstaggered). Useful identities for `a` on `I` (all covered by runtime tests in +the prototype): + +- `a(Staggered[I] + 1/2)(I + 1/2) == a(I + 1)` — two half shifts = one full. +- `a(Staggered[I] + 1/2)(I - 1/2) == a` — round trip. +- `a(Staggered[I] + 1/2) - a(Staggered[I] - 1/2)` — C-grid finite difference, + lives on `Staggered[I]`. + +A staggering shift is *pure relabeling + translation* (no interpolation); averaging is written explicitly, e.g. -`0.5 * (u(I + 1/2) + u(I - 1/2))`. Domain handling on bounded (non-periodic) +`0.5 * (u(Staggered[I] + 1/2) + u(Staggered[I] - 1/2))`. Domain handling on bounded (non-periodic) domains follows the same position arithmetic (result range = `{i : i + offset ∈ range(a)}`, i.e. half-open ranges shrink/shift by ceil/floor); the prototype sidesteps this with periodic `np.roll`, the real @@ -392,14 +432,25 @@ concrete typing handles it: - **FOAST**: `visit_BinOp` on `dimension ± literal` (today `type_deduction.py:604`, integers only) gains the fractional case, - producing `ts.OffsetType(source=I, target=(Staggered[I],))` resp. - `(source=Staggered[I], target=(I,))`. `return_type_field` already performs - the dims substitution for arbitrary source/target. The lowering of the - offset value keeps the integral part as today's cartesian offset; the grid - change is encoded in the (auto-registered) relocation offset provider, so - GTIR and the backends only ever see ordinary dimensions (e.g. `I½`) and - integer shifts — **backends are unaffected**. -- **Embedded**: `I + 1/2` constructs the relocation+translation + producing `ts.OffsetType(source=Staggered[I], target=(I,))` for `I ± 1/2` + resp. `(source=I, target=(Staggered[I],))` for `Staggered[I] ± 1/2` — i.e. + `source=conn.codomain, target=(conn.domain_dim,)`, exactly as the shipped + code already does. `return_type_field` already performs the dims substitution + for arbitrary source/target. The lowering of the offset value keeps the + integral part as today's cartesian offset; the grid change is encoded in the + (auto-registered) relocation offset provider, so GTIR sees only ordinary + dimensions (e.g. `I½`) and integer shifts. **Backends are not entirely + unaffected**, though: gtfn needed `_add_staggered_aliases` + (`itir_to_gtfn_ir.py`), emitting the staggered tag as a C++ `using` alias of + its base tag, because a gtfn shift can offset a SID but cannot *rename* its + axis — so base and staggered must collapse into one axis for the relocation to + become a plain offset. That is the concrete reason a staggering shift has to + stay expressible as integer offset + relabel. +- **A dimension and its staggered counterpart may not co-occur** in one field or + domain (`check_dims`, `common.py`). mypy cannot express that constraint, so + `Field[Dims[I, Staggered[I]]]` stays statically writable and the check remains + a runtime one — a known gap, not a solved case. +- **Embedded**: `Staggered[I] + 1/2` constructs the relocation+translation `CartesianConnectivity`, which `premap` supports today. - The `Field`/`Connectivity` overloads of §4.1/4.2 live in `common.py` (or a generated `.pyi`) and replace the current untyped @@ -612,7 +663,8 @@ all remaining named axes), `tag` the result. fails the runtime tick-equality check. Part II promotes exactly this metadata bit to a type (`Staggered[I]`, `Dual[X]`), making the same mistake a static + decoration-time error and giving shifts a *type-level* - effect (`a(I + 1/2)` changing the dims type has no coordax analogue). + effect (`a(Staggered[I] + 1/2)` changing the dims type has no coordax + analogue). **How coordax could fit.** Not as a dependency — gt4py's `Field`/`Domain` already cover its runtime role — but three concrete touchpoints: @@ -680,7 +732,7 @@ add ergonomics, not power: (self-contained, vendored next to this document; run with `pytest content/personal/havogt/dimension-generic-fields/` and `mypy --config-file content/personal/havogt/dimension-generic-fields/mypy.ini ...`; -mypy 1.19 without any plugin): +mypy 2.3 without any plugin): - `typed_dimensions.py`: `DimensionMeta`/`Dimension`/`Staggered`/ `Connectivity`/`Field` exactly as in §3–4 (~150 lines of typing surface). @@ -736,9 +788,11 @@ Stages 0–2 of the dtype plan are prerequisites for the *frontend* stages here ## 10. Risks and open questions 1. **The metaclass-overload mypy quirk** (§3.3): call-site behavior is - correct but the def-site suppression could break on a mypy upgrade; - pinned by tests, with a notation-only fallback. **pyright is untested** — - must be checked before committing to the `I + 1/2` notation (the + correct but the def-site suppression could break on a mypy upgrade; pinned + by tests, with a notation-only fallback. Re-verified on **mypy 2.3** + (2026-08-05) — the quirk and the suppression both still behave as described, + so the upgrade risk has not materialised so far. **pyright is untested** — + must be checked before committing to the `Staggered[I] + 1/2` notation (the fallback functions are checker-agnostic). 2. **Float-literal staggering offsets** are value-checked only at runtime (`I + 0.7`); decide literal notation vs. `HALF` token (§4.2) in the ADR. diff --git a/content/personal/havogt/dimension-generic-fields/mypy.ini b/content/personal/havogt/dimension-generic-fields/mypy.ini index fc31199..a36be59 100644 --- a/content/personal/havogt/dimension-generic-fields/mypy.ini +++ b/content/personal/havogt/dimension-generic-fields/mypy.ini @@ -2,7 +2,7 @@ # deliberately *without* the gt4py mypy plugin — the point of the prototype is # that the redesigned Dimension needs no plugin support. [mypy] -python_version = 3.10 +python_version = 3.12 strict_equality = True warn_redundant_casts = True warn_unused_ignores = True diff --git a/content/personal/havogt/dimension-generic-fields/static_checks.py b/content/personal/havogt/dimension-generic-fields/static_checks.py index 35d0ad7..5867f01 100644 --- a/content/personal/havogt/dimension-generic-fields/static_checks.py +++ b/content/personal/havogt/dimension-generic-fields/static_checks.py @@ -28,6 +28,7 @@ Dual, Field, Staggered, + dual, dual_operator, ) from typing_extensions import TypeVarTuple, Unpack, assert_type @@ -54,12 +55,13 @@ class K(Dimension, kind=DimensionKind.VERTICAL): ... def connectivity_types() -> None: assert_type(I + 1, Connectivity[I, I]) assert_type(I - 2, Connectivity[I, I]) - # half-integral offsets move to the dual grid ... - assert_type(I + 1 / 2, Connectivity[Staggered[I], I]) - assert_type(I - 1 / 2, Connectivity[Staggered[I], I]) + # A shift names the dimension the *result* lives on, so naming a staggered + # dimension consumes an unstaggered field ... + assert_type(Staggered[I] + 1 / 2, Connectivity[Staggered[I], I]) + assert_type(Staggered[I] - 3 / 2, Connectivity[Staggered[I], I]) # ... and back (involution: never `Staggered[Staggered[I]]`) - assert_type(Staggered[I] + 1 / 2, Connectivity[I, Staggered[I]]) - assert_type(Staggered[I] - 3 / 2, Connectivity[I, Staggered[I]]) + assert_type(I + 1 / 2, Connectivity[I, Staggered[I]]) + assert_type(I - 1 / 2, Connectivity[I, Staggered[I]]) # integer offsets on a staggered dimension stay on the staggered grid assert_type(Staggered[I] + 1, Connectivity[Staggered[I], Staggered[I]]) @@ -68,11 +70,11 @@ def connectivity_types() -> None: def staggering_roundtrip(a: Field[Dims[I], float]) -> None: - b = a(I + 1 / 2) + b = a(Staggered[I] + 1 / 2) assert_type(b, Field[Dims[Staggered[I]], float]) b_via_alias: Field[Dims[Ihalf], float] = b # the alias spelling is the same type - c = b(Staggered[I] + 1 / 2) + c = b(I + 1 / 2) assert_type(c, Field[Dims[I], float]) # ordinary shifts do not change the dimensions @@ -87,16 +89,16 @@ def staggering_roundtrip(a: Field[Dims[I], float]) -> None: def rank2_substitution(uv: Field[Dims[I, J], float]) -> None: assert_type(uv(I + 1), Field[Dims[I, J], float]) - assert_type(uv(J + 1 / 2), Field[Dims[I, Staggered[J]], float]) - assert_type(uv(I - 1 / 2), Field[Dims[Staggered[I], J], float]) + assert_type(uv(Staggered[J] + 1 / 2), Field[Dims[I, Staggered[J]], float]) + assert_type(uv(Staggered[I] - 1 / 2), Field[Dims[Staggered[I], J], float]) def rank3_substitution(w: Field[Dims[I, J, K], float]) -> None: - assert_type(w(K + 1 / 2), Field[Dims[I, J, Staggered[K]], float]) + assert_type(w(Staggered[K] + 1 / 2), Field[Dims[I, J, Staggered[K]], float]) assert_type(w(J + 2), Field[Dims[I, J, K], float]) # chaining: stagger in I, then in K assert_type( - w(I + 1 / 2)(K - 1 / 2), + w(Staggered[I] + 1 / 2)(Staggered[K] - 1 / 2), Field[Dims[Staggered[I], J, Staggered[K]], float], ) @@ -106,7 +108,7 @@ def rank3_substitution(w: Field[Dims[I, J, K], float]) -> None: def gradient_to_staggered(p: Field[Dims[I, J], float]) -> Field[Dims[Staggered[I], J], float]: """Pressure gradient at the staggered (velocity) points.""" - return p(I + 1 / 2) - p(I - 1 / 2) + return p(Staggered[I] + 1 / 2) - p(Staggered[I] - 1 / 2) # --- Dimension-generic operators --------------------------------------------- @@ -125,7 +127,9 @@ def to_staggered( f: Field[Dims[DimT], float], dim: type[DimT] ) -> Field[Dims[Staggered[DimT]], float]: """Staggering an operator that is *generic* in the dimension it staggers.""" - return f(dim + 1 / 2) + # The shift names the *result* grid, which here is the dual of `dim`; + # `dual` is the typed counterpart of gt4py's `flip_staggered`. + return f(dual(dim) + 1 / 2) def use_generic_operators( @@ -154,8 +158,8 @@ def avg_explicit(f: Field[Dims[DimT], float]) -> Field[Dims[Staggered[DimT]], fl def avg_explicit(f: Field[Any, float]) -> Field[Any, float]: # The body is dual-generic, so the dimension is only known as "the field's # dimension"; recover it at runtime (`Any`: precision lives in the overloads). - dim: Any = f.dims[0] - return f(dim - 1 / 2) + f(dim + 1 / 2) + to: Any = dual(f.dims[0]) + return f(to - 1 / 2) + f(to + 1 / 2) # (2) declarative: ONE natural signature with the `Dual[X]` marker; the overload @@ -167,8 +171,8 @@ def avg_explicit(f: Field[Any, float]) -> Field[Any, float]: @dual_operator def avg(f: Field[Dims[AnyDimT], float]) -> Field[Dims[Dual[AnyDimT]], float]: - dim: Any = f.dims[0] - return f(dim - 1 / 2) + f(dim + 1 / 2) + to: Any = dual(f.dims[0]) + return f(to - 1 / 2) + f(to + 1 / 2) # `Dual[X]` also works in *parameter* positions, e.g. a weight that already @@ -179,8 +183,8 @@ def avg(f: Field[Dims[AnyDimT], float]) -> Field[Dims[Dual[AnyDimT]], float]: def weighted_avg( a: Field[Dims[AnyDimT], float], weight: Field[Dims[Dual[AnyDimT]], float] ) -> Field[Dims[Dual[AnyDimT]], float]: - dim: Any = a.dims[0] - return weight * (a(dim + 1 / 2) + a(dim - 1 / 2)) + to: Any = dual(a.dims[0]) + return weight * (a(to + 1 / 2) + a(to - 1 / 2)) def use_avg(a: Field[Dims[I], float], b: Field[Dims[Staggered[I]], float]) -> None: diff --git a/content/personal/havogt/dimension-generic-fields/static_errors.py b/content/personal/havogt/dimension-generic-fields/static_errors.py index 65b214b..4b94928 100644 --- a/content/personal/havogt/dimension-generic-fields/static_errors.py +++ b/content/personal/havogt/dimension-generic-fields/static_errors.py @@ -28,7 +28,7 @@ class K(Dimension, kind=DimensionKind.VERTICAL): ... def staggered_is_not_unstaggered(a: Field[Dims[I], float]) -> None: - b: Field[Dims[I], float] = a(I + 1 / 2) # EXPECT-ERROR: result lives on Staggered[I] + b: Field[Dims[I], float] = a(Staggered[I] + 1 / 2) # EXPECT-ERROR: result lives on Staggered[I] c: Field[Dims[Staggered[I]], float] = a(I + 1) # EXPECT-ERROR: result lives on I del b, c @@ -53,8 +53,8 @@ def horizontal_only(f: Field[Dims[I], float]) -> Field[Dims[I], float]: def roundtrip_must_go_through_the_dual(a: Field[Dims[I], float]) -> None: - b = a(I + 1 / 2) - b(I + 1 / 2) # EXPECT-ERROR: b has no unstaggered I dimension anymore + b = a(Staggered[I] + 1 / 2) + b(Staggered[I] + 1 / 2) # EXPECT-ERROR: b is already on the staggered grid def doubly_staggered_dimensions_are_unrepresentable( diff --git a/content/personal/havogt/dimension-generic-fields/test_typed_dimensions.py b/content/personal/havogt/dimension-generic-fields/test_typed_dimensions.py index dbff73c..554e270 100644 --- a/content/personal/havogt/dimension-generic-fields/test_typed_dimensions.py +++ b/content/personal/havogt/dimension-generic-fields/test_typed_dimensions.py @@ -84,9 +84,9 @@ def test_staggering_is_an_involution_at_the_type_constructor_level(): def test_connectivity_construction(): conn = I + 1 assert (conn.domain_dim, conn.codomain, conn.offset) == (I, I, 1) - conn = I + 1 / 2 + conn = Staggered[I] + 1 / 2 assert (conn.domain_dim, conn.codomain, conn.offset) == (Staggered[I], I, 0.5) - conn = Staggered[I] - 1 / 2 + conn = I - 1 / 2 assert (conn.domain_dim, conn.codomain, conn.offset) == (I, Staggered[I], -0.5) with pytest.raises(ValueError, match="half-integral"): I + 0.7 # statically a (claimed) staggering shift; rejected at runtime @@ -108,11 +108,13 @@ def test_plain_shift(): def test_staggering_relabels_and_shifts(): a = make_ifield([0.0, 1.0, 2.0, 3.0]) - b = a(I + 1 / 2) # b[i] = a at position i + 1/2, i.e. a[i + 1] in index space + # the point of `b` at index i sits at position i - 1/2 and reads a half a + # cell up, i.e. a[i] in index space + b = a(Staggered[I] + 1 / 2) assert b.dimensions == (Staggered[I],) - np.testing.assert_array_equal(b.ndarray, [1.0, 2.0, 3.0, 0.0]) + np.testing.assert_array_equal(b.ndarray, [0.0, 1.0, 2.0, 3.0]) - c = b(Staggered[I] + 1 / 2) # back to the unstaggered grid + c = b(I + 1 / 2) # back to the unstaggered grid assert c.dimensions == (I,) # two +1/2 staggerings == one full shift np.testing.assert_array_equal(c.ndarray, a(I + 1).ndarray) @@ -120,37 +122,37 @@ def test_staggering_relabels_and_shifts(): def test_staggering_roundtrip_identity(): a = make_ifield([0.0, 1.0, 2.0, 3.0]) - back = a(I + 1 / 2)(Staggered[I] - 1 / 2) + back = a(Staggered[I] + 1 / 2)(I - 1 / 2) assert back.dimensions == (I,) np.testing.assert_array_equal(back.ndarray, a.ndarray) def avg(f: NdField) -> NdField: """Runtime twin of the dual-generic `avg` in `static_checks.py` (one body, both grids).""" - (dim,) = f.dims - return f(dim - 1 / 2) + f(dim + 1 / 2) + to = dual(f.dims[0]) # a shift names the result grid: the dual of the field's + return f(to - 1 / 2) + f(to + 1 / 2) def test_avg_to_staggered(): a = make_ifield([0.0, 1.0, 2.0, 3.0]) result = avg(a) assert result.dimensions == (Staggered[I],) - # at staggered point i+1/2: a[i] + a[i+1] (periodic) - np.testing.assert_array_equal(result.ndarray, [1.0, 3.0, 5.0, 3.0]) + # at staggered point i (position i-1/2): a[i-1] + a[i] (periodic) + np.testing.assert_array_equal(result.ndarray, [3.0, 1.0, 3.0, 5.0]) def test_avg_from_staggered(): g = NdField((Staggered[I],), np.asarray([0.0, 1.0, 2.0, 3.0])) result = avg(g) assert result.dimensions == (I,) - # at unstaggered point i: g[i-1] + g[i] (the staggered neighbors at i -+ 1/2) - np.testing.assert_array_equal(result.ndarray, [3.0, 1.0, 3.0, 5.0]) + # at unstaggered point i: g[i] + g[i+1] (the staggered neighbors at i -+ 1/2) + np.testing.assert_array_equal(result.ndarray, [1.0, 3.0, 5.0, 3.0]) def weighted_avg(a: NdField, w: NdField) -> NdField: """Runtime twin of the dual-generic `weighted_avg` in `static_checks.py`.""" - (dim,) = a.dims - return w * (a(dim + 1 / 2) + a(dim - 1 / 2)) + to = dual(a.dims[0]) + return w * (a(to + 1 / 2) + a(to - 1 / 2)) def test_weighted_avg(): @@ -158,8 +160,8 @@ def test_weighted_avg(): w = NdField((Staggered[I],), np.asarray([1.0, 2.0, 3.0, 4.0])) result = weighted_avg(a, w) assert result.dimensions == (Staggered[I],) - # at staggered point i+1/2: w[i] * (a[i+1] + a[i]) (periodic) - np.testing.assert_array_equal(result.ndarray, [1.0, 6.0, 15.0, 12.0]) + # at staggered point i (position i-1/2): w[i] * (a[i] + a[i-1]) (periodic) + np.testing.assert_array_equal(result.ndarray, [3.0, 2.0, 9.0, 20.0]) back = weighted_avg(result, a) assert back.dimensions == (I,) @@ -173,21 +175,21 @@ def test_avg_roundtrip_dims(): def test_c_grid_gradient(): a = make_ifield([0.0, 1.0, 4.0, 9.0]) - grad = a(I + 1 / 2) - a(I - 1 / 2) # grad[i+1/2] = a[i+1] - a[i] + grad = a(Staggered[I] + 1 / 2) - a(Staggered[I] - 1 / 2) # grad[i] = a[i] - a[i-1] assert grad.dimensions == (Staggered[I],) - np.testing.assert_array_equal(grad.ndarray, [1.0, 3.0, 5.0, -9.0]) + np.testing.assert_array_equal(grad.ndarray, [-9.0, 1.0, 3.0, 5.0]) def test_rank2_substitution_at_runtime(): a = NdField((I, J), np.zeros((3, 4))) - assert a(J + 1 / 2).dimensions == (I, Staggered[J]) - assert a(I - 1 / 2).dimensions == (Staggered[I], J) + assert a(Staggered[J] + 1 / 2).dimensions == (I, Staggered[J]) + assert a(Staggered[I] - 1 / 2).dimensions == (Staggered[I], J) def test_mixing_dual_grids_raises(): a = make_ifield([0.0, 1.0]) with pytest.raises(ValueError, match="different domains"): - a + a(I + 1 / 2) + a + a(Staggered[I] + 1 / 2) def test_shift_along_missing_dimension_raises(): diff --git a/content/personal/havogt/dimension-generic-fields/typed_dimensions.py b/content/personal/havogt/dimension-generic-fields/typed_dimensions.py index 1962df4..62197b9 100644 --- a/content/personal/havogt/dimension-generic-fields/typed_dimensions.py +++ b/content/personal/havogt/dimension-generic-fields/typed_dimensions.py @@ -20,10 +20,12 @@ the same value-level behavior dimension *instances* have today (``I + 1`` builds a connectivity, ``str(I)``, equality by name, ...). 3. **Staggering as a type-level function**: ``Staggered[I]`` is the dual - (staggered) counterpart of ``I``. ``I + 1/2`` is typed - ``Connectivity[Staggered[I], I]`` and ``Staggered[I] + 1/2`` is typed - ``Connectivity[I, Staggered[I]]`` — the dual-grid involution is encoded in a - pair of overloads, so no doubly-staggered types ever arise. + (staggered) counterpart of ``I``. A shift names the dimension the *result* + lives on, so ``Staggered[I] + 1/2`` is typed + ``Connectivity[Staggered[I], I]`` (applied to a field on ``I``) and + ``I + 1/2`` is typed ``Connectivity[I, Staggered[I]]`` — the dual-grid + involution is encoded in a pair of overloads, so no doubly-staggered types + ever arise. 4. **Shift typing as positional substitution**: ``field(conn)`` replaces exactly the matching dimension in the field's ``Dims[...]``, expressed with rank-bounded overloads (rank <= 3 here; generated code can extend this). @@ -98,9 +100,17 @@ def __hash__(cls) -> int: # The fractional-offset overloads encode staggering: # - integer offsets stay on the same grid, # - half-integral offsets move to the dual grid. + # + # A shift is precomposition: for `f: I -> value` and `phi: D -> I`, applying + # `phi` yields `f . phi`, a field on `D`. So a connectivity's domain is the + # *result's* dimension and its codomain the field's, and the index + # expression is written in the result's index space while evaluating to an + # index of the field's. Hence `D + 1/2` consumes a field on `Staggered[D]` + # and produces one on `D`. + # # The `Staggered[D]` overload must come before the generic `D` overload so - # that shifting *from* a staggered dimension resolves back to the base - # dimension (involution) instead of producing `Staggered[Staggered[D]]`. + # that a shift naming a staggered dimension consumes the base dimension + # (involution) instead of demanding a `Staggered[Staggered[D]]` field. # # The `type: ignore[misc]` comments silence a *definition-site* mypy # restriction ("Self argument missing for a non-static method"): mypy does @@ -112,11 +122,11 @@ def __add__(cls: type[AnyD], offset: int) -> Connectivity[AnyD, AnyD]: ... # ty @overload def __add__( # type: ignore[misc] cls: type[Staggered[D]], offset: float - ) -> Connectivity[D, Staggered[D]]: ... + ) -> Connectivity[Staggered[D], D]: ... @overload def __add__( # type: ignore[misc] cls: type[D], offset: float - ) -> Connectivity[Staggered[D], D]: ... + ) -> Connectivity[D, Staggered[D]]: ... def __add__(cls, offset: int | float) -> Connectivity[Any, Any]: if isinstance(offset, int): return Connectivity(cls, cls, offset) @@ -124,18 +134,18 @@ def __add__(cls, offset: int | float) -> Connectivity[Any, Any]: raise ValueError( f"Only integral or half-integral offsets are supported, got '{offset}'." ) - return Connectivity(dual(cls), cls, offset) + return Connectivity(cls, dual(cls), offset) @overload def __sub__(cls: type[AnyD], offset: int) -> Connectivity[AnyD, AnyD]: ... # type: ignore[misc] @overload def __sub__( # type: ignore[misc] cls: type[Staggered[D]], offset: float - ) -> Connectivity[D, Staggered[D]]: ... + ) -> Connectivity[Staggered[D], D]: ... @overload def __sub__( # type: ignore[misc] cls: type[D], offset: float - ) -> Connectivity[Staggered[D], D]: ... + ) -> Connectivity[D, Staggered[D]]: ... def __sub__(cls, offset: int | float) -> Connectivity[Any, Any]: return cast("Connectivity[Any, Any]", cast(Any, cls) + (-offset)) @@ -231,11 +241,21 @@ def _is_staggered(dim: DimensionMeta) -> bool: return issubclass(dim, Staggered) -def dual(dim: DimensionMeta) -> DimensionMeta: - """Return the dual-grid counterpart of a dimension class (involution).""" +@overload +def dual(dim: type[Staggered[D]]) -> type[D]: ... +@overload +def dual(dim: type[D]) -> type[Staggered[D]]: ... +def dual(dim: Any) -> Any: + """ + Return the dual-grid counterpart of a dimension class (involution). + + The typed counterpart of gt4py's value-level `flip_staggered`. A + dual-generic operator needs it because a shift names the *result* grid, + which for such an operator is the dual of the field's own grid. + """ if _is_staggered(dim): - return cast(DimensionMeta, dim.base) # type: ignore[attr-defined] # staggered dims carry `base` - return cast(DimensionMeta, Staggered[dim]) # type: ignore[valid-type] # value-level use + return dim.base + return Staggered[dim] @dataclasses.dataclass(frozen=True) @@ -391,11 +411,13 @@ class NdField(Generic[DimsT, DT]): """ Runtime implementation of the `Field` protocol (periodic boundaries). - Conventions (position of staggered point ``i`` is ``i + 1/2``): - given ``b = a(conn)``, ``b[p] = a[p + conn.offset]`` where positions are - grid points of the respective dimensions. In index space this is a shift by - ``ceil(offset)`` when reading from an unstaggered dimension and by - ``floor(offset)`` when reading from a staggered one. + Conventions follow gt4py as shipped (ADR 0026): staggered point ``i`` sits + at position ``i - 1/2``, so ``Staggered[I](0)`` is the point just *below* + ``I(0)``. Given ``b = a(conn)``, the point of ``b`` at index ``p`` reads + ``a`` at position ``pos(p) + frac(conn.offset)``. In index space that is a + shift by ``ceil(offset)`` when the shift names an unstaggered dimension and + by ``floor(offset)`` when it names a staggered one — keyed on the dimension + the *result* lives on, not the one the field is on. """ dimensions: tuple[DimensionMeta, ...] @@ -415,7 +437,7 @@ def __call__(self, conn: Connectivity[Any, Any]) -> NdField[Any, DT]: f"Field defined on '{self.dimensions}' has no dimension '{conn.codomain}'." ) axis = self.dimensions.index(conn.codomain) - if issubclass(conn.codomain, Staggered): + if issubclass(conn.domain_dim, Staggered): index_shift = math.floor(conn.offset) else: index_shift = math.ceil(conn.offset)