feat[next] high-level field storage API

src/gt4py/next/iterator/embedded.py

@@ -138,6 +138,8 @@ def mapped_index(self, primary: IntIndex, neighbor_idx: IntIndex) -> IntIndex:
 #: A ``None`` position flags invalid not-a-neighbor results in neighbor-table lookups
 MaybePosition: TypeAlias = Optional[Position]
 
+NamedFieldIndices: TypeAlias = Mapping[Tag, FieldIndex | SparsePositionEntry]
+
 
 def is_int_index(p: Any) -> TypeGuard[IntIndex]:
     return isinstance(p, (int, np.integer))
@@ -692,19 +694,18 @@ def _get_axes(
 
 
 def _single_vertical_idx(
-    indices: FieldIndices, column_axis_idx: int, column_index: IntIndex
-) -> tuple[FieldIndex, ...]:
-    transformed = tuple(
-        index if i != column_axis_idx else column_index for i, index in enumerate(indices)
-    )
+    indices: NamedFieldIndices, column_axis: Tag, column_index: IntIndex
+) -> NamedFieldIndices:
+    transformed = {
+        axis: (index if axis != column_axis else column_index) for axis, index in indices.items()
+    }
     return transformed
 
 
 @overload
 def _make_tuple(
     field_or_tuple: tuple[tuple | common.Field, ...],  # arbitrary nesting of tuples of Field
-    indices: FieldIndices,
-    *,
+    named_indices: NamedFieldIndices,
     column_axis: Tag,
 ) -> tuple[tuple | Column, ...]:
     ...
@@ -713,28 +714,32 @@ def _make_tuple(
 @overload
 def _make_tuple(
     field_or_tuple: tuple[tuple | Field, ...],  # arbitrary nesting of tuples of Field
-    indices: FieldIndices,
-    *,
+    named_indices: NamedFieldIndices,
     column_axis: Literal[None] = None,
 ) -> tuple[tuple | npt.DTypeLike, ...]:  # arbitrary nesting
     ...
 
 
 @overload
-def _make_tuple(field_or_tuple: common.Field, indices: FieldIndices, *, column_axis: Tag) -> Column:
+def _make_tuple(
+    field_or_tuple: common.Field, named_indices: NamedFieldIndices, *, column_axis: Tag
+) -> Column:
     ...
 
 
 @overload
 def _make_tuple(
-    field_or_tuple: common.Field, indices: FieldIndices, *, column_axis: Literal[None] = None
+    field_or_tuple: common.Field,
+    named_indices: NamedFieldIndices,
+    *,
+    column_axis: Literal[None] = None,
 ) -> npt.DTypeLike:
     ...
 
 
 def _make_tuple(
     field_or_tuple: common.Field | tuple[tuple | common.Field, ...],
-    indices: FieldIndices,
+    named_indices: NamedFieldIndices,
     *,
     column_axis: Optional[Tag] = None,
 ) -> Column | npt.DTypeLike | tuple[tuple | Column | npt.DTypeLike, ...]:
@@ -743,11 +748,8 @@ def _make_tuple(
         if column_axis is not None:
             assert column_range
             # construct a Column of tuples
-            column_axis_idx = _axis_idx(_get_axes(field_or_tuple), column_axis)
-            if column_axis_idx is None:
-                column_axis_idx = -1  # field doesn't have the column index, e.g. ContantField
             first = tuple(
-                _make_tuple(f, _single_vertical_idx(indices, column_axis_idx, column_range.start))
+                _make_tuple(f, _single_vertical_idx(named_indices, column_axis, column_range.start))
                 for f in field_or_tuple
             )
             col = Column(
@@ -756,14 +758,14 @@ def _make_tuple(
             col[0] = first
             for i in column_range[1:]:
                 col[i] = tuple(
-                    _make_tuple(f, _single_vertical_idx(indices, column_axis_idx, i))
+                    _make_tuple(f, _single_vertical_idx(named_indices, column_axis, i))
                     for f in field_or_tuple
                 )
             return col
         else:
-            return tuple(_make_tuple(f, indices) for f in field_or_tuple)
+            return tuple(_make_tuple(f, named_indices) for f in field_or_tuple)
     else:
-        data = field_or_tuple.field_getitem(indices)
+        data = field_or_tuple.field_getitem(named_indices)
         if column_axis is not None:
             # wraps a vertical slice of an input field into a `Column`
             assert column_range is not None
@@ -822,13 +824,10 @@ def deref(self) -> Any:
             slice_column[self.column_axis] = range(k_pos, k_pos + len(column_range))
 
         assert _is_concrete_position(shifted_pos)
-        ordered_indices = get_ordered_indices(
-            axes,
-            {**shifted_pos, **slice_column},
-        )
+        position = {**shifted_pos, **slice_column}
         return _make_tuple(
             self.field,
-            ordered_indices,
+            position,
             column_axis=self.column_axis,
         )
 
@@ -949,9 +948,7 @@ def _is_sparse_position_entry(
     return isinstance(pos, list)
 
 
-def get_ordered_indices(
-    axes: Iterable[Axis], pos: Mapping[Tag, FieldIndex | SparsePositionEntry]
-) -> tuple[FieldIndex, ...]:
+def get_ordered_indices(axes: Iterable[Axis], pos: NamedFieldIndices) -> tuple[FieldIndex, ...]:
     res: list[FieldIndex] = []
     sparse_position_tracker: dict[Tag, int] = {}
     for axis in axes:
@@ -1069,7 +1066,8 @@ def __init__(self, axis: common.Dimension, dtype: npt.DTypeLike) -> None:
         self.axis = axis
         self.dtype = np.dtype(dtype)
 
-    def field_getitem(self, index: FieldIndexOrIndices) -> Any:
+    def field_getitem(self, named_indices: NamedFieldIndices) -> Any:
+        index = get_ordered_indices(self.__gt_dims__, named_indices)
         if isinstance(index, int):
             return self.dtype.type(index)
         else:
@@ -1090,7 +1088,7 @@ def __init__(self, value: Any, dtype: npt.DTypeLike):
         self.value = value
         self.dtype = np.dtype(dtype).type
 
-    def field_getitem(self, _: FieldIndexOrIndices) -> Any:
+    def field_getitem(self, _: NamedFieldIndices) -> Any:
         return self.dtype(self.value)
 
     @property
@@ -1311,7 +1309,8 @@ def __init__(self, data):
         axeses = _get_axeses(data)
         self.__gt_dims__ = axeses[0]
 
-    def field_getitem(self, indices):
+    def field_getitem(self, named_indices: NamedFieldIndices):
+        indices = get_ordered_indices(self.__gt_dims__, named_indices)
         return _build_tuple_result(self.data, indices)
 
     def field_setitem(self, indices, value):
@@ -1417,15 +1416,13 @@ def _closure_runner():
 
                 if column is None:
                     assert _is_concrete_position(pos)
-                    ordered_indices = get_ordered_indices(_get_axes(out), pos)
-                    out.field_setitem(ordered_indices, res)
+                    out.field_setitem(pos, res)
                 else:
                     col_pos = pos.copy()
                     for k in column.col_range:
                         col_pos[column.axis] = k
                         assert _is_concrete_position(col_pos)
-                        ordered_indices = get_ordered_indices(_get_axes(out), col_pos)
-                        out.field_setitem(ordered_indices, res[k])
+                        out.field_setitem(col_pos, res[k])
 
         ctx = cvars.copy_context()
         ctx.run(_closure_runner)