mass solve bugfix

firedrakeproject · Dec 18, 2024 · 3e20ce7 · 3e20ce7
1 parent c5eb881
commit 3e20ce7
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 19 deletions.
diff --git a/gusto/solvers/parameters.py b/gusto/solvers/parameters.py
@@ -42,8 +42,9 @@ def mass_parameters(V, spaces=None, ignore_vertical=True):
     extruded = hasattr(V.mesh, "_base_mesh")
 
     continuous_fields = set()
-    for i, Vsub in V.subfunctions:
-        field = Vsub.name or str(i)
+    for i, Vsub in enumerate(V.subfunctions):
+        # field = Vsub.name or str(i)
+        field = str(i)
 
         if spaces is not None:
             continuous = spaces.continuity.get(field, is_cg(Vsub))
@@ -62,17 +63,25 @@ def mass_parameters(V, spaces=None, ignore_vertical=True):
         if continuous:
             continuous_fields.add(field)
 
-    parameters = {
-        'ksp_type': 'preonly',
-        'pc_type': 'fieldsplit',
-        'pc_fieldsplit_type': 'additive',
-        'pc_fieldsplit_0_fields': ','.join(continuous_fields),
-        'fieldsplit': {
-            'ksp_type': 'preonly',
+    if len(V.subfunctions) == 1:
+        parameters = {
+            'ksp_type': 'cg' if all(continuous_fields) else 'preonly',
             'pc_type': 'bjacobi',
-            'sub_pc_type': 'ilu'
-        },
-        'fieldsplit_0_ksp_type': 'cg',
-    }
+            'sub_pc_type': 'ilu',
+        }
+    else:
+
+        parameters = {
+            'ksp_type': 'preonly',
+            'pc_type': 'fieldsplit',
+            'pc_fieldsplit_type': 'additive',
+            'pc_fieldsplit_0_fields': ','.join(continuous_fields),
+            'fieldsplit': {
+                'ksp_type': 'preonly',
+                'pc_type': 'bjacobi',
+                'sub_pc_type': 'ilu'
+            },
+            'fieldsplit_0_ksp_type': 'cg',
+        }
 
     return parameters
diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py
@@ -487,12 +487,9 @@ def setup(self, equation, apply_bcs=True, *active_labels):
         super().setup(equation, apply_bcs, *active_labels)
 
         # get default solver options if none passed in
-        if solver_parameters is None:
-            self.solver_parameters = mass_parameters(
-                equation.function_space, equation.domain.spaces)
-            self.solver_parameters['snes_type'] = 'ksponly'
-        else:
-            self.solver_parameters = solver_parameters
+        self.solver_parameters.update(mass_parameters(
+            equation.function_space, 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