From 9ad537c51a72470b72cd33c5dd0158a0d58e01f4 Mon Sep 17 00:00:00 2001
From: Marcus Mohr <marcus.mohr@lmu.de>
Date: Mon, 2 Dec 2024 16:51:15 +0100
Subject: [PATCH] Add ability to generate affine mass and diffusion forms for
 P2PlusBubble space

Note that this is only for testing form generation and verification of the
result. The generated forms are, currently, not useable in HyTeG as there
is neither a base class for them nor any operator that could make use of
them.
---
 generate_all_hyteg_forms.py | 27 +++++++++++++++++++++++++++
 hog/forms.py                |  4 +++-
 hog/hyteg_form_template.py  |  6 +++++-
 3 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/generate_all_hyteg_forms.py b/generate_all_hyteg_forms.py
index d96ca5d..cf26ddb 100644
--- a/generate_all_hyteg_forms.py
+++ b/generate_all_hyteg_forms.py
@@ -34,6 +34,7 @@ from hog.element_geometry import (
 from hog.function_space import (
     LagrangianFunctionSpace,
     N1E1Space,
+    P2PlusBubbleSpace,
     TrialSpace,
     TestSpace,
 )
@@ -133,6 +134,8 @@ class FormInfo:
         """A compact representation of the function spaces."""
         if self.trial_family == "N1E1":
             return "n1e1"
+        elif self.trial_family == "P2 enhanced with Bubble":
+            return "p2_plus_bubble"
         elif self.trial_degree == self.test_degree:
             return f"p{self.trial_degree}"
         else:
@@ -163,6 +166,8 @@ 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_degree == self.test_degree:
             sub_dir = f"p{self.trial_degree}"
         else:
@@ -242,6 +247,15 @@ form_infos = [
         quad_schemes={2: 3, 3: 3},
         blending=ExternalMap(),
     ),
+    FormInfo(
+        "diffusion",
+        trial_degree=2,
+        test_degree=2,
+        trial_family="P2 enhanced with Bubble",
+        test_family="P2 enhanced with Bubble",
+        quad_schemes={2: "exact", 3: "exact"},
+        integrate_rows=[],
+    ),
     FormInfo(
         "mass",
         trial_degree=1,
@@ -296,6 +310,15 @@ form_infos = [
         blending=ExternalMap(),
         integrate_rows=[],
     ),
+    FormInfo(
+        "mass",
+        trial_degree=2,
+        test_degree=2,
+        trial_family="P2 enhanced with Bubble",
+        test_family="P2 enhanced with Bubble",
+        quad_schemes={2: "exact", 3: "exact"},
+        integrate_rows=[],
+    ),
     FormInfo(
         "curl_curl",
         trial_degree=1,
@@ -1077,6 +1100,8 @@ def main():
         trial: TrialSpace
         if form_info.trial_family == "N1E1":
             trial = TrialSpace(N1E1Space(symbolizer))
+        elif form_info.trial_family == "P2 enhanced with Bubble":
+            trial = TrialSpace(P2PlusBubbleSpace(symbolizer))
         else:
             trial = TrialSpace(
                 LagrangianFunctionSpace(form_info.trial_degree, symbolizer)
@@ -1085,6 +1110,8 @@ def main():
         test: TestSpace
         if form_info.test_family == "N1E1":
             test = TestSpace(N1E1Space(symbolizer))
+        elif form_info.test_family == "P2 enhanced with Bubble":
+            test = TestSpace(P2PlusBubbleSpace(symbolizer))
         else:
             test = TestSpace(LagrangianFunctionSpace(form_info.test_degree, symbolizer))
 
diff --git a/hog/forms.py b/hog/forms.py
index f67f078..ec630ce 100644
--- a/hog/forms.py
+++ b/hog/forms.py
@@ -36,6 +36,7 @@ from hog.fem_helpers import (
 from hog.function_space import (
     FunctionSpace,
     N1E1Space,
+    P2PlusBubbleSpace,
     TrialSpace,
     TestSpace,
     LagrangianFunctionSpace,
@@ -1097,7 +1098,7 @@ Weak formulation
 
     ∫ ((∇ρ / ρ) · u) v
 """
-    
+
     from hog.recipes.integrands.volume.frozen_velocity import integrand
 
     return process_integrand(
@@ -1114,6 +1115,7 @@ Weak formulation
         },
     )
 
+
 def zero_form(
     trial: TrialSpace,
     test: TestSpace,
diff --git a/hog/hyteg_form_template.py b/hog/hyteg_form_template.py
index f0d4308..b1c7c34 100644
--- a/hog/hyteg_form_template.py
+++ b/hog/hyteg_form_template.py
@@ -21,7 +21,7 @@ from hog.ast import Assignment, CodeBlock
 from hog.exception import HOGException
 from hog.quadrature import Quadrature
 from hog.symbolizer import Symbolizer
-from hog.function_space import N1E1Space, TrialSpace, TestSpace
+from hog.function_space import N1E1Space, P2PlusBubbleSpace, TrialSpace, TestSpace
 from hog.element_geometry import ElementGeometry
 from hog.code_generation import code_block_from_element_matrix
 from hog.multi_assignment import Member
@@ -292,6 +292,8 @@ class HyTeGFormClass:
 
         if isinstance(self.trial, N1E1Space):
             super_class = f"n1e1::N1E1Form"
+        elif isinstance(self.trial, P2PlusBubbleSpace):
+            super_class = f"P2PlusBubbleForm"
         elif self.trial.degree == self.test.degree:
             super_class = f"P{self.trial.degree}FormHyTeG"
         else:
@@ -358,6 +360,8 @@ class HyTeGForm:
 
         if isinstance(self.trial, N1E1Space):
             super_class = f"N1E1Form"
+        elif isinstance(self.trial, P2PlusBubbleSpace):
+            super_class = f"P2PlusBubbleForm"
         elif self.trial.degree == self.test.degree:
             super_class = f"form_hyteg_base/P{self.trial.degree}FormHyTeG"
         else:
-- 
GitLab