Solver parameters for conservative tracer transport

gusto/solvers/parameters.py

@@ -4,7 +4,10 @@
 """
 from gusto.core.function_spaces import is_cg
 
-__all__ = ['mass_parameters', 'hydrostatic_parameters']
+__all__ = [
+    'mass_parameters', 'hydrostatic_parameters',
+    'conservative_tracer_parameters'
+]
 
 
 def mass_parameters(V, spaces=None, ignore_vertical=True):
@@ -120,3 +123,56 @@ def mass_parameters(V, spaces=None, ignore_vertical=True):
         }
     }
 }
+
+
+def conservative_tracer_parameters(V, num_fields=2):
+    """
+    Returns PETSc solver settings for conservative tracer transport, in which
+    a tracer is transported simultaneously with a reference density field in a
+    discontinuous space.
+
+    As the density does not depend upon the tracer, we solve this with a
+    multiplicative fieldsplit, in which the density is solved for first.
+
+    Parameters
+    ----------
+    V : :class:`FunctionSpace`
+        The (sub)function space of the density/tracer.
+    num_fields : int, optional
+        The number of fields in the mixed function space. Default is 2.
+
+    Returns
+    -------
+    settings : dict
+        A dictionary containing the PETSc solver settings.
+    """
+    fs_name = V.name
+    settings = {
+        "mat_type": "aij",
+        "snes_type": "ksponly",
+        "ksp_type": "preonly",
+        "pc_type": "fieldsplit",
+        "ksp_monitor_true_residual": None,
+        "ksp_converged_reason": None,
+        "snes_monitor": None,
+        "snes_converged_reason": None,
+        "pc_fieldsplit_type": "multiplicative",
+        "pc_fieldsplit_0_fields": "0",
+        "pc_fieldsplit_1_fields": "1",
+
+        f"fieldsplit_{fs_name}_ksp_type": "preonly",
+        "fieldsplit_0_pc_type": "lu",
+        f"fieldsplit_{fs_name}_ksp_converged_reason": None,
+        f"fieldsplit_{fs_name}_ksp_monitor_true_residual": None,
+    }
+
+    if num_fields < 2:
+        raise ValueError(
+            "The number of fields used for conservative tracer transport must "
+            + f"be at least 2, but got {num_fields}."
+        )
+    elif num_fields > 2:
+        for i in range(2, num_fields):
+            settings[f"pc_fieldsplit_{i}_fields"] = str(i)
+
+    return settings

gusto/time_discretisation/time_discretisation.py

@@ -101,6 +101,8 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
         self.courant_max = None
         self.augmentation = augmentation
         self.subcycling_options = subcycling_options
+        self.solver_parameters = solver_parameters
+        self.default_solver_parameters = (solver_parameters is None)
 
         if self.subcycling_options is not None:
             self.subcycling_options.check_options()
@@ -136,14 +138,6 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
             self.wrapper = None
             self.wrapper_name = None
 
-        # get default solver options if none passed in
-        if solver_parameters is None:
-            self.solver_parameters = {'ksp_type': 'gmres',
-                                      'pc_type': 'bjacobi',
-                                      'sub_pc_type': 'ilu'}
-        else:
-            self.solver_parameters = solver_parameters
-
     def setup(self, equation, apply_bcs=True, *active_labels):
         """
         Set up the time discretisation based on the equation.
@@ -351,7 +345,17 @@ def replace_test_all_but_last(t):
                     )
 
                     self.residual = self.wrapper.label_terms(self.residual)
-                if self.solver_parameters is None:
+
+                # Use wrapper solver parameters if they are specified
+                use_wrapper_params = (
+                    self.solver_parameters is None
+                    and self.wrapper.solver_parameters is not None
+                )
+                if use_wrapper_params:
+                    logger.info(
+                        'Using default solver parameters for'
+                        + f'{self.wrapper_name} wrapper'
+                    )
                     self.solver_parameters = self.wrapper.solver_parameters
 
         # -------------------------------------------------------------------- #
@@ -384,6 +388,15 @@ def replace_test_all_but_last(t):
         self.x_out = Function(self.fs)
         self.x1 = Function(self.fs)
 
+        # Finally, set solver parameters to default
+        if self.solver_parameters is None:
+            logger.info('Using default TimeDiscretisation solver parameters')
+            self.solver_parameters = {
+                'ksp_type': 'gmres',
+                'pc_type': 'bjacobi',
+                'sub_pc_type': 'ilu'
+            }
+
     @property
     def nlevels(self):
         return 1
@@ -490,6 +503,7 @@ def __init__(self, domain, field_name=None, subcycling_options=None,
                 this time discretisation to be augmented, for instances with
                 extra terms of another auxiliary variable. Defaults to None.
         """
+
         super().__init__(domain, field_name,
                          subcycling_options=subcycling_options,
                          solver_parameters=solver_parameters,
@@ -509,10 +523,16 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         """
         super().setup(equation, apply_bcs, *active_labels)
 
-        # get default solver options if none passed in
-        self.solver_parameters.update(mass_parameters(
-            self.fs, equation.domain.spaces))
-        self.solver_parameters['snes_type'] = 'ksponly'
+        # Set default solver params for explicit schemes, if none were passed in
+        if self.default_solver_parameters:
+            logger.info(
+                'Actually using default ExplicitTimeDiscretisation '
+                + 'solver parameters'
+            )
+            self.solver_parameters = mass_parameters(
+                self.fs, equation.domain.spaces
+            )
+            self.solver_parameters['snes_type'] = 'ksponly'
 
         # if user has specified a number of fixed subcycles, then save this
         # and rescale dt accordingly; else perform just one cycle using dt