Actual source code: ex213.c
1: static char help[] = "Tests MatMPIBAIJSetPreallocationCSR()\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 petscvec.h - vectors
7: petscmat.h - matrices
8: petscis.h - index sets
9: petscviewer.h - viewers
10: */
11: #include <petscmat.h>
13: int main(int argc, char **args)
14: {
15: Mat A;
16: PetscInt *ia, *ja, bs = 2;
17: PetscInt N = 9, n;
18: PetscInt rstart, rend, row, col;
19: PetscInt i;
20: PetscMPIInt rank, size;
21: Vec v;
23: PetscFunctionBeginUser;
24: PetscCall(PetscInitialize(&argc, &args, NULL, help));
25: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
26: PetscCheck(size < 5, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Can only use at most 4 processors.");
27: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
29: /* Get a partition range based on the vector size */
30: PetscCall(VecCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, PETSC_DECIDE, N, &v));
31: PetscCall(VecGetLocalSize(v, &n));
32: PetscCall(VecGetOwnershipRange(v, &rstart, &rend));
33: PetscCall(VecDestroy(&v));
35: PetscCall(PetscMalloc1(n + 1, &ia));
36: PetscCall(PetscMalloc1(3 * n, &ja));
38: /* Construct a tri-diagonal CSR indexing */
39: i = 1;
40: ia[0] = 0;
41: for (row = rstart; row < rend; row++) {
42: ia[i] = ia[i - 1];
44: /* diagonal */
45: col = row;
46: {
47: ja[ia[i]] = col;
48: ia[i]++;
49: }
51: /* lower diagonal */
52: col = row - 1;
53: if (col >= 0) {
54: ja[ia[i]] = col;
55: ia[i]++;
56: }
58: /* upper diagonal */
59: col = row + 1;
60: if (col < N) {
61: ja[ia[i]] = col;
62: ia[i]++;
63: }
64: i++;
65: }
67: PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
68: PetscCall(MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE));
69: PetscCall(MatSetType(A, MATMPIAIJ));
70: PetscCall(MatMPIAIJSetPreallocationCSR(A, ia, ja, NULL));
71: PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
72: PetscCall(MatDestroy(&A));
74: PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
75: PetscCall(MatSetSizes(A, bs * n, bs * n, PETSC_DETERMINE, PETSC_DETERMINE));
76: PetscCall(MatSetType(A, MATMPIBAIJ));
77: PetscCall(MatMPIBAIJSetPreallocationCSR(A, bs, ia, ja, NULL));
78: PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
79: PetscCall(MatDestroy(&A));
81: PetscCall(PetscFree(ia));
82: PetscCall(PetscFree(ja));
83: PetscCall(PetscFinalize());
84: return 0;
85: }
87: /*TEST
89: test:
90: nsize: 4
92: TEST*/