tests: Start tests for petscsection constraining bcs in 1d

Author: ZoeLeibowitz

devito/petsc/equations.py

@@ -35,6 +35,8 @@ def constrain_essential_bcs(expressions, **kwargs):
 
             # obvs shouldn't be obtained from indexeds[0], but how should it be obtained?
             # USE e.lhs function -> the one that the BC is being applied to
+            # from IPython import embed; embed()
+            # f._size_nodomain.left
             halo_size_left = indexeds[0].function._size_halo[d].left
             halo_size_right = indexeds[0].function._size_halo[d].right
 
@@ -45,11 +47,11 @@ def constrain_essential_bcs(expressions, **kwargs):
 
             # in theory this class shoulod just take in d
             # TODO: use unique name
-            subdim_max = SubDimMax(d.name + '_max', subdim=d, thickness=d.thickness)
-            subdim_min = SubDimMin(d.name + '_min', subdim=d, thickness=d.thickness)
+            sregistry = kwargs.get('sregistry')
+            subdim_max = SubDimMax(sregistry.make_name(prefix=d.name + '_max'), subdim=d, thickness=d.thickness)
+            subdim_min = SubDimMin(sregistry.make_name(prefix=d.name + '_min'), subdim=d, thickness=d.thickness)
 
             # from IPython import embed; embed()
-
             new_dim = CustomBoundSubDimension(
             name=d.name,
             parent=d.parent,

devito/petsc/iet/callbacks.py

@@ -126,7 +126,7 @@ def _make_core(self):
         if self.field_data.initial_guess.exprs:
             self._make_initial_guess()
         # Make the callback used to constrain boundary nodes
-        if self.field_data.constrain_bc.exprs:
+        if self.field_data.constrain_bc:
             self._make_constrain_bc()
         self._make_user_struct_callback()
 
@@ -658,7 +658,6 @@ def _make_constrain_bc(self):
             exprs, options={'mpi': False}, sregistry=self.sregistry,
             concretize_mapper=self.concretize_mapper
         )
-        # from IPython import embed; embed()
         body = self._create_constrain_bc_body(
             List(body=irs.uiet.body)
         )
@@ -669,7 +668,60 @@ def _make_constrain_bc(self):
         self._efuncs[cb.name] = cb
 
     def _create_constrain_bc_body(self, body):
-        return body
+        linsolve_expr = self.inject_solve.expr.rhs
+        objs = self.objs
+        sobjs = self.solver_objs
+        target = self.target
+
+        dmda = sobjs['callbackdm']
+        ctx = objs['dummyctx']
+
+        x_arr = self.field_data.arrays[target]['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.time_dependence.uxreplace_time(body)
+
+        fields = get_user_struct_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_local_info
+        )
+
+        # Dereference function data in struct
+        derefs = dereference_funcs(ctx, fields)
+
+        # Force the struct definition to appear at the very start, since
+        # stacks, allocs etc may rely on its information
+        struct_definition = [Definition(ctx), dm_get_app_context]
+
+        body = self._make_callable_body(
+            body, standalones=struct_definition, stacks=stacks+derefs
+        )
+
+        # 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):
         """

devito/petsc/types/metadata.py

@@ -210,11 +210,12 @@ class MultipleFieldData(FieldData):
     arrays : dict
         A dictionary mapping the `targets` to their corresponding PETScArrays.
     """
-    def __init__(self, targets, arrays, jacobian=None, residual=None):
+    def __init__(self, targets, arrays, jacobian=None, residual=None, constrain_bc=None):
         self._targets = as_tuple(targets)
         self._arrays = arrays
         self._jacobian = jacobian
         self._residual = residual
+        self._constrain_bc = constrain_bc
 
     @cached_property
     def space_dimensions(self):