Merge branch 'pystencils_boundary' into 'master'

pystencils: Boundary generation for non-LB fields and for inner_or_boundary=False

Closes #111

See merge request walberla/walberla!266

python/lbmpy_walberla/boundary.py

@@ -2,13 +2,11 @@ import numpy as np
 from jinja2 import Environment, PackageLoader, StrictUndefined
 
 from lbmpy.boundaries.boundaryhandling import create_lattice_boltzmann_boundary_kernel
-from lbmpy_walberla.walberla_lbm_generation import KernelInfo
 from pystencils import Field, FieldType
-from pystencils.boundaries.createindexlist import (
-    boundary_index_array_coordinate_names, direction_member_name,
-    numpy_data_type_for_boundary_object)
+from pystencils.boundaries.createindexlist import numpy_data_type_for_boundary_object
 from pystencils.data_types import TypedSymbol, create_type
-from pystencils_walberla.codegen import default_create_kernel_parameters
+from pystencils_walberla.boundary import struct_from_numpy_dtype
+from pystencils_walberla.codegen import default_create_kernel_parameters, KernelInfo
 from pystencils_walberla.jinja_filters import add_pystencils_filters_to_jinja_env
 
 
@@ -41,7 +39,7 @@ def generate_boundary(generation_context, class_name, boundary_object, lb_method
     else:
         stencil = lb_method.stencil
 
-    stencil_info = [(i, ", ".join([str(e) for e in d])) for i, d in enumerate(stencil)]
+    stencil_info = [(i, d, ", ".join([str(e) for e in d])) for i, d in enumerate(stencil)]
 
     context = {
         'class_name': boundary_object.name,
@@ -52,9 +50,10 @@ def generate_boundary(generation_context, class_name, boundary_object, lb_method
         'dim': lb_method.dim,
         'target': target,
         'namespace': 'lbm',
+        'inner_or_boundary': boundary_object.inner_or_boundary
     }
 
-    env = Environment(loader=PackageLoader('lbmpy_walberla'), undefined=StrictUndefined)
+    env = Environment(loader=PackageLoader('pystencils_walberla'), undefined=StrictUndefined)
     add_pystencils_filters_to_jinja_env(env)
 
     header = env.get_template('Boundary.tmpl.h').render(**context)
@@ -63,30 +62,3 @@ def generate_boundary(generation_context, class_name, boundary_object, lb_method
     source_extension = "cpp" if create_kernel_params.get("target", "cpu") == "cpu" else "cu"
     generation_context.write_file("{}.h".format(class_name), header)
     generation_context.write_file("{}.{}".format(class_name, source_extension), source)
-
-
-def struct_from_numpy_dtype(struct_name, numpy_dtype):
-    result = "struct %s { \n" % (struct_name,)
-
-    equality_compare = []
-    constructor_params = []
-    constructor_initializer_list = []
-    for name, (sub_type, offset) in numpy_dtype.fields.items():
-        pystencils_type = create_type(sub_type)
-        result += "    %s %s;\n" % (pystencils_type, name)
-        if name in boundary_index_array_coordinate_names or name == direction_member_name:
-            constructor_params.append("%s %s_" % (pystencils_type, name))
-            constructor_initializer_list.append("%s(%s_)" % (name, name))
-        else:
-            constructor_initializer_list.append("%s()" % name)
-        if pystencils_type.is_float():
-            equality_compare.append("floatIsEqual(%s, o.%s)" % (name, name))
-        else:
-            equality_compare.append("%s == o.%s" % (name, name))
-
-    result += "    %s(%s) : %s {}\n" % \
-              (struct_name, ", ".join(constructor_params), ", ".join(constructor_initializer_list))
-    result += "    bool operator==(const %s & o) const {\n        return %s;\n    }\n" % \
-              (struct_name, " && ".join(equality_compare))
-    result += "};\n"
-    return result

python/pystencils_walberla/__init__.py

+from .boundary import generate_staggered_boundary, generate_staggered_flux_boundary
 from .cmake_integration import CodeGeneration
 from .codegen import (
     generate_pack_info, generate_pack_info_for_field, generate_pack_info_from_kernel,
@@ -5,4 +6,4 @@ from .codegen import (
 
 __all__ = ['CodeGeneration',
            'generate_sweep', 'generate_pack_info_from_kernel', 'generate_pack_info_for_field', 'generate_pack_info',
-           'generate_mpidtype_info_from_kernel']
+           'generate_mpidtype_info_from_kernel', 'generate_staggered_boundary', 'generate_staggered_flux_boundary']

python/pystencils_walberla/boundary.py

+import numpy as np
+from jinja2 import Environment, PackageLoader, StrictUndefined
+
+from pystencils import Field, FieldType
+from pystencils.boundaries.boundaryhandling import create_boundary_kernel
+from pystencils.boundaries.createindexlist import (
+    boundary_index_array_coordinate_names, direction_member_name,
+    numpy_data_type_for_boundary_object)
+from pystencils.data_types import TypedSymbol, create_type
+from pystencils_walberla.codegen import KernelInfo
+from pystencils_walberla.jinja_filters import add_pystencils_filters_to_jinja_env
+
+
+def generate_staggered_boundary(generation_context, class_name, boundary_object,
+                                dim, neighbor_stencil, index_shape, target='cpu'):
+    struct_name = "IndexInfo"
+    boundary_object.name = class_name
+
+    index_struct_dtype = numpy_data_type_for_boundary_object(boundary_object, dim)
+
+    staggered_field = Field.create_generic('field', dim,
+                                           np.float64 if generation_context.double_accuracy else np.float32,
+                                           index_dimensions=len(index_shape), layout='c', index_shape=index_shape,
+                                           field_type=FieldType.STAGGERED)
+
+    index_field = Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
+                        shape=(TypedSymbol("indexVectorSize", create_type(np.int64)), 1), strides=(1, 1))
+
+    kernel = create_boundary_kernel(staggered_field, index_field, neighbor_stencil, boundary_object, target=target,
+                                    openmp=generation_context.openmp)
+    kernel.function_name = "boundary_" + boundary_object.name
+
+    # waLBerla is a 3D framework. Therefore, a zero for the z index has to be added if we work in 2D
+    if dim == 2:
+        stencil = ()
+        for d in neighbor_stencil:
+            d = d + (0,)
+            stencil = stencil + (d,)
+    else:
+        stencil = neighbor_stencil
+
+    stencil_info = [(i, d, ", ".join([str(e) for e in d])) for i, d in enumerate(stencil)]
+
+    context = {
+        'class_name': boundary_object.name,
+        'StructName': struct_name,
+        'StructDeclaration': struct_from_numpy_dtype(struct_name, index_struct_dtype),
+        'kernel': KernelInfo(kernel),
+        'stencil_info': stencil_info,
+        'dim': dim,
+        'target': target,
+        'namespace': 'pystencils',
+        'inner_or_boundary': boundary_object.inner_or_boundary
+    }
+
+    env = Environment(loader=PackageLoader('pystencils_walberla'), undefined=StrictUndefined)
+    add_pystencils_filters_to_jinja_env(env)
+
+    header = env.get_template('Boundary.tmpl.h').render(**context)
+    source = env.get_template('Boundary.tmpl.cpp').render(**context)
+
+    source_extension = "cpp" if target == "cpu" else "cu"
+    generation_context.write_file("{}.h".format(class_name), header)
+    generation_context.write_file("{}.{}".format(class_name, source_extension), source)
+
+
+def generate_staggered_flux_boundary(generation_context, class_name, boundary_object,
+                                     dim, neighbor_stencil, index_shape, target='cpu'):
+    struct_name = "IndexInfo"
+    boundary_object.name = class_name
+
+    index_struct_dtype = numpy_data_type_for_boundary_object(boundary_object, dim)
+
+    staggered_field = Field.create_generic('flux', dim,
+                                           np.float64 if generation_context.double_accuracy else np.float32,
+                                           index_dimensions=len(index_shape), layout='c', index_shape=index_shape,
+                                           field_type=FieldType.STAGGERED_FLUX)
+
+    index_field = Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
+                        shape=(TypedSymbol("indexVectorSize", create_type(np.int64)), 1), strides=(1, 1))
+
+    kernel = create_boundary_kernel(staggered_field, index_field, neighbor_stencil, boundary_object, target=target,
+                                    openmp=generation_context.openmp)
+    kernel.function_name = "boundary_" + boundary_object.name
+
+    # waLBerla is a 3D framework. Therefore, a zero for the z index has to be added if we work in 2D
+    if dim == 2:
+        stencil = ()
+        for d in neighbor_stencil:
+            d = d + (0,)
+            stencil = stencil + (d,)
+    else:
+        stencil = neighbor_stencil
+
+    stencil_info = [(i, d, ", ".join([str(e) for e in d])) for i, d in enumerate(stencil)]
+
+    context = {
+        'class_name': boundary_object.name,
+        'StructName': struct_name,
+        'StructDeclaration': struct_from_numpy_dtype(struct_name, index_struct_dtype),
+        'kernel': KernelInfo(kernel),
+        'stencil_info': stencil_info,
+        'dim': dim,
+        'target': target,
+        'namespace': 'pystencils',
+        'inner_or_boundary': boundary_object.inner_or_boundary
+    }
+
+    env = Environment(loader=PackageLoader('pystencils_walberla'), undefined=StrictUndefined)
+    add_pystencils_filters_to_jinja_env(env)
+
+    header = env.get_template('Boundary.tmpl.h').render(**context)
+    source = env.get_template('Boundary.tmpl.cpp').render(**context)
+
+    source_extension = "cpp" if target == "cpu" else "cu"
+    generation_context.write_file("{}.h".format(class_name), header)
+    generation_context.write_file("{}.{}".format(class_name, source_extension), source)
+
+
+def struct_from_numpy_dtype(struct_name, numpy_dtype):
+    result = "struct %s { \n" % (struct_name,)
+
+    equality_compare = []
+    constructor_params = []
+    constructor_initializer_list = []
+    for name, (sub_type, offset) in numpy_dtype.fields.items():
+        pystencils_type = create_type(sub_type)
+        result += "    %s %s;\n" % (pystencils_type, name)
+        if name in boundary_index_array_coordinate_names or name == direction_member_name:
+            constructor_params.append("%s %s_" % (pystencils_type, name))
+            constructor_initializer_list.append("%s(%s_)" % (name, name))
+        else:
+            constructor_initializer_list.append("%s()" % name)
+        if pystencils_type.is_float():
+            equality_compare.append("floatIsEqual(%s, o.%s)" % (name, name))
+        else:
+            equality_compare.append("%s == o.%s" % (name, name))
+
+    result += "    %s(%s) : %s {}\n" % \
+              (struct_name, ", ".join(constructor_params), ", ".join(constructor_initializer_list))
+    result += "    bool operator==(const %s & o) const {\n        return %s;\n    }\n" % \
+              (struct_name, " && ".join(equality_compare))
+    result += "};\n"
+    return result

python/lbmpy_walberla/templates/Boundary.tmpl.cpp → python/pystencils_walberla/templates/Boundary.tmpl.cpp

@@ -14,8 +14,7 @@
 //  with waLBerla (see COPYING.txt). If not, see <http://www.gnu.org/licenses/>.
 //
 //! \\file {{class_name}}.cpp
-//! \\ingroup lbm
-//! \\author lbmpy
+//! \\author pystencils
 //======================================================================================================================
 
 #include <cmath>

python/lbmpy_walberla/templates/Boundary.tmpl.h → python/pystencils_walberla/templates/Boundary.tmpl.h

@@ -158,10 +158,14 @@ public:
             if( ! isFlagSet(it, domainFlag) )
                 continue;
 
-            {%- for dirIdx, offset in stencil_info %}
+            {%- for dirIdx, dirVec, offset in stencil_info %}
             if ( isFlagSet( it.neighbor({{offset}} {%if dim == 3%}, 0 {%endif %}), boundaryFlag ) )
             {
+                {% if inner_or_boundary -%}
                 auto element = {{StructName}}(it.x(), it.y(), {%if dim == 3%} it.z(), {%endif %} {{dirIdx}} );
+                {% else -%}
+                auto element = {{StructName}}(it.x() + cell_idx_c({{dirVec[0]}}), it.y() + cell_idx_c({{dirVec[1]}}), {%if dim == 3%} it.z() + cell_idx_c({{dirVec[2]}}), {%endif %} {{dirIdx}} );
+                {% endif -%}
                 indexVectorAll.push_back( element );
                 if( inner.contains( it.x(), it.y(), it.z() ) )
                     indexVectorInner.push_back( element );