misc: Add BaseJacobian

Author: ZoeLeibowitz

devito/petsc/types/types.py

@@ -137,6 +137,26 @@ def eval(cls, *args):
 
 
 class FieldData:
+    """
+    Metadata class passed to `LinearSolveExpr`. Encapsulates metadata for a single
+    `target` field needed to interface with PETSc SNES solvers.
+
+    Parameters
+    ----------
+
+    target : Function-like
+        The target field to solve into, which is a Function-like object.
+    jacobian : Jacobian
+        Defines the matrix-vector product for the linear system, where the vector is 
+        the PETScArray representing the `target`.
+    residual : Residual
+        Defines the nonlinear residual function F(target) = 0.
+    initialguess : InitialGuess
+        Defines the initial guess for the solution, which satisfies
+        essential boundary conditions.
+    arrays : dict
+        A dictionary mapping `target` to its corresponding PETScArrays.
+    """
     def __init__(self, target=None, jacobian=None, residual=None,
                  initialguess=None, arrays=None, **kwargs):
         self._target = target
@@ -190,6 +210,24 @@ def targets(self):
 
 
 class MultipleFieldData(FieldData):
+    """
+    Metadata class passed to `LinearSolveExpr`, for mixed-field problems,
+    where the solution vector spans multiple `targets`.
+
+    Parameters
+    ----------
+    targets : list of Function-like
+        The fields to solve into, each represented by a Function-like object.
+    jacobian : MixedJacobian
+        Defines the matrix-vector products for the full system Jacobian.
+    residual : MixedResidual
+        Defines the nonlinear residual function F(targets) = 0.
+    initialguess : InitialGuess
+        Defines the initial guess metadata, which satisfies
+        essential boundary conditions.
+    arrays : dict
+        A dictionary mapping the `targets` to their corresponding PETScArrays.
+    """
     def __init__(self, targets, arrays, jacobian=None, residual=None):
         self._targets = as_tuple(targets)
         self._arrays = arrays
@@ -208,6 +246,7 @@ def space_dimensions(self):
 
     @property
     def grid(self):
+        """The unique `Grid` associated with all targets."""
         grids = [t.grid for t in self.targets]
         if len(set(grids)) > 1:
             raise ValueError(
@@ -233,8 +272,70 @@ def space_order(self):
     def targets(self):
         return self._targets
 
+    
+class BaseJacobian:
+    def __init__(self, arrays, target=None):
+        self.arrays = arrays
+        self.target = target
+
+    def _scale_non_bcs(self, matvecs, target=None):
+        target = target or self.target
+        vol = target.grid.symbolic_volume_cell
+
+        return [
+            m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol)
+            for m in matvecs
+        ]
+
+    def _compute_scdiag(self, matvecs, col_target=None):
+        """
+        """
+        x = self.arrays[col_target or self.target]['x']
+
+        centres = {
+            centre_stencil(m.rhs, x, as_coeff=True)
+            for m in matvecs if not isinstance(m, EssentialBC)
+        }
+        return centres.pop() if len(centres) == 1 else 1.0
+
+    def _scale_bcs(self, matvecs, scdiag):
+        """
+        Scale the essential BCs
+        """
+        return [
+            m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m
+            for m in matvecs
+        ]
+    
+    def _build_matvec_expr(self, expr, **kwargs):
+        col_target = kwargs.get('col_target', self.target)
+        row_target = kwargs.get('row_target', self.target)
+
+        _, F_target, _, targets = separate_eqn(expr, col_target)
+        if F_target:
+            return self._make_matvec(
+                expr, F_target, targets, col_target, row_target
+            )
+        else:
+            return (None,)
+
+    def _make_matvec(self, expr, F_target, targets, col_target, row_target):
+        y = self.arrays[row_target]['y']
+        x = self.arrays[col_target]['x']
+
+        if isinstance(expr, EssentialBC):
+            # NOTE: Essential BCs are trivial equations in the solver.
+            # See `EssentialBC` for more details.
+            zero_row = ZeroRow(y, x, subdomain=expr.subdomain)
+            zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain)
+            return (zero_row, zero_column)
+        else:
+            rhs = F_target.subs(targets_to_arrays(x, targets))
+            rhs = rhs.subs(self.time_mapper)
+            return (Eq(y, rhs, subdomain=expr.subdomain),)
+
 
-class Jacobian:
+class Jacobian(BaseJacobian):
     """
     Represents a Jacobian matrix.
 
@@ -246,14 +347,14 @@ class Jacobian:
     require explicit symbolic differentiation.
     """
     def __init__(self, target, exprs, arrays, time_mapper):
-        self.target = target
+        super().__init__(arrays=arrays, target=target)
         self.exprs = exprs
-        self.arrays = arrays
         self.time_mapper = time_mapper
         self._build_matvecs()
 
     @property
     def matvecs(self):
+        # TODO: add shortcut explanation etc
         """
         Stores the expressions used to generate the `MatMult`
         callback generated at the IET level. This function is
@@ -279,7 +380,6 @@ 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)))
 
         matvecs = self._scale_non_bcs(matvecs)
@@ -289,83 +389,8 @@ def _build_matvecs(self):
         self._matvecs = matvecs
         self._scdiag = scdiag
 
-    def _build_matvec_expr(self, expr, col_target=None, row_target=None):
-        col_target = col_target or self.target
-        row_target = row_target or self.target
 
-        _, F_target, _, targets = separate_eqn(expr, col_target)
-        if F_target:
-            return self._make_matvec(
-                expr, F_target, targets, col_target, row_target
-            )
-        else:
-            return (None,)
-
-    def _make_matvec(self, expr, F_target, targets, col_target, row_target):
-        y = self.arrays[row_target]['y']
-        x = self.arrays[col_target]['x']
-
-        if isinstance(expr, EssentialBC):
-            # NOTE: Essential BCs are trivial equations in the solver.
-            # See `EssentialBC` for more details.
-            zero_row = ZeroRow(y, x, subdomain=expr.subdomain)
-            zero_column = ZeroColumn(x, 0., subdomain=expr.subdomain)
-            return (zero_row, zero_column)
-        else:
-            rhs = F_target.subs(targets_to_arrays(x, targets))
-            rhs = rhs.subs(self.time_mapper)
-            return (Eq(y, rhs, subdomain=expr.subdomain),)
-
-    def _scale_non_bcs(self, matvecs, target=None):
-        target = target or self.target
-        vol = target.grid.symbolic_volume_cell
-
-        return [
-            m if isinstance(m, EssentialBC) else m._rebuild(rhs=m.rhs * vol)
-            for m in matvecs
-        ]
-
-    def _compute_scdiag(self, matvecs, col_target=None):
-        """
-        """
-        x = self.arrays[col_target or self.target]['x']
-
-        centres = {
-            centre_stencil(m.rhs, x, as_coeff=True)
-            for m in matvecs if not isinstance(m, EssentialBC)
-        }
-        return centres.pop() if len(centres) == 1 else 1.0
-
-    def _scale_bcs(self, matvecs, scdiag):
-        """
-        Scale the essential BCs
-        """
-        return [
-            m._rebuild(rhs=m.rhs * scdiag) if isinstance(m, ZeroRow) else m
-            for m in matvecs
-        ]
-
-
-class SubMatrixBlock:
-    def __init__(self, name, matvecs, scdiag, row_target,
-                 col_target, row_idx, col_idx, linear_idx):
-        self.name = name
-        self.matvecs = matvecs
-        self.scdiag = scdiag
-        self.row_target = row_target
-        self.col_target = col_target
-        self.row_idx = row_idx
-        self.col_idx = col_idx
-        self.linear_idx = linear_idx
-
-    def is_diag(self):
-        return self.row_idx == self.col_idx
-
-    def __repr__(self):
-        return (f"<SubMatrixBlock {self.name}>")
-
-
-class MixedJacobian(Jacobian):
+class MixedJacobian(BaseJacobian):
     """
     Represents a Jacobian for a linear system with a solution vector
     composed of multiple fields (targets).
@@ -380,12 +405,12 @@ class MixedJacobian(Jacobian):
 
     # TODO: pcfieldsplit support for each block
     """
-    def __init__(self, target_eqns, arrays, time_mapper):
-        self.targets = tuple(target_eqns.keys())
-        self.arrays = arrays
+    def __init__(self, target_exprs, arrays, time_mapper):
+        super().__init__(arrays=arrays, target=None)
+        self.targets = tuple(target_exprs.keys())
         self.time_mapper = time_mapper
         self._submatrices = []
-        self._build_blocks(target_eqns)
+        self._build_blocks(target_exprs)
 
     @property
     def submatrices(self):
@@ -427,7 +452,9 @@ def _build_blocks(self, target_exprs):
                 matvecs = []
                 for expr in exprs:
                     matvecs.extend(
-                        e for e in self._build_matvec_expr(expr, col_target, row_target)
+                        e for e in self._build_matvec_expr(
+                            expr, col_target=col_target, row_target=row_target
+                        )
                     )
                 matvecs = [m for m in matvecs if m is not None]
 
@@ -469,6 +496,25 @@ def __repr__(self):
         return f"<MixedJacobian with {self.n_submatrices} submatrices: [{summary}]>"
 
 
+class SubMatrixBlock:
+    def __init__(self, name, matvecs, scdiag, row_target,
+                 col_target, row_idx, col_idx, linear_idx):
+        self.name = name
+        self.matvecs = matvecs
+        self.scdiag = scdiag
+        self.row_target = row_target
+        self.col_target = col_target
+        self.row_idx = row_idx
+        self.col_idx = col_idx
+        self.linear_idx = linear_idx
+
+    def is_diag(self):
+        return self.row_idx == self.col_idx
+
+    def __repr__(self):
+        return (f"<SubMatrixBlock {self.name}>")
+
+
 class Residual:
     """
     Gennerates the metadata needed to define the nonlinear residual function