diff --git a/generate_all_hyteg_forms.py b/generate_all_hyteg_forms.py
index b37446c88d8cea049feaafdbb64ba29ce217c5ed..2c49ddd8967f63013cdd3f4c4ec5cdfe9918196c 100644
--- a/generate_all_hyteg_forms.py
+++ b/generate_all_hyteg_forms.py
@@ -166,8 +166,11 @@ class FormInfo:
sub_dir = "p1"
if self.trial_family == "N1E1":
sub_dir = "n1e1"
- elif self.trial_family == "P2 enhanced with Bubble":
- return "p2_plus_bubble"
+ elif (
+ self.trial_family == "P2 enhanced with Bubble"
+ and self.test_family == "P2 enhanced with Bubble"
+ ):
+ sub_dir = "p2"
elif self.trial_degree == self.test_degree:
sub_dir = f"p{self.trial_degree}"
else:
@@ -836,7 +839,9 @@ def form_func(
quad, symbolizer
)
elif name.startswith("pspg"):
- return pspg(trial, test, geometry, quad, symbolizer, blending=blending)
+ return pspg(
+ trial, test, geometry, quad, symbolizer, blending=blending
+ ).integrate(quad, symbolizer)
elif name.startswith("linear_form"):
return linear_form(trial, test, geometry, quad, symbolizer, blending=blending)
elif name.startswith("divt"):
@@ -1135,30 +1140,28 @@ def main():
inline_values=form_info.inline_quad,
)
- mat = form_func(
- form_info.form_name,
- row,
- col,
- trial,
- test,
- geometry,
- quad,
- symbolizer,
- blending=form_info.blending,
- )
- form_codes.append(
- HyTeGIntegrator(
- form_info.class_name(row, col),
- mat,
+ if form_info.is_implemented(row, col, geometry.dimensions):
+ mat = form_func(
+ form_info.form_name,
+ row,
+ col,
+ trial,
+ test,
geometry,
quad,
symbolizer,
- not_implemented=not form_info.is_implemented(
- row, col, geometry.dimensions
- ),
- integrate_rows=form_info.integrate_rows,
+ blending=form_info.blending,
+ )
+ form_codes.append(
+ HyTeGIntegrator(
+ form_info.class_name(row, col),
+ mat,
+ geometry,
+ quad,
+ symbolizer,
+ integrate_rows=form_info.integrate_rows,
+ )
)
- )
form_classes.append(
HyTeGFormClass(
form_info.class_name(row, col),
diff --git a/hog/forms.py b/hog/forms.py
index ec630ce41999ca80a8396f1d0de0838f231fe435..5ed4b0eeea7aeecc044ad81ba042bb9ff4ad7e0f 100644
--- a/hog/forms.py
+++ b/hog/forms.py
@@ -579,6 +579,7 @@ Weak formulation
geometry,
symbolizer,
blending=blending,
+ component_index=component_index,
is_symmetric=False,
docstring=docstring,
rot_type=RotationType.POST_MULTIPLY
@@ -620,6 +621,7 @@ def gradient(
geometry,
symbolizer,
blending=blending,
+ component_index=component_index,
is_symmetric=False,
docstring=docstring,
rot_type=RotationType.PRE_MULTIPLY
diff --git a/hog/integrand.py b/hog/integrand.py
index 07313705041256cc2c0ffb8f2f35d62ee448dbbc..a95b19358ff6be220e022621941201cea9ac5149 100644
--- a/hog/integrand.py
+++ b/hog/integrand.py
@@ -220,6 +220,9 @@ class IntegrandSymbols:
#
tabulate: Callable[[Union[sp.Expr, sp.Matrix], str], sp.Matrix] | None = None
+ # For backward compatibility with (sub-)form generation this integer allows to select a component
+ component_index: int | None = None
+
def process_integrand(
integrand: Callable[..., Any],
@@ -230,6 +233,7 @@ def process_integrand(
blending: GeometryMap = IdentityMap(),
boundary_geometry: ElementGeometry | None = None,
fe_coefficients: Dict[str, Union[FunctionSpace, None]] | None = None,
+ component_index: int | None = None,
is_symmetric: bool = False,
rot_type: RotationType = RotationType.NO_ROTATION,
docstring: str = "",
@@ -293,7 +297,7 @@ def process_integrand(
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 rot_type: whether the operator has to be wrapped with rotation matrix and the type of rotation that needs
+ :param rot_type: whether the operator has to be wrapped with rotation matrix and the type of rotation that needs
to be applied, only applicable for Vectorial spaces
:param docstring: documentation of the integrand/bilinear form - will end up in the docstring of the generated code
"""
@@ -482,6 +486,9 @@ def process_integrand(
mat = create_empty_element_matrix(trial, test, volume_geometry)
it = element_matrix_iterator(trial, test, volume_geometry)
+ if component_index is not None:
+ s.component_index = component_index
+
for data in it:
s.u = data.trial_shape
s.grad_u = data.trial_shape_grad
diff --git a/hog/recipes/integrands/volume/divergence.py b/hog/recipes/integrands/volume/divergence.py
index 75ae7517998e3c5e4a53920530817760ef70238c..f2ea6530cb567f2b2874f35582f080692884e18e 100644
--- a/hog/recipes/integrands/volume/divergence.py
+++ b/hog/recipes/integrands/volume/divergence.py
@@ -26,10 +26,21 @@ def integrand(
grad_u,
v,
tabulate,
+ component_index,
**_,
):
- return (
- -(jac_b_inv.T * tabulate(jac_a_inv.T * grad_u)).trace()
- * tabulate(v * jac_a_abs_det)
- * jac_b_abs_det
- )
+ # working with vector-valued functions
+ if grad_u.is_square:
+ return (
+ -(jac_b_inv.T * tabulate(jac_a_inv.T * grad_u)).trace()
+ * tabulate(v * jac_a_abs_det)
+ * jac_b_abs_det
+ )
+
+ # working with scalar-valued functions (backward compatibility)
+ else:
+ return (
+ -(jac_b_inv.T * tabulate(jac_a_inv.T * grad_u))[component_index]
+ * tabulate(v * jac_a_abs_det)
+ * jac_b_abs_det
+ )
diff --git a/hog/recipes/integrands/volume/gradient.py b/hog/recipes/integrands/volume/gradient.py
index 982ba2329ecdc6b5aee3491e6e8a5c1703f19d5f..134271d7ebaadc30e44d913dd69cbfbc3d274ad5 100644
--- a/hog/recipes/integrands/volume/gradient.py
+++ b/hog/recipes/integrands/volume/gradient.py
@@ -26,10 +26,21 @@ def integrand(
u,
grad_v,
tabulate,
+ component_index,
**_,
):
- return (
- -(jac_b_inv.T * tabulate(jac_a_inv.T * grad_v)).trace()
- * tabulate(u * jac_a_abs_det)
- * jac_b_abs_det
- )
+ # working with vector-valued functions
+ if grad_v.is_square:
+ return (
+ -(jac_b_inv.T * tabulate(jac_a_inv.T * grad_v)).trace()
+ * tabulate(u * jac_a_abs_det)
+ * jac_b_abs_det
+ )
+
+ # working with scalar-valued functions (backward compatibility)
+ else:
+ return (
+ -(jac_b_inv.T * tabulate(jac_a_inv.T * grad_v))[component_index]
+ * tabulate(u * jac_a_abs_det)
+ * jac_b_abs_det
+ )