misc: Move vecset to function inside iet/utils.py

Author: ZoeLeibowitz

devito/petsc/iet/routines.py

@@ -16,7 +16,7 @@
 from devito.petsc.types import PETScArray, PetscBundle
 from devito.petsc.iet.nodes import (PETScCallable, FormFunctionCallback,
                                     MatShellSetOp, PetscMetaData)
-from devito.petsc.iet.utils import petsc_call, petsc_struct
+from devito.petsc.iet.utils import petsc_call, petsc_struct, zero_vector
 from devito.petsc.utils import solver_mapper
 from devito.petsc.types import (DM, Mat, CallbackVec, Vec, KSP, PC, SNES,
                                 PetscInt, StartPtr, PointerIS, PointerDM, VecScatter,
@@ -99,13 +99,6 @@ def initialguesses(self):
     def user_struct_callback(self):
         return self._user_struct_callback
 
-    @property
-    def zero_memory(self):
-        """Indicates whether the memory of the output
-        vector should be set to zero before the computation
-        in the callback."""
-        return True
-
     @property
     def fielddata(self):
         return self.injectsolve.expr.rhs.fielddata
@@ -169,7 +162,7 @@ def _create_matvec_body(self, body, jacobian):
             'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
         )
 
-        zero_y_memory = self.zero_vector(objs['Y'])
+        zero_y_memory = zero_vector(objs['Y']) if jacobian.zero_memory else None
 
         dm_get_local_xvec = petsc_call(
             'DMGetLocalVector', [dmda, Byref(xlocal)]
@@ -188,9 +181,7 @@ def _create_matvec_body(self, body, jacobian):
             'DMGetLocalVector', [dmda, Byref(ylocal)]
         )
 
-        zero_ylocal_memory = petsc_call(
-            'VecSet', [ylocal, 0.0]
-        )
+        zero_ylocal_memory = zero_vector(ylocal)
 
         vec_get_array_y = petsc_call(
             'VecGetArray', [ylocal, Byref(y_matvec._C_symbol)]
@@ -320,7 +311,7 @@ def _create_formfunc_body(self, body):
             'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
         )
 
-        zero_f_memory = self.zero_vector(objs['F'])
+        zero_f_memory = zero_vector(objs['F'])
 
         dm_get_local_xvec = petsc_call(
             'DMGetLocalVector', [dmda, Byref(objs['xloc'])]
@@ -652,12 +643,6 @@ 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
-
     def dereference_funcs(self, struct, fields):
         return tuple(
             [Dereference(i, struct) for i in
@@ -691,13 +676,6 @@ def main_matvec_callback(self):
     def main_formfunc_callback(self):
         return self._main_formfunc_callback
 
-    @property
-    def zero_memory(self):
-        """Indicates whether the memory of the output
-        vector should be set to zero before the computation
-        in the callback."""
-        return False
-
     def _make_core(self):
         for sm in self.fielddata.jacobian.nonzero_submatrices:
             self._make_matvec(sm, prefix=f'{sm.name}_MatMult')
@@ -730,9 +708,7 @@ def _whole_matvec_body(self):
 
         nonzero_submats = self.jacobian.nonzero_submatrices
 
-        zero_y_memory = petsc_call(
-            'VecSet', [objs['Y'], 0.0]
-        )
+        zero_y_memory = zero_vector(objs['Y'])
 
         calls = ()
         for sm in nonzero_submats:
@@ -815,9 +791,7 @@ def _whole_formfunc_body(self, body):
             'DMGetApplicationContext', [dmda, Byref(ctx._C_symbol)]
         )
 
-        zero_f_memory = petsc_call(
-            'VecSet', [objs['F'], 0.0]
-        )
+        zero_f_memory = zero_vector(objs['F'])
 
         dm_get_local_xvec = petsc_call(
             'DMGetLocalVector', [dmda, Byref(objs['xloc'])]

devito/petsc/iet/utils.py

@@ -18,6 +18,10 @@ def petsc_struct(name, fields, pname, liveness='lazy', modifier=None):
                        modifier=modifier)
 
 
+def zero_vector(vec):
+    return petsc_call('VecSet', [vec, 0.0])
+
+
 # Mapping special Eq operations to their corresponding IET Expression subclass types.
 # These operations correspond to subclasses of Eq utilised within PETScSolve.
 petsc_iet_mapper = {OpPetsc: PetscMetaData}

devito/petsc/types/types.py

@@ -382,6 +382,10 @@ def row_target(self):
     def col_target(self):
         return self.target
 
+    @property
+    def zero_memory(self):
+        return True
+
     def _build_matvecs(self):
         matvecs = []
         for eq in self.exprs:
@@ -516,8 +520,9 @@ def __init__(self, name, matvecs, scdiag, row_target,
         self.col_idx = col_idx
         self.linear_idx = linear_idx
 
-    def is_diag(self):
-        return self.row_idx == self.col_idx
+    @property
+    def zero_memory(self):
+        return False
 
     def __repr__(self):
         return (f"<SubMatrixBlock {self.name}>")

examples/petsc/cfd/01_navierstokes.py

@@ -297,3 +297,8 @@ def neumann_right(eq, subdomain):
 # Pressure norm check
 tol = 1e-3
 assert np.sum((pn1.data[0]-pn1.data[1])**2/np.maximum(pn1.data[0]**2, 1e-10)) < tol
+
+from devito import norm
+print(norm(u1))
+print(norm(v1))
+print(norm(pn1))