misc: Address comments

Author: ZoeLeibowitz

devito/passes/iet/engine.py

@@ -530,7 +530,6 @@ 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

@@ -114,6 +114,10 @@ def fielddata(self):
     def arrays(self):
         return self.fielddata.arrays
 
+    @property
+    def target(self):
+        return self.fielddata.target
+
     def _make_core(self):
         self._make_matvec(self.fielddata.jacobian)
         self._make_formfunc()
@@ -163,9 +167,7 @@ def _create_matvec_body(self, body, jacobian):
             'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
         )
 
-        zero_y_memory = petsc_call(
-            'VecSet', [objs['Y'], 0.0]
-        ) if self.zero_memory else None
+        zero_y_memory = self.zero_vector(objs['Y'])
 
         dm_get_local_xvec = petsc_call(
             'DMGetLocalVector', [dmda, Byref(xlocal)]
@@ -297,7 +299,8 @@ def _create_formfunc_body(self, body):
         linsolve_expr = self.injectsolve.expr.rhs
         objs = self.objs
         sobjs = self.solver_objs
-        target = self.fielddata.target
+        arrays = self.arrays
+        target = self.target
 
         dmda = sobjs['callbackdm']
         ctx = objs['dummyctx']
@@ -307,18 +310,16 @@ def _create_formfunc_body(self, body):
         fields = self._dummy_fields(body)
         self._struct_params.extend(fields)
 
-        f_formfunc = self.fielddata.arrays[target]['f']
-        x_formfunc = self.fielddata.arrays[target]['x']
+        f_formfunc = arrays[target]['f']
+        x_formfunc = arrays[target]['x']
 
         dm_cast = DummyExpr(dmda, DMCast(objs['dummyptr']), init=True)
 
         dm_get_app_context = petsc_call(
             'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
         )
 
-        zero_f_memory = petsc_call(
-            'VecSet', [objs['F'], 0.0]
-        ) if self.zero_memory else None
+        zero_f_memory = self.zero_vector(objs['F'])
 
         dm_get_local_xvec = petsc_call(
             'DMGetLocalVector', [dmda, Byref(objs['xloc'])]
@@ -437,7 +438,7 @@ def _create_form_rhs_body(self, body):
         linsolve_expr = self.injectsolve.expr.rhs
         objs = self.objs
         sobjs = self.solver_objs
-        target = self.fielddata.target
+        target = self.target
 
         dmda = sobjs['callbackdm']
         ctx = objs['dummyctx']
@@ -508,13 +509,15 @@ def _create_form_rhs_body(self, body):
         )
 
         # Dereference function data in struct
-        dereference_funcs = [Dereference(i, ctx) for i in
-                             fields if isinstance(i.function, AbstractFunction)]
+        dereference_funcs = tuple(
+            [Dereference(i, ctx) for i in
+             fields if isinstance(i.function, AbstractFunction)]
+        )
 
         formrhs_body = CallableBody(
             List(body=[body]),
             init=(objs['begin_user'],),
-            stacks=stacks+tuple(dereference_funcs),
+            stacks=stacks+dereference_funcs,
             retstmt=(Call('PetscFunctionReturn', arguments=[0]),)
         )
 
@@ -550,7 +553,7 @@ def _create_initial_guess_body(self, body):
         linsolve_expr = self.injectsolve.expr.rhs
         objs = self.objs
         sobjs = self.solver_objs
-        target = self.fielddata.target
+        target = self.target
 
         dmda = sobjs['callbackdm']
         ctx = objs['dummyctx']
@@ -653,6 +656,12 @@ def _uxreplace_efuncs(self):
             mapper.update({k: visitor.visit(v)})
         return mapper
 
+    def zero_vector(self, vec):
+        """
+        Zeros the memory of the output vector before computation
+        """
+        return petsc_call('VecSet', [vec, 0.0]) if self.zero_memory else None
+
 
 class CCBBuilder(CBBuilder):
     def __init__(self, **kwargs):
@@ -1059,7 +1068,7 @@ def residual_bundle(self, body, bundles):
             if i.base in mapper:
                 bundle = mapper[i.base]
                 index = bundles['target_indices'][i.function.target]
-                index = (index,)+i.indices
+                index = (index,) + i.indices
                 subs[i] = bundle.__getitem__(index)
 
         body = Uxreplace(subs).visit(body)
@@ -1178,9 +1187,9 @@ def _extend_build(self, base_dict):
             base_dict[f'{name}F'] = CallbackVec(f'{name}F')
 
         # Bundle objects/metadata required by the coupled residual callback
-        f_components = []
-        x_components = []
+        f_components, x_components = [], []
         bundle_mapper = {}
+        pname = sreg.make_name(prefix='Field')
 
         target_indices = {t: i for i, t in enumerate(targets)}
 
@@ -1190,18 +1199,15 @@ def _extend_build(self, base_dict):
             f_components.append(f_arr)
             x_components.append(x_arr)
 
-        bundle_pname = sreg.make_name(prefix='Field')
         fbundle = PetscBundle(
-            name='f_bundle', components=f_components, pname=bundle_pname
+            name='f_bundle', components=f_components, pname=pname
         )
         xbundle = PetscBundle(
-            name='x_bundle', components=x_components, pname=bundle_pname
+            name='x_bundle', components=x_components, pname=pname
         )
 
         # Build the bundle mapper
-        for i, t in enumerate(targets):
-            f_arr = arrays[t]['f']
-            x_arr = arrays[t]['x']
+        for f_arr, x_arr in zip(f_components, x_components):
             bundle_mapper[f_arr.base] = fbundle
             bundle_mapper[x_arr.base] = xbundle
 

devito/petsc/solve.py

@@ -83,7 +83,7 @@ def linear_solve_args(self):
         self.time_mapper = generate_time_mapper(funcs)
         arrays = self.generate_arrays_combined(target)
 
-        eqns = sorted(eqns, key=lambda e: 0 if isinstance(e, EssentialBC) else 1)
+        eqns = sorted(eqns, key=lambda e: not isinstance(e, EssentialBC))
 
         jacobian = Jacobian(target, eqns, arrays, self.time_mapper)
         residual = Residual(target, eqns, arrays, self.time_mapper, jacobian.scdiag)

devito/petsc/types/array.py

@@ -126,6 +126,34 @@ def symbolic_shape(self):
 
 class PetscBundle(Bundle):
     """
+    Tensor symbol representing an unrolled vector of PETScArrays.
+
+    This class declares a struct in the generated ccode to represent the
+    fields defined at each node of the grid. For example:
+
+            typedef struct {
+                PetscScalar u,v,omega,temperature;
+            } Field;
+
+    Residual evaluations are then written using:
+
+    f[i][j].omega = ...
+
+    Reference - https://petsc.org/release/manual/vec/#sec-struct
+
+    Parameters
+    ----------
+    name : str
+        Name of the symbol.
+    components : tuple of PETScArray
+        The PETScArrays of the Bundle.
+    pname : str, optional
+        The name of the struct in the generated C code. Defaults to "Field".
+
+    Warnings
+    --------
+    PetscBundles are created and managed directly by Devito (IOW, they are not
+    expected to be used directly in user code).
     """
     is_Bundle = True
     _data_alignment = False
@@ -136,11 +164,11 @@ def __init__(self, *args, pname="Field", **kwargs):
         super().__init__(*args, **kwargs)
         self._pname = pname
 
-    @property
+    @cached_property
     def _C_ctype(self):
         fields = [(i.target.name, dtype_to_ctype(i.dtype)) for i in self.components]
         return POINTER(type(self.pname, (Structure,), {'_fields_': fields}))
-    
+
     @cached_property
     def symbolic_shape(self):
         return self.c0.symbolic_shape
@@ -176,8 +204,10 @@ def __getitem__(self, index):
                 component_names=names
             )
         else:
-            raise ValueError("Expected %d or %d indices, got %d instead"
-                             % (self.ndim, self.ndim + 1, len(index)))
+            raise ValueError(
+                f"Expected {self.ndim} or {self.ndim + 1} indices, "
+                f"got {len(index)} instead"
+            )
 
     @property
     def pname(self):
@@ -187,7 +217,7 @@ def pname(self):
 class PetscComponentAccess(ComponentAccess):
     def __new__(cls, arg, index=0, component_names=None, **kwargs):
         if not arg.is_Indexed:
-            raise ValueError("Expected Indexed, got `%s` instead" % type(arg))
+            raise ValueError(f"Expected Indexed, got `{type(arg)}` instead")
         names = component_names or cls._default_component_names
 
         obj = Expr.__new__(cls, arg)

devito/petsc/types/types.py

@@ -2,7 +2,7 @@
 
 from itertools import chain
 
-from devito.tools import Reconstructable, sympy_mutex, as_tuple
+from devito.tools import Reconstructable, sympy_mutex, as_tuple, frozendict
 from devito.tools.dtypes_lowering import dtype_mapper
 from devito.petsc.utils import petsc_variables
 from devito.symbolics.extraction import separate_eqn, generate_targets, centre_stencil
@@ -269,11 +269,11 @@ def col_target(self):
         return self.target
 
     def _build_matvecs(self):
-        matvecs = [
-            e for eq in self.eqns for e in
-            self._build_matvec_eq(eq)
-            if e is not None
-        ]
+        matvecs = []
+        for eq in self.eqns:
+            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)
@@ -287,12 +287,13 @@ def _build_matvec_eq(self, eq, col_target=None, row_target=None):
         col_target = col_target or self.target
         row_target = row_target or self.target
 
-        b, F_target, _, targets = separate_eqn(eq, col_target)
+        _, F_target, _, targets = separate_eqn(eq, col_target)
         if F_target:
             return self._make_matvec(
                 eq, F_target, targets, col_target, row_target
             )
-        return (None,)
+        else:
+            return (None,)
 
     def _make_matvec(self, eq, F_target, targets, col_target, row_target):
         y = self.arrays[row_target]['y']
@@ -307,7 +308,7 @@ def _make_matvec(self, eq, F_target, targets, col_target, row_target):
         else:
             rhs = F_target.subs(targets_to_arrays(x, targets))
             rhs = rhs.subs(self.time_mapper)
-        return as_tuple(Eq(y, rhs, subdomain=eq.subdomain))
+            return (Eq(y, rhs, subdomain=eq.subdomain),)
 
     def _scale_non_bcs(self, matvecs, target=None):
         target = target or self.target
@@ -376,7 +377,7 @@ class MixedJacobian(Jacobian):
     def __init__(self, target_eqns, arrays, time_mapper):
         """
         """
-        self.targets = as_tuple(target_eqns.keys())
+        self.targets = tuple(target_eqns.keys())
         self.arrays = arrays
         self.time_mapper = time_mapper
         self._submatrices = []
@@ -391,7 +392,7 @@ def submatrices(self):
         return self._submatrices
 
     @property
-    def no_submatrices(self):
+    def n_submatrices(self):
         """
         Return the number of submatrix blocks.
         """
@@ -421,11 +422,13 @@ def _build_blocks(self, target_eqns):
         for i, row_target in enumerate(self.targets):
             eqns = target_eqns[row_target]
             for j, col_target in enumerate(self.targets):
-                matvecs = [
-                    e for eq in eqns for e in
-                    self._build_matvec_eq(eq, col_target, row_target)
-                ]
+                matvecs = []
+                for eq in eqns:
+                    matvecs.extend(
+                        e for e in self._build_matvec_eq(eq, col_target, row_target)
+                    )
                 matvecs = [m for m in matvecs if m is not None]
+
                 # Sort to put EssentialBC first if any
                 matvecs = tuple(
                     sorted(matvecs, key=lambda e: not isinstance(e, EssentialBC))
@@ -461,7 +464,7 @@ def __repr__(self):
             f"{sm.name} (row={sm.row_idx}, col={sm.col_idx})"
             for sm in self.submatrices
         )
-        return f"<MixedJacobian with {self.no_submatrices} submatrices: [{summary}]>"
+        return f"<MixedJacobian with {self.n_submatrices} submatrices: [{summary}]>"
 
 
 class Residual:
@@ -499,12 +502,13 @@ def _build_equations(self):
             # TODO: If b is zero then don't need a rhs vector+callback
             rhs.extend(self._make_b(eq, b))
 
-        self._formfuncs = [self._scale_bcs(eq) for eq in funcs]
-        self._formrhs = rhs
+        self._formfuncs = tuple([self._scale_bcs(eq) for eq in funcs])
+        self._formrhs = tuple(rhs)
 
     def _make_F_target(self, eq, F_target, targets):
         arrays = self.arrays[self.target]
         volume = self.target.grid.symbolic_volume_cell
+
         if isinstance(eq, EssentialBC):
             # The initial guess satisfies the essential BCs, so this term is zero.
             # Still included to support Jacobian testing via finite differences.
@@ -513,19 +517,20 @@ def _make_F_target(self, eq, F_target, targets):
             # Move essential boundary condition to the right-hand side
             zero_col = ZeroColumn(arrays['x'], eq.rhs, subdomain=eq.subdomain)
             return (zero_row, zero_col)
+
         else:
             if isinstance(F_target, (int, float)):
                 rhs = F_target * volume
             else:
                 rhs = F_target.subs(targets_to_arrays(arrays['x'], targets))
                 rhs = rhs.subs(self.time_mapper) * volume
-        return as_tuple(Eq(arrays['f'], rhs, subdomain=eq.subdomain))
+        return (Eq(arrays['f'], rhs, subdomain=eq.subdomain),)
 
     def _make_b(self, eq, b):
         b_arr = self.arrays[self.target]['b']
         rhs = 0. if isinstance(eq, EssentialBC) else b.subs(self.time_mapper)
         rhs = rhs * self.target.grid.symbolic_volume_cell
-        return as_tuple(Eq(b_arr, rhs, subdomain=eq.subdomain))
+        return (Eq(b_arr, rhs, subdomain=eq.subdomain),)
 
     def _scale_bcs(self, eq, scdiag=None):
         """
@@ -617,11 +622,11 @@ def _build_equations(self):
         """
         Return a list of initial guess equations.
         """
-        self._eqs = [
+        self._eqs = tuple([
             eq for eq in
             (self._make_initial_guess(e) for e in self.eqns)
             if eq is not None
-        ]
+        ])
 
     def _make_initial_guess(self, eq):
         if isinstance(eq, EssentialBC):
@@ -655,4 +660,4 @@ def targets_to_arrays(array, targets):
     array_targets = [
         array.subs(dict(zip(array.indices, i))) for i in space_indices
     ]
-    return dict(zip(targets, array_targets))
+    return frozendict(zip(targets, array_targets))