diff --git a/hog/forms.py b/hog/forms.py index 8ba9b5571dc6526acedd8a0bf4d8814a09454578..113f035fcb8f8a2bf7f29e23a192bbd7a1efb5f4 100644 --- a/hog/forms.py +++ b/hog/forms.py @@ -18,7 +18,7 @@ import logging import sympy as sp from typing import Optional, Callable, Any -from hog.element_geometry import ElementGeometry, TetrahedronElement +from hog.element_geometry import ElementGeometry from hog.exception import HOGException from hog.fem_helpers import ( trafo_ref_to_affine, @@ -30,12 +30,10 @@ from hog.fem_helpers import ( element_matrix_iterator, scalar_space_dependent_coefficient, vector_space_dependent_coefficient, - fem_function_on_element, - fem_function_gradient_on_element, ) from hog.function_space import FunctionSpace, N1E1Space from hog.math_helpers import dot, inv, abs, det, double_contraction -from hog.quadrature import Quadrature, Tabulation +from hog.quadrature import Quadrature from hog.symbolizer import Symbolizer from hog.logger import TimedLogger from hog.blending import GeometryMap, ExternalMap, IdentityMap @@ -163,7 +161,7 @@ Weak formulation geometry, symbolizer, blending=blending, - fe_coefficients=[("k", coefficient_function_space)], + fe_coefficients={"k": coefficient_function_space}, is_symmetric=trial == test, docstring=docstring, ) @@ -230,7 +228,7 @@ Note: :math:`a(c) = 1/8 + u^2` is currently hard-coded and the form is intended blending=blending, is_symmetric=trial == test, docstring=docstring, - fe_coefficients=[("u", coefficient_function_space)], + fe_coefficients={"u": coefficient_function_space}, ) @@ -307,7 +305,7 @@ Note: :math:`a(k) = 1/8 + k^2` is currently hard-coded and the form is intended blending=blending, is_symmetric=False, docstring=docstring, - fe_coefficients=[("k", coefficient_function_space)], + fe_coefficients={"k": coefficient_function_space}, ) @@ -363,7 +361,7 @@ where blending=blending, is_symmetric=trial == test, docstring=docstring, - fe_coefficients=[("mu", coefficient_function_space)], + fe_coefficients={"mu": coefficient_function_space}, ) @@ -658,7 +656,7 @@ where blending=blending, is_symmetric=trial == test, docstring=docstring, - fe_coefficients=[("mu", coefficient_function_space)], + fe_coefficients={"mu": coefficient_function_space}, ) @@ -726,12 +724,12 @@ The resulting matrix must be multiplied with a vector of ones to be used as the ): """First function: mu, other functions: ux, uy, uz.""" - mu = k[0] + mu = k["mu"] # grad_k[0] is grad_mu_ref - grad_wx = jac_b_inv.T * jac_a_inv.T * grad_k[1] - grad_wy = jac_b_inv.T * jac_a_inv.T * grad_k[2] - grad_wz = jac_b_inv.T * jac_a_inv.T * grad_k[3] + grad_wx = jac_b_inv.T * jac_a_inv.T * grad_k["wx"] + grad_wy = jac_b_inv.T * jac_a_inv.T * grad_k["wy"] + grad_wz = jac_b_inv.T * jac_a_inv.T * grad_k["wz"] grad_w = grad_wx.row_join(grad_wy) dim = volume_geometry.dimensions @@ -758,12 +756,12 @@ The resulting matrix must be multiplied with a vector of ones to be used as the geometry, symbolizer, blending=blending, - fe_coefficients=[ - ("mu", viscosity_function_space), - ("wx", velocity_function_space), - ("wy", velocity_function_space), - ("wz", velocity_function_space), - ], + fe_coefficients={ + "mu": viscosity_function_space, + "wx": velocity_function_space, + "wy": velocity_function_space, + "wz": velocity_function_space, + }, is_symmetric=trial == test, docstring=docstring, ) @@ -863,10 +861,10 @@ Weak formulation ): """First function: k𝛿, other functions: ux, uy, uz.""" - k_times_delta = k[0] - wx = k[1] - wy = k[2] - wz = k[3] + k_times_delta = k["diffusivity_times_delta"] + wx = k["wx"] + wy = k["wy"] + wz = k["wz"] dim = volume_geometry.dimensions if dim == 2: @@ -917,12 +915,12 @@ Weak formulation geometry, symbolizer, blending=blending, - fe_coefficients=[ - ("diffusivity_times_delta", diffusivityXdelta_function_space), - ("wx", velocity_function_space), - ("wy", velocity_function_space), - ("wz", velocity_function_space), - ], + fe_coefficients={ + "diffusivity_times_delta": diffusivityXdelta_function_space, + "wx": velocity_function_space, + "wy": velocity_function_space, + "wz": velocity_function_space, + }, ) def grad_rho_by_rho_dot_u( diff --git a/hog/integrand.py b/hog/integrand.py index f6c54db5aaa530d8cdc932752704ea153610f228..e7c52bfec9c322b4033cce757a00bec0885bdff8 100644 --- a/hog/integrand.py +++ b/hog/integrand.py @@ -46,8 +46,8 @@ This module contains various functions and data structures to formulate the inte The actual integration and code generation is handled somewhere else. """ -from typing import Callable, List, Union, Tuple, Any -from dataclasses import dataclass, asdict, fields +from typing import Callable, List, Union, Tuple, Any, Dict +from dataclasses import dataclass, asdict, fields, field import sympy as sp @@ -64,6 +64,7 @@ from hog.fem_helpers import ( fem_function_gradient_on_element, scalar_space_dependent_coefficient, jac_affine_to_physical, + trafo_ref_to_physical, ) from hog.math_helpers import inv, det @@ -78,6 +79,7 @@ class Form: integrand: sp.MatrixBase tabulation: Tabulation symmetric: bool + free_symbols: List[sp.Symbol] docstring: str = "" def integrate(self, quad: Quadrature, symbolizer: Symbolizer) -> sp.Matrix: @@ -144,13 +146,15 @@ class IntegrandSymbols: # The Hessian of the test shape function (reference space!). hessian_v: sp.Matrix | None = None - # A list of scalar constants. - c: List[sp.Symbol] | None = None + # The physical coordinates. + x: sp.Matrix | None = None - # A list of finite element functions that can be used as function parameters. - k: List[sp.Symbol] | None = None + # A dict of finite element functions that can be used as function parameters. + # The keys are specified by the strings that are passed to process_integrand. + k: Dict[str, sp.Symbol] | None = None # A list of the gradients of the parameter finite element functions. - grad_k: List[sp.Matrix] | None = None + # The keys are specified by the strings that are passed to process_integrand. + grad_k: Dict[str, sp.Matrix] | None = None # The geometry of the volume element. volume_geometry: ElementGeometry | None = None @@ -165,6 +169,32 @@ class IntegrandSymbols: # element. jac_a_boundary: sp.Matrix | None = None + # A callback to generate free symbols that can be chosen by the user later. + # + # To get (and register new symbols) simply pass a list of symbol names to this function. + # It returns sympy symbols that can be safely used in the integrand: + # + # n_x, n_y, n_z = scalars(["n_x", "n_y", "n_z"]) + # + # or for a single symbol + # + # a = scalars("a") + # + # or use the sympy-like space syntax + # + # a, b = scalars("a b") + # + # Simply using sp.Symbol will not work since the symbols must be registered in the generator. + scalars: Callable[[str | List[str]], sp.Symbol | List[sp.Symbol]] | None = None + + # Same as scalars, but returns the symbols arranged as matrices. + # + # For a matrix with three rows and two columns: + # + # A = matrix("A", 3, 2) + # + matrix: Callable[[str, int, int], sp.Matrix] | None = None + # A callback to tabulate (aka precompute) terms that are identical on all elements of the same type. # # Simply enclose such a factor with this function, e.g., replace @@ -192,8 +222,7 @@ def process_integrand( symbolizer: Symbolizer, blending: GeometryMap = IdentityMap(), boundary_geometry: ElementGeometry | None = None, - scalar_coefficients: List[str] | None = None, - fe_coefficients: List[Tuple[str, Union[FunctionSpace, None]]] | None = None, + fe_coefficients: Dict[str, Union[FunctionSpace, None]] | None = None, is_symmetric: bool = False, docstring: str = "", ) -> Form: @@ -252,20 +281,15 @@ def process_integrand( :param blending: an optional blending map e.g., for curved geometries :param boundary_geometry: the geometry to integrate over for boundary integrals - passed through to the callable via the IntegrandSymbols object - :param scalar_coefficients: a list of strings that are names for scalar coefficients, they will be available to the - callable as `c` - :param fe_coefficients: a list of tuples of type (str, FunctionSpace) that are names and spaces for scalar + :param fe_coefficients: a dictionary of type (str, FunctionSpace) that are names and spaces for scalar finite-element function coefficients, they will be available to the callable as `k` supply None as the FunctionSpace for a std::function-type coeff (only works for old forms) :param is_symmetric: whether the bilinear form is symmetric - this is exploited by the generator :param docstring: documentation of the integrand/bilinear form - will end up in the docstring of the generated code """ - if scalar_coefficients is None: - scalar_coefficients = [] - if fe_coefficients is None: - fe_coefficients = [] + fe_coefficients = {} s = IntegrandSymbols() @@ -281,6 +305,30 @@ def process_integrand( s.tabulate = _tabulate + free_symbols = set() + + def _scalars(symbol_names: str | List[str]) -> sp.Symbol | List[sp.Symbol]: + nonlocal free_symbols + symbs = sp.symbols(symbol_names) + if isinstance(symbs, list): + free_symbols |= set(symbs) + elif isinstance(symbs, sp.Symbol): + free_symbols.add(symbs) + else: + raise HOGException( + f"I did not expect sp.symbols() to return whatever this is: {type(symbs)}" + ) + return symbs + + def _matrix(base_name: str, rows: int, cols: int) -> sp.Matrix: + symbs = _scalars( + [f"{base_name}_{row}_{col}" for row in range(rows) for col in range(cols)] + ) + return sp.Matrix(symbs).reshape(rows, cols) + + s.scalars = _scalars + s.matrix = _matrix + s.volume_geometry = volume_geometry s.jac_a = symbolizer.jac_ref_to_affine(volume_geometry) @@ -317,11 +365,11 @@ def process_integrand( s.boundary_geometry = boundary_geometry s.jac_a_boundary = symbolizer.jac_ref_to_affine(boundary_geometry) - s.c = [sp.Symbol(s) for s in scalar_coefficients] + s.x = trafo_ref_to_physical(volume_geometry, symbolizer, blending) - s.k = [] - s.grad_k = [] - for name, coefficient_function_space in fe_coefficients: + s.k = dict() + s.grad_k = dict() + for name, coefficient_function_space in fe_coefficients.items(): if coefficient_function_space is None: k = scalar_space_dependent_coefficient( name, volume_geometry, symbolizer, blending=blending @@ -350,8 +398,8 @@ def process_integrand( function_id=f"grad_{name}", dof_symbols=dof_symbols, ) - s.k.append(k) - s.grad_k.append(grad_k) + s.k[name] = k + s.grad_k[name] = grad_k mat = create_empty_element_matrix(trial, test, volume_geometry) it = element_matrix_iterator(trial, test, volume_geometry) @@ -367,4 +415,12 @@ def process_integrand( mat[data.row, data.col] = integrand(**asdict(s)) - return Form(mat, tabulation, symmetric=is_symmetric, docstring=docstring) + free_symbols = sorted(list(free_symbols), key=lambda x: str(x)) + + return Form( + mat, + tabulation, + symmetric=is_symmetric, + free_symbols=free_symbols, + docstring=docstring, + ) diff --git a/hog/operator_generation/operators.py b/hog/operator_generation/operators.py index 1052857a20166b5bb241f62f5508a886b51b5b27..0dec0085b599f874bc146dab9ac6836558f2fce3 100644 --- a/hog/operator_generation/operators.py +++ b/hog/operator_generation/operators.py @@ -164,9 +164,12 @@ class IntegrationInfo: optimizations: Set[Opts] + free_symbols: List[sp.Symbol] + name: str = "unknown_integral" docstring: str = "" boundary_uid_name: str = "" + integrand_name: str = "unknown_integrand" def _str_(self): return f"Integration Info: {self.name}, {self.geometry}, {self.integration_domain}, mat shape {self.mat.shape}, quad degree {self.quad.degree}, blending {self.blending}" @@ -311,6 +314,7 @@ class HyTeGElementwiseOperator: loop_strategy: LoopStrategy, boundary_uid_name: str, optimizations: Set[Opts], + integrand_name: str | None = None, ) -> None: """ Use this method to add integrals to the operator if you know what you are doing. There are helper methods for @@ -327,6 +331,12 @@ class HyTeGElementwiseOperator: :param boundary_uid_name: string that defines the name of the boundary UID if this is a boundary integral the parameter is ignored for volume integrals :param optimizations: optimizations that shall be applied to this integral + :param integrand_name: while each integral is a separate kernel with a name, for some types of integrals (e.g., + boundary integrals) more than one kernel with the same integrand is added - for some + features (e.g., symbol naming) it is convenient to be able to identify all those + integrals by a string - this (optional) integrand_name does not have to be unique (other + than the name parameter which has to be unique) but can be the same for all integrals + with different domains but the same integrand """ if "".join(name.split()) != name: @@ -334,6 +344,9 @@ class HyTeGElementwiseOperator: "Please give the integral an identifier without white space." ) + if integrand_name is None: + integrand_name = name + if volume_geometry.space_dimension in self.integration_infos: if name in [ ii.name @@ -401,6 +414,8 @@ class HyTeGElementwiseOperator: loop_strategy=loop_strategy, boundary_uid_name=boundary_uid_name, optimizations=optimizations, + free_symbols=form.free_symbols, + integrand_name=integrand_name, ) ) @@ -491,7 +506,7 @@ class HyTeGElementwiseOperator: # Since we will only integrate over the reference facet that lies on the x-line (2D) or xy-plane (3D) we need to # set the last reference coordinate to zero since it will otherwise appear as a free, uninitialized variable. # - # This has to be repeated later before the qudrature is applied in case we are working with symbols. + # This has to be repeated later before the quadrature is applied in case we are working with symbols. # form.integrand = form.integrand.subs( self.symbolizer.ref_coords_as_list(volume_geometry.dimensions)[-1], 0 @@ -508,6 +523,7 @@ class HyTeGElementwiseOperator: BOUNDARY(facet_id=facet_id), name + "_boundary_uid", optimizations, + integrand_name=name, ) def coefficients(self) -> List[FunctionSpaceImpl]: @@ -648,6 +664,27 @@ class HyTeGElementwiseOperator: ) ) + # Free symbols that shall be settable through the ctor. + free_symbol_vars = set() + for integration_infos in self.integration_infos.values(): + for integration_info in integration_infos: + for fs in integration_info.free_symbols: + free_symbol_vars.add(f"{str(fs)}_{integration_info.integrand_name}") + + free_symbol_vars = [ + CppVariable( + name=fs, + type=str(self._type_descriptor.pystencils_type), + ) + for fs in free_symbol_vars + ] + + free_symbol_vars = sorted(free_symbol_vars, key=lambda x: x.name) + + free_symbol_vars_members = [ + CppVariable(name=fsv.name + "_", type=fsv.type) for fsv in free_symbol_vars + ] + # Let's now check whether we need ctor arguments and member variables for boundary integrals. boundary_condition_vars = [] for integration_infos in self.integration_infos.values(): @@ -696,9 +733,14 @@ class HyTeGElementwiseOperator: ) for coeff in self.coefficients() ] + + free_symbol_vars + boundary_condition_vars, initializer_list=["Operator( storage, minLevel, maxLevel )"] + [f"{coeff.name}( _{coeff.name} )" for coeff in self.coefficients()] + + [ + f"{fsv[0].name}( {fsv[1].name} )" + for fsv in zip(free_symbol_vars_members, free_symbol_vars) + ] + [ f"{bcv[0].name}( {bcv[1].name} )" for bcv in zip( @@ -720,6 +762,9 @@ class HyTeGElementwiseOperator: ) ) + for fsv in free_symbol_vars_members: + operator_cpp_class.add(CppMemberVariable(fsv, visibility="private")) + for bcv in boundary_condition_vars_members: operator_cpp_class.add(CppMemberVariable(bcv, visibility="private")) @@ -1219,15 +1264,6 @@ class HyTeGElementwiseOperator: for a in ass.atoms(DoFSymbol) } - if ( - integration_info.integration_domain != MacroIntegrationDomain.VOLUME - and len(dof_symbols_set) > 0 - ): - raise HOGException( - "Boundary integrals and FE coefficients cannot be combined at the moment." - "Dev note: the coefficients need to be sorted." - ) - dof_symbols = sorted(dof_symbols_set, key=lambda ds: ds.name) coeffs = dict( ( @@ -1633,8 +1669,36 @@ class HyTeGElementwiseOperator: kernel_parameters = kernel_function.get_parameters() + if any( + [ + str(free_symbol) + not in [str(kp) for kp in kernel_parameters] + for free_symbol in integration_info.free_symbols + ] + ): + raise HOGException( + "Hmm, some free symbols from the integrand have been lost...\n" + f"free symbols: {[str(fs) for fs in integration_info.free_symbols]}\n" + f"kernel parameters: {[str(kp) for kp in kernel_parameters]}" + ) + + # We prepend the name of the kernel to the free symbols we found in the integrand to make sure that + # two integrands (e.g., a boundary and a volume integrand) that use the same symbol name do not + # clash. + + kernel_parameters_updated = [] + for prm in kernel_parameters: + if str(prm) in [ + str(fs) for fs in integration_info.free_symbols + ]: + kernel_parameters_updated.append( + f"{str(prm)}_{integration_info.integrand_name}_" + ) + else: + kernel_parameters_updated.append(str(prm)) + kernel_function_call_parameter_string = ",\n".join( - [str(prm) for prm in kernel_parameters] + kernel_parameters_updated ) kernel_function_call_strings.append( diff --git a/hog/recipes/integrands/volume/div_k_grad.py b/hog/recipes/integrands/volume/div_k_grad.py index 6e0aca147186b717796682e8751e026489ec2ce0..f0ca46a64ad6d83a2c66fed6500b43ccf8fb8cfb 100644 --- a/hog/recipes/integrands/volume/div_k_grad.py +++ b/hog/recipes/integrands/volume/div_k_grad.py @@ -29,7 +29,7 @@ def integrand( tabulate, **_, ): - return k[0] * ( + return k["k"] * ( double_contraction( jac_b_inv.T * tabulate(jac_a_inv.T * grad_u), jac_b_inv.T * tabulate(jac_a_inv.T * grad_v), diff --git a/hog/recipes/integrands/volume/div_k_grad_affine.py b/hog/recipes/integrands/volume/div_k_grad_affine.py index 49e2028200f7913bd5d2708ebe987eb1c0d13482..2a78278e5a1f19747b1b80ccf354417c52c204da 100644 --- a/hog/recipes/integrands/volume/div_k_grad_affine.py +++ b/hog/recipes/integrands/volume/div_k_grad_affine.py @@ -27,7 +27,7 @@ def integrand( tabulate, **_, ): - return k[0] * tabulate( + return k["k"] * tabulate( double_contraction( jac_a_inv.T * grad_u, jac_a_inv.T * grad_v, diff --git a/hog/recipes/integrands/volume/epsilon.py b/hog/recipes/integrands/volume/epsilon.py index 2a76bb4d3470e1808064d0f0fb0265a77bcab24a..ab7888e56fa474658e72208089d050b8d4fe02b5 100644 --- a/hog/recipes/integrands/volume/epsilon.py +++ b/hog/recipes/integrands/volume/epsilon.py @@ -40,7 +40,7 @@ def integrand( symm_grad_v = symm_grad(grad_v_chain) return ( - double_contraction(2 * k[0] * symm_grad_u, symm_grad_v) + double_contraction(2 * k["mu"] * symm_grad_u, symm_grad_v) * tabulate(jac_a_abs_det) * jac_b_abs_det ) diff --git a/hog/recipes/integrands/volume/epsilon_affine.py b/hog/recipes/integrands/volume/epsilon_affine.py index eadd0ff7c73da0fb6dbc5d0e4fcb359b68116a90..049f334a027fa363c38be1f18f59da494cb6c12c 100644 --- a/hog/recipes/integrands/volume/epsilon_affine.py +++ b/hog/recipes/integrands/volume/epsilon_affine.py @@ -37,6 +37,6 @@ def integrand( symm_grad_u = symm_grad(grad_u_chain) symm_grad_v = symm_grad(grad_v_chain) - return k[0] * tabulate( + return k["mu"] * tabulate( double_contraction(2 * symm_grad_u, symm_grad_v) * jac_a_abs_det ) diff --git a/hog/recipes/integrands/volume/full_stokes.py b/hog/recipes/integrands/volume/full_stokes.py index ed8e1824a1c53a93fb85d919a715655ac04790d0..d8e40a1d0e94c278278de8d0f340c28d0ca34c51 100644 --- a/hog/recipes/integrands/volume/full_stokes.py +++ b/hog/recipes/integrands/volume/full_stokes.py @@ -42,7 +42,7 @@ def integrand( div_u = (jac_b_inv.T * tabulate(jac_a_inv.T * grad_u)).trace() div_v = (jac_b_inv.T * tabulate(jac_a_inv.T * grad_v)).trace() - return k[0] * ( + return k["mu"] * ( ( double_contraction(2 * symm_grad_u, symm_grad_v) * tabulate(jac_a_abs_det) diff --git a/hog/recipes/integrands/volume/k_mass.py b/hog/recipes/integrands/volume/k_mass.py index f43de9479a88b5fc5653ac55db5f666836a0e157..bd2f997473aadeb9181c3d9c23205cc24f528daa 100644 --- a/hog/recipes/integrands/volume/k_mass.py +++ b/hog/recipes/integrands/volume/k_mass.py @@ -18,4 +18,4 @@ from hog.recipes.common import * def integrand(*, u, v, jac_a_abs_det, jac_b_abs_det, k, tabulate, **_): - return k[0] * tabulate(u * v * jac_a_abs_det) * jac_b_abs_det + return k["k"] * tabulate(u * v * jac_a_abs_det) * jac_b_abs_det