`create_kernel` API Update

To reduce the kwargs hell, we introduced dataclasses, which handle the settings of create_kernel and similar functions. In addition we introduced type-hints for the API functions to increase usability and simplify development.

pystencils/kernelcreation.py

@@ -43,6 +43,8 @@ class CreateKernelConfig:
         gpu_indexing_params: dict with indexing parameters (constructor parameters of indexing class)
                              e.g. for 'block' one can specify '{'block_size': (20, 20, 10) }'
         cpu_prepend_optimizations: list of extra optimizations to perform first on the AST
+        index_fields: list of index fields, i.e. 1D fields with struct data type
+        coordinate_names: name of the coordinate fields in the struct data type
     """
     target: Target = Target.CPU
     function_name: str = 'kernel'
@@ -56,11 +58,13 @@ class CreateKernelConfig:
     omp_single_loop: bool = True
     gpu_indexing: str = 'block'
     gpu_indexing_params: MappingProxyType = field(default=MappingProxyType({}))
-    use_textures_for_interpolation: bool = True
+    use_textures_for_interpolation: bool = True  # TODO docstring
     cpu_prepend_optimizations: List[Callable] = field(default_factory=list)
     use_auto_for_assignments: bool = False
     opencl_queue = None  # TODO docstring, typing
     opencl_ctx = None  # TODO docstring, typing
+    index_fields: List[Field] = None
+    coordinate_names: Tuple[str, ...] = ('x', 'y', 'z')
 
 
 def create_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCollection], *,
@@ -112,6 +116,43 @@ def create_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCol
     if isinstance(assignments, Assignment):
         assignments = [assignments]
     assert assignments, "Assignments must not be empty!"
+
+    if config.index_fields:
+        return create_indexed_kernel(assignments, config=config)
+    else:
+        return create_domain_kernel(assignments, config=config)
+
+
+def create_domain_kernel(assignments: List[Assignment], *, config: CreateKernelConfig):
+    """
+    Creates abstract syntax tree (AST) of kernel, using a list of update equations.
+
+    Args:
+        assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection`
+        config: CreateKernelConfig which includes the needed configuration
+
+    Returns:
+        abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
+        can be compiled with through its 'compile()' member
+
+    Example:
+        >>> import pystencils as ps
+        >>> import numpy as np
+        >>> s, d = ps.fields('s, d: [2D]')
+        >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
+        >>> config = ps.CreateKernelConfig(cpu_openmp=True)
+        >>> kernel_ast = ps.kernelcreation.create_domain_kernel([assignment], config=config)
+        >>> kernel = kernel_ast.compile()
+        >>> d_arr = np.zeros([5, 5])
+        >>> kernel(d=d_arr, s=np.ones([5, 5]))
+        >>> d_arr
+        array([[0., 0., 0., 0., 0.],
+               [0., 4., 4., 4., 0.],
+               [0., 4., 4., 4., 0.],
+               [0., 4., 4., 4., 0.],
+               [0., 0., 0., 0., 0.]])
+    """
+    # ----  Normalizing parameters
     split_groups = ()
     if isinstance(assignments, AssignmentCollection):
         if 'split_groups' in assignments.simplification_hints:
@@ -167,7 +208,7 @@ def create_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCol
             ast._backend = 'opencl'
         return ast
     else:
-        raise ValueError(f"Unknown {config.target=}. Has to be one of 'cpu', 'gpu' or 'llvm' ")
+        raise NotImplementedError(f'{config.target} is not supported by `create_domain_kernel`')
 
     if config.use_auto_for_assignments:
         for a in ast.atoms(SympyAssignment):
@@ -176,17 +217,7 @@ def create_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCol
     return ast
 
 
-def create_indexed_kernel(assignments: Union[Assignment, List[Assignment], AssignmentCollection],
-                          index_fields,
-                          target='cpu',
-                          data_type="double",
-                          coordinate_names=('x', 'y', 'z'),
-                          cpu_openmp=True,
-                          gpu_indexing='block',
-                          gpu_indexing_params=MappingProxyType({}),
-                          use_textures_for_interpolation=True,
-                          opencl_queue=None,
-                          opencl_ctx=None):
+def create_indexed_kernel(assignments: List[Assignment], *, config: CreateKernelConfig):
     """
     Similar to :func:`create_kernel`, but here not all cells of a field are updated but only cells with
     coordinates which are stored in an index field. This traversal method can e.g. be used for boundary handling.
@@ -196,8 +227,13 @@ def create_indexed_kernel(assignments: Union[Assignment, List[Assignment], Assig
     'coordinate_names' parameter. The struct can have also other fields that can be read and written in the kernel, for
     example boundary parameters.
 
-    index_fields: list of index fields, i.e. 1D fields with struct data type
-    coordinate_names: name of the coordinate fields in the struct data type
+    Args:
+        assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection`
+        config: CreateKernelConfig which includes the needed configuration
+
+    Returns:
+        abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
+        can be compiled with through its 'compile()' member
 
     Example:
         >>> import pystencils as ps
@@ -211,7 +247,8 @@ def create_indexed_kernel(assignments: Union[Assignment, List[Assignment], Assig
         >>> # Additional values  stored in index field can be accessed in the kernel as well
         >>> s, d = ps.fields('s, d: [2D]')
         >>> assignment = ps.Assignment(d[0,0], 2 * s[0, 1] + 2 * s[1, 0] + idx_field('val'))
-        >>> kernel_ast = ps.create_indexed_kernel(assignment, [idx_field], coordinate_names=('x', 'y'))
+        >>> config = ps.CreateKernelConfig(index_fields=[idx_field], coordinate_names=('x', 'y'))
+        >>> kernel_ast = ps.create_indexed_kernel([assignment], config=config)
         >>> kernel = kernel_ast.compile()
         >>> d_arr = np.zeros([5, 5])
         >>> kernel(s=np.ones([5, 5]), d=d_arr, idx=index_arr)
@@ -222,37 +259,31 @@ def create_indexed_kernel(assignments: Union[Assignment, List[Assignment], Assig
                [0. , 0. , 0. , 4.3, 0. ],
                [0. , 0. , 0. , 0. , 0. ]])
     """
-    if isinstance(assignments, Assignment):
-        assignments = [assignments]
-    elif isinstance(assignments, AssignmentCollection):
-        assignments = assignments.all_assignments
-    if target == 'cpu':
+    if config.target == Target.CPU:
         from pystencils.cpu import add_openmp, create_indexed_kernel
-        ast = create_indexed_kernel(assignments, index_fields=index_fields, type_info=data_type,
-                                    coordinate_names=coordinate_names)
-        if cpu_openmp:
-            add_openmp(ast, num_threads=cpu_openmp)
+        ast = create_indexed_kernel(assignments, index_fields=config.index_fields, type_info=config.data_type,
+                                    coordinate_names=config.coordinate_names)
+        if config.cpu_openmp:
+            add_openmp(ast, num_threads=config.cpu_openmp)
         return ast
-    elif target == 'llvm':
-        raise NotImplementedError("Indexed kernels are not yet supported in LLVM backend")
-    elif target == 'gpu' or target == 'opencl':
+    elif config.target == Target.GPU or config.target == Target.OPENCL:
         from pystencils.gpucuda import created_indexed_cuda_kernel
-        idx_creator = indexing_creator_from_params(gpu_indexing, gpu_indexing_params)
+        idx_creator = indexing_creator_from_params(config.gpu_indexing, config.gpu_indexing_params)
         ast = created_indexed_cuda_kernel(assignments,
-                                          index_fields,
-                                          type_info=data_type,
-                                          coordinate_names=coordinate_names,
+                                          config.index_fields,
+                                          type_info=config.data_type,
+                                          coordinate_names=config.coordinate_names,
                                           indexing_creator=idx_creator,
-                                          use_textures_for_interpolation=use_textures_for_interpolation)
-        if target == 'opencl':
+                                          use_textures_for_interpolation=config.use_textures_for_interpolation)
+        if config.target == Target.OPENCL:
             from pystencils.opencl.opencljit import make_python_function
             ast._backend = 'opencl'
-            ast.compile = functools.partial(make_python_function, ast, opencl_queue, opencl_ctx)
+            ast.compile = functools.partial(make_python_function, ast, config.opencl_queue, config.opencl_ctx)
             ast._target = 'opencl'
             ast._backend = 'opencl'
         return ast
     else:
-        raise ValueError(f"Unknown target {target}. Has to be either 'cpu' or 'gpu'")
+        raise NotImplementedError(f'Indexed kernels are not yet supported in {config.target} backend')
 
 
 def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=False, **kwargs):