diff --git a/src/pystencils_autodiff/walberla.py b/src/pystencils_autodiff/walberla.py
index 0f4ac2e348898c33a05adda4c0ca2d44aba3008b..69f939d8180f00041ef1b80cc56b4b483c77bfa7 100644
--- a/src/pystencils_autodiff/walberla.py
+++ b/src/pystencils_autodiff/walberla.py
@@ -12,6 +12,7 @@ from abc import ABC
from os.path import dirname, join
import jinja2
+import numpy as np
import sympy as sp
from stringcase import camelcase, pascalcase
@@ -208,7 +209,7 @@ BlockDataID {{ field_name }}_data_gpu = cuda::addGPUFieldToStorage<{{ field_type
{{ usePitchedMem }} );
{%- else -%}
BlockDataID {{ field_name }}_data = field::addToStorage<{{ field_type }}>( {{ block_forest }},
- "{{ field_name }}"
+ "{{ field_name }}"
{%- if init_value -%} , {{ init_value }}{% endif %}
{%- if layout -%} , {{ layout }}{% endif %}
{%- if num_ghost_layers -%}, {{ num_ghost_layers }} {% endif %}
@@ -216,7 +217,7 @@ BlockDataID {{ field_name }}_data = field::addToStorage<{{ field_type }}>( {{ bl
{%- endif %}
""") # noqa
- def __init__(self, block_forest, field, on_gpu=False, usePitchedMem=True, num_ghost_layers=1):
+ 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_forest': block_forest,
@@ -247,6 +248,13 @@ BlockDataID {{ field_name }}_data = field::addToStorage<{{ field_type }}>( {{ bl
else ['"field/AddToStorage.h"'])
+class FlagFieldAllocation(FieldAllocation):
+ """We need this special class to pass waLBerla's runtime type checks"""
+ TEMPLATE = jinja2.Template("""BlockDataID {{ field_name }}_data = field::addFlagFieldToStorage<FlagField_T>({{ block_forest }},
+ "{{ field_name }}");
+""") # noqa
+
+
class WalberlaVector(JinjaCppFile):
TEMPLATE = jinja2.Template("""math::Vector{{ndim}}<{{dtype}}>({{offsets}})""")
@@ -263,65 +271,6 @@ class WalberlaVector(JinjaCppFile):
return sp.Matrix(self.ast_dict.offset)
-class PdfFieldAllocation(FieldAllocation):
- """
- .. code: : cpp
-
- BlockDataID addPdfFieldToStorage(const shared_ptr < BlockStorage_T > & blocks, const std:: string & identifier,
- const LatticeModel_T & latticeModel,
- const Vector3 < real_t > & initialVelocity, const real_t initialDensity,
- const uint_t ghostLayers,
- const field:: Layout & layout = field: : zyxf,
- 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_forest }},
- {{field_name}},
- {{lattice_model}}
- {%- if initial_velocity -%} , {{initial_velocity }} {% endif %}
- {%- if initial_density -%} , {{initial_density }} {% endif %}
- {%- if num_ghost_layers -%}, {{num_ghost_layers }} {% endif %});
-""") # noqa
-
- 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)
-
- self.ast_dict.update({
- 'initial_density': initial_density,
- 'lattice_model': lattice_model,
- 'initial_velocity': initial_velocity,
- })
-
- headers = ['"lbm/field/AddToStorage.h"']
-
-
-class FlagFieldAllocation(FieldAllocation):
- """
- .. code: : cpp
-
- template< typename FlagField_T, typename BlockStorage_T >
- BlockDataID addFlagFieldToStorage( const shared_ptr< BlockStorage_T > & blocks,
- const std::string & identifier,
- const uint_t nrOfGhostLayers = uint_t(1),
- const bool alwaysInitialize = false,
- const std::function< void ( FlagField_T * field, IBlock * const block ) > & initFunction =
- std::function< void ( FlagField_T * field, IBlock * const block ) >(),
- 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_forest }},
- {{field_name}}
- {%- if num_ghost_layers -%}, {{num_ghost_layers }} {% endif %});
-""") # noqa
-
- 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 }}");')
@@ -693,6 +642,15 @@ class DefineKernelObjects(JinjaCppFile):
JinjaCppFile.__init__(self, ast_dict)
+def _allocate_field(block_forest, field, on_gpu):
+ # TODO: hack!
+ if np.issubdtype(field.dtype.numpy_dtype, np.integer) and 'flag' in field.name.lower():
+ assert on_gpu is False
+ return FlagFieldAllocation(block_forest, field)
+ else:
+ return FieldAllocation(block_forest, field, on_gpu=on_gpu)
+
+
class AllocateAllFields(JinjaCppFile):
TEMPLATE = jinja2.Template("""
@@ -731,13 +689,13 @@ class AllocateAllFields(JinjaCppFile):
@property
def _cpu_allocations(self):
- return {s: FieldAllocation(self.block_forest, self.data_handling.fields[s])
+ return {s: _allocate_field(self.block_forest, self.data_handling.fields[s], on_gpu=False)
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)
+ return {s: _allocate_field(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)}
@@ -914,7 +872,7 @@ static inline auto sweep(walberla::shared_ptr<BlockStorage_T> blocks, Functor_T
class FieldCopy(JinjaCppFile):
- TEMPLATE = jinja2.Template("""cuda::fieldCpy < {{src_type }}, {{dst_type }} > ({{block_forest }}, {{src_id }}, {{dst_id }}); """) # noqa
+ TEMPLATE = jinja2.Template("""cuda::fieldCpy< {{src_type }}, {{dst_type }} > ({{block_forest }}, {{src_id }}, {{dst_id }}); """) # noqa
def __init__(self, block_forest, src_id, src_field, src_gpu, dst_id, dst_field, dst_gpu):
src_type = _make_field_type(src_field, src_gpu)
diff --git a/tests/test_walberla.py b/tests/test_walberla.py
index 58707632d9b82e8ea726c6c6084aba3152674664..13e5c8d041e160a16ecfc38e9906e163b28dbdba 100644
--- a/tests/test_walberla.py
+++ b/tests/test_walberla.py
@@ -19,7 +19,7 @@ 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, FlagFieldAllocation, GetParameter, PdfFieldAllocation,
+ DefinitionsHeader, FieldAllocation, FlagFieldAllocation, GetParameter,
ResolveUndefinedSymbols, UniformBlockforestFromConfig, WalberlaMain, WalberlaModule)
from pystencils_autodiff.wald_und_wiesen_simulation import WaldUndWiesenSimulation
from pystencils_walberla.cmake_integration import ManualCodeGenerationContext
@@ -44,8 +44,6 @@ def test_walberla():
SympyAssignment(crazy_plus_one, number_symbol + 1),
EmptyLine(),
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)
])