Insert zero blocks on matrix diagonals

OP2 · Aug 22, 2023 · e62d57b · e62d57b
1 parent d230953
commit e62d57b
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/pyop2/sparsity.pyx b/pyop2/sparsity.pyx
@@ -185,14 +185,14 @@ def fill_with_zeros(PETSc.Mat mat not None, dims, maps, iteration_regions, set_d
     cdef:
         PetscInt rdim, cdim
         PetscScalar *values
+        PetscScalar *diag_values
         int set_entry
         int set_size
         int region_selector
         bint constant_layers, extruded_periodic
         PetscInt layer_start, layer_end, layer_bottom, num_layers, effective_offset, layer
         PetscInt[:, ::1] layers
         PetscInt i, k, irem
-        PetscScalar zero = 0.0
         PetscInt nrow, ncol
         PetscInt rarity, carity, tmp_rarity, tmp_carity
         PetscInt[:, ::1] rmap, cmap, tempmap
@@ -211,9 +211,11 @@ def fill_with_zeros(PETSc.Mat mat not None, dims, maps, iteration_regions, set_d
     # Always allocate space for diagonal
     nrow, ncol = mat.getLocalSize()
     if set_diag:
-        for i in range(nrow):
-            if i < ncol:
-                CHKERR(MatSetValuesLocal(mat.mat, 1, &i, 1, &i, &zero, PETSC_INSERT_VALUES))
+        CHKERR(PetscCalloc1(rdim*cdim, &diag_values))
+        for i in range(nrow // rdim):
+            if i < ncol  // cdim:
+                CHKERR(MatSetValuesBlockedLocal(mat.mat, 1, &i, 1, &i, diag_values, PETSC_INSERT_VALUES))
+        CHKERR(PetscFree(diag_values))
     extruded = maps[0][0].iterset._extruded
     for iteration_region, pair in zip(iteration_regions, maps):
         # Iterate over row map values including value entries