misc: Clean up

Author: ZoeLeibowitz

.github/workflows/pytest-petsc.yml

@@ -72,7 +72,7 @@ jobs:
       run: |
         ${{ env.RUN_CMD }} python3 -c "from devito import configuration; print(''.join(['%s: %s \n' % (k, v) for (k, v) in configuration.items()]))"
 
-    - name: Test with pytest - non parallel
+    - name: Test with pytest - serial
       run: |
         ${{ env.RUN_CMD }} mpiexec -n 1 pytest -m "not parallel" --cov --cov-config=.coveragerc --cov-report=xml ${{ env.TESTS }}
 

devito/petsc/iet/logging.py

@@ -22,9 +22,12 @@ def __init__(self, level, **kwargs):
 
         if level <= PERF:
             self.function_list.extend([
+                # KSP specific
                 'kspgetiterationnumber',
+                'kspgettolerances',
+                'kspgetconvergedreason',
+                # SNES specific
                 'snesgetiterationnumber',
-                'kspgettolerances'
             ])
 
         # TODO: To be extended with if level <= DEBUG: ...

devito/petsc/iet/passes.py

@@ -69,20 +69,21 @@ def lower_petsc(iet, **kwargs):
     # Map PETScSolve to its Section (for logging)
     section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet)
 
-    # Prefixes within a single Operator should not be duplicated
+    # Prefixes within the same Operator should not be duplicated
     prefixes = [d.expr.rhs.user_prefix for d in data if d.expr.rhs.user_prefix]
     duplicates = {p for p in prefixes if prefixes.count(p) > 1}
 
-    # TODO: Avoid the other exception raised given it is an iet_pass?
+    # TODO: Avoid the other exception raised - think due to exception being
+    # raised inside the @iet_pass?
     if duplicates:
         dup_list = ", ".join(repr(p) for p in sorted(duplicates))
         raise ValueError(
             f"The following `options_prefix` values are duplicated "
             f"among your PETScSolves. Ensure each one is unique: {dup_list}"
         )
 
-    # List of calls to clear options from the global PETSc options database.
-    # These are executed at the end of the Operator.
+    # List of `Call`s to clear options from the global PETSc options database,
+    # executed at the end of the Operator.
     clear_options = []
 
     for iters, (inject_solve,) in inject_solve_mapper.items():

devito/petsc/iet/routines.py

@@ -136,7 +136,7 @@ def _make_callable_body(self, body, stacks=(), casts=()):
 
     def _make_options_callback(self):
         """
-        Create two callbacks: one to set PETSc options and one for
+        Create two callbacks: one to set PETSc options and one
         to clear them.
 
         Options are only set/cleared if they were not specifed via
@@ -153,6 +153,7 @@ def _make_options_callback(self):
                 # Ensures that the command line args take priority
                 continue
             option_name = String(option)
+            # For options without a value e.g `ksp_view`, pass Null
             option_value = Null if v is None else String(str(v))
             set_body.append(
                 petsc_call('PetscOptionsSetValue', [Null, option_name, option_value])

devito/petsc/logging.py

@@ -134,10 +134,12 @@ def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper,
 
         mapper = {v: k for k, v in petsc_type_mappings.items()}
 
-        fields = [
-            (str(ptype), mapper[str(ptype._C_ctype)])
-            for option in petsc_option_mapper.values() for ptype in option.values()
-        ]
+        fields = []
+        for obj_mapper in petsc_option_mapper.values():
+            for petsc_option in obj_mapper.values():
+                ctype = mapper[str(petsc_option._C_ctype)]
+                fields.append((petsc_option.name, ctype))
+
         super().__init__(name, pname, fields)
 
     @property
@@ -164,24 +166,22 @@ def __getattr__(self, attr):
         # Helper to get the value from the profiling struct
         get_val = lambda v: getattr(self.value._obj, v.name)
 
-        # Return the value(s) for the given PETSc attribute:
-        # - If there's only one output (e.g., for KSPGetIterationNumber),
-        #   return it directly.
-        # - If there are multiple outputs (e.g., for KSPGetTolerances),
-        #   return a dictionary mapping each output name to
-        #   its value, e.g., {'rtol': val0, 'abstol': val1, ...}.
+        # - If the function returns a single value (e.g., KSPGetIterationNumber),
+        #   return that value directly.
+        # - If the function returns multiple values (e.g., KSPGetTolerances),
+        #   return a dictionary mapping each output name to its value,
+        #   e.g., {'rtol': val0, 'abstol': val1, ...}.
         if len(obj_mapper) == 1:
             return get_val(next(iter(obj_mapper.values())))
         return {k: get_val(v) for k, v in obj_mapper.items()}
 
 
-# TODO: change the lists to tuples
 @dataclass
 class PetscReturnVariable:
     name: str
-    variable_type: list
+    variable_type: tuple
     input_params: str
-    output_param: list[str]
+    output_param: tuple[str]
 
 
 # NOTE:
@@ -190,23 +190,32 @@ class PetscReturnVariable:
 # If any of the PETSc function signatures change (e.g., names, input/output parameters),
 # this dictionary must be updated accordingly.
 
+# TODO: To be extended
 petsc_return_variable_dict = {
+    # KSP specific
     'kspgetiterationnumber': PetscReturnVariable(
         name='KSPGetIterationNumber',
-        variable_type=[PetscInt],
+        variable_type=(PetscInt,),
+        input_params='ksp',
+        output_param=('kspits',)
+    ),
+    'kspgettolerances': PetscReturnVariable(
+        name='KSPGetTolerances',
+        variable_type=(PetscScalar, PetscScalar, PetscScalar, PetscInt),
         input_params='ksp',
-        output_param=['kspits']
+        output_param=('rtol', 'abstol', 'dtol', 'maxits'),
     ),
+    'kspgetconvergedreason': PetscReturnVariable(
+        name='KSPGetConvergedReason',
+        variable_type=(PetscInt,),
+        input_params='ksp',
+        output_param=('reason',),
+    ),
+    # SNES specific
     'snesgetiterationnumber': PetscReturnVariable(
         name='SNESGetIterationNumber',
-        variable_type=[PetscInt],
+        variable_type=(PetscInt,),
         input_params='snes',
-        output_param=['snesits'],
+        output_param=('snesits',),
     ),
-    'kspgettolerances': PetscReturnVariable(
-        name='KSPGetTolerances',
-        variable_type=[PetscScalar, PetscScalar, PetscScalar, PetscInt],
-        input_params='ksp',
-        output_param=['rtol', 'abstol', 'dtol', 'maxits'],
-    )
 }

devito/petsc/solver_parameters.py

@@ -39,6 +39,7 @@
 
 
 def linear_solver_parameters(solver_parameters):
+    # Flatten parameters to support nested dictionaries
     flattened = flatten_parameters(solver_parameters or {})
     processed = linear_solve_defaults.copy()
     processed.update(flattened)

tests/test_petsc.py

@@ -1476,8 +1476,7 @@ def test_logging_multiple_solves(self, log_level):
     @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG'])
     def test_logging_user_prefixes(self, log_level):
         """
-        Test that the PetscSummary uses the user provided options prefix
-        if provided, otherwise it uses the default one provided by Devito.
+        Verify that `PetscSummary` uses the user provided `options_prefix` when given.
         """
         grid = Grid(shape=(11, 11), dtype=np.float64)
 
@@ -1505,8 +1504,8 @@ def test_logging_user_prefixes(self, log_level):
     @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG'])
     def test_logging_default_prefixes(self, log_level):
         """
-        Test that the PetscSummary uses the default options prefix
-        # generated by Devito
+        Verify that `PetscSummary` uses the default options prefix
+        provided by Devito if no user `options_prefix` is specified.
         """
         grid = Grid(shape=(11, 11), dtype=np.float64)
 
@@ -1526,7 +1525,8 @@ def test_logging_default_prefixes(self, log_level):
 
         petsc_summary = summary.petsc
 
-        # Note, if users want to use logging, they should really set the options_prefix
+        # Users should set a custom options_prefix if they want logging; otherwise,
+        # the default automatically generated prefix is used in the `PetscSummary`.
         assert all(re.fullmatch(r"devito_\d+_", k.options_prefix) for k in petsc_summary)
 
 
@@ -1594,7 +1594,7 @@ def test_options_prefix(self):
     def test_options_no_value(self):
         """
         Test solver parameters that do not require a value, such as
-        `snes_view` and `ksp_view`
+        `snes_view` and `ksp_view`.
         """
         solver = PETScSolve(
             self.eq1, target=self.e, solver_parameters={'snes_view': None},
@@ -1629,6 +1629,8 @@ def test_tolerances(self, log_level):
         entry = petsc_summary.get_entry('section0', 'solver')
         tolerances = entry.KSPGetTolerances
 
+        # Test that the tolerances have been set correctly and therefore
+        # appear as expected in the `PetscSummary`.
         assert tolerances['rtol'] == params['ksp_rtol']
         assert tolerances['abstol'] == params['ksp_atol']
         assert tolerances['dtol'] == params['ksp_divtol']
@@ -1672,7 +1674,7 @@ def test_error_if_same_prefix(self):
 
     @skipif('petsc')
     @pytest.mark.parametrize('log_level', ['PERF', 'DEBUG'])
-    def test_multiple_operators(self):
+    def test_multiple_operators(self, log_level):
         """
         Verify that solver parameters are set correctly when multiple Operators
         are created with PETScSolve instances sharing the same options_prefix.
@@ -1690,7 +1692,7 @@ def test_multiple_operators(self):
             self.eq2, target=self.g, options_prefix='poisson',
             solver_parameters={'ksp_rtol': '1e-12'}
         )
-        with switchconfig(language='petsc'):
+        with switchconfig(language='petsc', log_level=log_level):
             op1 = Operator(solver1)
             op2 = Operator(solver2)
             summary1 = op1.apply()
@@ -1708,7 +1710,6 @@ def test_multiple_operators(self):
     # TODO: Add test to check that the command line args override anything set
     # in the solver_parameters dictionary
 
-    # TODO: update names of all of these tests
     # @skipif('petsc')
     # def test_command_line_priority(self):
     #     """