diff --git a/src/pystencils_autodiff/framework_integration/astnodes.py b/src/pystencils_autodiff/framework_integration/astnodes.py index 131816316810cf4987b65cb1b3c92c28441a681f..99f1da8ad8845b4394b2ef9854f441485ecbf5b4 100644 --- a/src/pystencils_autodiff/framework_integration/astnodes.py +++ b/src/pystencils_autodiff/framework_integration/astnodes.py @@ -214,7 +214,7 @@ class JinjaCppFile(Node): 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)) and a is not None else a) for a in v] for k, v in self.ast_dict.items()} diff --git a/src/pystencils_autodiff/walberla.py b/src/pystencils_autodiff/walberla.py index 328a8b32624807299b875f7563117e0cf60f7e8d..8716a7e78211c70c2a5031605b9be1c46f0f2d9f 100644 --- a/src/pystencils_autodiff/walberla.py +++ b/src/pystencils_autodiff/walberla.py @@ -11,6 +11,7 @@ from os.path import dirname, join import jinja2 +import sympy as sp from pystencils.data_types import TypedSymbol from pystencils_autodiff._file_io import read_template_from_file @@ -160,8 +161,8 @@ class FieldAllocation(JinjaCppFile): {%- if init_value -%} , {{ init_value }} {% endif %} {%- if layout_str -%} , {{ layout_str }} {% endif %} {%- if num_ghost_layers -%}, {{ num_ghost_layers }} {% endif %} - {%- if always_initialize -%}, {{ always_initialize }} {% endif %}) - """) # noqa + {%- if always_initialize -%}, {{ always_initialize }} {% endif %}); +""") # noqa def __init__(self, block_forrest, field): self._symbol = TypedSymbol(field.name + '_data', 'BlockDataID') @@ -181,3 +182,78 @@ class FieldAllocation(JinjaCppFile): @property def symbols_defined(self): return {self.symbol} + + +class WalberlaVector(JinjaCppFile): + TEMPLATE = jinja2.Template("""math::Vector{{ndim}}<{{dtype}}>({{offsets}})""") + + def __init__(self, offset, dtype='real_t'): + ast_dict = { + 'offset': offset, + 'offsets': ', '.join(str(o) for o in offset), + 'dtype': dtype, + 'ndim': len(offset), + } + super().__init__(ast_dict) + + def __sympy__(self): + 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_forrest }}, + {{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_forrest, field, lattice_model, initial_velocity=None, initial_density=None): + super().__init__(block_forrest, field) + 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_forrest }}, + {{field_name}} + {%- if num_ghost_layers -%}, {{num_ghost_layers }} {% endif %}); +""") # noqa + + def __init__(self, block_forrest, field): + super().__init__(block_forrest, field) + + headers = ['field/AddToStorage.h'] diff --git a/tests/test_walberla.py b/tests/test_walberla.py index 9147f41495050b4844108e0cab40dea32fd5679b..5f2859e6e8327d04b6a1dc2bb0bc8e8005bf86af 100644 --- a/tests/test_walberla.py +++ b/tests/test_walberla.py @@ -14,12 +14,16 @@ from pystencils.astnodes import Block, EmptyLine, SympyAssignment from pystencils.data_types import TypedSymbol from pystencils_autodiff._file_io import write_file from pystencils_autodiff.walberla import ( - DefinitionsHeader, FieldAllocation, GetParameter, UniformBlockForrestFromConfig, - WalberlaMain, WalberlaModule) + DefinitionsHeader, FieldAllocation, GetParameter, PdfFieldAllocation,FlagFieldAllocation, + UniformBlockForrestFromConfig, WalberlaMain, WalberlaModule) def test_walberla(): x, y = pystencils.fields('x, y: float32[3d]') + pdf = pystencils.fields('pdf(27): float32[3d]') + pdf2 = pystencils.fields('pdf2(27): float32[3d]') + flags = pystencils.fields('flag_field: float32[3d]') + foo_symbol = TypedSymbol('foo', np.bool) number_symbol = TypedSymbol('number', np.float32) crazy_plus_one = TypedSymbol('crazy', np.float32) @@ -32,7 +36,10 @@ def test_walberla(): SympyAssignment(number_symbol, GetParameter('parameters', number_symbol, 1.2)), SympyAssignment(crazy_plus_one, number_symbol + 1), EmptyLine(), - FieldAllocation(block_forrest.blocks, x) + 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) ]) module = WalberlaModule(WalberlaMain(block))