tests: Add more petsc tests for linear coupled

Author: ZoeLeibowitz

.github/workflows/examples-mpi.yml

@@ -17,11 +17,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   build:

.github/workflows/examples.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   tutorials:

.github/workflows/flake8.yml

@@ -10,11 +10,11 @@ on:
    push:
      branches:
        - main
-       - master
+       - petsc
    pull_request:
      branches:
        - main
-       - master
+       - petsc
 
 jobs:
   flake8:

.github/workflows/pytest-core-mpi.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   test-mpi-basic:

.github/workflows/pytest-core-nompi.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   pytest:

.github/workflows/pytest-petsc.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   pytest:

.github/workflows/tutorials.yml

@@ -10,11 +10,11 @@ on:
   push:
     branches:
       - main
-      - master
+      - petsc
   pull_request:
     branches:
       - main
-      - master
+      - petsc
 
 jobs:
   tutorials:

devito/passes/iet/engine.py

@@ -1,6 +1,7 @@
 from collections import OrderedDict, defaultdict
 from functools import partial, singledispatch, wraps
 
+from devito.finite_differences.differentiable import Differentiable
 from devito.ir.iet import (Call, ExprStmt, Iteration, SyncSpot, AsyncCallable,
                            FindNodes, FindSymbols, MapNodes, MetaCall, Transformer,
                            EntryFunction, FixedArgsCallable, Uxreplace,
@@ -10,13 +11,13 @@
 from devito.passes import needs_transfer
 from devito.symbolics import FieldFromComposite, FieldFromPointer
 from devito.tools import DAG, as_tuple, filter_ordered, sorted_priority, timed_pass
-from devito.types import (Bundle, CompositeObject, Lock, IncrDimension,
+from devito.types import (Array, Bundle, CompositeObject, Lock, IncrDimension,
                           ModuloDimension, Indirection, Pointer, SharedData,
                           ThreadArray, Temp, NPThreads, NThreadsBase, Wildcard)
-from devito.types.array import ArrayBasic
 from devito.types.args import ArgProvider
 from devito.types.dense import DiscreteFunction
 from devito.types.dimension import AbstractIncrDimension, BlockDimension
+from devito.types.array import ArrayBasic
 
 __all__ = ['Graph', 'iet_pass', 'iet_visit']
 
@@ -414,7 +415,8 @@ def abstract_objects(objects0, sregistry=None):
 
     # Precedence rules make it possible to reconstruct objects that depend on
     # higher priority objects
-    keys = [Bundle, ArrayBasic, DiscreteFunction, AbstractIncrDimension, BlockDimension]
+    keys = [Bundle, Array, Differentiable, DiscreteFunction,
+            AbstractIncrDimension, BlockDimension]
     priority = {k: i for i, k in enumerate(keys, start=1)}
     objects = sorted_priority(objects, priority)
 
@@ -449,6 +451,7 @@ def _(i, mapper, sregistry):
     })
 
 
+@abstract_object.register(Array)
 @abstract_object.register(ArrayBasic)
 def _(i, mapper, sregistry):
     if isinstance(i, Lock):

devito/petsc/iet/routines.py

@@ -688,15 +688,11 @@ def zero_memory(self):
         return False
 
     def _make_core(self):
-        injectsolve = self.injectsolve
-        all_fielddata = injectsolve.expr.rhs.fielddata
-
-        for sm in all_fielddata.jacobian.nonzero_submatrices:
-            # arrays = all_fielddata.arrays[sm.row_target]
+        for sm in self.fielddata.jacobian.nonzero_submatrices:
             self._make_matvec(sm, prefix=f'{sm.name}_MatMult')
 
         self._make_whole_matvec()
-        self._make_whole_formfunc(all_fielddata)
+        self._make_whole_formfunc()
         self._make_user_struct_callback()
         self._create_submatrices()
         self._efuncs['PopulateMatContext'] = self.objs['dummyefunc']
@@ -764,15 +760,15 @@ def _whole_matvec_body(self):
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
-    def _make_whole_formfunc(self, fielddata):
-        formfuncs = fielddata.residual.formfuncs
+    def _make_whole_formfunc(self):
+        formfuncs = self.fielddata.residual.formfuncs
         # Compile formfunc `eqns` into an IET via recursive compilation
         irs_formfunc, _ = self.rcompile(
             formfuncs, options={'mpi': False}, sregistry=self.sregistry,
             concretize_mapper=self.concretize_mapper
         )
-        body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body),
-                                                  fielddata)
+        body_formfunc = self._whole_formfunc_body(List(body=irs_formfunc.uiet.body))
+
         objs = self.objs
         cb = PETScCallable(
             self.sregistry.make_name(prefix='WholeFormFunc'),
@@ -783,7 +779,7 @@ def _make_whole_formfunc(self, fielddata):
         self._main_formfunc_callback = cb
         self._efuncs[cb.name] = cb
 
-    def _whole_formfunc_body(self, body, fielddata):
+    def _whole_formfunc_body(self, body):
         linsolve_expr = self.injectsolve.expr.rhs
         objs = self.objs
         sobjs = self.solver_objs
@@ -796,7 +792,7 @@ def _whole_formfunc_body(self, body, fielddata):
         fields = self._dummy_fields(body)
         self._struct_params.extend(fields)
 
-        # Process body for residual callback, including generating bundles etc
+        # Process body with bundles for residual callback
         bundles = sobjs['bundles']
         fbundle = bundles['f']
         xbundle = bundles['x']
@@ -1143,7 +1139,6 @@ def _extend_build(self, base_dict):
 
 class CoupledObjectBuilder(BaseObjectBuilder):
     def _extend_build(self, base_dict):
-        injectsolve = self.injectsolve
         sreg = self.sregistry
         objs = self.objs
         targets = self.fielddata.targets
@@ -1164,8 +1159,7 @@ def _extend_build(self, base_dict):
             dim_labels[i]: PetscInt(dim_labels[i]) for i in range(space_dims)
         })
 
-        jacobian = injectsolve.expr.rhs.fielddata.jacobian
-        submatrices = jacobian.nonzero_submatrices
+        submatrices = self.fielddata.jacobian.nonzero_submatrices
 
         base_dict['jacctx'] = JacobianStruct(
             name=sreg.make_name(prefix=objs['ljacctx'].name),

devito/petsc/solve.py

@@ -27,12 +27,12 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None):
     target_eqns : Eq or list of Eq, or dict of Function-like -> Eq or list of Eq
         The targets and symbolic equations defining the system to be solved.
 
-        - **Single-field problem**:
+        - Single-field problem:
             Pass a single Eq or list of Eq, and specify `target` separately:
                 PETScSolve(Eq1, target)
                 PETScSolve([Eq1, Eq2], target)
 
-        - **Multi-field (mixed) problem**:
+        - Multi-field (mixed) problem:
             Pass a dictionary mapping each target field to its Eq(s):
                 PETScSolve({u: Eq1, v: Eq2})
                 PETScSolve({u: [Eq1, Eq2], v: [Eq3, Eq4]})
@@ -47,8 +47,8 @@ def PETScSolve(target_eqns, target=None, solver_parameters=None):
     Returns
     -------
     Eq
-        A symbolic equation that wraps the system, solver metadata,
-        and boundary conditions. This can be passed directly to a Devito Operator.
+        A symbolic equation that wraps the linear solver.
+        This can be passed directly to a Devito Operator.
     """
     if target is not None:
         return InjectSolve(solver_parameters, {target: target_eqns}).build_eq()
@@ -143,9 +143,6 @@ def linear_solve_args(self):
 
         return coupled_targets[0], tuple(funcs), all_data
 
-    # def generate_arrays_combined(self, *targets):
-    #     return {target: self.generate_arrays(target) for target in targets}
-
 
 def generate_time_mapper(funcs):
     """