fieldsplit SIQN forcing mass solves

firedrakeproject · Nov 7, 2024 · 2e5158f · 2e5158f
1 parent c5fd4d8
commit 2e5158f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 19 deletions.
diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py
@@ -398,20 +398,18 @@ def setup(self, equation, apply_bcs=True, *active_labels):
             logger.warning(message)
             self.solver_parameters['snes_type'] = 'newtonls'
 
+            # TODO: move this outside the if-else and just set the fieldsplit_type here
             if len(self.x0.subfunctions) > 1:
                 self.solver_parameters.update({
                     'ksp_type': 'preonly',
                     'pc_type': 'fieldsplit',
                     'pc_fieldsplit_type': 'multiplicative',
+                    **{f'fieldsplit_{fs.name}': {
+                        'ksp_type': 'cg',
+                        'pc_type': 'bjacobi',
+                        'sub_pc_type': 'ilu'
+                    } for fs in self.fs.subfunctions}
                 })
-                for fs in self.fs.subfunctions:
-                    self.solver_parameters.update({
-                        f'fieldsplit_{fs.name}': {
-                            'ksp_type': 'cg',
-                            'pc_type': 'bjacobi',
-                            'sub_pc_type': 'ilu',
-                        }
-                    })
         else:
             self.solver_parameters.setdefault('snes_lag_jacobian', -2)
             self.solver_parameters.setdefault('snes_lag_jacobian_persists', None)
@@ -425,15 +423,12 @@ def setup(self, equation, apply_bcs=True, *active_labels):
                     'ksp_type': 'preonly',
                     'pc_type': 'fieldsplit',
                     'pc_fieldsplit_type': 'additive',
+                    **{f'fieldsplit_{fs.name}': {
+                        'ksp_type': 'cg',
+                        'pc_type': 'bjacobi',
+                        'sub_pc_type': 'ilu'
+                    } for fs in self.fs.subfunctions}
                 })
-                for fs in self.fs.subfunctions:
-                    self.solver_parameters.update({
-                        f'fieldsplit_{fs.name}': {
-                            'ksp_type': 'cg',
-                            'pc_type': 'bjacobi',
-                            'sub_pc_type': 'ilu',
-                        }
-                    })
 
     @cached_property
     def lhs(self):

diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py
@@ -541,9 +541,14 @@ def __init__(self, equation, alpha):
         )
 
         self.solver_parameters = {
-            'ksp_type': 'cg',
-            'pc_type': 'bjacobi',
-            'sub_pc_type': 'ilu'
+            'ksp_type': 'preonly',
+            'pc_type': 'fieldsplit',
+            'pc_fieldsplit_type': 'additive',
+            **{f'fieldsplit_{fs.name}': {
+                'ksp_type': 'cg',
+                'pc_type': 'bjacobi',
+                'sub_pc_type': 'ilu'
+            } for fs in W.subfunctions}
         }
 
         self.solvers = {}