Generate parameters also for boundary conditions

src/pystencils_autodiff/walberla.py

@@ -526,8 +526,7 @@ else {
 }
 {%- else %}
 timeloop.run();
-{%- endif %}
-    """)  # noqa
+{%- endif %}""")  # noqa
 
     def __init__(self, block_forest, timeloop, with_gui=False, use_gui_default='false'):
         self.density_name = "DensityAdaptor"
@@ -684,7 +683,7 @@ class InitBoundaryHandling(JinjaCppFile):
 {% endfor %}
 """)  # noqa
 
-    def __init__(self, block_forest, flag_field_id, pdf_field_id, boundary_conditions):
+    def __init__(self, block_forest, flag_field_id, pdf_field_id, boundary_conditions, boundary_kernel: dict, field_allocations):
         self.fluid = FlagUidDefinition("fluid")
         ast_dict = {'fluid_uid_definition': self.fluid,
                     'geometry_initialization': BoundaryHandlingFromConfig(block_forest,
@@ -694,20 +693,51 @@ class InitBoundaryHandling(JinjaCppFile):
                                                                              b,
                                                                              pdf_field_id,
                                                                              flag_field_id,
-                                                                             self.fluid.symbol)
+                                                                             self.fluid.symbol,
+                                                                             boundary_kernel[b.name],
+                                                                             field_allocations)
                                              for b in boundary_conditions]
                     }
         super().__init__(ast_dict)
 
     headers = ['"cuda/FieldCopy.h"', '"geometry/InitBoundaryHandling.h"']
+    @property
+    def undefined_symbols(self):
+        rtn = super().undefined_symbols
+        for b in self.ast_dict.generated_boundaries:
+            rtn = rtn | b.undefined_symbols
+        return rtn
 
 
 class GeneratedBoundaryInitialization(JinjaCppFile):
-    TEMPLATE = jinja2.Template("""lbm::{{ boundary_condition }} {{ identifier }}( {{ block_forest }}, {{ pdf_field_id }} );
+    TEMPLATE = jinja2.Template("""lbm::{{ boundary_condition }} {{ identifier }}( {{ block_forest }}, {{ pdf_field_id }}{{ parameter_str }} );
 {{ 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):
+    def __init__(self,
+                 block_forest,
+                 boundary_condition,
+                 pdf_field_id,
+                 flag_field_id,
+                 fluid_uid,
+                 kernel,
+                 field_allocations):
+        def resolve_parameter(p):
+            if kernel.target == 'cpu':
+                dict = field_allocations._cpu_allocations
+            else:
+                dict = field_allocations._gpu_allocations
+
+            return dict.get(p.symbol.name.replace('_data_', ''), p).symbol
+        parameters = kernel.get_parameters()
+        parameter_ids = [resolve_parameter(p)
+                         for p in parameters
+                         if (p.is_field_pointer or not p.is_field_parameter) and
+                         p.symbol.name not in ('_data_indexVector', '_data_pdfs', 'indexVectorSize',)]
+        parameter_str = ', '.join(p.name for p in parameter_ids)
+        if parameter_str:
+            parameter_str = ', ' + parameter_str
+
         self.fluid = FlagUidDefinition("fluid")
         ast_dict = {'block_forest': block_forest,
                     'boundary_condition': pascalcase(boundary_condition.name),
@@ -715,6 +745,8 @@ class GeneratedBoundaryInitialization(JinjaCppFile):
                     'pdf_field_id': pdf_field_id,
                     'fluid_uid': fluid_uid,
                     'flag_field_id': flag_field_id,
+                    'parameter_ids': parameter_ids,
+                    'parameter_str': parameter_str
                     }
         super().__init__(ast_dict)
 
@@ -728,28 +760,41 @@ class GeneratedBoundaryInitialization(JinjaCppFile):
         # TypedSymbol(self.ast_dict.identifier, 'FlagUID'),
         return {TypedSymbol(self.ast_dict.identifier, f'lbm::{self.ast_dict.boundary_condition}')}
 
+    @property
+    def undefined_symbols(self):
+        return super().undefined_symbols | set(self.ast_dict.parameter_ids)
+
 
 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):
+        def resolve_parameter(p):
+            if ast.target == 'cpu':
+                dict = field_allocation._cpu_allocations
+            else:
+                dict = field_allocation._gpu_allocations
+
+            return dict.get(p.symbol.name.replace('_data_', ''), p).symbol
+
         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
+        parameter_ids = [resolve_parameter(p)
                          for p in parameters
-                         if p.is_field_pointer or not p.is_field_parameter
-                         ]
+                         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)}
+                    'parameter_str': ', '.join(p.name for p in parameter_ids)}
         super().__init__(ast_dict)
 
     @property
     def headers(self):
         return [f'"{self.ast_dict.sweep_class_name}.h"']
 
+    @property
+    def undefined_symbols(self):
+        return set(self.ast_dict.parameter_ids)
+
 
 class SweepOverAllBlocks(JinjaCppFile):
     # TEMPLATE = jinja2.Template("""std::for_each({{block_forest}}->begin(), {{block_forest}}->end(), {{functor}});""")  # noqa

src/pystencils_autodiff/wald_und_wiesen_simulation.py

@@ -48,6 +48,7 @@ class WaldUndWiesenSimulation():
         self._kernel_class_generator = WaldUndWiesenSimulation._get_sweep_class_name()
         self._with_gui = False
         self._with_gui_default = False
+        self._boundary_kernels = {}
 
     def _create_helper_files(self) -> Dict[str, str]:
         if self._lb_rule:
@@ -56,8 +57,8 @@ class WaldUndWiesenSimulation():
                                                   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)
+                    self._boundary_kernels.update({bc.name: lbmpy_walberla.generate_boundary(
+                        self._codegen_context, pascalcase(bc.name), bc, self._lb_rule.method)})
 
     def _create_module(self):
         if self._lb_rule:
@@ -87,8 +88,12 @@ class WaldUndWiesenSimulation():
             ResolveUndefinedSymbols(
                 Block([
                     field_allocations,
-                    InitBoundaryHandling(self._block_forest.blocks, flag_field_id,
-                                         pdf_field_id, self.boundary_conditions)
+                    InitBoundaryHandling(self._block_forest.blocks,
+                                         flag_field_id,
+                                         pdf_field_id,
+                                         self.boundary_conditions,
+                                         self._boundary_kernels,
+                                         self._field_allocations)
                     if self._boundary_handling else EmptyLine(),
                     LbCommunicationSetup(self._lb_model_name,
                                          pdf_field_id)

tests/test_walberla.py

@@ -78,7 +78,10 @@ def test_wald_wiesen_lbm():
     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)
+        import sympy as sp
+        lid_velocity = sp.symbols('lid_velocity')
+        lbm_step = ldc_setup(domain_size=(30, 30), optimization=opt_params,
+                             fixed_loop_sizes=False, lid_velocity=lid_velocity)
 
         sim = WaldUndWiesenSimulation(lbm_step.data_handling,
                                       ctx,
@@ -108,4 +111,3 @@ def test_resolve_parameters():
     ])))
 
     print(module)
-