First autogenerated walberla app that compiles without errors

src/pystencils_autodiff/framework_integration/astnodes.py

@@ -210,11 +210,11 @@ class JinjaCppFile(Node):
     def __str__(self):
         assert self.TEMPLATE, f"Template of {self.__class__} must be set"
         render_dict = {k: (self._print(v)
-                           if not isinstance(v, (pystencils.Field, pystencils.TypedSymbol)) and v is not None
+                           if not isinstance(v, (pystencils.Field, pystencils.TypedSymbol, bool)) and v is not None
                            else v)
                        if not isinstance(v, Iterable) or isinstance(v, str)
                        else [(self._print(a)
-                              if not isinstance(a, (pystencils.Field, pystencils.TypedSymbol)) and a is not None
+                              if not isinstance(a, (pystencils.Field, pystencils.TypedSymbol, bool)) and a is not None
                               else a)
                              for a in v]
                        for k, v in self.ast_dict.items()}

src/pystencils_autodiff/framework_integration/printer.py

+import sympy as sp
+
 import pystencils.backends.cbackend
 from pystencils.kernelparameters import FieldPointerSymbol
 
@@ -12,14 +14,17 @@ class FrameworkIntegrationPrinter(pystencils.backends.cbackend.CBackend):
     """
 
     def __init__(self):
-        super().__init__(sympy_printer=None,
-                         dialect='c')
+        super().__init__(dialect='c')
 
     def _print(self, node):
         from pystencils_autodiff.framework_integration.astnodes import JinjaCppFile
         if isinstance(node, JinjaCppFile):
             node.printer = self
-        return super()._print(node)
+
+        if isinstance(node, sp.Expr):
+            return self.sympy_printer._print(node)
+        else:
+            return super()._print(node)
 
     def _print_WrapperFunction(self, node):
         super_result = super()._print_KernelFunction(node)

src/pystencils_autodiff/walberla.py

@@ -8,16 +8,44 @@
 """
 
 
+from abc import ABC
+from itertools import chain
 from os.path import dirname, join
 
 import jinja2
 import sympy as sp
+from stringcase import camelcase, pascalcase
 
+import pystencils
+from pystencils.astnodes import SympyAssignment
 from pystencils.data_types import TypedSymbol
 from pystencils_autodiff._file_io import read_template_from_file
 from pystencils_autodiff.framework_integration.astnodes import JinjaCppFile
 
 
+class FieldType(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("{{ field_type }}")
+
+    def __init__(self, field: pystencils.Field, on_gpu: bool):
+        from pystencils_walberla.jinja_filters import make_field_type, get_field_fsize
+
+        f_size = get_field_fsize(field)
+        field_type = make_field_type(pystencils.data_types.get_base_type(field.dtype), f_size, on_gpu)
+
+        ast_dict = {'on_gpu': on_gpu,
+                    'field_type': field_type
+                    }
+        JinjaCppFile.__init__(self, ast_dict)
+
+    @property
+    def headers(self):
+        if self.ast_dict.on_gpu:
+            return ['"field/GhostLayerField.h"']
+        else:
+            return ['"cuda/GPUField.h"']
+
+
 class WalberlaModule(JinjaCppFile):
     TEMPLATE = read_template_from_file(join(dirname(__file__), 'walberla_main.tmpl.cpp'))
 
@@ -32,6 +60,9 @@ class WalberlaMain(JinjaCppFile):
 int main( int argc, char ** argv )
 {
    using namespace walberla;
+   using namespace walberla_user;
+   using namespace walberla::pystencils;
+
    Environment {{ walberl_env }}( argc, argv );
 
    {{ body | indent(3) }}
@@ -58,12 +89,12 @@ int main( int argc, char ** argv )
         return {self.ast_dict.walberl_env}
 
 
-class BlockForrestCreation(JinjaCppFile):
-    TEMPLATE = jinja2.Template("""auto {{ blocks }} = walberla_user::createBlockForrest(walberlaEnv);""")
-    def __init__(self, block_forrest_name):
+class BlockForestCreation(JinjaCppFile):
+    TEMPLATE = jinja2.Template("""auto {{ blocks }} = walberla_user::createBlockForest(walberlaEnv);""")
+    def __init__(self, block_forest_name):
 
         ast_dict = {
-            'blocks': TypedSymbol(block_forrest_name, "auto")
+            'blocks': TypedSymbol(block_forest_name, "auto")
         }
 
         super().__init__(ast_dict)
@@ -73,7 +104,7 @@ class BlockForrestCreation(JinjaCppFile):
         return {self.ast_dict.blocks}
 
 
-class UniformBlockForrestFromConfig(BlockForrestCreation):
+class UniformBlockforestFromConfig(BlockForestCreation):
     TEMPLATE = jinja2.Template(
         """auto {{ blocks }} = blockforest::createUniformBlockGridFromConfig( walberlaEnv.config() );""")
     headers = ['"blockforest/Initialization.h"']
@@ -93,13 +124,9 @@ class UniformBlockForrestFromConfig(BlockForrestCreation):
 class DefinitionsHeader(JinjaCppFile):
     TEMPLATE = read_template_from_file(join(dirname(__file__), 'walberla_user_defintions.tmpl.hpp'))
 
-    def __init__(self, walberla_module):
-        self.main_module = walberla_module
-        super().__init__({})
-
-    @property
-    def config_required(self):
-        return {"DomainSetup": {"blocks": [1, 1, 1], "cellsPerBlock": [100, 100, 100]}}
+    def __init__(self, lb_model_name, flag_field_type):
+        self.headers = ['<cstdint>', '"lbm/field/PdfField.h"', f'"{lb_model_name}.h"', '"field/FlagField.h"']
+        super().__init__({'lb_model_name': lb_model_name, 'flag_field_type': flag_field_type})
 
 
 class Using(JinjaCppFile):
@@ -120,7 +147,7 @@ class Using(JinjaCppFile):
 class GetParameter(JinjaCppFile):
 
     TEMPLATE = jinja2.Template(
-        'walberlaEnv.config()->getOneBlock("{{ block }}").getParameter<{{ key.dtype }}>("{{ key }}"{% if default %}, {{ default }}{% endif %})'  # noqa
+        'walberlaEnv.config()->getOneBlock("{{ block }}").getParameter<{{ key.dtype }}>("{{ key }}"{% if default %}, static_cast<{{ key.dtype }}>({{ default }}){% endif %})'  # noqa
         )
 
     def __init__(self, block: str, key, default_value=None):
@@ -138,6 +165,9 @@ class GetParameter(JinjaCppFile):
     def __sympy__(self):
         return TypedSymbol(self.ast_dict.key.name + "Config", str(self.ast_dict.key.dtype))
 
+    def __getattr__(self, name):
+        return getattr(self.__sympy__(), name)
+
 
 class FieldAllocation(JinjaCppFile):
     """
@@ -156,25 +186,40 @@ class FieldAllocation(JinjaCppFile):
 
     """  # noqa
 
-    TEMPLATE = jinja2.Template("""BlockDataID {{ field_name }}_data = field::addToStorage<{{ field_type }}>( {{ block_forrest }},
-                              {{ field_name }}
-     {%- if init_value -%}      , {{ init_value }} {% endif %}
-     {%- if layout_str -%}      , {{ layout_str }} {% endif %}
+    TEMPLATE = jinja2.Template("""
+{%- if on_gpu -%}
+BlockDataID {{ field_name }}_data_gpu = cuda::addGPUFieldToStorage<{{ field_type }}>({{ block_forest }},
+                                    "{{ field_name }}",
+                                    {{ f_size }},
+                                    {{ layout }},
+                                    {{ num_ghost_layers }},
+                                    {{ usePitchedMem }} );
+{%- else -%}
+BlockDataID {{ field_name }}_data = field::addToStorage<{{ field_type }}>( {{ block_forest }},
+                              "{{ field_name }}"
+     {%- if init_value -%}      , {{ init_value }}{% endif %}
+     {%- if layout -%}      , {{ layout }}{% endif %}
      {%- if num_ghost_layers -%}, {{ num_ghost_layers }} {% endif %}
      {%- if always_initialize -%}, {{ always_initialize }} {% endif %});
+{%- endif %}
 """)  # noqa
 
-    def __init__(self, block_forrest, field):
-        self._symbol = TypedSymbol(field.name + '_data', 'BlockDataID')
+    def __init__(self, block_forest, field, on_gpu=False, usePitchedMem=True, num_ghost_layers=1):
+        self._symbol = TypedSymbol(field.name + ('_data_gpu' if on_gpu else '_data'), 'BlockDataID')
         ast_dict = {
-            'block_forrest': block_forrest,
+            'block_forest': block_forest,
             'field_name': field.name,
-            'field_type': f'GhostLayerField< {field.dtype}, {field.index_shape[0] if field.index_shape else 1} >'
+            # f'GhostLayerField< {field.dtype}, {field.index_shape[0] if field.index_shape else 1} >'
+            'field_type': FieldType(field, on_gpu),
+            'on_gpu': on_gpu,
+            'f_size': field.index_shape[0] if field.index_shape else 1,
+            'init_value': f'{field.dtype}{{}}',
+            'num_ghost_layers': num_ghost_layers,
+            'layout': 'field::zyxf',
+            'usePitchedMem': 'true' if usePitchedMem else 'false',
         }
         super().__init__(ast_dict)
 
-    headers = ['"field/GhostLayerField.h"']
-
     @property
     def symbol(self):
         return self._symbol
@@ -183,6 +228,12 @@ class FieldAllocation(JinjaCppFile):
     def symbols_defined(self):
         return {self.symbol}
 
+    @property
+    def headers(self):
+        return (['"cuda/AddGPUFieldToStorage.h"', '"field/GhostLayerField.h"']
+                if self.ast_dict.on_gpu
+                else ['"field/AddToStorage.h"'])
+
 
 class WalberlaVector(JinjaCppFile):
     TEMPLATE = jinja2.Template("""math::Vector{{ndim}}<{{dtype}}>({{offsets}})""")
@@ -212,7 +263,7 @@ class PdfFieldAllocation(FieldAllocation):
                                           const Set < SUID > & requiredSelectors     = Set < SUID > : : emptySet(),
                                           const Set < SUID > & incompatibleSelectors = Set < SUID > : : emptySet() )
     """
-    TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = lbm::field::addPdfFieldToStorage < {{ field_type }} > ({{ block_forrest }},
+    TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = lbm::field::addPdfFieldToStorage < {{ field_type }} > ({{ block_forest }},
                               {{field_name}},
                               {{lattice_model}}
      {%- if initial_velocity -%}      , {{initial_velocity }} {% endif %}
@@ -220,8 +271,8 @@ class PdfFieldAllocation(FieldAllocation):
      {%- if num_ghost_layers -%}, {{num_ghost_layers }} {% endif %});
 """)  # noqa
 
-    def __init__(self, block_forrest, field, lattice_model, initial_velocity=None, initial_density=None):
-        super().__init__(block_forrest, field)
+    def __init__(self, block_forest, field, lattice_model, initial_velocity=None, initial_density=None, on_gpu=False):
+        super().__init__(block_forest, field, on_gpu)
         if initial_velocity and not isinstance(initial_velocity, WalberlaVector):
             initial_velocity = WalberlaVector(initial_velocity)
 
@@ -231,7 +282,7 @@ class PdfFieldAllocation(FieldAllocation):
             'initial_velocity': initial_velocity,
         })
 
-    headers = ['lbm/field/AddToStorage.h']
+    headers = ['"lbm/field/AddToStorage.h"']
 
 
 class FlagFieldAllocation(FieldAllocation):
@@ -248,12 +299,480 @@ class FlagFieldAllocation(FieldAllocation):
                                            const Set<SUID> & requiredSelectors = Set<SUID>::emptySet(),
                                            const Set<SUID> & incompatibleSelectors = Set<SUID>::emptySet() )
     """  # noqa
-    TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = field::addFlagFieldToStorage < {{ field_type }} > ({{ block_forrest }},
+    TEMPLATE = jinja2.Template("""BlockDataID {{field_name}}_data = field::addFlagFieldToStorage < {{ field_type }} > ({{ block_forest }},
                               {{field_name}}
      {%- if num_ghost_layers -%}, {{num_ghost_layers }} {% endif %});
 """)  # noqa
 
-    def __init__(self, block_forrest, field):
-        super().__init__(block_forrest, field)
+    def __init__(self, block_forest, field, on_gpu=False):
+        super().__init__(block_forest, field, on_gpu)
+
+    headers = ['"field/AddToStorage.h"']
+
+
+class FlagUidDefinition(JinjaCppFile):
+    TEMPLATE = jinja2.Template('const FlagUID {{ name }}FlagUID("{{ name }}");')
+
+    def __init__(self, name):
+        self._symbol = TypedSymbol(name + 'FlagUID', 'FlagUID')
+        ast_dict = {
+            'name': name,
+        }
+        super().__init__(ast_dict)
+
+    headers = ['"field/FlagUID.h"']
+
+    @property
+    def symbol(self):
+        return self._symbol
+
+    @property
+    def symbols_defined(self):
+        return {self.symbol}
+
+
+class BoundaryHandling(ABC):
+    pass
+
+
+class BoundaryHandlingFromConfig(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""auto {{  boundaries_config }} = walberlaEnv.config()->getOneBlock( "Boundaries" );
+geometry::initBoundaryHandling<FlagField_T>(*{{ block_forest }}, {{ flag_field_id }}, {{ boundaries_config }});
+geometry::setNonBoundaryCellsToDomain<FlagField_T>(*{{ block_forest }}, {{ flag_field_id }}, {{ fluid_uid }});""")
+
+    def __init__(self, block_forest, flag_field_id, fluid_uid):
+        self._symbol = boundaries_config = TypedSymbol('boundariesConfig', 'auto')
+        ast_dict = {
+            'flag_field_id': flag_field_id,
+            'block_forest': block_forest,
+            'fluid_uid': fluid_uid,
+            'boundaries_config': boundaries_config,
+        }
+        super().__init__(ast_dict)
+
+    @property
+    def symbol(self):
+        return self._symbol
+
+    @property
+    def symbols_defined(self):
+        return {self.symbol}
+
+
+class FillFromFlagField(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""
+    {{ boundary_condition }}.fillFromFlagField<{{ flag_field_type }}>( {{ block_forest }}, {{ flag_field_id }}, FlagUID("{{ boundary_name }}"), {{ flag_field_id }});
+    """)  # noqa
+
+    def __init__(self, flag_field_id, fluid_uid):
+        self._symbol = boundaries_config = TypedSymbol('boundariesConfig', 'auto')
+        ast_dict = {
+            'flag_field_id': flag_field_id,
+            'fluid_uid': fluid_uid,
+            'boundaries_config': boundaries_config,
+        }
+        super().__init__(ast_dict)
+
+    @property
+    def symbol(self):
+        return self._symbol
+
+    @property
+    def symbols_defined(self):
+        return {self.symbol}
+
+
+class LbCommunicationSetup(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""blockforest::communication::UniformBufferedScheme<lbm::{{ lb_model_type }}::CommunicationStencil> {{ communication }}( blocks );
+{{ communication }}.addPackInfo( make_shared< lbm::PdfFieldPackInfo<lbm::{{ lb_model_type }}> >( {{ pdf_id }} ) );
+ """)  # noqa
+
+    def __init__(self, lb_model_type, pdf_id):
+        self._symbol = TypedSymbol('communication', 'auto')
+        ast_dict = {
+            'lb_model_type': lb_model_type,
+            'pdf_id': pdf_id,
+            'communication': self._symbol,
+        }
+        super().__init__(ast_dict)
+
+    @property
+    def symbol(self):
+        return self._symbol
+
+    @property
+    def symbols_defined(self):
+        return {self.symbol}
+
+    headers = ['"blockforest/communication/UniformBufferedScheme.h"', '"lbm/communication/PdfFieldPackInfo.h"']
+
+
+class BeforeFunction(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""BeforeFunction( {{ f }}, "{{ f }}" )""")  # noqa
+
+    def __init__(self, function):
+        ast_dict = {
+            'f': function,
+        }
+        super().__init__(ast_dict)
+
+
+class AfterFunction(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""AfterFunction( {{ f }}, "{{ f }}" )""")  # noqa
+
+    def __init__(self, function):
+        ast_dict = {
+            'f': function,
+        }
+        super().__init__(ast_dict)
+
+
+class Sweep(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""Sweep( {{ f }}, "{{ f }}" )""")  # noqa
+
+    def __init__(self, function):
+        ast_dict = {
+            'f': function,
+        }
+        super().__init__(ast_dict)
+
+
+class TimeLoop(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""
+{{ timeloop_definition }}
+
+{% for f in before_functions -%}
+{{timeloop_name}}.add() << {{ f }};
+{% endfor %}
+{% for s in sweeps -%}
+{{timeloop_name}}.add() << {{ s }};
+{% endfor %}
+{% for f in after_functions -%}
+{{timeloop_name}}.add() << {{ f }};
+{% endfor %}
+    """)  # noqa
+
+    def __init__(self, block_forest, before_functions, sweeps, after_functions, timesteps, timeloop_name='timeloop'):
+        self._symbol = TypedSymbol(timeloop_name, 'SweepTimeloop')
+        ast_dict = {
+            'timeloop_definition': CppObjectConstruction(self._symbol,
+                                                         [f'{block_forest}->getBlockStorage()',
+                                                          timesteps]),
+            'before_functions': [BeforeFunction(f) for f in before_functions],
+            'sweeps': [Sweep(f) for f in sweeps],
+            'after_functions': [AfterFunction(f) for f in after_functions],
+            'timeloop_name': timeloop_name
+        }
+        super().__init__(ast_dict)
+
+    @property
+    def symbol(self):
+        return self._symbol
+
+    @property
+    def symbols_defined(self):
+        return {self.symbol}
+
+    headers = ['"timeloop/all.h"']
+
+
+class U_Rho_Adaptor(JinjaCppFile):
+    """Docstring for U_Rho_Adaptor. """
+
+    TEMPLATE = jinja2.Template("""
+field::addFieldAdaptor<lbm::Adaptor<lbm::{{ lb_model_type }}>::Density>       ( {{ block_forest }}, {{ pdf_id }}, "{{ density_name }}" );
+field::addFieldAdaptor<lbm::Adaptor<lbm::{{ lb_model_type }}>::VelocityVector>( {{ block_forest }}, {{ pdf_id }}, "{{ velocity_name }}" );
+    """)  # noqa
+
+    def __init__(self, block_forest, lb_model_type, pdf_id):
+        self.density_name = "DensityAdaptor"
+        self.velocity_name = "VelocityAdaptor"
+        ast_dict = {
+            'lb_model_type': lb_model_type,
+            'pdf_id': pdf_id,
+            'block_forest': block_forest,
+            'density_name': self.density_name,
+            'velocity_name': self.velocity_name,
+        }
+        super().__init__(ast_dict)
+
+    headers = ['"timeloop/all.h"']
+
+
+class RunTimeLoop(JinjaCppFile):
+    TEMPLATE = jinja2.Template("""{%- if with_gui == 'true' %}
+if( parameters.getParameter<bool>( "useGui", {{ use_gui_default }}) )
+{
+   GUI gui ( {{ timeloop }}, {{ block_forest }}, argc, argv );
+   lbm::connectToGui<LatticeModel_T> ( gui );
+   gui.run();
+}
+else {
+   {{ timeloop }}.run();
+}
+{%- else %}
+timeloop.run();
+{%- endif %}
+    """)  # noqa
+
+    def __init__(self, block_forest, timeloop, with_gui=False, use_gui_default='false'):
+        self.density_name = "DensityAdaptor"
+        self.velocity_name = "VelocityAdaptor"
+        ast_dict = {
+            'block_forest': block_forest,
+            'timeloop': timeloop.symbol,
+            'with_gui': with_gui,
+            'use_gui_default': use_gui_default,
+        }
+        super().__init__(ast_dict)
+
+    @property
+    def headers(self):
+        if self.ast_dict.with_gui:
+            return [
+                "gui/all.h",
+                "lbm/gui/Connection.h"
+            ]
+        else:
+            return []
+
+
+class CppObjectConstruction(JinjaCppFile, pystencils.astnodes.SympyAssignment):
+
+    TEMPLATE = jinja2.Template("""{{  symbol.dtype  }} {{ symbol }}({{ arg_list }});""")  # noqa
+
+    def __init__(self, symbol, args):
+        JinjaCppFile.__init__(self, {})
+        self.ast_dict.update({
+            'symbol': symbol,
+            'args': args,
+            'arg_list': ', '.join(self.printer(a) for a in args)
+        })
+        pystencils.astnodes.SympyAssignment.__init__(self, symbol, 1)
+
+
+class ResolveUndefinedSymbols(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""
+{% for p in parameters %}
+{{ p }}
+{% endfor %}
+{{ block  }}
+""")  # noqa
+
+    def __init__(self, block, config_block):
+        self.block = block
+        JinjaCppFile.__init__(self, {})
+        self.ast_dict.update({
+            'block': block,
+            'config_block': config_block,
+            'parameters': self.parameter_definitions
+        })
+
+    @property
+    def symbols_defined(self):
+        return self.block.undefined_symbols
+
+    def __repr__(self):
+        self.ast_dict.parameters = self.parameter_definitions
+        return super().__repr__()
+
+    def __str__(self):
+        self.ast_dict.parameters = self.parameter_definitions
+        return super().__str__()
+
+    @property
+    def parameter_definitions(self):
+        parameter_definitions = [SympyAssignment(s, GetParameter(self.ast_dict.config_block, s))
+                                 for s in self.block.undefined_symbols]
+        return parameter_definitions
+
+    @property
+    def config_required(self):
+        return {self.ast_dict.config_block, {s.name: None for s in self.symbols_defined}}
+
+
+class FieldCopyFunctor(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""cuda::fieldCpyFunctor<{{from_type}}, {{ to_type }} >({{from_id}}, {{to_id}})""")  # noqa
+
+    def __init__(self, from_id, from_type, to_id, to_type):
+        ast_dict = {'from_id': from_id,
+                    'from_type': from_type,
+                    'to_id': to_id,
+                    'to_type': to_type
+                    }
+        super().__init__(ast_dict)
+
+    headers = ['"cuda/FieldCopy.h"']
+
+
+class AllocateAllFields(JinjaCppFile):
+
+    TEMPLATE = jinja2.Template("""
+// CPU Arrays
+{% for c in cpu_arrays %}
+{{- c }}
+{% endfor %}
+// GPU Arrays
+{% for g in gpu_arrays %}
+{{- g }}
+{% endfor %}
+{{ block  }}
+""")  # noqa
+
+    def __init__(self, block_forest, data_handling, lb_index_shape=(-1,), lb_model_name=None):
+        self.data_handling = data_handling
+        self.block_forest = block_forest
+        self.lb_index_shape = lb_index_shape
+        self.lb_model_name = lb_model_name
+        JinjaCppFile.__init__(self, {})
+        self._update()
+
+    def __repr__(self):
+        self._update()
+        return super().__repr__()
+
+    def __str__(self):
+        self._update()
+        return super().__str__()
+
+    def _update(self):
+        self.ast_dict.update({
+            'cpu_arrays': self._cpu_allocations.values(),
+            'gpu_arrays': self._gpu_allocations.values()
+        })
+
+    @property
+    def _cpu_allocations(self):
+        return {s: FieldAllocation(self.block_forest, self.data_handling.fields[s])
+                for s in self.data_handling.cpu_arrays.keys()}
+        # if not (self.data_handling.fields[s].index_shape == self.lb_index_shape)}
+
+    @property
+    def _gpu_allocations(self):
+        return {s: FieldAllocation(self.block_forest, self.data_handling.fields[s], on_gpu=True)
+                for s in self.data_handling.gpu_arrays.keys()}
+        # if not (self.data_handling.fields[s].index_shape == self.lb_index_shape)}
+
+
+class InitBoundaryHandling(JinjaCppFile):
+    TEMPLATE = jinja2.Template("""
+// This is my fluid 🏊. It's special... 🤔 but for the computer just a 0
+{{ fluid_uid_definition }}
+
+// Initialize geometry 🎲
+{{ geometry_initialization }}
+
+// Here are the generated boundaries. They are not so special... 👎
+{%- for b in generated_boundaries %}
+{{ b }}
+{% endfor %}
+""")  # noqa
+
+    def __init__(self, block_forest, flag_field_id, pdf_field_id, boundary_conditions):
+        self.fluid = FlagUidDefinition("fluid")
+        ast_dict = {'fluid_uid_definition': self.fluid,
+                    'geometry_initialization': BoundaryHandlingFromConfig(block_forest,
+                                                                          flag_field_id,
+                                                                          self.fluid.symbol),
+                    'generated_boundaries': [GeneratedBoundaryInitialization(block_forest,
+                                                                             b,
+                                                                             pdf_field_id,
+                                                                             flag_field_id,
+                                                                             self.fluid.symbol)
+                                             for b in boundary_conditions]
+                    }
+        super().__init__(ast_dict)
+
+    headers = ['"cuda/FieldCopy.h"', '"geometry/InitBoundaryHandling.h"']
+
+
+class GeneratedBoundaryInitialization(JinjaCppFile):
+    TEMPLATE = jinja2.Template("""lbm::{{ boundary_condition }} {{ identifier }}( {{ block_forest }}, {{ pdf_field_id }} );
+{{ identifier }}.fillFromFlagField<FlagField_T>( {{ block_forest }}, {{ flag_field_id }}, FlagUID("{{ boundary_condition }}"), {{ fluid_uid }} );
+""")  # noqa
+
+    def __init__(self, block_forest, boundary_condition, pdf_field_id, flag_field_id, fluid_uid):
+        self.fluid = FlagUidDefinition("fluid")
+        ast_dict = {'block_forest': block_forest,
+                    'boundary_condition': pascalcase(boundary_condition.name),
+                    'identifier': camelcase(boundary_condition.name),
+                    'pdf_field_id': pdf_field_id,
+                    'fluid_uid': fluid_uid,
+                    'flag_field_id': flag_field_id,
+                    }
+        super().__init__(ast_dict)
+
+    @property
+    def headers(self):
+        return [f'"{pascalcase(self.ast_dict.boundary_condition)}.h"']
+
+    @property
+    def symbols_defined(self):
+        # Could also be defined
+        # TypedSymbol(self.ast_dict.identifier, 'FlagUID'),
+        return {TypedSymbol(self.ast_dict.identifier, f'lbm::{self.ast_dict.boundary_condition}')}
+
+
+class SweepCreation(JinjaCppFile):
+    TEMPLATE = jinja2.Template("""{{ sweep_class_name }}( {{ parameter_str }} )""")  # noqa
+
+    def __init__(self, sweep_class_name: str, field_allocation: AllocateAllFields, ast, parameters_to_ignore=None):
+        parameters = ast.get_parameters()
+        parameter_ids = [field_allocation._cpu_allocations[p.symbol.name.replace('_data_', '')].symbol.name
+                         if ast.target == 'cpu'
+                         else field_allocation._gpu_allocations[p.symbol.name.replace('_data_', '')].symbol.name
+                         for p in parameters
+                         if p.is_field_pointer or not p.is_field_parameter
+                         ]
+
+        ast_dict = {'sweep_class_name': sweep_class_name,
+                    'parameter_ids': parameter_ids,
+                    'parameter_str': ', '.join(parameter_ids)}
+        super().__init__(ast_dict)
+
+    @property
+    def headers(self):
+        return [f'"{self.ast_dict.sweep_class_name}.h"']
+
+
+class SweepOverAllBlocks(JinjaCppFile):
+    # TEMPLATE = jinja2.Template("""std::for_each({{block_forest}}->begin(), {{block_forest}}->end(), {{functor}});""")  # noqa
+    TEMPLATE = jinja2.Template("""auto {{sweep_class_name | lower() }} = {{functor}};
+for( auto& block : *{{block_forest}}) {{sweep_class_name | lower() }}(&block);""")  # noqa
+
+    def __init__(self, functor: SweepCreation, block_forest):
+        ast_dict = {'functor': functor,
+                    'sweep_class_name': functor.ast_dict.sweep_class_name,
+                    'block_forest': block_forest}
+        super().__init__(ast_dict)
+
+    headers = ['<algorithm>']
+
+
+# class GetSweep(JinjaCppFile):
+    # TEMPLATE = jinja2.Template("""std::function<void(*domain_decomposition::IBlock)>({{ sweep_class }})""")  # noqa
+
+    # def __init__(self, sweep_class: SweepCreation):
+
+    # ast_dict = {'sweep_class': sweep_class, }
+    # super().__init__(ast_dict)
+
+    # headers = ['"domain_decomposition/IBlock.h"']
+
+
+# class MakeSharedFunctor(JinjaCppFile):
+    # TEMPLATE = jinja2.Template("""makeSharedSweep{{ functor }}(&{{block_forest}})""")  # noqa
 
-    headers = ['field/AddToStorage.h']
+    # def __init__(self, functor, block_forest):
+    # ast_dict = {'functor': functor,
+    # 'block_forest': block_forest}
+    # super().__init__(ast_dict)

src/pystencils_autodiff/walberla_user_defintions.tmpl.hpp

@@ -7,22 +7,21 @@
 
 #pragma once
 
-#include "lbm/communication/PdfFieldPackInfo.h"
-#include "lbm/field/AddToStorage.h"
-#include "lbm/field/PdfField.h"
-#include "lbm/gui/Connection.h"
-#include "lbm/vtk/VTKOutput.h"
+
+{% for header in headers -%}
+#include {{ header }}
+{% endfor %}
 
 namespace walberla_user {
 using namespace walberla;
-using LatticeModel_T = lbm::LbCodeGenerationExample_LatticeModel;
+using LatticeModel_T = lbm::{{ lb_model_name }};
 using Stencil_T = LatticeModel_T::Stencil;
-using CommunicationStencil_T = LatticeModel_T::CommunicationStencil_T;
-using lPdfField_T = bm::PdfField<LatticeModel_T>;
-using flag_t = walberla::uint8_t;
-using FlagField_T FlagField<flag_t>;
+using CommunicationStencil_T = LatticeModel_T::CommunicationStencil;
+using PdfField_T = lbm::PdfField<LatticeModel_T>;
+using flag_t = {{ flag_field_type }};
+using FlagField_T = FlagField<flag_t>;
 
 using PdfField_T = lbm::PdfField<LatticeModel_T>;
-typedef VectorField_T = GhostLayerField<real_t, LatticeModel_T::Stencil::D>;
-typedef ScalarField_T GhostLayerField<real_t, 1>;
+using VectorField_T = GhostLayerField<real_t, LatticeModel_T::Stencil::D>;
+using ScalarField_T = GhostLayerField<real_t, 1>;
 } // namespace walberla_user

src/pystencils_autodiff/wald_und_wiesen_simulation.py

+#
+# Copyright © 2020 Stephan Seitz <stephan.seitz@fau.de>
+#
+# Distributed under terms of the GPLv3 license.
+
+"""
+
+"""
+
+from typing import Dict
+
+import sympy as sp
+from stringcase import pascalcase
+
+import lbmpy_walberla
+import pystencils
+import pystencils_walberla.codegen
+from pystencils.astnodes import Block, EmptyLine
+from pystencils_autodiff.walberla import (
+    AllocateAllFields, DefinitionsHeader, InitBoundaryHandling, LbCommunicationSetup,
+    ResolveUndefinedSymbols, RunTimeLoop, SweepCreation, SweepOverAllBlocks, TimeLoop,
+    UniformBlockforestFromConfig, WalberlaMain, WalberlaModule)
+
+
+class WaldUndWiesenSimulation():
+    def _get_sweep_class_name(prefix='Kernel'):
+
+        ctr = 0
+        while True:
+            yield f'{prefix}{ctr}'
+            ctr += 1
+
+    def __init__(self,
+                 graph_data_handling,
+                 codegen_context,
+                 boundary_handling: pystencils.boundaries.BoundaryHandling = None,
+                 lb_rule=None,
+                 refinement_scaling=None):
+        self._data_handling = graph_data_handling
+        self._lb_rule = lb_rule
+        self._refinement_scaling = refinement_scaling
+        self._block_forest = UniformBlockforestFromConfig()
+        self.parameter_config_block = 'parameters'
+        self._codegen_context = codegen_context
+        self._boundary_handling = boundary_handling
+        self._lb_model_name = 'GeneratedLatticeModel'
+        self._flag_field_dtype = 'uint32_t'
+        self._kernel_class_generator = WaldUndWiesenSimulation._get_sweep_class_name()
+        self._with_gui = False
+        self._with_gui_default = False
+
+    def _create_helper_files(self) -> Dict[str, str]:
+        if self._lb_rule:
+            lbmpy_walberla.generate_lattice_model(self._codegen_context, self._lb_model_name,
+                                                  self._lb_rule,
+                                                  refinement_scaling=self._refinement_scaling)
+            if self._boundary_handling:
+                for bc in self.boundary_conditions:
+                    lbmpy_walberla.generate_boundary(
+                        self._codegen_context, pascalcase(bc.name), bc, self._lb_rule.method)
+
+    def _create_module(self):
+        if self._lb_rule:
+            lb_shape = (len(self._lb_rule.method.stencil),)
+        else:
+            lb_shape = (-1,)
+
+        self._field_allocations = field_allocations = AllocateAllFields(self._block_forest.blocks,
+                                                                        self._data_handling,
+                                                                        lb_shape,
+                                                                        self._lb_model_name)
+
+        if self._boundary_handling:
+            flag_field_id = field_allocations._cpu_allocations[
+                self._boundary_handling.flag_interface.flag_field_name].symbol
+
+        if self._lb_rule:
+            pdf_field_id = field_allocations._gpu_allocations.get(
+                'ldc_pdfSrc', field_allocations._cpu_allocations['ldc_pdfSrc']).symbol
+        else:
+            pdf_field_id = None
+
+        call_nodes = filter(lambda x: x, [self._graph_to_sweep(c) for c in self._data_handling.call_queue])
+
+        module = WalberlaModule(WalberlaMain(Block([
+            self._block_forest,
+            ResolveUndefinedSymbols(
+                Block([
+                    field_allocations,
+                    InitBoundaryHandling(self._block_forest.blocks, flag_field_id,
+                                         pdf_field_id, self.boundary_conditions)
+                    if self._boundary_handling else EmptyLine(),
+                    LbCommunicationSetup(self._lb_model_name,
+                                         pdf_field_id)
+                    if self._lb_rule else EmptyLine(),
+                    *call_nodes
+                ]), self.parameter_config_block)
+        ])))
+
+        self._codegen_context.write_file("main.cpp", str(module))
+        return module
+
+    def _create_defintions_header(self):
+        self._codegen_context.write_file("UserDefinitions.h",
+                                         str(DefinitionsHeader(self._lb_model_name, self._flag_field_dtype)))
+
+    def write_files(self):
+        self._create_helper_files()
+        self._create_module()
+        self._create_defintions_header()
+
+    @property
+    def boundary_conditions(self):
+        return self._boundary_handling._boundary_object_to_boundary_info.keys()
+
+    def _graph_to_sweep(self, c):
+        from pystencils_autodiff.graph_datahandling import KernelCall, TimeloopRun
+
+        if isinstance(c, KernelCall):
+            sweep_class_name = next(self._kernel_class_generator)
+            pystencils_walberla.codegen.generate_sweep(
+                self._codegen_context, sweep_class_name, c.kernel.ast)
+            rtn = SweepOverAllBlocks(SweepCreation(sweep_class_name, self._field_allocations,
+                                                   c.kernel.ast), self._block_forest.blocks)
+
+        elif isinstance(c, TimeloopRun):
+            sweeps = []
+            for a in c.timeloop._single_step_asts:
+                if 'indexField' in [f.name for f in a.fields_accessed]:
+                    continue
+                sweep_class_name = next(self._kernel_class_generator)
+                pystencils_walberla.codegen.generate_sweep(
+                    self._codegen_context, sweep_class_name, a)
+                sweeps.append(SweepCreation(sweep_class_name, self._field_allocations, a))
+
+            loop = TimeLoop(self._block_forest.blocks, [], sweeps, [], sp.S(c.time_steps))
+            rtn = Block([loop, RunTimeLoop(self._block_forest.blocks, loop, self._with_gui, self._with_gui_default)])
+
+        else:
+            return None
+        return rtn

tests/test_walberla.py

@@ -6,16 +6,24 @@
 """
 
 """
+import os
+import sys
+from os.path import dirname, join
+
 import numpy as np
-import sympy as sp
 
+import lbmpy_walberla
 import pystencils
+from lbmpy.creationfunctions import create_lb_collision_rule, create_lbm_kernel
 from pystencils.astnodes import Block, EmptyLine, SympyAssignment
-from pystencils.data_types import TypedSymbol
+from pystencils.data_types import TypedSymbol, create_type
 from pystencils_autodiff._file_io import write_file
+from pystencils_autodiff.graph_datahandling import GraphDataHandling
 from pystencils_autodiff.walberla import (
-    DefinitionsHeader, FieldAllocation, GetParameter, PdfFieldAllocation,FlagFieldAllocation,
-    UniformBlockForrestFromConfig, WalberlaMain, WalberlaModule)
+    DefinitionsHeader, FieldAllocation, FlagFieldAllocation, GetParameter, PdfFieldAllocation,
+    ResolveUndefinedSymbols, UniformBlockforestFromConfig, WalberlaMain, WalberlaModule)
+from pystencils_autodiff.wald_und_wiesen_simulation import WaldUndWiesenSimulation
+from pystencils_walberla.cmake_integration import ManualCodeGenerationContext
 
 
 def test_walberla():
@@ -28,18 +36,18 @@ def test_walberla():
     number_symbol = TypedSymbol('number', np.float32)
     crazy_plus_one = TypedSymbol('crazy', np.float32)
 
-    block_forrest = UniformBlockForrestFromConfig()
+    block_forest = UniformBlockforestFromConfig()
 
     block = Block([
-        block_forrest,
+        block_forest,
         SympyAssignment(foo_symbol, GetParameter('parameters', foo_symbol)),
         SympyAssignment(number_symbol, GetParameter('parameters', number_symbol, 1.2)),
         SympyAssignment(crazy_plus_one, number_symbol + 1),
         EmptyLine(),
-        FieldAllocation(block_forrest.blocks, x),
-        PdfFieldAllocation(block_forrest.blocks, pdf, 'LbModel_T'),
-        PdfFieldAllocation(block_forrest.blocks, pdf2, 'LbModel_T', [0, 0, 0], 1),
-        FlagFieldAllocation(block_forrest.blocks, flags)
+        FieldAllocation(block_forest.blocks, x, on_gpu=False),
+        PdfFieldAllocation(block_forest.blocks, pdf, 'LbModel_T', on_gpu=True),
+        PdfFieldAllocation(block_forest.blocks, pdf2, 'LbModel_T', [0, 0, 0], 1, on_gpu=True),
+        FlagFieldAllocation(block_forest.blocks, flags)
     ])
 
     module = WalberlaModule(WalberlaMain(block))
@@ -48,5 +56,56 @@ def test_walberla():
 
     write_file('/localhome/seitz_local/projects/walberla/apps/autogen/main.cpp', code)
 
-    definitions = DefinitionsHeader(module)
+    definitions = DefinitionsHeader(module, 'uint8_t')
     write_file('/localhome/seitz_local/projects/walberla/apps/autogen/UserDefinitions.h', str(definitions))
+
+
+def test_wald_wiesen_simulation():
+    with ManualCodeGenerationContext() as ctx:
+        dh = GraphDataHandling((30, 30),
+                               periodicity=False,
+                               default_ghost_layers=1,
+                               default_target='cpu')
+        dh.add_arrays('x, y')
+        dh.add_arrays('w, z', gpu=True)
+
+        sim = WaldUndWiesenSimulation(dh, ctx)
+        print(sim._create_module())
+
+
+def test_wald_wiesen_lbm():
+    sys.path.append(dirname(__file__))
+    with ManualCodeGenerationContext() as ctx:
+        from test_graph_datahandling import ldc_setup
+        opt_params = {'target': 'gpu'}
+        lbm_step = ldc_setup(domain_size=(30, 30), optimization=opt_params, fixed_loop_sizes=False)
+
+        sim = WaldUndWiesenSimulation(lbm_step.data_handling,
+                                      ctx,
+                                      lbm_step.boundary_handling,
+                                      create_lb_collision_rule(lbm_step.method, optimization=opt_params))
+        sim.write_files()
+
+        dir = '/localhome/seitz_local/projects/walberla/apps/autogen/'
+        os.makedirs(dir, exist_ok=True)
+        for k, v in ctx.files.items():
+            with open(join(dir, k), 'w') as file:
+                file.write(v)
+
+
+def test_resolve_parameters():
+    sym = TypedSymbol('s', create_type('double'))
+    sym2 = TypedSymbol('t', create_type('double'))
+
+    block_forest = UniformBlockforestFromConfig()
+
+    module = WalberlaModule(WalberlaMain(Block([
+        block_forest,
+        ResolveUndefinedSymbols(
+            Block([
+                SympyAssignment(sym, 1 + sym2),
+            ]), 'parameters')
+    ])))
+
+    print(module)
+