make original_space a required argument for wrapper.setup()

gusto/time_discretisation.py

@@ -184,7 +184,7 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                         raise ValueError(f"The option defined for {field} is for a field that does not exist in the equation set")
 
                     field_idx = equation.field_names.index(field)
-                    subwrapper.setup(self.equation.spaces[field_idx])
+                    subwrapper.setup(equation.spaces[field_idx])
 
                     # Update the function space to that needed by the wrapper
                     self.wrapper.wrapper_spaces[field_idx] = subwrapper.function_space
@@ -200,7 +200,10 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     map_if_true=replace_test_function(new_test_mixed))
 
             else:
-                self.wrapper.setup()
+                if self.wrapper_name == "supg":
+                    self.wrapper.setup()
+                else:
+                    self.wrapper.setup(self.fs)
                 self.fs = self.wrapper.function_space
                 if self.solver_parameters is None:
                     self.solver_parameters = self.wrapper.solver_parameters

gusto/wrappers.py

@@ -36,7 +36,7 @@ def __init__(self, time_discretisation, wrapper_options):
         self.original_space = None
 
     @abstractmethod
-    def setup(self, original_space=None):
+    def setup(self, original_space):
         """
         Store the original function space of the prognostic variable.
 
@@ -47,13 +47,9 @@ def setup(self, original_space=None):
         Args:
             original_space (:class:`FunctionSpace`): the space that the
         prognostic variable is defined on. This is a subset space of
-        a mixed function space when using a MixedFSWrapper. Defaults
-        to None, when not using a MixedFSWrapper.
+        a mixed function space when using a MixedFSWrapper. 
         """
-        if original_space is not None:
-            self.original_space = original_space
-        else:
-            self.original_space = self.time_discretisation.fs
+        self.original_space = original_space
 
     @abstractmethod
     def pre_apply(self):
@@ -89,7 +85,7 @@ class EmbeddedDGWrapper(Wrapper):
     the original space.
     """
 
-    def setup(self, original_space=None):
+    def setup(self, original_space):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, EmbeddedDGOptions), \
@@ -173,7 +169,7 @@ class RecoveryWrapper(Wrapper):
     field is then returned to the original space.
     """
 
-    def setup(self, original_space=None):
+    def setup(self, original_space):
         """Sets up function spaces and fields needed for this wrapper."""
 
         assert isinstance(self.options, RecoveryOptions), \