misc: Rename petsc/utils.py

Author: ZoeLeibowitz

conftest.py

@@ -14,7 +14,7 @@
 from devito.ir.iet import (FindNodes, FindSymbols, Iteration, ParallelBlock,
                            retrieve_iteration_tree)
 from devito.tools import as_tuple
-from devito.petsc.utils import PetscOSError, get_petsc_dir
+from devito.petsc.internals import PetscOSError, get_petsc_dir
 
 try:
     from mpi4py import MPI  # noqa

devito/passes/iet/languages/C.py

@@ -8,7 +8,7 @@
 from devito.symbolics import c_complex, c_double_complex
 from devito.tools import dtype_to_cstr
 
-from devito.petsc.utils import petsc_type_mappings
+from devito.petsc.internals import petsc_type_mappings
 
 __all__ = ['CBB', 'CDataManager', 'COrchestrator']
 

devito/petsc/iet/callback_builder.py

@@ -3,7 +3,9 @@
 from devito.ir.iet import (Call, FindSymbols, List, Uxreplace, CallableBody,
                            Dereference, DummyExpr, BlankLine, Callable, Iteration,
                            PointerCast, Definition)
-from devito.symbolics import (Byref, FieldFromPointer, IntDiv, Deref, Mod, String, Null, VOID)
+from devito.symbolics import (
+    Byref, FieldFromPointer, IntDiv, Deref, Mod, String, Null, VOID
+)
 from devito.symbolics.unevaluation import Mul
 from devito.types.basic import AbstractFunction
 from devito.types import Dimension, Temp, TempArray

devito/petsc/iet/passes.py

@@ -17,7 +17,7 @@
 )
 from devito.petsc.types.macros import petsc_func_begin_user
 from devito.petsc.iet.nodes import PetscMetaData, petsc_call
-from devito.petsc.utils import core_metadata, petsc_languages
+from devito.petsc.internals import core_metadata, petsc_languages
 from devito.petsc.iet.callback_builder import (
     BaseCallback, CoupledCallback, populate_matrix_context, get_user_struct_fields
 )

devito/petsc/iet/pre_solver.py

@@ -83,11 +83,13 @@ def _setup(self):
                                   [sobjs['snes'], Byref(sobjs['ksp'])])
 
         matvec = self.cbbuilder.main_matvec_callback
-        matvec_operation = petsc_call('MatShellSetOperation',
+        matvec_operation = petsc_call(
+            'MatShellSetOperation',
             [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)]
         )
         formfunc = self.cbbuilder._F_efunc
-        formfunc_operation = petsc_call('SNESSetFunction',
+        formfunc_operation = petsc_call(
+            'SNESSetFunction',
             [sobjs['snes'], Null, FormFunctionCallback(formfunc.name, void, void),
              self.snes_ctx]
         )

devito/petsc/utils.py devito/petsc/internals.pydevito/petsc/utils.py renamed to devito/petsc/internals.py

@@ -2,10 +2,7 @@
 import ctypes
 from pathlib import Path
 
-from devito.tools import memoized_func, filter_ordered, as_tuple
-from devito.types import Symbol, SteppingDimension, TimeDimension
-from devito.operations.solve import eval_time_derivatives
-from devito.symbolics import retrieve_functions, retrieve_dimensions
+from devito.tools import memoized_func
 
 
 class PetscOSError(OSError):
@@ -106,43 +103,3 @@ def get_petsc_type_mappings():
 
 
 petsc_languages = ['petsc']
-
-
-def get_funcs(exprs):
-    funcs = [
-        f for e in exprs
-        for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs))
-    ]
-    return as_tuple(filter_ordered(funcs))
-
-
-def generate_time_mapper(exprs):
-    """
-    Replace time indices with `Symbols` in expressions used within
-    PETSc callback functions. These symbols are Uxreplaced at the IET
-    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 `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.
-    Examples
-    --------
-    >>> exprs = (Eq(f1(t + dt, x, y), g1(t + dt, x, y) + g2(t, x, y)*f1(t, x, y)),)
-    >>> generate_time_mapper(exprs)
-    {t + dt: tau0, t: tau1}
-    """
-    # First, map any actual TimeDimensions
-    time_indices = [d for d in retrieve_dimensions(exprs) if isinstance(d, TimeDimension)]
-
-    funcs = get_funcs(exprs)
-
-    time_indices.extend(list({
-        i if isinstance(d, SteppingDimension) else d
-        for f in funcs
-        for i, d in zip(f.indices, f.dimensions)
-        if d.is_Time
-    }))
-    tau_symbs = [Symbol('tau%d' % i) for i in range(len(time_indices))]
-    return dict(zip(time_indices, tau_symbs))

devito/petsc/logging.py

@@ -7,7 +7,7 @@
 from devito.petsc.types import (
     PetscInt, PetscScalar, KSPType, KSPConvergedReason, KSPNormType
 )
-from devito.petsc.utils import petsc_type_to_ctype
+from devito.petsc.internals import petsc_type_to_ctype
 
 
 class PetscEntry:

devito/petsc/solve.py

@@ -1,12 +1,15 @@
 from devito.types.equation import PetscEq
-from devito.tools import as_tuple
+from devito.tools import filter_ordered, as_tuple
+from devito.types import Symbol, SteppingDimension, TimeDimension
+from devito.operations.solve import eval_time_derivatives
+from devito.symbolics import retrieve_functions, retrieve_dimensions
+
 from devito.petsc.types import (LinearSolverMetaData, PETScArray, DMDALocalInfo,
                                 FieldData, MultipleFieldData, Jacobian, Residual,
                                 MixedResidual, MixedJacobian, InitialGuess)
 from devito.petsc.types.equation import EssentialBC
 from devito.petsc.solver_parameters import (linear_solver_parameters,
                                             format_options_prefix)
-from devito.petsc.utils import get_funcs, generate_time_mapper
 
 
 __all__ = ['PETScSolve']
@@ -186,5 +189,45 @@ def linear_solve_args(self):
         return targets[0], funcs, all_data
 
 
+def get_funcs(exprs):
+    funcs = [
+        f for e in exprs
+        for f in retrieve_functions(eval_time_derivatives(e.lhs - e.rhs))
+    ]
+    return as_tuple(filter_ordered(funcs))
+
+
+def generate_time_mapper(exprs):
+    """
+    Replace time indices with `Symbols` in expressions used within
+    PETSc callback functions. These symbols are Uxreplaced at the IET
+    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 `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.
+    Examples
+    --------
+    >>> exprs = (Eq(f1(t + dt, x, y), g1(t + dt, x, y) + g2(t, x, y)*f1(t, x, y)),)
+    >>> generate_time_mapper(exprs)
+    {t + dt: tau0, t: tau1}
+    """
+    # First, map any actual TimeDimensions
+    time_indices = [d for d in retrieve_dimensions(exprs) if isinstance(d, TimeDimension)]
+
+    funcs = get_funcs(exprs)
+
+    time_indices.extend(list({
+        i if isinstance(d, SteppingDimension) else d
+        for f in funcs
+        for i, d in zip(f.indices, f.dimensions)
+        if d.is_Time
+    }))
+    tau_symbs = [Symbol('tau%d' % i) for i in range(len(time_indices))]
+    return dict(zip(time_indices, tau_symbs))
+
+
 localinfo = DMDALocalInfo(name='info', liveness='eager')
 prefixes = ['y', 'x', 'f', 'b']

devito/petsc/types/metadata.py

@@ -8,7 +8,7 @@
 from devito.types.equation import Eq
 from devito.operations.solve import eval_time_derivatives
 
-from devito.petsc.utils import petsc_variables
+from devito.petsc.internals import petsc_variables
 from devito.petsc.types.equation import EssentialBC, ZeroRow, ZeroColumn