Change SolveExpr name to SolverMetaData and optimise logger

Author: ZoeLeibowitz

devito/petsc/clusters.py

@@ -1,5 +1,5 @@
 from devito.tools import timed_pass
-from devito.petsc.types import SolveExpr
+from devito.petsc.types import SolverMetaData
 
 
 @timed_pass()
@@ -19,7 +19,7 @@ def petsc_lift(clusters):
     """
     processed = []
     for c in clusters:
-        if isinstance(c.exprs[0].rhs, SolveExpr):
+        if isinstance(c.exprs[0].rhs, SolverMetaData):
             ispace = c.ispace.lift(c.exprs[0].rhs.field_data.space_dimensions)
             processed.append(c.rebuild(ispace=ispace))
         else:

devito/petsc/iet/logging.py

@@ -13,9 +13,10 @@ class PetscLogger:
     """
     Class for PETSc loggers that collect solver related statistics.
     """
-    def __init__(self, level, get_info=[], **kwargs):
+    # TODO: Update docstring with kwargs
+    def __init__(self, level, **kwargs):
 
-        self.query_functions = get_info
+        self.query_functions = kwargs.get('get_info', [])
         self.sobjs = kwargs.get('solver_objs')
         self.sreg = kwargs.get('sregistry')
         self.section_mapper = kwargs.get('section_mapper', {})
@@ -32,9 +33,9 @@ def __init__(self, level, get_info=[], **kwargs):
                 # SNES specific
                 'snesgetiterationnumber',
             ]
-            self.query_functions.extend(
-                [f for f in funcs if f not in self.query_functions]
-            )
+            self.query_functions = set(self.query_functions)
+            self.query_functions.update(funcs)
+            self.query_functions = sorted(list(self.query_functions))
 
         # TODO: To be extended with if level <= DEBUG: ...
 

devito/petsc/iet/passes.py

@@ -203,7 +203,7 @@ def solve(self):
     def logger(self):
         log_level = devito.logger.logger.level
         return PetscLogger(
-            log_level, self.get_info, **self.common_kwargs
+            log_level, get_info=self.get_info, **self.common_kwargs
         )
 
     @cached_property

devito/petsc/iet/routines.py

@@ -1726,10 +1726,10 @@ class TimeDependent(NonTimeDependent):
       for each `SNESSolve` at every time step, don't require the time loop, but
       may still need access to data from other time steps.
     - All `Function` objects are passed through the initial lowering via the
-      `SolveExpr` object, ensuring the correct time loop is generated
+      `SolverMetaData` object, ensuring the correct time loop is generated
       in the main kernel.
     - Another mapper is created based on the modulo dimensions
-      generated by the `SolveExpr` object in the main kernel
+      generated by the `SolverMetaData` object in the main kernel
       (e.g., {time: time, t: t0, t + 1: t1}).
     - These two mappers are used to generate a final mapper `symb_to_moddim`
       (e.g. {tau0: t0, tau1: t1}) which is used at the IET level to

devito/petsc/solve.py

@@ -1,6 +1,6 @@
 from devito.types.equation import PetscEq
 from devito.tools import as_tuple
-from devito.petsc.types import (LinearSolveExpr, PETScArray, DMDALocalInfo,
+from devito.petsc.types import (LinearSolverMetaData, PETScArray, DMDALocalInfo,
                                 FieldData, MultipleFieldData, Jacobian, Residual,
                                 MixedResidual, MixedJacobian, InitialGuess)
 from devito.petsc.types.equation import EssentialBC
@@ -103,7 +103,7 @@ def build_expr(self):
         target, funcs, field_data = self.linear_solve_args()
         # Placeholder expression for inserting calls to the solver
 
-        linear_solve = LinearSolveExpr(
+        linear_solve = LinearSolverMetaData(
             funcs,
             solver_parameters=self.solver_parameters,
             field_data=field_data,

devito/petsc/types/types.py

@@ -32,7 +32,7 @@ class GetArgs(MetaData):
     pass
 
 
-class SolveExpr(MetaData):
+class SolverMetaData(MetaData):
     """
     A symbolic expression passed through the Operator, containing the metadata
     needed to execute the PETSc solver.
@@ -59,7 +59,7 @@ def __new__(cls, expr, solver_parameters=None,
         obj.localinfo = localinfo
         obj.user_prefix = user_prefix
         obj.formatted_prefix = formatted_prefix
-        obj.get_info = get_info
+        obj.get_info = get_info if get_info is not None else []
         return obj
 
     def __repr__(self):
@@ -76,7 +76,7 @@ def _hashable_content(self):
         return (self.expr, self.formatted_prefix, self.solver_parameters)
 
     def __eq__(self, other):
-        return (isinstance(other, SolveExpr) and
+        return (isinstance(other, SolverMetaData) and
                 self.expr == other.expr and
                 self.formatted_prefix == other.formatted_prefix
                 and self.solver_parameters == other.solver_parameters)
@@ -92,15 +92,15 @@ def eval(cls, *args):
     func = Reconstructable._rebuild
 
 
-class LinearSolveExpr(SolveExpr):
+class LinearSolverMetaData(SolverMetaData):
     """
     Linear problems are handled by setting the SNESType to 'ksponly',
     enabling a unified interface for both linear and nonlinear solvers.
     """
     pass
 
 
-class NonLinearSolveExpr(SolveExpr):
+class NonLinearSolverMetaData(SolverMetaData):
     """
     TODO: Non linear solvers are not yet supported.
     """
@@ -109,7 +109,7 @@ class NonLinearSolveExpr(SolveExpr):
 
 class FieldData:
     """
-    Metadata for a single `target` field passed to `SolveExpr`.
+    Metadata for a single `target` field passed to `SolverMetaData`.
     Used to interface with PETSc SNES solvers at the IET level.
 
     Parameters
@@ -134,8 +134,8 @@ def __init__(self, target=None, jacobian=None, residual=None,
         petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']]
         if self._target.dtype != petsc_precision:
             raise TypeError(
-                f"Your target dtype must match the precision of your "
-                f"PETSc configuration. "
+                "Your target dtype must match the precision of your "
+                "PETSc configuration. "
                 f"Expected {petsc_precision}, but got {self._target.dtype}."
             )
         self._jacobian = jacobian
@@ -182,7 +182,7 @@ def targets(self):
 
 class MultipleFieldData(FieldData):
     """
-    Metadata class passed to `SolveExpr`, for mixed-field problems,
+    Metadata class passed to `SolverMetaData`, for mixed-field problems,
     where the solution vector spans multiple `targets`. Used to interface
     with PETSc SNES solvers at the IET level.
 

devito/petsc/utils.py

@@ -123,7 +123,7 @@ def generate_time_mapper(funcs):
     level to align with the `TimeDimension` and `ModuloDimension` objects
     present in the initial lowering.
     NOTE: All functions used in PETSc callback functions are attached to
-    the `SolveExpr` object, which is passed through the initial lowering
+    the `SolverMetaData` object, which is passed through the initial lowering
     (and subsequently dropped and replaced with calls to run the solver).
     Therefore, the appropriate time loop will always be correctly generated inside
     the main kernel.