misc: Clean up and docstrings

Author: ZoeLeibowitz

devito/petsc/iet/routines.py

@@ -42,14 +42,12 @@ def __init__(self, **kwargs):
         self._struct_params = []
 
         self._main_matvec_callback = None
-        self._main_formfunc_callback = None
         self._user_struct_callback = None
-        # TODO: Test pickling. The mutability of these lists
-        # could cause issues when pickling?
+        self._F_efunc = None
+        self._b_efunc = None
+
         self._J_efuncs = []
-        self._F_efuncs = []
-        self._b_efuncs = []
-        self._initialguesses = []
+        self._initial_guesses = []
 
         self._make_core()
         self._efuncs = self._uxreplace_efuncs()
@@ -69,31 +67,26 @@ def filtered_struct_params(self):
     @property
     def main_matvec_callback(self):
         """
-        This is the matvec callback associated with the whole Jacobian i.e
-        is set in the main kernel via
-        `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))...));`
+        The matrix-vector callback for the full Jacobian.
+        This is the function set in the main Kernel via:
+            PetscCall(MatShellSetOperation(J, MATOP_MULT, (void (*)(void))...));
+        The callback has the signature `(Mat, Vec, Vec)`.
         """
         return self._J_efuncs[0]
 
-    @property
-    def main_formfunc_callback(self):
-        return self._F_efuncs[0]
-
     @property
     def J_efuncs(self):
+        """
+        List of matrix-vector callbacks.
+        Each callback has the signature `(Mat, Vec, Vec)`. Typically, this list
+        contains a single element, but in mixed systems it can include multiple
+        callbacks, one for each subblock.
+        """
         return self._J_efuncs
 
     @property
-    def F_efuncs(self):
-        return self._F_efuncs
-
-    @property
-    def b_efuncs(self):
-        return self._b_efuncs
-
-    @property
-    def initialguesses(self):
-        return self._initialguesses
+    def initial_guesses(self):
+        return self._initial_guesses
 
     @property
     def user_struct_callback(self):
@@ -284,7 +277,7 @@ def _make_formfunc(self):
             retval=objs['err'],
             parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr'])
         )
-        self._F_efuncs.append(cb)
+        self._F_efunc = cb
         self._efuncs[cb.name] = cb
 
     def _create_formfunc_body(self, body):
@@ -422,7 +415,7 @@ def _make_formrhs(self):
             retval=objs['err'],
             parameters=(sobjs['callbackdm'], objs['B'])
         )
-        self._b_efuncs.append(cb)
+        self._b_efunc = cb
         self._efuncs[cb.name] = cb
 
     def _create_form_rhs_body(self, body):
@@ -534,7 +527,7 @@ def _make_initialguess(self):
             retval=objs['err'],
             parameters=(sobjs['callbackdm'], objs['xloc'])
         )
-        self._initialguesses.append(cb)
+        self._initial_guesses.append(cb)
         self._efuncs[cb.name] = cb
 
     def _create_initial_guess_body(self, body):
@@ -660,16 +653,12 @@ def jacobian(self):
     @property
     def main_matvec_callback(self):
         """
-        This is the matvec callback associated with the whole Jacobian i.e
+        This is the matrix-vector callback associated with the whole Jacobian i.e
         is set in the main kernel via
         `PetscCall(MatShellSetOperation(J,MATOP_MULT,(void (*)(void))MyMatShellMult));`
         """
         return self._main_matvec_callback
 
-    @property
-    def main_formfunc_callback(self):
-        return self._main_formfunc_callback
-
     def _make_core(self):
         for sm in self.fielddata.jacobian.nonzero_submatrices:
             self._make_matvec(sm, prefix=f'{sm.name}_MatMult')
@@ -757,7 +746,7 @@ def _make_whole_formfunc(self):
             retval=objs['err'],
             parameters=(objs['snes'], objs['X'], objs['F'], objs['dummyptr'])
         )
-        self._main_formfunc_callback = cb
+        self._F_efunc = cb
         self._efuncs[cb.name] = cb
 
     def _whole_formfunc_body(self, body):
@@ -1310,7 +1299,7 @@ def _setup(self):
             'MatShellSetOperation',
             [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)]
         )
-        formfunc = self.cbbuilder.main_formfunc_callback
+        formfunc = self.cbbuilder._F_efunc
         formfunc_operation = petsc_call(
             'SNESSetFunction',
             [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void),
@@ -1477,7 +1466,7 @@ def _setup(self):
             'MatShellSetOperation',
             [sobjs['Jac'], 'MATOP_MULT', MatShellSetOp(matvec.name, void, void)]
         )
-        formfunc = self.cbbuilder.main_formfunc_callback
+        formfunc = self.cbbuilder._F_efunc
         formfunc_operation = petsc_call(
             'SNESSetFunction',
             [sobjs['snes'], objs['Null'], FormFunctionCallback(formfunc.name, void, void),
@@ -1593,16 +1582,16 @@ def _execute_solve(self):
 
         struct_assignment = self.timedep.assign_time_iters(sobjs['userctx'])
 
-        b_efunc = self.cbbuilder.b_efuncs[0]
+        b_efunc = self.cbbuilder._b_efunc
 
         dmda = sobjs['dmda']
 
         rhs_call = petsc_call(b_efunc.name, [sobjs['dmda'], sobjs['bglobal']])
 
         vec_place_array = self.timedep.place_array(target)
 
-        if self.cbbuilder.initialguesses:
-            initguess = self.cbbuilder.initialguesses[0]
+        if self.cbbuilder.initial_guesses:
+            initguess = self.cbbuilder.initial_guesses[0]
             initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']])
         else:
             initguess_call = None

devito/petsc/types/array.py

@@ -169,7 +169,7 @@ def _C_ctype(self):
         fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components]
         return POINTER(type(self.pname, (Structure,), {'_fields_': fields}))
 
-    @cached_property
+    @property
     def symbolic_shape(self):
         return self.c0.symbolic_shape
 

devito/petsc/types/types.py

@@ -1,6 +1,7 @@
 import sympy
 
 from itertools import chain
+from functools import cached_property
 
 from devito.tools import Reconstructable, sympy_mutex, as_tuple, frozendict
 from devito.tools.dtypes_lowering import dtype_mapper
@@ -241,7 +242,7 @@ def __init__(self, targets, arrays, jacobian=None, residual=None):
         self._jacobian = jacobian
         self._residual = residual
 
-    @property
+    @cached_property
     def space_dimensions(self):
         space_dims = {t.space_dimensions for t in self.targets}
         if len(space_dims) > 1:
@@ -251,7 +252,7 @@ def space_dimensions(self):
             )
         return space_dims.pop()
 
-    @property
+    @cached_property
     def grid(self):
         """The unique `Grid` associated with all targets."""
         grids = [t.grid for t in self.targets]
@@ -261,7 +262,7 @@ def grid(self):
             )
         return grids.pop()
 
-    @property
+    @cached_property
     def space_order(self):
         # NOTE: since we use DMDA to create vecs for the coupled solves,
         # all fields must have the same space order
@@ -398,7 +399,8 @@ def _build_matvecs(self):
             matvecs.extend(
                 e for e in self._build_matvec_expr(eq) if e is not None
             )
-        matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)))
+        key = lambda e: not isinstance(e, EssentialBC)
+        matvecs = tuple(sorted(matvecs, key=key))
 
         matvecs = self._scale_non_bcs(matvecs)
         scdiag = self._compute_scdiag(matvecs)
@@ -425,7 +427,7 @@ class MixedJacobian(BaseJacobian):
     """
     def __init__(self, target_exprs, arrays, time_mapper):
         super().__init__(arrays=arrays, target=None)
-        self.targets = tuple(target_exprs.keys())
+        self.targets = tuple(target_exprs)
         self.time_mapper = time_mapper
         self._submatrices = []
         self._build_blocks(target_exprs)
@@ -443,12 +445,12 @@ def n_submatrices(self):
         """Return the number of submatrix blocks."""
         return len(self._submatrices)
 
-    @property
+    @cached_property
     def nonzero_submatrices(self):
         """Return SubMatrixBlock objects that have non-empty matvecs."""
         return [m for m in self.submatrices if m.matvecs]
 
-    @property
+    @cached_property
     def target_scaler_mapper(self):
         """
         Map each row target to the scdiag of its corresponding
@@ -467,13 +469,8 @@ def _build_blocks(self, target_exprs):
         for i, row_target in enumerate(self.targets):
             exprs = target_exprs[row_target]
             for j, col_target in enumerate(self.targets):
-                matvecs = []
-                for expr in exprs:
-                    matvecs.extend(
-                        e for e in self._build_matvec_expr(
-                            expr, col_target=col_target, row_target=row_target
-                        )
-                    )
+
+                matvecs = self._build_submatrix_matvecs(exprs, row_target, col_target)
                 matvecs = [m for m in matvecs if m is not None]
 
                 # Sort to put EssentialBC first if any
@@ -497,6 +494,16 @@ def _build_blocks(self, target_exprs):
                 )
                 self._submatrices.append(block)
 
+    def _build_submatrix_matvecs(self, exprs, row_target, col_target):
+        matvecs = []
+        for expr in exprs:
+            matvecs.extend(
+                e for e in self._build_matvec_expr(
+                    expr, col_target=col_target, row_target=row_target
+                )
+            )
+        return matvecs
+
     def get_submatrix(self, row_idx, col_idx):
         """
         Return the SubMatrixBlock at (row_idx, col_idx), or None if not found.

devito/symbolics/extraction.py

@@ -79,10 +79,13 @@ def centre_stencil(expr, target, as_coeff=False):
     Extract the centre stencil from an expression. Its coefficient is what
     would appear on the diagonal of the matrix system if the matrix were
     formed explicitly.
+
     Parameters
     ----------
-    expr : the expression to extract the centre stencil from
-    target : the target function whose centre stencil we want
+    expr : expr-like
+        The expression to extract the centre stencil from
+    target : Function
+        The target function whose centre stencil we want
     as_coeff : bool, optional
         If True, return the coefficient of the centre stencil
     """

devito/types/object.py

@@ -252,10 +252,10 @@ class LocalCompositeObject(CompositeObject, LocalType):
     __rkwargs__ = ('modifier', 'liveness')
 
     def __init__(self, name, pname, fields, modifier=None, liveness='lazy'):
-        self.modifier = modifier
         dtype = CustomDtype(f"struct {pname}", modifier=modifier)
         Object.__init__(self, name, dtype, None)
         self._pname = pname
+        self.modifier = modifier
         assert liveness in ['eager', 'lazy']
         self._liveness = liveness
         self._fields = fields