Add back in MatInvertVariableBlockDiagonal_Ceed

examples/fluids/src/mat-ceed.c

@@ -209,6 +209,58 @@ static PetscErrorCode MatGetBlockDiagonal_Ceed(Mat mat_ceed, Mat *mat_block) {
   PetscFunctionReturn(PETSC_SUCCESS);
 }
 
+/**
+  @brief Invert `MATCEED` diagonal block for Jacobi.
+
+  Collective across MPI processes.
+
+  @param[in]   mat_ceed  `MATCEED` to invert
+  @param[out]  values    The block inverses in column major order
+
+  @return An error code: 0 - success, otherwise - failure
+**/
+static PetscErrorCode MatInvertBlockDiagonal_Ceed(Mat mat_ceed, const PetscScalar **values) {
+  Mat            mat_inner = NULL;
+  MatCeedContext ctx;
+
+  PetscFunctionBeginUser;
+  PetscCall(MatShellGetContext(mat_ceed, &ctx));
+
+  // Assemble inner mat if needed
+  PetscCall(MatCeedAssembleInnerBlockDiagonalMat(mat_ceed, ctx->is_ceed_pbd_valid, &mat_inner));
+
+  // Invert PB diagonal
+  PetscCall(MatInvertBlockDiagonal(mat_inner, values));
+  PetscFunctionReturn(PETSC_SUCCESS);
+}
+
+/**
+  @brief Invert `MATCEED` variable diagonal block for Jacobi.
+
+  Collective across MPI processes.
+
+  @param[in]   mat_ceed     `MATCEED` to invert
+  @param[in]   num_blocks   The number of blocks on the process
+  @param[in]   block_sizes  The size of each block on the process
+  @param[out]  values       The block inverses in column major order
+
+  @return An error code: 0 - success, otherwise - failure
+**/
+static PetscErrorCode MatInvertVariableBlockDiagonal_Ceed(Mat mat_ceed, PetscInt num_blocks, const PetscInt *block_sizes, PetscScalar *values) {
+  Mat            mat_inner = NULL;
+  MatCeedContext ctx;
+
+  PetscFunctionBeginUser;
+  PetscCall(MatShellGetContext(mat_ceed, &ctx));
+
+  // Assemble inner mat if needed
+  PetscCall(MatCeedAssembleInnerBlockDiagonalMat(mat_ceed, ctx->is_ceed_vpbd_valid, &mat_inner));
+
+  // Invert PB diagonal
+  PetscCall(MatInvertVariableBlockDiagonal(mat_inner, num_blocks, block_sizes, values));
+  PetscFunctionReturn(PETSC_SUCCESS);
+}
+
 /**
   @brief Get on-process diagonal block of `MATCEED`
 
@@ -472,6 +524,8 @@ PetscErrorCode MatCreateCeed(DM dm_x, DM dm_y, CeedOperator op_mult, CeedOperato
   if (op_mult_transpose) PetscCall(MatShellSetOperation(*mat, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_Ceed));
   PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiagonal_Ceed));
   PetscCall(MatShellSetOperation(*mat, MATOP_GET_DIAGONAL_BLOCK, (void (*)(void))MatGetDiagonalBlock_Ceed));
+  PetscCall(MatShellSetOperation(*mat, MATOP_INVERT_BLOCK_DIAGONAL, (void (*)(void))MatInvertBlockDiagonal_Ceed));
+  PetscCall(MatShellSetOperation(*mat, MATOP_INVERT_VBLOCK_DIAGONAL, (void (*)(void))MatInvertVariableBlockDiagonal_Ceed));
   PetscCall(MatShellSetOperation(*mat, MATOP_GET_BLOCK_DIAGONAL, (void (*)(void))MatGetBlockDiagonal_Ceed));
   PetscCall(MatShellSetOperation(*mat, MATOP_GET_VBLOCK_DIAGONAL, (void (*)(void))MatGetVariableBlockDiagonal_Ceed));
   PetscCall(MatShellSetVecType(*mat, vec_type));
@@ -538,6 +592,8 @@ PetscErrorCode MatCeedCopy(Mat mat_ceed, Mat mat_other) {
     if (ctx->op_mult_transpose) PetscCall(MatShellSetOperation(mat_other, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_Ceed));
     PetscCall(MatShellSetOperation(mat_other, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiagonal_Ceed));
     PetscCall(MatShellSetOperation(mat_other, MATOP_GET_DIAGONAL_BLOCK, (void (*)(void))MatGetDiagonalBlock_Ceed));
+    PetscCall(MatShellSetOperation(mat_other, MATOP_INVERT_BLOCK_DIAGONAL, (void (*)(void))MatInvertBlockDiagonal_Ceed));
+    PetscCall(MatShellSetOperation(mat_other, MATOP_INVERT_VBLOCK_DIAGONAL, (void (*)(void))MatInvertVariableBlockDiagonal_Ceed));
     PetscCall(MatShellSetOperation(mat_other, MATOP_GET_BLOCK_DIAGONAL, (void (*)(void))MatGetBlockDiagonal_Ceed));
     PetscCall(MatShellSetOperation(mat_other, MATOP_GET_VBLOCK_DIAGONAL, (void (*)(void))MatGetVariableBlockDiagonal_Ceed));
     {