From 50e6324d76c6e35c41c45a2f8496ae30a1002396 Mon Sep 17 00:00:00 2001
From: Ponsuganth Ilangovan Ponkumar Ilango
 <pponkumar@geophysik.uni-muenchen.de>
Date: Fri, 16 Aug 2024 20:08:03 +0200
Subject: [PATCH] Modify generate.py to support updates and folder/op-name
 changes

Missed to adapt changes for general case
---
 generate/generate.py | 44 ++++++++++++++++++------
 operators.toml       | 82 ++++++++++++++++++++++++++++++++++++++------
 2 files changed, 106 insertions(+), 20 deletions(-)

diff --git a/generate/generate.py b/generate/generate.py
index 613c3155..3a90089f 100644
--- a/generate/generate.py
+++ b/generate/generate.py
@@ -71,13 +71,16 @@ def unfold_toml_dict(toml_dict):
                         )
                         operators_unfolded.append(operator_cleared)
             elif "components-trial" in operator or "components-test" in operator:
-                comp_str = "components-trial" if "components-trial" in operator else "components-test"
+                comp_str = (
+                    "components-trial"
+                    if "components-trial" in operator
+                    else "components-test"
+                )
                 for comp in operator[comp_str]:
                     operator_cleared = operator.copy()
                     operator_cleared.pop(comp_str)
-                    operator_cleared.update({comp_str.replace('s-', '_'): comp})
+                    operator_cleared.update({comp_str.replace("s-", "_"): comp})
                     operators_unfolded.append(operator_cleared)
-                print(operators_unfolded)
             else:
                 operators_unfolded.append(operator)
         toml_dict_unfolded[form_str] = operators_unfolded
@@ -313,7 +316,9 @@ def generate_operator(
     trial_space = fe_spaces[spec["trial-space"]]
     test_space = fe_spaces[spec["test-space"]]
 
-    name = elementwise_operator_name(form_str, spec) # Will use operator-name from spec if provided
+    name = elementwise_operator_name(
+        form_str, spec
+    )  # Will use operator-name from spec if provided
     optimizations = {
         operator_generation.optimizer.opts_arg_mapping[opt.upper()]
         for opt in spec["optimizations"]
@@ -350,12 +355,24 @@ def generate_operator(
                 == spec["form-args"]["component_trial"]
             )
 
-    if "component_test" in spec and "component_trial" not in spec:
+    if "form-args" in spec:
+        if (
+            "component_test" in spec["form-args"]
+            and "component_trial" in spec["form-args"]
+        ):
+            test_space = function_space.TensorialVectorFunctionSpace(
+                test_space, single_component=spec["form-args"]["component_test"]
+            )
+            trial_space = function_space.TensorialVectorFunctionSpace(
+                trial_space, single_component=spec["form-args"]["component_trial"]
+            )
+
+    elif "component_test" in spec:
         test_space = function_space.TensorialVectorFunctionSpace(
             test_space, single_component=spec["component_test"]
         )
 
-    if "component_trial" in spec and "component_test" not in spec:
+    elif "component_trial" in spec:
         trial_space = function_space.TensorialVectorFunctionSpace(
             trial_space, single_component=spec["component_trial"]
         )
@@ -415,7 +432,10 @@ def generate_operator(
                 optimizations=optimizations,
             )
 
-        dir_path = os.path.join(args.output, spec["folder-name"] if "folder-name" in spec.keys() else form_str)
+        dir_path = os.path.join(
+            args.output,
+            spec["folder-name"] if "folder-name" in spec.keys() else form_str,
+        )
         operator.generate_class_code(
             dir_path,
             class_files=operator_generation.operators.CppClassFiles.HEADER_IMPL_AND_VARIANTS,
@@ -437,7 +457,11 @@ def generate_operator(
 
 
 def elementwise_operator_name(form_str: str, spec: Dict[str, Any]) -> str:
-    operator_name = spec["operator-name"] if "operator-name" in spec else form_str.title().replace("_", "")
+    operator_name = (
+        spec["operator-name"]
+        if "operator-name" in spec
+        else form_str.title().replace("_", "")
+    )
 
     if spec["trial-space"] == spec["test-space"]:
         space_mapping = spec["trial-space"]
@@ -455,10 +479,10 @@ def elementwise_operator_name(form_str: str, spec: Dict[str, Any]) -> str:
     # I do not like this, but should do the trick until we have actual vector function spaces in the HOG.
     if "component_test" in spec and "component_trial" not in spec:
         component = f"_{spec['component_test']}_0"
-    
+
     if "component_trial" in spec and "component_test" not in spec:
         component = f"_0_{spec['component_trial']}"
-        
+
     blending = ""
     if spec.get("blending", "IdentityMap") != "IdentityMap":
         blending = spec["blending"]
diff --git a/operators.toml b/operators.toml
index 064d423e..a956d4ef 100644
--- a/operators.toml
+++ b/operators.toml
@@ -200,15 +200,7 @@ optimizations = ["moveconstants", "vectorize", "quadloops", "tabulate"]
 blending      = "IcosahedralShellMap"
 
 [[divergence]]
-trial-space   = "P2Vector"
-test-space    = "P1"
-dimensions    = [2, 3]
-quadrature    = 2
-loop-strategy = "cubes"
-optimizations = ["moveconstants", "vectorize"]
-
-[[divergence]]
-trial-space   = "P2Vector"
+trial-space   = "P2"
 test-space    = "P1"
 components-trial = [0, 1, 2]
 dimensions    = [2, 3]
@@ -227,7 +219,7 @@ optimizations = ["moveconstants", "vectorize"]
 blending      = "AnnulusMap"
 
 [[divergence]]
-trial-space   = "P2Vector"
+trial-space   = "P2"
 test-space    = "P1"
 components-trial = [0, 1, 2]
 dimensions    = [3]
@@ -317,6 +309,41 @@ loop-strategy = "sawtooth"
 optimizations = ["moveconstants", "vectorize"]
 blending      = "IcosahedralShellMap"
 
+[[epsilon]]
+trial-space               = "P2"
+test-space                = "P2"
+components-trial = [0, 1, 2]
+components-test  = [0, 1, 2]
+form-space-args.coefficient_function_space = "P2"
+dimensions      = [2, 3]
+quadrature      = 2
+loop-strategy   = "sawtooth"
+optimizations   = ["moveconstants", "vectorize", "quadloops", "tabulate"]
+
+[[epsilon]]
+trial-space               = "P2"
+test-space                = "P2"
+components-trial = [0, 1]
+components-test  = [0, 1]
+form-space-args.coefficient_function_space = "P2"
+dimensions      = [2]
+quadrature      = 3
+loop-strategy   = "sawtooth"
+optimizations   = ["moveconstants", "vectorize", "quadloops", "tabulate"]
+blending        = "AnnulusMap"
+
+[[epsilon]]
+trial-space               = "P2"
+test-space                = "P2"
+components-trial = [0, 1, 2]
+components-test  = [0, 1, 2]
+form-space-args.coefficient_function_space = "P2"
+dimensions      = [3]
+quadrature      = 3
+loop-strategy   = "sawtooth"
+optimizations   = ["moveconstants", "vectorize", "quadloops", "tabulate"]
+blending        = "IcosahedralShellMap"
+
 [[epsilon]]
 trial-space               = "P2Vector"
 test-space                = "P2Vector"
@@ -346,6 +373,41 @@ loop-strategy   = "sawtooth"
 optimizations   = ["moveconstants", "vectorize", "quadloops", "tabulate"]
 blending        = "IcosahedralShellMap"
 
+[[full_stokes]]
+trial-space               = "P2"
+test-space                = "P2"
+components-trial = [0, 1, 2]
+components-test  = [0, 1, 2]
+form-space-args.coefficient_function_space = "P2"
+dimensions      = [2, 3]
+quadrature      = 2
+loop-strategy   = "sawtooth"
+optimizations   = ["moveconstants", "vectorize", "quadloops", "tabulate"]
+
+[[full_stokes]]
+trial-space               = "P2"
+test-space                = "P2"
+components-trial = [0, 1]
+components-test  = [0, 1]
+form-space-args.coefficient_function_space = "P2"
+dimensions      = [2]
+quadrature      = 3
+loop-strategy   = "sawtooth"
+optimizations   = ["moveconstants", "vectorize", "quadloops", "tabulate"]
+blending        = "AnnulusMap"
+
+[[full_stokes]]
+trial-space               = "P2"
+test-space                = "P2"
+components-trial = [0, 1, 2]
+components-test  = [0, 1, 2]
+form-space-args.coefficient_function_space = "P2"
+dimensions      = [3]
+quadrature      = 3
+loop-strategy   = "sawtooth"
+optimizations   = ["moveconstants", "vectorize", "quadloops", "tabulate"]
+blending        = "IcosahedralShellMap"
+
 [[full_stokes]]
 trial-space               = "P2Vector"
 test-space                = "P2Vector"
-- 
GitLab