Actual source code: ex139.c

  1: const char help[] = "Test MatCreateLocalRef()\n\n";

  3: #include <petscmat.h>

  5: static PetscErrorCode GetLocalRef(Mat A, IS isrow, IS iscol, Mat *B)
  6: {
  7:   IS istmp;

  9:   PetscFunctionBegin;
 10:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Extracting LocalRef with isrow:\n"));
 11:   PetscCall(ISOnComm(isrow, PETSC_COMM_WORLD, PETSC_COPY_VALUES, &istmp));
 12:   PetscCall(ISView(istmp, PETSC_VIEWER_STDOUT_WORLD));
 13:   PetscCall(ISDestroy(&istmp));
 14:   PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Extracting LocalRef with iscol (only rank=0 shown):\n"));
 15:   PetscCall(ISOnComm(iscol, PETSC_COMM_WORLD, PETSC_COPY_VALUES, &istmp));
 16:   PetscCall(ISView(istmp, PETSC_VIEWER_STDOUT_WORLD));
 17:   PetscCall(ISDestroy(&istmp));

 19:   PetscCall(MatCreateLocalRef(A, isrow, iscol, B));
 20:   PetscFunctionReturn(PETSC_SUCCESS);
 21: }

 23: int main(int argc, char *argv[])
 24: {
 25:   MPI_Comm               comm;
 26:   Mat                    J, B;
 27:   PetscInt               i, j, k, l, rstart, rend, m, n, top_bs, row_bs, col_bs, nlocblocks, *idx, nrowblocks, ncolblocks, *ridx, *cidx, *icol, *irow;
 28:   PetscScalar           *vals;
 29:   ISLocalToGlobalMapping brmap;
 30:   IS                     is0, is1;
 31:   PetscBool              diag, blocked;

 33:   PetscFunctionBeginUser;
 34:   PetscCall(PetscInitialize(&argc, &argv, 0, help));
 35:   comm = PETSC_COMM_WORLD;

 37:   PetscOptionsBegin(comm, NULL, "LocalRef Test Options", NULL);
 38:   {
 39:     top_bs  = 2;
 40:     row_bs  = 2;
 41:     col_bs  = 2;
 42:     diag    = PETSC_FALSE;
 43:     blocked = PETSC_FALSE;
 44:     PetscCall(PetscOptionsInt("-top_bs", "Block size of top-level matrix", 0, top_bs, &top_bs, NULL));
 45:     PetscCall(PetscOptionsInt("-row_bs", "Block size of row map", 0, row_bs, &row_bs, NULL));
 46:     PetscCall(PetscOptionsInt("-col_bs", "Block size of col map", 0, col_bs, &col_bs, NULL));
 47:     PetscCall(PetscOptionsBool("-diag", "Extract a diagonal black", 0, diag, &diag, NULL));
 48:     PetscCall(PetscOptionsBool("-blocked", "Use block insertion", 0, blocked, &blocked, NULL));
 49:   }
 50:   PetscOptionsEnd();

 52:   PetscCall(MatCreate(comm, &J));
 53:   PetscCall(MatSetSizes(J, 6, 6, PETSC_DETERMINE, PETSC_DETERMINE));
 54:   PetscCall(MatSetBlockSize(J, top_bs));
 55:   PetscCall(MatSetFromOptions(J));
 56:   PetscCall(MatSeqBAIJSetPreallocation(J, top_bs, PETSC_DECIDE, 0));
 57:   PetscCall(MatMPIBAIJSetPreallocation(J, top_bs, PETSC_DECIDE, 0, PETSC_DECIDE, 0));
 58:   PetscCall(MatSetUp(J));
 59:   PetscCall(MatGetSize(J, &m, &n));
 60:   PetscCall(MatGetOwnershipRange(J, &rstart, &rend));

 62:   nlocblocks = (rend - rstart) / top_bs + 2;
 63:   PetscCall(PetscMalloc1(nlocblocks, &idx));
 64:   for (i = 0; i < nlocblocks; i++) idx[i] = (rstart / top_bs + i - 1 + m / top_bs) % (m / top_bs);
 65:   PetscCall(ISLocalToGlobalMappingCreate(comm, top_bs, nlocblocks, idx, PETSC_OWN_POINTER, &brmap));
 66:   PetscCall(PetscPrintf(comm, "Block ISLocalToGlobalMapping:\n"));
 67:   PetscCall(ISLocalToGlobalMappingView(brmap, PETSC_VIEWER_STDOUT_WORLD));

 69:   PetscCall(MatSetLocalToGlobalMapping(J, brmap, brmap));
 70:   PetscCall(ISLocalToGlobalMappingDestroy(&brmap));

 72:   /* Create index sets for local submatrix */
 73:   nrowblocks = (rend - rstart) / row_bs;
 74:   ncolblocks = (rend - rstart) / col_bs;
 75:   PetscCall(PetscMalloc2(nrowblocks, &ridx, ncolblocks, &cidx));
 76:   for (i = 0; i < nrowblocks; i++) ridx[i] = i + ((i > nrowblocks / 2) ^ !rstart);
 77:   for (i = 0; i < ncolblocks; i++) cidx[i] = i + ((i < ncolblocks / 2) ^ !!rstart);
 78:   PetscCall(ISCreateBlock(PETSC_COMM_SELF, row_bs, nrowblocks, ridx, PETSC_COPY_VALUES, &is0));
 79:   PetscCall(ISCreateBlock(PETSC_COMM_SELF, col_bs, ncolblocks, cidx, PETSC_COPY_VALUES, &is1));
 80:   PetscCall(PetscFree2(ridx, cidx));

 82:   if (diag) {
 83:     PetscCall(ISDestroy(&is1));
 84:     PetscCall(PetscObjectReference((PetscObject)is0));
 85:     is1        = is0;
 86:     ncolblocks = nrowblocks;
 87:   }

 89:   PetscCall(GetLocalRef(J, is0, is1, &B));

 91:   PetscCall(MatZeroEntries(J));

 93:   PetscCall(PetscMalloc3(row_bs, &irow, col_bs, &icol, row_bs * col_bs, &vals));
 94:   for (i = 0; i < nrowblocks; i++) {
 95:     for (j = 0; j < ncolblocks; j++) {
 96:       for (k = 0; k < row_bs; k++) {
 97:         for (l = 0; l < col_bs; l++) vals[k * col_bs + l] = i * 100000 + j * 1000 + k * 10 + l;
 98:         irow[k] = i * row_bs + k;
 99:       }
100:       for (l = 0; l < col_bs; l++) icol[l] = j * col_bs + l;
101:       if (blocked) {
102:         PetscCall(MatSetValuesBlockedLocal(B, 1, &i, 1, &j, vals, ADD_VALUES));
103:       } else {
104:         PetscCall(MatSetValuesLocal(B, row_bs, irow, col_bs, icol, vals, ADD_VALUES));
105:       }
106:     }
107:   }
108:   PetscCall(PetscFree3(irow, icol, vals));

110:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
111:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));

113:   PetscCall(MatView(J, PETSC_VIEWER_STDOUT_WORLD));

115:   PetscCall(ISDestroy(&is0));
116:   PetscCall(ISDestroy(&is1));
117:   PetscCall(MatDestroy(&B));
118:   PetscCall(MatDestroy(&J));
119:   PetscCall(PetscFinalize());
120:   return 0;
121: }

123: /*TEST

125:    test:
126:       nsize: 2
127:       filter: grep -v "type: mpi"
128:       args: -blocked {{0 1}} -mat_type {{aij baij}}

130: TEST*/