diff --git a/src/sfg_walberla/api.py b/src/sfg_walberla/api.py
index 482ec88fb4e10bcc732d0e8f04dfd389857d89a0..f5d391a9728ff4a57c7cab5e6a8999bcef0a2c3e 100644
--- a/src/sfg_walberla/api.py
+++ b/src/sfg_walberla/api.py
@@ -15,6 +15,19 @@ from pystencilssfg.ir import SfgHeaderInclude
real_t = PsCustomType("walberla::real_t")
+class Vector2(SrcVector):
+ def __init__(self, val_type: UserTypeSpec):
+ self._value_type = create_type(val_type)
+ val_type_str = self._value_type.c_string()
+ super().__init__(PsCustomType(f"walberla::Vector2< {val_type_str} >"))
+
+ def extract_component(self, coordinate: int) -> AugExpr:
+ if coordinate > 1:
+ raise ValueError(f"Cannot extract component {coordinate} from Vector2")
+
+ return AugExpr(self._value_type).bind("{}[{}]", self, coordinate)
+
+
class Vector3(SrcVector):
def __init__(self, val_type: UserTypeSpec):
self._value_type = create_type(val_type)
@@ -22,6 +35,9 @@ class Vector3(SrcVector):
super().__init__(PsCustomType(f"walberla::Vector3< {val_type_str} >"))
def extract_component(self, coordinate: int) -> AugExpr:
+ if coordinate > 2:
+ raise ValueError(f"Cannot extract component {coordinate} from Vector3")
+
return AugExpr(self._value_type).bind("{}[{}]", self, coordinate)
diff --git a/src/sfg_walberla/boundaries/hbb.py b/src/sfg_walberla/boundaries/hbb.py
index 27626306239007bdb88469094c925474b617c767..97db11db6c82e07111d6e93d6848aa15d5632b0d 100644
--- a/src/sfg_walberla/boundaries/hbb.py
+++ b/src/sfg_walberla/boundaries/hbb.py
@@ -13,7 +13,7 @@ from lbmpy.boundaries.boundaryconditions import LbBoundary
from lbmpy.boundaries.boundaryhandling import create_lattice_boltzmann_boundary_kernel
from pystencilssfg.lang.expressions import IFieldExtraction
-from ..sweep import FieldInfo, SweepClassProperties
+from ..sweep import FieldInfo, SweepClassProperties, combine_vectors
from ..api import GhostLayerFieldPtr, BlockDataID, IBlockPtr, StructuredBlockForest
BoundaryIndexType = create_type("int32")
@@ -107,14 +107,15 @@ class SimpleHbbBoundary(CustomGenerator):
key=lambda f: f.field.name,
)
- scalar_params = khandle.scalar_parameters
+ parameters = khandle.scalar_parameters
+ vector_groups = combine_vectors(parameters)
props = HbbBoundaryProperties()
for fi in domain_fields:
props.add_property(fi.data_id, setter=False, getter=True)
- for s in sorted(scalar_params, key=lambda p: p.name):
+ for s in sorted(parameters, key=lambda p: p.name):
props.add_property(s, setter=True, getter=True)
block = IBlockPtr().var("block")
@@ -142,14 +143,16 @@ class SimpleHbbBoundary(CustomGenerator):
),
# Map GhostLayerFields to pystencils fields
*(sfg.map_field(fi.field, fi.glfield_ptr) for fi in domain_fields),
- # Extract scalar parameters
- *(sfg.init(param)(props.get(param)) for param in scalar_params),
# Get index vector
sfg.init(idx_vector)(
block.getData(idx_vector.vector_type, idx_vector_id)
),
# Map index vector to pystencils field
sfg.map_field(self._index_field, idx_vector),
+ # Extract parameters
+ *(sfg.init(param)(props.get(param)) for param in parameters),
+ # Extract vector components
+ *(sfg.map_vector(comps, vec) for vec, comps in vector_groups.items()),
# Call kernel
sfg.call(khandle),
)
diff --git a/src/sfg_walberla/sweep.py b/src/sfg_walberla/sweep.py
index 4e48c515a0f61e8bf44bc98bf95320de35d064d7..59f7b387074921c572e8b6b161aa92a3c25e7a7d 100644
--- a/src/sfg_walberla/sweep.py
+++ b/src/sfg_walberla/sweep.py
@@ -2,6 +2,7 @@ from typing import Sequence, Iterable, Callable
from dataclasses import dataclass, replace
from itertools import chain
from collections import defaultdict
+import warnings
import sympy as sp
@@ -19,8 +20,8 @@ from pystencils import (
from pystencils.types import PsType, constify, deconstify, PsCustomType
from pystencilssfg import SfgComposer
-from pystencilssfg.lang import VarLike, asvar, SfgVar, AugExpr, Ref
-from .api import GhostLayerFieldPtr, IBlockPtr, BlockDataID, Vector3
+from pystencilssfg.lang import VarLike, asvar, SfgVar, AugExpr, Ref, SrcVector
+from .api import GhostLayerFieldPtr, IBlockPtr, BlockDataID, Vector2, Vector3
class SweepClassProperties:
@@ -188,18 +189,68 @@ class ShadowFieldCache:
return f"shadow_{str(orig)}_"
-# TODO: Happens outside of Properties class
-# def combine_vectors():
-# vector_components: defaultdict[str, list[SweepClassProperties.Property]] = (
-# defaultdict(lambda _: [None, None, None])
-# )
+def combine_vectors(scalars: set[SfgVar]) -> dict[SrcVector, tuple[SfgVar, ...]]:
+ """Attempt to combine vector component symbols into vectors.
-# for prop_name, prop in self._properties.items():
-# tokens = prop_name.split("_")
-# if tokens[-1] in ("0", "1", "2"):
-# vec_name = "_".join(tokens[:-1])
-# vec_comp = int(tokens[-1])
-# vector_components[vec_name][vec_comp] = prop
+ This function modifies the `scalars` parameter in-place by removing grouped scalar
+ components and replacing them by their vectors.
+ """
+ potential_vectors: defaultdict[str, list[tuple[int, SfgVar]]] = defaultdict(list)
+
+ for s in scalars:
+ tokens = s.name.split("_")
+ try:
+ coord = int(tokens[-1])
+ vec_name = "_".join(tokens[:-1])
+ potential_vectors[vec_name].append((coord, s))
+ except ValueError:
+ pass
+
+ all_vectors: dict[SrcVector, tuple[SfgVar, ...]] = dict()
+
+ for name, entries in potential_vectors.items():
+ entries = sorted(entries, key=lambda t: t[0])
+ coords = [e[0] for e in entries]
+ components = tuple(e[1] for e in entries)
+ scalar_types = set(e[1].dtype for e in entries)
+
+ if set(coords) == {0, 1}:
+ vec_type = Vector2
+ elif set(coords) == {0, 1, 2}:
+ vec_type = Vector3
+ else:
+ continue
+
+ if len(scalar_types) > 1:
+ warnings.warn(
+ f"Skipping vector grouping for vector {name}: "
+ f"Encountered vector components {components} with conflicting data types {scalar_types}.",
+ UserWarning,
+ )
+ continue
+
+ name_conflict = False
+ for s in scalars:
+ if s.name == name:
+ warnings.warn(
+ f"Skipping vector grouping for vector {name}: "
+ f"Vector name conflicts with existing scalar parameter {s}.",
+ UserWarning,
+ )
+ name_conflict = True
+
+ if name_conflict:
+ continue
+
+ scalar_type = deconstify(scalar_types.pop())
+ vector = vec_type(scalar_type).var(name)
+ all_vectors[vector] = components
+
+ for c in components:
+ scalars.remove(c)
+ scalars.add(asvar(vector))
+
+ return all_vectors
class Sweep(CustomGenerator):
@@ -225,25 +276,9 @@ class Sweep(CustomGenerator):
self._kernel = kernel
self._gen_config = config
- # Vectors
- self._vectors: dict[Vector3, tuple[sp.Symbol, ...]] = dict()
-
# Map from shadow field to shadowed field
self._shadow_fields: dict[Field, Field] = dict()
- def vector_param(self, vector: Vector3, vector_entries: tuple[sp.Symbol]):
- if len(vector_entries) not in (2, 3):
- raise ValueError("Can only combine 2D or 3D vectors")
-
- if vector in self._vectors:
- raise KeyError(f"Duplicate vector parameter: {vector}")
-
- if not vector.is_bound():
- v_name = vector_entries[0].name.split("_")[0]
- vector = vector.var(v_name)
-
- self._vectors[vector] = vector_entries
-
def swap_fields(self, field: Field, shadow_field: Field):
if field in self._shadow_fields:
raise ValueError(f"Field swap for {field} was already registered.")
@@ -286,14 +321,7 @@ class Sweep(CustomGenerator):
parameters = khandle.scalar_parameters
- # Remove vector entries from parameter list, add vectors instead
- vector_entry_names: set[str] = (
- set.union(*(set(e.name for e in vec) for vec in self._vectors.values()))
- if self._vectors
- else set()
- )
- parameters = set(filter(lambda p: p.name not in vector_entry_names, parameters))
- parameters |= set(asvar(vec) for vec in self._vectors)
+ vector_groups = combine_vectors(parameters)
props = SweepClassProperties()
@@ -336,7 +364,7 @@ class Sweep(CustomGenerator):
# Extract scalars from vectors
*(
sfg.map_vector(components, vector)
- for vector, components in self._vectors.items()
+ for vector, components in vector_groups.items()
),
sfg.call(khandle),
# Perform field swaps