diff --git a/.gitignore b/.gitignore
index d50d80ff99631648f03dd5ef6d62664e1643e697..41879ab794f8d9527543dfa833f6206672e82326 100644
--- a/.gitignore
+++ b/.gitignore
@@ -20,7 +20,9 @@ pystencils/boundaries/createindexlistcython.*.so
pystencils_tests/tmp
pystencils_tests/kerncraft_inputs/.2d-5pt.c_kerncraft/
pystencils_tests/kerncraft_inputs/.3d-7pt.c_kerncraft/
+report.xml
+coverage_report/
# macOS
-**/.DS_Store
\ No newline at end of file
+**/.DS_Store
diff --git a/pystencils/__init__.py b/pystencils/__init__.py
index 3dc1abcd4d4956ce23bda209b540c2d73f0ea770..a10acb8f6e721cfa71fff1224379114b06f7e933 100644
--- a/pystencils/__init__.py
+++ b/pystencils/__init__.py
@@ -7,7 +7,7 @@ from .data_types import TypedSymbol
from .datahandling import create_data_handling
from .display_utils import get_code_obj, get_code_str, show_code, to_dot
from .field import Field, FieldType, fields
-from .kernel_decorator import kernel
+from .kernel_decorator import kernel, kernel_config
from .kernelcreation import (
CreateKernelConfig, create_domain_kernel, create_indexed_kernel, create_kernel, create_staggered_kernel)
from .simp import AssignmentCollection
@@ -34,7 +34,7 @@ __all__ = ['Field', 'FieldType', 'fields',
'assignment_from_stencil',
'SymbolCreator',
'create_data_handling',
- 'kernel',
+ 'kernel', 'kernel_config',
'x_', 'y_', 'z_',
'x_staggered', 'y_staggered', 'z_staggered',
'x_vector', 'x_staggered_vector',
diff --git a/pystencils/kernel_decorator.py b/pystencils/kernel_decorator.py
index 92974a40f3dc9473eb3c586f11167c30cecdd968..19938e1507fd82970f15a2c2926a01775d18519d 100644
--- a/pystencils/kernel_decorator.py
+++ b/pystencils/kernel_decorator.py
@@ -1,38 +1,25 @@
import ast
import inspect
import textwrap
-from typing import Callable, Union, List, Dict
+from typing import Callable, Union, List, Dict, Tuple
import sympy as sp
from pystencils.assignment import Assignment
from pystencils.sympyextensions import SymbolCreator
+from pystencils.kernelcreation import CreateKernelConfig
-__all__ = ['kernel']
+__all__ = ['kernel', 'kernel_config']
-def kernel(func: Callable[..., None], return_config: bool = False, **kwargs) -> Union[List[Assignment], Dict]:
- """Decorator to simplify generation of pystencils Assignments.
-
- Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
- in the result list. Furthermore the meaning of the ternary inline 'if-else' changes meaning to denote a
- sympy Piecewise.
-
- The decorated function may not receive any arguments, with exception of an argument called 's' that specifies
- a SymbolCreator()
- func: the decorated function
- return_config: Specify whether to return the list with assignments, or a dictionary containing additional settings
- like func_name
-
- Examples:
- >>> import pystencils as ps
- >>> @kernel
- ... def my_kernel(s):
- ... f, g = ps.fields('f, g: [2D]')
- ... s.neighbors @= f[0,1] + f[1,0]
- ... g[0,0] @= s.neighbors + f[0,0] if f[0,0] > 0 else 0
- >>> f, g = ps.fields('f, g: [2D]')
- >>> assert my_kernel[0].rhs == f[0,1] + f[1,0]
+def _kernel(func: Callable[..., None], **kwargs) -> Tuple[List[Assignment], str]:
+ """
+ Convenient function for kernel decorator to prevent code duplication
+ Args:
+ func: decorated function
+ **kwargs: kwargs for the function
+ Returns:
+ assignments, function_name
"""
source = inspect.getsource(func)
source = textwrap.dedent(source)
@@ -55,10 +42,74 @@ def kernel(func: Callable[..., None], return_config: bool = False, **kwargs) ->
if 's' in args and 's' not in kwargs:
kwargs['s'] = SymbolCreator()
func(**kwargs)
- if return_config:
- return {'assignments': assignments, 'function_name': func.__name__}
- else:
- return assignments
+ return assignments, func.__name__
+
+
+def kernel(func: Callable[..., None], **kwargs) -> List[Assignment]:
+ """Decorator to simplify generation of pystencils Assignments.
+
+ Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
+ in the result list. Furthermore the meaning of the ternary inline 'if-else' changes meaning to denote a
+ sympy Piecewise.
+
+ The decorated function may not receive any arguments, with exception of an argument called 's' that specifies
+ a SymbolCreator()
+ Args:
+ func: decorated function
+ **kwargs: kwargs for the function
+
+ Examples:
+ >>> import pystencils as ps
+ >>> @kernel
+ ... def my_kernel(s):
+ ... f, g = ps.fields('f, g: [2D]')
+ ... s.neighbors @= f[0,1] + f[1,0]
+ ... g[0,0] @= s.neighbors + f[0,0] if f[0,0] > 0 else 0
+ >>> f, g = ps.fields('f, g: [2D]')
+ >>> assert my_kernel[0].rhs == f[0,1] + f[1,0]
+ """
+ assignments, _ = _kernel(func, **kwargs)
+ return assignments
+
+
+def kernel_config(config: CreateKernelConfig, **kwargs) -> Callable[..., Dict]:
+ """Decorator to simplify generation of pystencils Assignments, which takes a configuration
+ and updates the function name accordingly.
+
+ Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
+ in the result list. Furthermore the meaning of the ternary inline 'if-else' changes meaning to denote a
+ sympy Piecewise.
+
+ The decorated function may not receive any arguments, with exception of an argument called 's' that specifies
+ a SymbolCreator()
+ Args:
+ config: Specify whether to return the list with assignments, or a dictionary containing additional settings
+ like func_name
+ Returns:
+ decorator with config
+
+ Examples:
+ >>> import pystencils as ps
+ >>> config = ps.CreateKernelConfig()
+ >>> @kernel_config(config)
+ ... def my_kernel(s):
+ ... f, g = ps.fields('f, g: [2D]')
+ ... s.neighbors @= f[0,1] + f[1,0]
+ ... g[0,0] @= s.neighbors + f[0,0] if f[0,0] > 0 else 0
+ >>> f, g = ps.fields('f, g: [2D]')
+ >>> assert my_kernel['assignments'][0].rhs == f[0,1] + f[1,0]
+ """
+ def decorator(func: Callable[..., None]) -> Union[List[Assignment], Dict]:
+ """
+ Args:
+ func: decorated function
+ Returns:
+ Dict for unpacking into create_kernel
+ """
+ assignments, func_name = _kernel(func, **kwargs)
+ config.function_name = func_name
+ return {'assignments': assignments, 'config': config}
+ return decorator
# noinspection PyMethodMayBeStatic
diff --git a/pystencils_tests/test_create_kernel_config.py b/pystencils_tests/test_create_kernel_config.py
index c36ba483ca9fd9718f284731f648a48777de29e3..836569a9382626b808a1bfb9f75c94ba3722e9ac 100644
--- a/pystencils_tests/test_create_kernel_config.py
+++ b/pystencils_tests/test_create_kernel_config.py
@@ -1,3 +1,4 @@
+import numpy as np
import pystencils as ps
@@ -15,3 +16,14 @@ def test_create_kernel_config():
c = ps.CreateKernelConfig(backend=ps.Backend.CUDA)
assert c.target == ps.Target.CPU
assert c.backend == ps.Backend.CUDA
+
+
+def test_kernel_decorator_config():
+ config = ps.CreateKernelConfig()
+ a, b, c = ps.fields(a=np.ones(100), b=np.ones(100), c=np.ones(100))
+
+ @ps.kernel_config(config)
+ def test():
+ a[0] @= b[0] + c[0]
+
+ ps.create_kernel(**test)