more changes to test functions for multiple wrappers

firedrakeproject · Jan 11, 2024 · 5f05367 · 5f05367
1 parent 793768e
commit 5f05367
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py
@@ -216,14 +216,14 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = self.wrapper.subwrappers[field].function_space
 
-                    # Replace with new test function
-                    # STILL TO SORT OUT.
-                    #new_test = TestFunction(self.wrapper.subwrappers[field].test_space)
-
-                    #self.residual = self.residual.label_map(
-                    #    all_terms,
-                    #    map_if_true=replace_test_function(new_test, old_idx=field_idx))
 
+                    # Replace the test function space
+                    # This currently won't work: supg test 
+                    # is a function, not a space
+                    if self.wrapper.suboptions[field].name == "supg":
+                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test
+                    else:
+                        self.wrapper.test_spaces[field_idx] = self.wrapper.subwrappers[field].test_space
 
                 self.wrapper.setup()
 
@@ -244,7 +244,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
 
                 # Only call this if with SUPG, so should put this into the 
                 # previous section
-                self.residual = self.wrapper.label_terms(self.residual)
+                # self.residual = self.wrapper.label_terms(self.residual)
 
 
             else:

diff --git a/gusto/wrappers.py b/gusto/wrappers.py
@@ -423,6 +423,7 @@ def __init__(self, equation, suboptions):
 
         #self.wrapper_spaces = equation.space_names
         self.wrapper_spaces = equation.spaces
+        self.test_spaces = equation.spaces
         print(len(equation.spaces))
 
         #self.x_in = Function(equation.function_space)
@@ -475,8 +476,7 @@ def setup(self):
         self.function_space = MixedFunctionSpace(self.wrapper_spaces)
         self.x_in = Function(self.function_space)
         self.x_out = Function(self.function_space) 
-        self.test_space = self.function_space
-        #pass
+        self.test_space = MixedFunctionSpace(self.test_spaces)
 
     def pre_apply(self, x_in):
         """