misc: Add tests and clean up

ZoeLeibowitz · ZoeLeibowitz · commit 586b77ef87a3 · 2025-08-14T15:16:46.000+01:00
diff --git a/devito/petsc/iet/logging.py b/devito/petsc/iet/logging.py
@@ -29,9 +29,6 @@ def __init__(self, level, **kwargs):
 
         # TODO: To be extended with if level <= DEBUG: ...
 
-        # from IPython import embed; embed()
-        # if str(self.inject_solve.expr.rhs.solver_parameters['ksp_rtol']) == '1e-15':
-
         name = self.sreg.make_name(prefix='petscinfo')
         pname = self.sreg.make_name(prefix='petscprofiler')
 
@@ -40,36 +37,31 @@ def __init__(self, level, **kwargs):
             self.section_mapper, self.inject_solve,
             self.function_list
         )
-        # else:
-        #     name = self.sreg.make_name(prefix='petscinfooo')
-        #     pname = self.sreg.make_name(prefix='petscprofilerrrr')
-
-        #     self.statstruct = PetscInfo(
-        #         name, pname, self.petsc_option_mapper, self.sobjs,
-        #         self.section_mapper, self.inject_solve,
-        #         self.function_list
-        #     )
-
-        # from IPython import embed; embed()  # noqa: E402
-
-    # @property
-    # def statstruct(self):
-    #     return self._statstruct
 
     @cached_property
     def petsc_option_mapper(self):
         """
-        Create PETSc objects specifically needed for logging solver statistics.
-
-        ADD EXTENDED DOCSTRING 
+        For each function in `self.function_list`, look up its metadata in
+        `petsc_return_variable_dict` and instantiate the corresponding PETSc logging
+        variables with names from the symbol registry.
+
+        Example:
+        --------
+        >>> self.function_list
+        ['kspgetiterationnumber', 'snesgetiterationnumber', 'kspgettolerances']
+
+        >>> self.petsc_option_mapper
+        {
+            'KSPGetIterationNumber': {'kspits': kspits0},
+            'KSPGetTolerances': {'rtol': rtol0, 'abstol': abstol0, ...}
+        }
         """
         opts = {}
         for func_name in self.function_list:
             info = petsc_return_variable_dict[func_name]
             opts[info.name] = {}
             for vtype, out in zip(info.variable_type, info.output_param, strict=True):
                 opts[info.name][out] = vtype(self.sreg.make_name(prefix=out))
-        
         return opts
 
     @cached_property
@@ -86,7 +78,7 @@ def calls(self):
             input = self.sobjs[return_variable.input_params]
             output_params = self.petsc_option_mapper[return_variable.name].values()
             outputs = [Byref(i) for i in output_params]
-            # from IPython import embed; embed()
+
             calls.append(
                 petsc_call(return_variable.name, [input] + outputs)
             )
diff --git a/devito/petsc/iet/passes.py b/devito/petsc/iet/passes.py
@@ -1,12 +1,11 @@
 import cgen as c
 import numpy as np
 from functools import cached_property
-from sympy import Not
 
 from devito.passes.iet.engine import iet_pass
 from devito.ir.iet import (Transformer, MapNodes, Iteration, BlankLine,
                            DummyExpr, CallableBody, List, Call, Callable,
-                           FindNodes, Section, Conditional)
+                           FindNodes, Section)
 from devito.symbolics import Byref, FieldFromPointer, Macro
 from devito.types import Symbol, Scalar
 from devito.types.basic import DataSymbol
@@ -17,10 +16,9 @@
                                 PointerIS, Mat, CallbackVec, Vec, CallbackMat, SNES,
                                 DummyArg, PetscInt, PointerDM, PointerMat, MatReuse,
                                 CallbackPointerIS, CallbackPointerDM, JacobianStruct,
-                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol,
-                                CharPtr, PetscBool)
+                                SubMatrixStruct, Initialize, Finalize, ArgvSymbol)
 from devito.petsc.types.macros import petsc_func_begin_user, Null
-from devito.petsc.iet.nodes import PetscMetaData, PETScCallable
+from devito.petsc.iet.nodes import PetscMetaData
 from devito.petsc.utils import core_metadata, petsc_languages
 from devito.petsc.iet.routines import (CBBuilder, CCBBuilder, BaseObjectBuilder,
                                        CoupledObjectBuilder, BaseSetup, CoupledSetup,
@@ -71,16 +69,19 @@ def lower_petsc(iet, **kwargs):
     # Map PETScSolve to its Section (for logging)
     section_mapper = MapNodes(Section, PetscMetaData, 'groupby').visit(iet)
 
-    # Generate a shared callback used by all PETScSolve instances to set
-    # individual PetscOptions
-    # set_solver_option(efuncs)
-    # List of all callbacks that clear PetscOptions
+    # prefixes within a single 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: throw a warning/error if the user passes a solver in with the same options_prefix
-    # it's going to lead to weird solver option behaviour. Note, if you use the options_prefix across
-    # different Operator runs, it will not be an issue
-    clear_options = []
+    # TODO: How to avoid the other exception raised given it is an iet_pass?
+    if duplicates:
+        dup_list = ", ".join(sorted(duplicates))
+        raise ValueError(
+            f"The following `options_prefix` values are duplicated "
+            f"among your PETScSolves. Ensure each one is unique: {dup_list}"
+        )
 
+    clear_options = []
     for iters, (inject_solve,) in inject_solve_mapper.items():
 
         builder = Builder(inject_solve, iters, comm, section_mapper, **kwargs)
@@ -99,7 +100,6 @@ def lower_petsc(iet, **kwargs):
     populate_matrix_context(efuncs)
 
     iet = Transformer(subs).visit(iet)
-    # from IPython import embed; embed()
     body = core + tuple(setup) + iet.body.body + tuple(clear_options)
     body = iet.body._rebuild(body=body)
     iet = iet._rebuild(body=body)
@@ -231,32 +231,6 @@ def populate_matrix_context(efuncs):
     )
 
 
-def set_solver_option(efuncs):
-
-    option = CharPtr(name='option', is_const=True)
-    value = CharPtr(name='value', is_const=True)
-    set = PetscBool(name='set')
-
-    body = List(body=[
-        petsc_call('PetscOptionsHasName', [Null, Null, option, Byref(set)]),
-        Conditional(Not(set), petsc_call('PetscOptionsSetValue', [Null, option, value]))
-    ])
-
-    body = CallableBody(
-        body,
-        init=(petsc_func_begin_user,),
-        retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
-    )
-
-    cb = PETScCallable(
-        'SetPetscOption',
-        body,
-        retval=objs['err'],
-        parameters=(option, value)
-    )
-    efuncs[cb.name] = cb
-
-
 subdms = PointerDM(name='subdms')
 fields = PointerIS(name='fields')
 submats = PointerMat(name='submats')
diff --git a/devito/petsc/iet/routines.py b/devito/petsc/iet/routines.py
@@ -132,8 +132,12 @@ def _make_options_callback(self):
                 continue
             option_name = String(option)
             option_value = Null if v is None else String(str(v))
-            set_body.append(petsc_call('PetscOptionsSetValue', [Null, option_name, option_value]))
-            clear_body.append(petsc_call('PetscOptionsClearValue', [Null, option_name]))
+            set_body.append(
+                petsc_call('PetscOptionsSetValue', [Null, option_name, option_value])
+            )
+            clear_body.append(
+                petsc_call('PetscOptionsClearValue', [Null, option_name])
+            )
 
         set_body = CallableBody(
             List(body=set_body),
diff --git a/devito/petsc/logging.py b/devito/petsc/logging.py
@@ -63,9 +63,7 @@ def petsc_entry(self, petscinfo):
         containing the values for each PETSc function call.
         """
         funcs = self._functions
-        # from IPython import embed; embed()
         values = tuple(getattr(petscinfo, c) for c in funcs)
-        # from IPython import embed; embed()
         return PetscEntry(**{k: v for k, v in zip(funcs, values)})
 
     def _add_properties(self):
@@ -133,16 +131,14 @@ def __init__(self, name, pname, petsc_option_mapper, sobjs, section_mapper,
         self.section_mapper = section_mapper
         self.inject_solve = inject_solve
         self.function_list = function_list
+        self.formatted_prefix = inject_solve.expr.rhs.formatted_prefix
 
         mapper = {v: k for k, v in petsc_type_mappings.items()}
 
-        self.formatted_prefix = inject_solve.expr.rhs.formatted_prefix
-
         fields = [
             (str(ptype), mapper[str(ptype._C_ctype)])
             for option in petsc_option_mapper.values() for ptype in option.values()
         ]
-
         super().__init__(name, pname, fields)
 
     @property
@@ -152,45 +148,29 @@ def section(self):
 
     @property
     def summary_key(self):
+        # If users provide an options prefix, use it in the summary;
+        # otherwise, use the default one generated by Devito
         user_prefix = self.inject_solve.expr.rhs.user_prefix
-        # TODO: this will be the case when using the default options prefix provided by Devito
-        # if user_prefix is None:
-        #     user_prefix = self.formatted_prefix
+        if user_prefix is None:
+            user_prefix = self.formatted_prefix
         return (self.section, user_prefix)
 
     def __getattr__(self, attr):
-        if attr in self.petsc_option_mapper.keys():
-            if len(self.petsc_option_mapper[attr].values()) > 1:
-                tmp = {}
-                for i, j in self.petsc_option_mapper[attr].items():
-                    tmp2 = getattr(self.value._obj, j.name)
-                    tmp[i] = tmp2
-                return tmp
-            else:
-                # TODO: CLEANNN
-                first_val = list(self.petsc_option_mapper[attr].values())[0]
-                return getattr(self.value._obj, first_val.name)
-        # from IPython import embed; embed()  # noqa: E402
-        raise AttributeError(f"{attr} not found in PETSc return variables")
-    
-    # TODO: maybe just overrider _hashable_content??
-
-    # def _hashable_content(self):
-    #     """
-    #     Return a tuple of the formatted prefix and section for hashing.
-    #     This is used to ensure that two PetscInfo objects with the same
-    #     formatted prefix and section are considered equal.
-    #     """
-    #     return (self.name, self.dtype, self.inject_solve.expr.rhs)
-    
-    # def __eq__(self, other):
-    #     if not isinstance(other, PetscInfo):
-    #         return NotImplemented
-    #     # return self.formatted_prefix == other.formatted_prefix
-    #     return self.inject_solve.expr.rhs == other.inject_solve.expr.rhs
-
-    # def __hash__(self):
-    #     return hash(self.inject_solve.expr.rhs)
+        if attr not in self.petsc_option_mapper:
+            raise AttributeError(f"{attr} not found in PETSc return variables")
+
+        # Maps the petsc_option to its generated variable name e.g {'its': its0}
+        obj_mapper = self.petsc_option_mapper[attr]
+
+        # Helper to get the value from the petsc profiling struct
+        get_val = lambda v: getattr(self.value._obj, v.name)
+
+        # If there's only one value to retrieve for the given attribute, for example, KSPGetIterationNumber
+        # we return it directly, otherwise we return a dictionary of all values e.g for KSPGetTolerances
+        # we return {'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
@@ -213,13 +193,13 @@ class PetscReturnVariable:
         name='KSPGetIterationNumber',
         variable_type=[PetscInt],
         input_params='ksp',
-        output_param=['kspiter']
+        output_param=['kspits']
     ),
     'snesgetiterationnumber': PetscReturnVariable(
         name='SNESGetIterationNumber',
         variable_type=[PetscInt],
         input_params='snes',
-        output_param=['snesiter'],
+        output_param=['snesits'],
     ),
     'kspgettolerances': PetscReturnVariable(
         name='KSPGetTolerances',
diff --git a/devito/petsc/solve.py b/devito/petsc/solve.py
@@ -52,10 +52,9 @@ def PETScSolve(target_exprs, target=None, solver_parameters=None, options_prefix
         This can be passed directly to a Devito Operator.
     """
     if target is not None:
-        tmp = InjectSolve(
+        return InjectSolve(
             solver_parameters, {target: target_exprs}, options_prefix
-        )
-        return tmp.build_expr()
+        ).build_expr()
     else:
         return InjectMixedSolve(
             solver_parameters, target_exprs, options_prefix
@@ -83,7 +82,6 @@ def build_expr(self):
             user_prefix=self.user_prefix,
             formatted_prefix=self.formatted_prefix
         )
-        # from IPython import embed; embed()
         return PetscEq(target, linear_solve)
 
     def linear_solve_args(self):
diff --git a/devito/petsc/solver_parameters.py b/devito/petsc/solver_parameters.py
@@ -47,7 +47,7 @@ def linear_solver_parameters(solver_parameters):
 
 _options_prefix_counter = itertools.count()
 
-# TODO: add a default options prefix if not provided
+
 def format_options_prefix(options_prefix):
     # NOTE: Modified from the `OptionsManager` inside petsctools
     if options_prefix is None:
diff --git a/devito/petsc/types/types.py b/devito/petsc/types/types.py
@@ -49,7 +49,7 @@ def __new__(cls, expr, solver_parameters=None,
             obj = sympy.Basic.__new__(cls, expr)
 
         obj.expr = expr
-        obj.solver_parameters = solver_parameters
+        obj.solver_parameters = frozendict(solver_parameters)
         obj.field_data = field_data if field_data else FieldData()
         obj.time_mapper = time_mapper
         obj.localinfo = localinfo
@@ -68,12 +68,13 @@ def _sympystr(self, printer):
     __hash__ = sympy.Basic.__hash__
 
     def _hashable_content(self):
-            return (self.expr, self.formatted_prefix)
+            return (self.expr, self.formatted_prefix, self.solver_parameters)
 
     def __eq__(self, other):
         return (isinstance(other, SolveExpr) and
                 self.expr == other.expr and
-                self.formatted_prefix == other.formatted_prefix)
+                self.formatted_prefix == other.formatted_prefix
+                and self.solver_parameters == other.solver_parameters)
 
     @property
     def grid(self):
diff --git a/examples/petsc/solver_options.py b/examples/petsc/solver_options.py
@@ -3,6 +3,7 @@
 
 from devito import (Grid, Function, Eq, Operator, configuration,
                     switchconfig, Constant)
+from devito.tools import frozendict
 from devito.petsc import PETScSolve
 from devito.petsc.initialize import PetscInitialize
 configuration['compiler'] = 'custom'
@@ -51,21 +52,29 @@
 eq = Eq(e.laplace, f)
 
 params1 = {'ksp_view': None, 'ksp_rtol': 1e-15}
-params2 = {'ksp_view': None, 'ksp_rtol': 1e-11}
+params2 = {'ksp_rtol': 1e-12, 'ksp_view': None}
 petsc1 = PETScSolve(eq, target=e, solver_parameters=params1, options_prefix='pde1')
-petsc2 = PETScSolve(eq, target=e, solver_parameters=params2, options_prefix='pde1')
+petsc2 = PETScSolve(eq, target=e, solver_parameters=params2, options_prefix='pde2')
+
+
+# frozen1 = frozendict(params1)
+# frozen2 = frozendict(params2)
+
+# assert hash(frozen1) == hash(frozen2)
 
 # from IPython import embed; embed()
 
-# with switchconfig(language='petsc'):
+with switchconfig(language='petsc', log_level='DEBUG'):
 
-#     op1 = Operator([petsc1, petsc2])
-#     # op2 = Operator(petsc2)
-#     # summary1 = op1.apply()
-#     # summary2 = op2.apply()
+    op1 = Operator([petsc1])
+    # op2 = Operator(petsc2)
+    summary1 = op1.apply()
+    # summary2 = op2.apply()
     
-#     print(op1.ccode)
-#     # print(op2.ccode)
+    # print(op1.ccode)
+    # print(op2.ccode)
+
+petsc_summary = summary1.petsc
 
 
 
diff --git a/tests/test_petsc.py b/tests/test_petsc.py
