pc (facet_split): not use vec.permute()

firedrakeproject · Jan 30, 2024 · ba5f780 · ba5f780
1 parent 15fb8e7
commit ba5f780
Showing 1 changed file with 16 additions and 9 deletions.
diff --git a/firedrake/preconditioners/facet_split.py b/firedrake/preconditioners/facet_split.py
@@ -45,6 +45,8 @@ def initialize(self, pc):
         _, P = pc.getOperators()
         appctx = self.get_appctx(pc)
         fcp = appctx.get("form_compiler_parameters")
+        self.work_vec_x = P.createVecLeft()
+        self.work_vec_y = P.createVecRight()
 
         prefix = pc.getOptionsPrefix()
         options_prefix = prefix + self._prefix
@@ -102,7 +104,8 @@ def _permute_nullspace(nsp):
                     return nsp
                 vecs = [vec.duplicate() for vec in nsp.getVecs()]
                 for vec in vecs:
-                    vec.permute(self.iperm)
+                    self.work_vec_x.getArray(False)[:] = vec.getArray(True)
+                    self._vec_permute(self.work_vec_x, vec, self.iperm)
                 return PETSc.NullSpace().create(constant=nsp.hasConstant(), vectors=vecs, comm=nsp.getComm())
 
             mixed_opmat.setNullSpace(_permute_nullspace(P.getNullSpace()))
@@ -152,23 +155,27 @@ def update(self, pc):
 
     def apply(self, pc, x, y):
         if self.perm:
-            x.permute(self.iperm)
+            self._vec_permute(x, self.work_vec_x, self.iperm)
         dm = self._dm
         with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
-            self.pc.apply(x, y)
+            self.pc.apply(self.work_vec_x, self.work_vec_y)
         if self.perm:
-            x.permute(self.perm)
-            y.permute(self.perm)
+            self._vec_permute(self.work_vec_y, y, self.perm)
 
     def applyTranspose(self, pc, x, y):
         if self.perm:
-            x.permute(self.iperm)
+            self._vec_permute(x, self.work_vec_y, self.iperm)
         dm = self._dm
         with dmhooks.add_hooks(dm, self, appctx=self._ctx_ref):
-            self.pc.applyTranspose(x, y)
+            self.pc.applyTranspose(self.work_vec_y, self.work_vec_x)
         if self.perm:
-            x.permute(self.perm)
-            y.permute(self.perm)
+            self._vec_permute(self.work_vec_x, y, self.perm)
+
+    def _vec_permute(self, vec_in, vec_out, perm):
+        array_in = vec_in.getArray(True)
+        array_out = vec_out.getArray(False)
+        with perm as perm_indices:
+            array_out[:] = array_in[perm_indices]
 
     def view(self, pc, viewer=None):
         super(FacetSplitPC, self).view(pc, viewer)