misc: Address more comments and add docstrings

Author: ZoeLeibowitz

devito/petsc/iet/routines.py

@@ -258,13 +258,12 @@ def _create_matvec_body(self, body, jacobian):
         )
 
         # Dereference function data in struct
-        dereference_funcs = [Dereference(i, ctx) for i in
-                             fields if isinstance(i.function, AbstractFunction)]
+        derefs = self.dereference_funcs(ctx, fields)
 
         matvec_body = CallableBody(
             List(body=body),
             init=(objs['begin_user'],),
-            stacks=stacks+tuple(dereference_funcs),
+            stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
@@ -398,13 +397,12 @@ def _create_formfunc_body(self, body):
         )
 
         # Dereference function data in struct
-        dereference_funcs = [Dereference(i, ctx) for i in
-                             fields if isinstance(i.function, AbstractFunction)]
+        derefs = self.dereference_funcs(ctx, fields)
 
         formfunc_body = CallableBody(
             List(body=body),
             init=(objs['begin_user'],),
-            stacks=stacks+tuple(dereference_funcs),
+            stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),))
 
         # Replace non-function data with pointer to data in struct
@@ -509,15 +507,12 @@ def _create_form_rhs_body(self, body):
         )
 
         # Dereference function data in struct
-        dereference_funcs = tuple(
-            [Dereference(i, ctx) for i in
-             fields if isinstance(i.function, AbstractFunction)]
-        )
+        derefs = self.dereference_funcs(ctx, fields)
 
         formrhs_body = CallableBody(
             List(body=[body]),
             init=(objs['begin_user'],),
-            stacks=stacks+dereference_funcs,
+            stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
@@ -590,13 +585,12 @@ def _create_initial_guess_body(self, body):
         )
 
         # Dereference function data in struct
-        dereference_funcs = [Dereference(i, ctx) for i in
-                             fields if isinstance(i.function, AbstractFunction)]
+        derefs = self.dereference_funcs(ctx, fields)
 
         body = CallableBody(
             List(body=[body]),
             init=(objs['begin_user'],),
-            stacks=stacks+tuple(dereference_funcs),
+            stacks=stacks+derefs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
@@ -662,6 +656,12 @@ def zero_vector(self, vec):
         """
         return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None
 
+    def dereference_funcs(self, struct, fields):
+        return tuple(
+            [Dereference(i, struct) for i in
+             fields if isinstance(i.function, AbstractFunction)]
+        )
+
 
 class CCBBuilder(CBBuilder):
     def __init__(self, **kwargs):
@@ -894,16 +894,15 @@ def _whole_formfunc_body(self, body):
         )
 
         # Dereference function data in struct
-        dereference_funcs = [Dereference(i, ctx) for i in
-                             fields if isinstance(i.function, AbstractFunction)]
+        derefs = self.dereference_funcs(ctx, fields)
 
         f_soa = PointerCast(fbundle)
         x_soa = PointerCast(xbundle)
 
         formfunc_body = CallableBody(
             List(body=body),
             init=(objs['begin_user'],),
-            stacks=stacks+tuple(dereference_funcs),
+            stacks=stacks+derefs,
             casts=(f_soa, x_soa),
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )

devito/petsc/types/types.py

@@ -241,7 +241,7 @@ class Jacobian:
     This Jacobian is defined implicitly via matrix-vector products
     derived from the symbolic equations provided in `matvecs`.
 
-    It assumes the problem is linear, meaning the Jacobian
+    The class assumes the problem is linear, meaning the Jacobian
     corresponds to a constant coefficient matrix and does not
     require explicit symbolic differentiation.
     """
@@ -274,6 +274,7 @@ def _build_matvecs(self):
             matvecs.extend(
                 e for e in self._build_matvec_eq(eq) if e is not None
             )
+
         matvecs = tuple(sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC)))
 
         matvecs = self._scale_non_bcs(matvecs)
@@ -469,6 +470,20 @@ def __repr__(self):
 
 class Residual:
     """
+    Gennerates the metadata needed to define the nonlinear residual function
+    F(target) = 0 for use with PETSc's SNES interface.
+
+    PETSc's SNES interface includes methods for solving nonlinear systems of
+    equations using Newton-type methods. For linear problems, `SNESKSPONLY`
+    is used to perform a single Newton iteration, unifying the
+    interface for both linear and nonlinear problems.
+
+    This class encapsulates the symbolic equations used to construct the
+    residual function F(target) = F_(target) - b, where b contains all
+    terms independent of the solution `target`.
+
+    References:
+        - https://petsc.org/main/manual/snes/
     """
     def __init__(self, target, eqns, arrays, time_mapper, scdiag):
         self.target = target
@@ -481,12 +496,18 @@ def __init__(self, target, eqns, arrays, time_mapper, scdiag):
     @property
     def formfuncs(self):
         """
+        Stores the equations used to build the `FormFunction`
+        callback generated at the IET level. This function is
+        passed to PETSc via `SNESSetFunction(..., FormFunction, ...)`.
         """
         return self._formfuncs
 
     @property
     def formrhs(self):
         """
+        Stores the equations used to generate the RHS
+        vector `b` through the `FormRHS` callback generated at the IET level.
+        The SNES solver is then called via `SNESSolve(..., b, target)`.
         """
         return self._formrhs
 
@@ -544,7 +565,7 @@ class MixedResidual(Residual):
     """
     """
     def __init__(self, target_eqns, arrays, time_mapper, scdiag):
-        self.targets = as_tuple(target_eqns.keys())
+        self.targets = tuple(target_eqns.keys())
         self.arrays = arrays
         self.time_mapper = time_mapper
         self.scdiag = scdiag
@@ -592,14 +613,15 @@ def _build_function_eq(self, eq, target):
                 self.arrays[target]['x'], eq.rhs, subdomain=eq.subdomain
             )
             return (zero_row, zero_col)
+
         else:
             if isinstance(zeroed, (int, float)):
                 rhs = zeroed * volume
             else:
                 rhs = zeroed.subs(mapper)
                 rhs = rhs.subs(self.time_mapper)*volume
 
-        return as_tuple(Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain))
+        return (Eq(self.arrays[target]['f'], rhs, subdomain=eq.subdomain),)
 
 
 class InitialGuess:

devito/symbolics/extraction.py

@@ -20,9 +20,10 @@ def separate_eqn(eqn, target):
 
     from devito.operations.solve import eval_time_derivatives
     zeroed_eqn = eval_time_derivatives(zeroed_eqn.lhs)
-    target_funcs = set(generate_targets(zeroed_eqn, target))
 
+    target_funcs = set(generate_targets(zeroed_eqn, target))
     b, F_target = remove_targets(zeroed_eqn, target_funcs)
+
     return -b, F_target, zeroed_eqn, target_funcs