diff --git a/src/sfg_walberla/api.py b/src/sfg_walberla/api.py
index 8e087b2530b717f741d1b3fabc85d1bdca18d822..2fec8a000dc80f5755e503a13194de4ba5b1d23c 100644
--- a/src/sfg_walberla/api.py
+++ b/src/sfg_walberla/api.py
@@ -8,11 +8,20 @@ from pystencils.types import (
PsPointerType,
PsType,
)
-from pystencilssfg.lang import IFieldExtraction, AugExpr, SrcField, SrcVector, Ref, ExprLike
+from pystencilssfg.lang import (
+ IFieldExtraction,
+ AugExpr,
+ SrcField,
+ SrcVector,
+ Ref,
+ ExprLike,
+)
from pystencilssfg.ir import SfgHeaderInclude
real_t = PsCustomType("walberla::real_t")
+cell_idx_t = PsCustomType("walberla::cell_idx_t")
+uint_t = PsCustomType("walberla::uint_t")
class Vector2(SrcVector):
@@ -50,10 +59,10 @@ class AABB(AugExpr):
super().__init__(dtype)
def min(self) -> Vector3:
- return Vector3().bind("{}.min()", self)
+ return Vector3(real_t).bind("{}.min()", self)
def max(self) -> Vector3:
- return Vector3().bind("{}.max()", self)
+ return Vector3(real_t).bind("{}.max()", self)
class BlockDataID(AugExpr):
@@ -79,10 +88,25 @@ class IBlockPtr(AugExpr):
def required_includes(self) -> set[SfgHeaderInclude]:
return {SfgHeaderInclude.parse("domain_decomposition/IBlock.h")}
+ def deref(self) -> AugExpr:
+ return AugExpr.format("*{}", self)
+
+
+class StructuredBlockForest(AugExpr):
+ def __init__(self, ref: bool = True, const: bool = False):
+ dtype = self.ref_type(const) if ref else self.cpp_type(const)
+ super().__init__(dtype)
+
+ @staticmethod
+ def shared_ptr(name: str = "blocks", const: bool = False):
+ const_str = "const " if const else ""
+ dtype = PsCustomType(
+ f"std::shared_ptr< {const_str}walberla::StructuredBlockForest >", const=True
+ )
+ return AugExpr(dtype).var(name)
-class StructuredBlockForest:
@staticmethod
- def shared_ptr_ref(name, const: bool = False):
+ def shared_ptr_ref(name: str = "blocks", const: bool = False):
const_str = "const " if const else ""
dtype = PsCustomType(
f"std::shared_ptr< {const_str}walberla::StructuredBlockForest >", const=True
@@ -90,12 +114,36 @@ class StructuredBlockForest:
return AugExpr(Ref(dtype)).var(name)
@staticmethod
- def ref_type(const: bool = False):
- return Ref(PsCustomType("walberla::StructuredBlockForest", const=const))
+ def cpp_type(const: bool = False) -> PsType:
+ return PsCustomType("walberla::StructuredBlockForest", const=const)
+
+ @staticmethod
+ def ref_type(const: bool = False) -> Ref:
+ return Ref(StructuredBlockForest.cpp_type(const))
@staticmethod
- def ref(name, const: bool = False):
- return AugExpr(StructuredBlockForest.ref_type(const)).var(name)
+ def ref(name: str = "blocks", const: bool = False):
+ return StructuredBlockForest(ref=True).var(name)
+
+ def dx(self, level: int | AugExpr = 0) -> AugExpr:
+ return AugExpr(real_t).bind("{}.dx({})", self, level)
+
+ def dy(self, level: int | AugExpr = 0) -> AugExpr:
+ return AugExpr(real_t).bind("{}.dy({})", self, level)
+
+ def dz(self, level: int | AugExpr = 0) -> AugExpr:
+ return AugExpr(real_t).bind("{}.dz({})", self, level)
+
+ def cell_extents(self, coord: int, level: int | AugExpr = 0) -> AugExpr:
+ dx_getter = [self.dx, self.dy, self.dz][coord]
+ return dx_getter(level)
+
+ def cellsPerBlock(self, coord: int, block: AugExpr) -> AugExpr:
+ x = ["X", "Y", "Z"][coord]
+ return AugExpr(uint_t).bind("{}.getNumberOf{}Cells({})", self, x, block)
+
+ def getLevel(self, block: AugExpr) -> AugExpr:
+ return AugExpr(uint_t).bind("{}.getLevel({})", self, block)
class GhostLayerFieldPtr(SrcField):
diff --git a/src/sfg_walberla/sweep.py b/src/sfg_walberla/sweep.py
index c1c2c1f403c677d913578f77922f59fd578077c7..4e73446b52e0fae63271481048ba4f9a906857a6 100644
--- a/src/sfg_walberla/sweep.py
+++ b/src/sfg_walberla/sweep.py
@@ -11,17 +11,23 @@ from pystencilssfg.composer.class_composer import SfgClassComposer
from pystencilssfg.ir import SfgMethod
from pystencils import (
- KernelFunction,
Assignment,
AssignmentCollection,
CreateKernelConfig,
Field,
- TypedSymbol,
)
from pystencils.types import PsType, constify, deconstify, PsCustomType
from pystencilssfg import SfgComposer
-from pystencilssfg.lang import VarLike, asvar, SfgVar, AugExpr, Ref, SrcVector
+from pystencilssfg.lang import (
+ VarLike,
+ asvar,
+ SfgVar,
+ AugExpr,
+ Ref,
+ SrcVector,
+ strip_ptr_ref,
+)
from .api import (
StructuredBlockForest,
GhostLayerFieldPtr,
@@ -29,6 +35,7 @@ from .api import (
BlockDataID,
Vector2,
Vector3,
+ real_t,
)
@@ -39,7 +46,7 @@ class SweepClassProperties:
dtype: PsType
getter: bool
setter: bool
- initializer: AugExpr | None = None
+ initializer: tuple[AugExpr, ...] | None = None
def __post_init__(self):
if self.dtype.const and self.setter:
@@ -86,7 +93,12 @@ class SweepClassProperties:
return self._properties.values()
def add_property(
- self, prop: VarLike, const=False, setter: bool = True, getter: bool = True
+ self,
+ prop: VarLike,
+ const=False,
+ setter: bool = True,
+ getter: bool = True,
+ initializer: tuple[AugExpr, ...] | None = None,
):
prop = asvar(prop)
if prop.name in self._properties:
@@ -95,7 +107,7 @@ class SweepClassProperties:
dtype = constify(prop.dtype) if const else deconstify(prop.dtype)
self._properties[prop.name] = SweepClassProperties.Property(
- prop.name, dtype, getter, setter
+ prop.name, dtype, getter, setter, initializer
)
def get(self, prop: VarLike) -> AugExpr:
@@ -115,16 +127,19 @@ class SweepClassProperties:
if p.initializer is None:
ctor.add_param(p.var)
else:
- for var in p.initializer.depends:
+ for var in chain.from_iterable(e.depends for e in p.initializer):
if var not in ctor.parameters:
- if deconstify(var.dtype) == StructuredBlockForest.ref_type():
+ if (
+ deconstify(strip_ptr_ref(var.dtype))
+ == StructuredBlockForest.cpp_type()
+ ):
ctor.add_param(var, 0)
else:
ctor.add_param(var)
for p in self._properties.values():
if p.initializer is not None:
- ctor.init(p.member_var)(p.initializer)
+ ctor.init(p.member_var)(*p.initializer)
else:
ctor.init(p.member_var)(p.var)
@@ -148,13 +163,14 @@ class BlockforestParameters:
self._block = block
self._ci = cell_interval
self._extractions: dict[SfgVar, AugExpr] = dict()
- self._properties: list[SweepClassProperties.Property] = []
+
+ self._blockforest: StructuredBlockForest | None = None
@staticmethod
def process(asms: AssignmentCollection) -> AssignmentCollection:
from sympy.core.function import AppliedUndef
- expandable_appls: set[AppliedUndef] = filter(
+ expandable_appls: set[AppliedUndef] = filter( # type: ignore
lambda expr: hasattr(expr, "expansion_func"), asms.atoms(AppliedUndef)
)
@@ -167,6 +183,22 @@ class BlockforestParameters:
return asms.new_with_substitutions(subs)
+ def blockforest(self) -> StructuredBlockForest:
+ if self._blockforest is None:
+ blocks = StructuredBlockForest.shared_ptr()
+ self._property_cache.add_property(
+ blocks,
+ getter=False,
+ setter=False,
+ initializer=(StructuredBlockForest.shared_ptr_ref(),),
+ )
+ self._blockforest = StructuredBlockForest().bind(
+ "*({})", self._property_cache.get(blocks)
+ )
+ return self._blockforest
+ else:
+ return self._blockforest
+
def filter_params(self, params: set[SfgVar]) -> set[SfgVar]:
from .symbolic import block, cell
@@ -180,12 +212,34 @@ class BlockforestParameters:
elif param.name == block.aabb_max[coord].name:
self._extractions[param] = self._block.getAABB().max()[coord]
break
- # TODO: ci
+ elif param.name == block.ci_min[coord].name:
+ if self._ci is None:
+ self._extractions[param] = AugExpr.format("0")
+ else:
+ raise NotImplementedError()
+ break
+ elif param.name == block.ci_max[coord].name:
+ if self._ci is None:
+ self._extractions[param] = AugExpr.format(
+ "{} - 1",
+ self.blockforest().cellsPerBlock(coord, self._block),
+ )
+ else:
+ raise NotImplementedError()
+ break
elif param.name == cell.cell_extents[coord].name:
- pass # TODO
+ self._extractions[param] = self.blockforest().cell_extents(
+ coord, self.blockforest().getLevel(self._block.deref())
+ )
+ break
else:
params_filtered.add(param)
+ return params_filtered
+
+ def render_extractions(self, sfg: SfgComposer):
+ return (sfg.set_param(p, expr) for p, expr in self._extractions.items())
+
@dataclass
class FieldInfo:
@@ -329,7 +383,7 @@ class Sweep(CustomGenerator):
def __init__(
self,
name: str,
- kernel: Sequence[Assignment] | AssignmentCollection | KernelFunction,
+ assignments: Sequence[Assignment] | AssignmentCollection,
config: CreateKernelConfig | None = None,
):
if config is not None:
@@ -343,7 +397,10 @@ class Sweep(CustomGenerator):
config = CreateKernelConfig(ghost_layers=0)
self._name = name
- self._kernel = kernel
+ if isinstance(assignments, AssignmentCollection):
+ self._assignments = assignments
+ else:
+ self._assignments = AssignmentCollection(assignments)
self._gen_config = config
# Map from shadow field to shadowed field
@@ -364,10 +421,8 @@ class Sweep(CustomGenerator):
def generate(self, sfg: SfgComposer) -> None:
knamespace = sfg.kernel_namespace(f"{self._name}_kernels")
- if isinstance(self._kernel, KernelFunction):
- khandle = knamespace.add(self._kernel, self._name)
- else:
- khandle = knamespace.create(self._kernel, self._name, self._gen_config)
+ assignments = BlockforestParameters.process(self._assignments)
+ khandle = knamespace.create(assignments, self._name, self._gen_config)
all_fields: dict[str, FieldInfo] = {
f.name: FieldInfo(
@@ -389,11 +444,15 @@ class Sweep(CustomGenerator):
key=lambda fi: fi.field.name,
)
+ props = SweepClassProperties()
+
parameters = khandle.scalar_parameters
- vector_groups = combine_vectors(parameters)
+ blockforest_params = BlockforestParameters(props, block, None)
- props = SweepClassProperties()
+ parameters = blockforest_params.filter_params(parameters)
+
+ vector_groups = combine_vectors(parameters)
for fi in block_fields:
props.add_property(fi.data_id, setter=False, getter=True)
@@ -436,6 +495,8 @@ class Sweep(CustomGenerator):
sfg.map_vector(components, vector)
for vector, components in vector_groups.items()
),
+ # Extract geometry information
+ *(blockforest_params.render_extractions(sfg)),
sfg.call(khandle),
# Perform field swaps
*(
diff --git a/src/sfg_walberla/symbolic.py b/src/sfg_walberla/symbolic.py
index 8f6c0763ff4d203429d213250984fa173f772202..5349d00584b295729b4b34da9a8c959ad3f97e8c 100644
--- a/src/sfg_walberla/symbolic.py
+++ b/src/sfg_walberla/symbolic.py
@@ -28,7 +28,9 @@ block = BlockCoordinates
def expandable(name: str):
def wrap(expansion_func):
nargs = len(inspect.signature(expansion_func).parameters)
- return sp.Function(name, nargs=nargs, expansion_func=staticmethod(expansion_func))
+ return sp.Function(
+ name, nargs=nargs, expansion_func=staticmethod(expansion_func)
+ )
return wrap
@@ -56,13 +58,35 @@ class CellCoordinates:
2
] * CastFunc.as_numeric(block.ci_min[2] + DEFAULTS.spatial_counters[2])
- local_x = sp.Function("cell.local_x", nargs=0, walberla_info="cell.global_center.x")
- local_y = sp.Function("cell.local_y", nargs=0, walberla_info="cell.global_center.y")
- local_z = sp.Function("cell.local_z", nargs=0, walberla_info="cell.global_center.z")
-
- dx = sp.Function("cell.dx", nargs=0, walberla_info="cell.dx")
- dy = sp.Function("cell.dy", nargs=0, walberla_info="cell.dy")
- dz = sp.Function("cell.dz", nargs=0, walberla_info="cell.dz")
+ @expandable("cell.local_x")
+ def local_x():
+ return CellCoordinates.cell_extents[0] * CastFunc.as_numeric(
+ block.ci_min[0] + DEFAULTS.spatial_counters[0]
+ )
+
+ @expandable("cell.local_y")
+ def local_y():
+ return CellCoordinates.cell_extents[1] * CastFunc.as_numeric(
+ block.ci_min[1] + DEFAULTS.spatial_counters[1]
+ )
+
+ @expandable("cell.local_z")
+ def local_z():
+ return CellCoordinates.cell_extents[2] * CastFunc.as_numeric(
+ block.ci_min[2] + DEFAULTS.spatial_counters[2]
+ )
+
+ @expandable("cell.dx")
+ def dx():
+ return CellCoordinates.cell_extents[0]
+
+ @expandable("cell.dy")
+ def dy():
+ return CellCoordinates.cell_extents[1]
+
+ @expandable("cell.dz")
+ def dz():
+ return CellCoordinates.cell_extents[2]
cell = CellCoordinates
diff --git a/tests/BlockforestGeometry.py b/tests/BlockforestGeometry.py
new file mode 100644
index 0000000000000000000000000000000000000000..09445e9e8be29f59556150ff6f9b0754ff040d21
--- /dev/null
+++ b/tests/BlockforestGeometry.py
@@ -0,0 +1,20 @@
+from pystencils.session import *
+from pystencilssfg import SourceFileGenerator, SfgConfiguration
+from sfg_walberla import Sweep
+from sfg_walberla.symbolic import cell
+
+cfg = SfgConfiguration(
+ output_directory="out"
+)
+
+with SourceFileGenerator(cfg) as sfg:
+ f = ps.fields("f(3): [3D]")
+
+ asms = [
+ ps.Assignment(f.center(0), cell.x()),
+ ps.Assignment(f.center(1), cell.y()),
+ ps.Assignment(f.center(2), cell.z())
+ ]
+
+ sweep = Sweep("coordinates", asms)
+ sfg.generate(sweep)