Actual source code: ex213.c
2: static char help[] = "Tests MatMPIBAIJSetPreallocationCSR()\n\n";
4: /*
5: Include "petscmat.h" so that we can use matrices. Note that this file
6: automatically includes:
7: petscsys.h - base PETSc routines petscvec.h - vectors
8: petscmat.h - matrices
9: petscis.h - index sets
10: petscviewer.h - viewers
11: */
12: #include <petscmat.h>
14: int main(int argc, char **args)
15: {
16: Mat A;
17: PetscInt *ia, *ja, bs = 2;
18: PetscInt N = 9, n;
19: PetscInt rstart, rend, row, col;
20: PetscInt i;
21: PetscMPIInt rank, size;
22: Vec v;
24: PetscFunctionBeginUser;
25: PetscCall(PetscInitialize(&argc, &args, (char *)0, help));
26: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
27: PetscCheck(size < 5, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Can only use at most 4 processors.");
28: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
30: /* Get a partition range based on the vector size */
31: PetscCall(VecCreateMPI(PETSC_COMM_WORLD, PETSC_DECIDE, N, &v));
32: PetscCall(VecGetLocalSize(v, &n));
33: PetscCall(VecGetOwnershipRange(v, &rstart, &rend));
34: PetscCall(VecDestroy(&v));
36: PetscCall(PetscMalloc1(n + 1, &ia));
37: PetscCall(PetscMalloc1(3 * n, &ja));
39: /* Construct a tri-diagonal CSR indexing */
40: i = 1;
41: ia[0] = 0;
42: for (row = rstart; row < rend; row++) {
43: ia[i] = ia[i - 1];
45: /* diagonal */
46: col = row;
47: {
48: ja[ia[i]] = col;
49: ia[i]++;
50: }
52: /* lower diagonal */
53: col = row - 1;
54: if (col >= 0) {
55: ja[ia[i]] = col;
56: ia[i]++;
57: }
59: /* upper diagonal */
60: col = row + 1;
61: if (col < N) {
62: ja[ia[i]] = col;
63: ia[i]++;
64: }
65: i++;
66: }
68: PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
69: PetscCall(MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE));
70: PetscCall(MatSetType(A, MATMPIAIJ));
71: PetscCall(MatMPIAIJSetPreallocationCSR(A, ia, ja, NULL));
72: PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
73: PetscCall(MatDestroy(&A));
75: PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
76: PetscCall(MatSetSizes(A, bs * n, bs * n, PETSC_DETERMINE, PETSC_DETERMINE));
77: PetscCall(MatSetType(A, MATMPIBAIJ));
78: PetscCall(MatMPIBAIJSetPreallocationCSR(A, bs, ia, ja, NULL));
79: PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
80: PetscCall(MatDestroy(&A));
82: PetscCall(PetscFree(ia));
83: PetscCall(PetscFree(ja));
84: PetscCall(PetscFinalize());
85: return 0;
86: }
88: /*TEST
90: test:
91: nsize: 4
93: TEST*/