1: static char help[] = "Partition tiny grid.\n\n";

  3: /*
  4:   Include "petscmat.h" so that we can use matrices.  Note that this file
  5:   automatically includes:
  6:      petscsys.h - base PETSc routines
  7:      petscvec.h - vectors
  8:      petscis.h  - index sets
  9: */
 10: #include <petscmat.h>

 12: int main(int argc, char **args)
 13: {
 14:   Mat             A, At;
 15:   PetscMPIInt     rank, size;
 16:   PetscInt       *ia, *ja, row;
 17:   MatPartitioning part;
 18:   IS              is, isn;
 19:   PetscBool       equal;

 21:   PetscFunctionBeginUser;
 22:   PetscCall(PetscInitialize(&argc, &args, NULL, help));
 23:   PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
 24:   PetscCheck(size == 4, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Must run with 4 processors");
 25:   PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));

 27:   PetscCall(PetscMalloc1(5, &ia));
 28:   PetscCall(PetscMalloc1(16, &ja));
 29:   if (rank == 0) {
 30:     ja[0] = 1;
 31:     ja[1] = 4;
 32:     ja[2] = 0;
 33:     ja[3] = 2;
 34:     ja[4] = 5;
 35:     ja[5] = 1;
 36:     ja[6] = 3;
 37:     ja[7] = 6;
 38:     ja[8] = 2;
 39:     ja[9] = 7;
 40:     ia[0] = 0;
 41:     ia[1] = 2;
 42:     ia[2] = 5;
 43:     ia[3] = 8;
 44:     ia[4] = 10;
 45:   } else if (rank == 1) {
 46:     ja[0]  = 0;
 47:     ja[1]  = 5;
 48:     ja[2]  = 8;
 49:     ja[3]  = 1;
 50:     ja[4]  = 4;
 51:     ja[5]  = 6;
 52:     ja[6]  = 9;
 53:     ja[7]  = 2;
 54:     ja[8]  = 5;
 55:     ja[9]  = 7;
 56:     ja[10] = 10;
 57:     ja[11] = 3;
 58:     ja[12] = 6;
 59:     ja[13] = 11;
 60:     ia[0]  = 0;
 61:     ia[1]  = 3;
 62:     ia[2]  = 7;
 63:     ia[3]  = 11;
 64:     ia[4]  = 14;
 65:   } else if (rank == 2) {
 66:     ja[0]  = 4;
 67:     ja[1]  = 9;
 68:     ja[2]  = 12;
 69:     ja[3]  = 5;
 70:     ja[4]  = 8;
 71:     ja[5]  = 10;
 72:     ja[6]  = 13;
 73:     ja[7]  = 6;
 74:     ja[8]  = 9;
 75:     ja[9]  = 11;
 76:     ja[10] = 14;
 77:     ja[11] = 7;
 78:     ja[12] = 10;
 79:     ja[13] = 15;
 80:     ia[0]  = 0;
 81:     ia[1]  = 3;
 82:     ia[2]  = 7;
 83:     ia[3]  = 11;
 84:     ia[4]  = 14;
 85:   } else {
 86:     ja[0] = 8;
 87:     ja[1] = 13;
 88:     ja[2] = 9;
 89:     ja[3] = 12;
 90:     ja[4] = 14;
 91:     ja[5] = 10;
 92:     ja[6] = 13;
 93:     ja[7] = 15;
 94:     ja[8] = 11;
 95:     ja[9] = 14;
 96:     ia[0] = 0;
 97:     ia[1] = 2;
 98:     ia[2] = 5;
 99:     ia[3] = 8;
100:     ia[4] = 10;
101:   }

103:   PetscCall(MatCreateMPIAdj(PETSC_COMM_WORLD, 4, 16, ia, ja, NULL, &A));
104:   PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));

106:   /* Create the same matrix but using MatSetValues() */
107:   PetscCall(MatCreate(PETSC_COMM_WORLD, &At));
108:   PetscCall(MatSetSizes(At, 4, 4, 16, 16));
109:   PetscCall(MatSetType(At, MATMPIADJ));
110:   for (PetscInt i = 0; i < 4; i++) {
111:     row = i + 4 * rank;
112:     PetscCall(MatSetValues(At, 1, &row, ia[i + 1] - ia[i], ja + ia[i], NULL, INSERT_VALUES));
113:   }
114:   PetscCall(MatAssemblyBegin(At, MAT_FINAL_ASSEMBLY));
115:   PetscCall(MatAssemblyEnd(At, MAT_FINAL_ASSEMBLY));
116:   PetscCall(MatEqual(A, At, &equal));
117:   PetscCheck(equal, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "Matrices are not equal that should be equal");
118:   PetscCall(MatDestroy(&At));

120:   /*
121:        Partition the graph of the matrix
122:   */
123:   PetscCall(MatPartitioningCreate(PETSC_COMM_WORLD, &part));
124:   PetscCall(MatPartitioningSetAdjacency(part, A));
125:   PetscCall(MatPartitioningSetFromOptions(part));
126:   /* get new processor owner number of each vertex */
127:   PetscCall(MatPartitioningApply(part, &is));
128:   /* get new global number of each old global number */
129:   PetscCall(ISPartitioningToNumbering(is, &isn));
130:   PetscCall(ISView(isn, PETSC_VIEWER_STDOUT_WORLD));
131:   PetscCall(ISDestroy(&is));

133:   PetscCall(ISDestroy(&isn));
134:   PetscCall(MatPartitioningDestroy(&part));

136:   /*
137:        Free work space.  All PETSc objects should be destroyed when they
138:        are no longer needed.
139:   */
140:   PetscCall(MatDestroy(&A));

142:   PetscCall(PetscFinalize());
143:   return 0;
144: }
145: /*
146:    test:
147:      requires: parmetis
148:      args: -mat_view
149:      nsize: 4
150: */