Compiler: Add basic initguess callback (#72)

* Compiler: Add basic initguess callback - to be improved

Author: ZoeLeibowitz

devito/petsc/iet/routines.py

@@ -47,6 +47,7 @@ def __init__(self, **kwargs):
         self._matvecs = []
         self._formfuncs = []
         self._formrhs = []
+        self._initialguesses = []
 
         self._make_core()
         self._efuncs = self._uxreplace_efuncs()
@@ -88,6 +89,10 @@ def formfuncs(self):
     def formrhs(self):
         return self._formrhs
 
+    @property
+    def initialguesses(self):
+        return self._initialguesses
+
     @property
     def user_struct_callback(self):
         return self._user_struct_callback
@@ -97,6 +102,8 @@ def _make_core(self):
         self._make_matvec(fielddata, fielddata.matvecs)
         self._make_formfunc(fielddata)
         self._make_formrhs(fielddata)
+        if fielddata.initialguess:
+            self._make_initialguess(fielddata)
         self._make_user_struct_callback()
 
     def _make_matvec(self, fielddata, matvecs, prefix='MatMult'):
@@ -483,6 +490,84 @@ def _create_form_rhs_body(self, body, fielddata):
 
         return Uxreplace(subs).visit(formrhs_body)
 
+    def _make_initialguess(self, fielddata):
+        initguess = fielddata.initialguess
+        sobjs = self.solver_objs
+
+        # Compile initital guess `eqns` into an IET via recursive compilation
+        irs, _ = self.rcompile(
+            initguess, options={'mpi': False}, sregistry=self.sregistry,
+            concretize_mapper=self.concretize_mapper
+        )
+        body_init_guess = self._create_initial_guess_body(
+            List(body=irs.uiet.body), fielddata
+        )
+        objs = self.objs
+        cb = PETScCallable(
+            self.sregistry.make_name(prefix='FormInitialGuess'),
+            body_init_guess,
+            retval=objs['err'],
+            parameters=(sobjs['callbackdm'], objs['xloc'])
+        )
+        self._initialguesses.append(cb)
+        self._efuncs[cb.name] = cb
+
+    def _create_initial_guess_body(self, body, fielddata):
+        linsolve_expr = self.injectsolve.expr.rhs
+        objs = self.objs
+        sobjs = self.solver_objs
+
+        dmda = sobjs['callbackdm']
+        ctx = objs['dummyctx']
+
+        x_arr = fielddata.arrays['x']
+
+        vec_get_array = petsc_call(
+            'VecGetArray', [objs['xloc'], Byref(x_arr._C_symbol)]
+        )
+
+        dm_get_local_info = petsc_call(
+            'DMDAGetLocalInfo', [dmda, Byref(linsolve_expr.localinfo)]
+        )
+
+        body = self.timedep.uxreplace_time(body)
+
+        fields = self._dummy_fields(body)
+        self._struct_params.extend(fields)
+
+        dm_get_app_context = petsc_call(
+            'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
+        )
+
+        vec_restore_array = petsc_call(
+            'VecRestoreArray', [objs['xloc'], Byref(x_arr._C_symbol)]
+        )
+
+        body = body._rebuild(body=body.body + (vec_restore_array,))
+
+        stacks = (
+            vec_get_array,
+            dm_get_app_context,
+            dm_get_local_info
+        )
+
+        # Dereference function data in struct
+        dereference_funcs = [Dereference(i, ctx) for i in
+                             fields if isinstance(i.function, AbstractFunction)]
+
+        body = CallableBody(
+            List(body=[body]),
+            init=(objs['begin_user'],),
+            stacks=stacks+tuple(dereference_funcs),
+            retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
+        )
+
+        # Replace non-function data with pointer to data in struct
+        subs = {i._C_symbol: FieldFromPointer(i._C_symbol, ctx) for
+                i in fields if not isinstance(i.function, AbstractFunction)}
+
+        return Uxreplace(subs).visit(body)
+
     def _make_user_struct_callback(self):
         """
         This is the struct initialised inside the main kernel and
@@ -1250,6 +1335,12 @@ def _execute_solve(self):
 
         vec_replace_array = self.timedep.replace_array(target)
 
+        if self.cbbuilder.initialguesses:
+            initguess = self.cbbuilder.initialguesses[0]
+            initguess_call = petsc_call(initguess.name, [dmda, sobjs['xlocal']])
+        else:
+            initguess_call = None
+
         dm_local_to_global_x = petsc_call(
             'DMLocalToGlobal', [dmda, sobjs['xlocal'], insert_vals,
                                 sobjs['xglobal']]
@@ -1268,6 +1359,7 @@ def _execute_solve(self):
         run_solver_calls = (struct_assignment,) + (
             rhs_call,
         ) + vec_replace_array + (
+            initguess_call,
             dm_local_to_global_x,
             snes_solve,
             dm_global_to_local_x,

devito/petsc/solve.py

@@ -57,11 +57,18 @@ def generate_field_data(self, eqns, target, arrays):
         )
         matvecs = [self.build_matvec_eqns(eq, target, arrays) for eq in eqns]
 
+        initialguess = [
+            eq for eq in
+            (self.make_initial_guess(e, target, arrays) for e in eqns)
+            if eq is not None
+        ]
+
         return FieldData(
             target=target,
             matvecs=matvecs,
             formfuncs=formfuncs,
             formrhs=formrhs,
+            initialguess=initialguess,
             arrays=arrays
         )
 
@@ -99,6 +106,22 @@ def make_rhs(self, eq, b, arrays):
         rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper)
         return Eq(arrays['b'], rhs, subdomain=eq.subdomain)
 
+    def make_initial_guess(self, eq, target, arrays):
+        """
+        Enforce initial guess to satisfy essential BCs.
+        # TODO: For time-stepping, only enforce these once outside the time loop
+        and use the previous time-step solution as the initial guess for next time step.
+        # TODO: Extend this to "coupled".
+        """
+        if isinstance(eq, EssentialBC):
+            assert eq.lhs == target
+            return Eq(
+                arrays['x'], eq.rhs,
+                subdomain=eq.subdomain
+            )
+        else:
+            return None
+
     def generate_arrays(self, target):
         return {
             p: PETScArray(name=f'{p}_{target.name}',

devito/petsc/types/types.py

@@ -132,7 +132,7 @@ def eval(cls, *args):
 
 class FieldData:
     def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None,
-                 arrays=None, **kwargs):
+                 initialguess=None, arrays=None, **kwargs):
         self._target = kwargs.get('target', target)
 
         petsc_precision = dtype_mapper[petsc_variables['PETSC_PRECISION']]
@@ -145,6 +145,7 @@ def __init__(self, target=None, matvecs=None, formfuncs=None, formrhs=None,
         self._matvecs = matvecs
         self._formfuncs = formfuncs
         self._formrhs = formrhs
+        self._initialguess = initialguess
         self._arrays = arrays
 
     @property
@@ -163,6 +164,10 @@ def formfuncs(self):
     def formrhs(self):
         return self._formrhs
 
+    @property
+    def initialguess(self):
+        return self._initialguess
+
     @property
     def arrays(self):
         return self._arrays

examples/seismic/tutorials/13_LSRTM_acoustic.ipynb

@@ -304,7 +304,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -329,12 +329,13 @@
     "    dm_true =  (solver.model.vp.data**(-2) - model0.vp.data**(-2))\n",
     "    \n",
     "    objective = 0.\n",
+    "    u0 = None\n",
     "    for i in range(nshots):\n",
     "        \n",
     "        #Observed Data using Born's operator\n",
     "        geometry.src_positions[0, :] = source_locations[i, :]\n",
     "\n",
-    "        _, u0, _ = solver.forward(vp=model0.vp, save=True)\n",
+    "        _, u0, _ = solver.forward(vp=model0.vp, save=True, u=u0)\n",
     "        \n",
     "        _, _, _,_ = solver.jacobian(dm_true, vp=model0.vp, rec = d_obs)\n",
     "        \n",

tests/test_petsc.py

@@ -4,15 +4,16 @@
 
 from conftest import skipif
 from devito import (Grid, Function, TimeFunction, Eq, Operator,
-                    configuration, norm, switchconfig)
+                    configuration, norm, switchconfig, SubDomain)
 from devito.ir.iet import (Call, ElementalFunction,
                            FindNodes, retrieve_iteration_tree)
 from devito.types import Constant, LocalCompositeObject
 from devito.passes.iet.languages.C import CDataManager
 from devito.petsc.types import (DM, Mat, Vec, PetscMPIInt, KSP,
                                 PC, KSPConvergedReason, PETScArray,
                                 LinearSolveExpr, FieldData, MultipleFieldData)
-from devito.petsc.solve import PETScSolve, separate_eqn, centre_stencil
+from devito.petsc.solve import (PETScSolve, separate_eqn, centre_stencil,
+                                EssentialBC)
 from devito.petsc.iet.nodes import Expression
 from devito.petsc.initialize import PetscInitialize
 
@@ -786,6 +787,74 @@ def test_solve_output():
     assert np.allclose(u.data, v.data)
 
 
+@skipif('petsc')
+def test_essential_bcs():
+    """
+    Verify that PETScSolve returns the correct output with
+    essential boundary conditions.
+    """
+    class SubTop(SubDomain):
+        name = 'subtop'
+
+        def define(self, dimensions):
+            x, y = dimensions
+            return {x: x, y: ('right', 1)}
+    sub1 = SubTop()
+
+    class SubBottom(SubDomain):
+        name = 'subbottom'
+
+        def define(self, dimensions):
+            x, y = dimensions
+            return {x: x, y: ('left', 1)}
+    sub2 = SubBottom()
+
+    class SubLeft(SubDomain):
+        name = 'subleft'
+
+        def define(self, dimensions):
+            x, y = dimensions
+            return {x: ('left', 1), y: y}
+    sub3 = SubLeft()
+
+    class SubRight(SubDomain):
+        name = 'subright'
+
+        def define(self, dimensions):
+            x, y = dimensions
+            return {x: ('right', 1), y: y}
+    sub4 = SubRight()
+
+    subdomains = (sub1, sub2, sub3, sub4)
+    grid = Grid(shape=(11, 11), subdomains=subdomains, dtype=np.float64)
+
+    u = Function(name='u', grid=grid, space_order=2)
+    v = Function(name='v', grid=grid, space_order=2)
+
+    # Solving Ax=b where A is the identity matrix
+    v.data[:] = 5.0
+    eqn = Eq(u, v)
+
+    bcs = [EssentialBC(u, 1., subdomain=sub1)]  # top
+    bcs += [EssentialBC(u, 2., subdomain=sub2)]  # bottom
+    bcs += [EssentialBC(u, 3., subdomain=sub3)]  # left
+    bcs += [EssentialBC(u, 4., subdomain=sub4)]  # right
+
+    petsc = PETScSolve([eqn]+bcs, target=u)
+
+    with switchconfig(language='petsc'):
+        op = Operator(petsc)
+        op.apply()
+
+    # Check u is equal to v on the interior
+    assert np.allclose(u.data[1:-1, 1:-1], v.data[1:-1, 1:-1])
+    # Check u satisfies the boundary conditions
+    assert np.allclose(u.data[1:-1, -1], 1.0)  # top
+    assert np.allclose(u.data[1:-1, 0], 2.0)  # bottom
+    assert np.allclose(u.data[0, 1:-1], 3.0)  # left
+    assert np.allclose(u.data[-1, 1:-1], 4.0)  # right
+
+
 class TestCoupledLinear:
     # The coupled interface can be used even for uncoupled problems, meaning
     # the equations will be solved within a single matrix system.