extend dyn-type support for fields. Extend user guide on typing.

docs/source/api/sympyextensions.rst

@@ -71,6 +71,9 @@ Typed Expressions
 .. autoclass:: pystencils.DynamicType
     :members:
 
+.. autoclass:: pystencils.sympyextensions.typed_sympy.TypeCast
+    :members:
+
 .. autoclass:: pystencils.sympyextensions.tcast
 
 

docs/source/reference/WorkingWithTypes.md

 ---
 file_format: mystnb
 kernelspec:
-    name: python3
+  display_name: Python 3 (ipykernel)
+  language: python
+  name: python3
+mystnb:
+  execution_mode: cache
 ---
 
 # Working with Data Types
@@ -13,54 +17,129 @@ Individual fields and symbols,
 single subexpressions,
 or the entire kernel.
 
-```{code-cell}
+```{code-cell} ipython3
 :tags: [remove-cell]
 import pystencils as ps
+import sympy as sp
 ```
 
-## Understanding the pystencils Type Inference Algorithm
+## Setting the Types of Fields and Symbols
 
-To correctly apply varying data types to pystencils kernels, it is important to understand
-how pystencils computes and propagates the data types of expressions.
-These are the rules by which untyped symbols learn their data type:
-
- - All *free symbols* (that is, symbols not defined by an assignment in the kernel) receive the
-  *default data type* from the code generator configuration.
- - All symbols defined using a *constant expression* also receive the default data type.
- - All other symbols receive the data type computed for the right-hand side expression of their
-   defining assignment.
-
-To determine the type of right hand-side expressions, pystencils looks for any subexpressions
-nested inside them which already have a known type.
-These might be symbols whose type was already determined,
-or expressions with a fixed type, such as field accesses (which get the element type of their field),
-explicitly typed symbols, or type casts (see [](explicit_expression_types)).
-
-:::{attention}
-Expressions must always have a unique and unambiguous data type,
-and pystencils will not introduce any implicit casts.
-If pystencils finds subexpressions with conflicting types inside one expression,
-type inference will fail and `create_kernel` will raise an error.
-:::
+### Untyped Symbols
+
+Symbols used inside a kernel are most commonly created using
+{any}`sp.symbols <sympy.core.symbol.symbols>` or
+{any}`sp.Symbol <sympy.core.symbol.Symbol>`.
+These symbols are *untyped*; they will receive a type during code generation
+according to these rules:
+ - Free untyped symbols (i.e. symbols not defined by an assignment inside the kernel) receive the 
+   {any}`default data type <CreateKernelConfig.default_dtype>` specified in the code generator configuration.
+ - Bound untyped symbols (i.e. symbols that *are* defined in an assignment)
+   receive the data type that was computed for the right-hand side expression of their defining assignment.
 
-Through these rules, there are multiple ways to modify the data types used inside a kernel.
-These are highlighted in the following sections.
+If you are working on kernels with homogenous data types, using untyped symbols will mostly be enough.
 
-## Changing the Default Data Type
+### Explicitly Typed Symbols and Fields
 
-The *default data type* is the fallback type assigned by pystencils to all free symbols and symbols with constant
-definitions.
-It can be modified by setting the {any}`default_dtype <CreateKernelConfig.default_dtype>` option
-of the code generator configuration:
+If you need more control over the data types in (parts of) your kernel,
+you will have to explicitly specify them.
+To set an explicit data type for a symbol, use the {any}`TypedSymbol` class of pystencils:
 
 ```{code-cell} ipython3
-cfg = ps.CreateKernelConfig()
-cfg.default_dtype = "float32"
+x_typed = ps.TypedSymbol("x", "uint32")
+x_typed, str(x_typed.dtype)
 ```
 
-:::{admonition} Developers To Do
-Fields should use DynamicType by default!
-:::
+You can set a `TypedSymbol` to any data type provided by [the type system](#page_type_system),
+which will then be enforced by the code generator.
+
+The same holds for fields:
+When creating fields through the {any}`fields <pystencils.field.fields>` function,
+add the type to the descriptor string; for instance:
+
+```{code-cell} ipython3
+f, g = ps.fields("f(1), g(3): float32[3D]")
+str(f.dtype), str(g.dtype)
+```
+
+When using `Field.create_generic` or `Field.create_fixed_size`, on the other hand,
+you can set the data type via the `dtype` keyword argument.
+
+### Dynamically Typed Symbols and Fields
+
+Apart from explicitly setting data types,
+`TypedSymbol`s and fields can also receive a *dynamic data type* (see {any}`DynamicType`).
+There are two options:
+ - Symbols or fields annotated with {any}`DynamicType.NUMERIC_TYPE` will always receive
+   the {any}`default numeric type <CreateKernelConfig.default_dtype>` configured for the
+   code generator.
+   This is the default setting for fields
+   created through `fields`, `Field.create_generic` or `Field.create_fixed_size`.
+ - When annotated with {any}`DynamicType.INDEX_TYPE`, on the other hand, they will receive
+   the {any}`index data type <CreateKernelConfig.index_dtype>` configured for the kernel.
+
+Using dynamic typing, you can enforce symbols to receive either the standard numeric or
+index type without explicitly stating it, such that your kernel definition becomes
+independent from the code generator configuration.
+
+## Mixing Types Inside Expressions
+
+Pystencils enforces that all symbols, constants, and fields occuring inside an expression
+have the same data type.
+The code generator will never introduce implicit casts --
+if any type conflicts arise, it will terminate with an error.
+
+Still, there are cases where you want to combine subexpressions of different types;
+maybe you need to compute geometric information from loop counters or other integers,
+or you are doing mixed-precision numerical computations.
+In these cases, you might have to
+ 1. Introduce explicit type casts when values move from one type context to another;
+ 2. Annotate expressions with a specific data type to ensure computations are performed in that type.
+
+(type_casts)=
+### Type Casts
+
+Type casts can be introduced into kernels using the {any}`tcast` symbolic function.
+It takes an expression and a data type, which is either an explicit type (see [the type system](#page_type_system))
+or a dynamic type ({any}`DynamicType`):
+
+```{code-cell} ipython3
+x, y = sp.symbols("x, y")
+expr1 = ps.tcast(x, "float32")
+expr2 = ps.tcast(3 + y, ps.DynamicType.INDEX_TYPE)
+
+str(expr1.dtype), str(expr2.dtype)
+```
+
+When a type cast occurs, pystencils will compute the type of its argument independently
+and then introduce a runtime cast to the target type.
+That target type must comply with the type computed for the outer expression,
+which the cast is embedded in.
 
 (explicit_expression_types)=
-## Setting Explicit Types for Expressions
+### Setting Explicit Types for Expressions
+
+
+
+## Understanding the pystencils Type Inference System
+
+To correctly apply varying data types to pystencils kernels, it is important to understand
+how pystencils computes and propagates the data types of symbols and expressions.
+
+Type inference happens on the level of assignments.
+For each assignment $x := \mathrm{calc}(y_1, \dots, y_n)$,
+the system first attempts to compute a *unique* type for the right-hand side (RHS) $\mathrm{calc}(y_1, \dots, y_n)$.
+It searches for any subexpression inside the RHS for which a type is already known --
+these might be typed symbols
+(whose types are either set explicitly by the user,
+or have been determined from their defining assignment),
+field accesses,
+or explicitly typed expressions.
+It then attempts to apply that data type to the entire expression.
+If type conflicts occur, the process fails and the code generator raises an error.
+Otherwise, the resulting type is assigned to the left-hand side symbol $x$.
+
+:::{admonition} Developer's To Do
+It would be great to illustrate this using a GraphViz-plot of an AST,
+with nodes colored according to their data types
+:::

src/pystencils/field.py

@@ -1123,11 +1123,30 @@ def fields(
     layout=None,
     field_type=FieldType.GENERIC,
     **kwargs,
-) -> Union[Field, List[Field]]:
+) -> Field | list[Field]:
     """Creates pystencils fields from a string description.
 
+    The description must be a string of the form
+    ``"name(index-shape) [name(index-shape) ...] <data-type>[<dimension-or-shape>]"``,
+    where:
+
+    - ``name`` is the name of the respective field
+    - ``(index-shape)`` is a tuple of integers describing the shape of the tensor on each field node
+      (can be omitted for scalar fields)
+    - ``<data-type>`` is the numerical data type of the field's entries;
+      this can be any type parseable by `create_type`,
+      as well as ``dyn`` for `DynamicType.NUMERIC_TYPE`
+      and ``dynidx`` for `DynamicType.INDEX_TYPE`.
+    - ``<dimension-or-shape>`` can be a dimensionality (e.g. ``1D``, ``2D``, ``3D``)
+      or a tuple of integers defining the spatial shape of the field.
+
     Examples:
-        Create a 2D scalar and vector field:
+        Create a 3D scalar field of default numeric type:
+            >>> f = fields("f(1): [2D]")
+            >>> str(f.dtype)
+            'DynamicType.NUMERIC_TYPE'
+
+        Create a 2D scalar and vector field of 64-bit float type:
             >>> s, v = fields("s, v(2): double[2D]")
             >>> assert s.spatial_dimensions == 2 and s.index_dimensions == 0
             >>> assert (v.spatial_dimensions, v.index_dimensions, v.index_shape) == (2, 1, (2,))
@@ -1153,6 +1172,9 @@ def fields(
             >>> f = fields("pdfs(19) : float32[3D]", layout='fzyx')
             >>> f.layout
             (2, 1, 0)
+
+    Returns:
+        Sequence of fields created from the description
     """
     result = []
     if description:
@@ -1429,7 +1451,10 @@ def _parse_description(description):
                 data_type_str = data_type_str.lower().strip()
 
             if data_type_str:
-                dtype = create_type(data_type_str)
+                match data_type_str:
+                    case "dyn": dtype = DynamicType.NUMERIC_TYPE
+                    case "dynidx": dtype = DynamicType.INDEX_TYPE
+                    case _: dtype = create_type(data_type_str)
             else:
                 dtype = DynamicType.NUMERIC_TYPE
 

tests/frontend/test_field.py

@@ -84,6 +84,12 @@ def test_field_description_parsing():
     assert f.index_shape == (1,)
     assert g.index_shape == (3,)
 
+    f = fields("f: dyn[3D]")
+    assert f.dtype == DynamicType.NUMERIC_TYPE
+
+    idx = fields("idx: dynidx[3D]")
+    assert idx.dtype == DynamicType.INDEX_TYPE
+
     h = fields("h: float32[3D]")
 
     assert h.index_shape == ()