1c4762a1bSJed Brown static char help[] = "Tests MatMPIBAIJSetPreallocationCSR()\n\n";
2c4762a1bSJed Brown
3c4762a1bSJed Brown /*
4c4762a1bSJed Brown Include "petscmat.h" so that we can use matrices. Note that this file
5c4762a1bSJed Brown automatically includes:
6c4762a1bSJed Brown petscsys.h - base PETSc routines petscvec.h - vectors
7c4762a1bSJed Brown petscmat.h - matrices
8c4762a1bSJed Brown petscis.h - index sets
9c4762a1bSJed Brown petscviewer.h - viewers
10c4762a1bSJed Brown */
11c4762a1bSJed Brown #include <petscmat.h>
12c4762a1bSJed Brown
main(int argc,char ** args)13d71ae5a4SJacob Faibussowitsch int main(int argc, char **args)
14d71ae5a4SJacob Faibussowitsch {
15c4762a1bSJed Brown Mat A;
16c4762a1bSJed Brown PetscInt *ia, *ja, bs = 2;
17c4762a1bSJed Brown PetscInt N = 9, n;
18c4762a1bSJed Brown PetscInt rstart, rend, row, col;
19c4762a1bSJed Brown PetscInt i;
20c4762a1bSJed Brown PetscMPIInt rank, size;
21c4762a1bSJed Brown Vec v;
22c4762a1bSJed Brown
23327415f7SBarry Smith PetscFunctionBeginUser;
24*c8025a54SPierre Jolivet PetscCall(PetscInitialize(&argc, &args, NULL, help));
259566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
26be096a46SBarry Smith PetscCheck(size < 5, PETSC_COMM_WORLD, PETSC_ERR_WRONG_MPI_SIZE, "Can only use at most 4 processors.");
279566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
28c4762a1bSJed Brown
29c4762a1bSJed Brown /* Get a partition range based on the vector size */
3077433607SBarry Smith PetscCall(VecCreateFromOptions(PETSC_COMM_WORLD, NULL, 1, PETSC_DECIDE, N, &v));
319566063dSJacob Faibussowitsch PetscCall(VecGetLocalSize(v, &n));
329566063dSJacob Faibussowitsch PetscCall(VecGetOwnershipRange(v, &rstart, &rend));
339566063dSJacob Faibussowitsch PetscCall(VecDestroy(&v));
34c4762a1bSJed Brown
359566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(n + 1, &ia));
369566063dSJacob Faibussowitsch PetscCall(PetscMalloc1(3 * n, &ja));
37c4762a1bSJed Brown
38c4762a1bSJed Brown /* Construct a tri-diagonal CSR indexing */
39c4762a1bSJed Brown i = 1;
40c4762a1bSJed Brown ia[0] = 0;
419371c9d4SSatish Balay for (row = rstart; row < rend; row++) {
42c4762a1bSJed Brown ia[i] = ia[i - 1];
43c4762a1bSJed Brown
44c4762a1bSJed Brown /* diagonal */
45c4762a1bSJed Brown col = row;
46c4762a1bSJed Brown {
47c4762a1bSJed Brown ja[ia[i]] = col;
48c4762a1bSJed Brown ia[i]++;
49c4762a1bSJed Brown }
50c4762a1bSJed Brown
51c4762a1bSJed Brown /* lower diagonal */
52c4762a1bSJed Brown col = row - 1;
539371c9d4SSatish Balay if (col >= 0) {
54c4762a1bSJed Brown ja[ia[i]] = col;
55c4762a1bSJed Brown ia[i]++;
56c4762a1bSJed Brown }
57c4762a1bSJed Brown
58c4762a1bSJed Brown /* upper diagonal */
59c4762a1bSJed Brown col = row + 1;
609371c9d4SSatish Balay if (col < N) {
61c4762a1bSJed Brown ja[ia[i]] = col;
62c4762a1bSJed Brown ia[i]++;
63c4762a1bSJed Brown }
64c4762a1bSJed Brown i++;
65c4762a1bSJed Brown }
66c4762a1bSJed Brown
679566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
689566063dSJacob Faibussowitsch PetscCall(MatSetSizes(A, n, n, PETSC_DETERMINE, PETSC_DETERMINE));
699566063dSJacob Faibussowitsch PetscCall(MatSetType(A, MATMPIAIJ));
709566063dSJacob Faibussowitsch PetscCall(MatMPIAIJSetPreallocationCSR(A, ia, ja, NULL));
719566063dSJacob Faibussowitsch PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
729566063dSJacob Faibussowitsch PetscCall(MatDestroy(&A));
73c4762a1bSJed Brown
749566063dSJacob Faibussowitsch PetscCall(MatCreate(PETSC_COMM_WORLD, &A));
759566063dSJacob Faibussowitsch PetscCall(MatSetSizes(A, bs * n, bs * n, PETSC_DETERMINE, PETSC_DETERMINE));
769566063dSJacob Faibussowitsch PetscCall(MatSetType(A, MATMPIBAIJ));
779566063dSJacob Faibussowitsch PetscCall(MatMPIBAIJSetPreallocationCSR(A, bs, ia, ja, NULL));
789566063dSJacob Faibussowitsch PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
799566063dSJacob Faibussowitsch PetscCall(MatDestroy(&A));
80c4762a1bSJed Brown
819566063dSJacob Faibussowitsch PetscCall(PetscFree(ia));
829566063dSJacob Faibussowitsch PetscCall(PetscFree(ja));
839566063dSJacob Faibussowitsch PetscCall(PetscFinalize());
84b122ec5aSJacob Faibussowitsch return 0;
85c4762a1bSJed Brown }
86c4762a1bSJed Brown
87c4762a1bSJed Brown /*TEST
88c4762a1bSJed Brown
89c4762a1bSJed Brown test:
90c4762a1bSJed Brown nsize: 4
91c4762a1bSJed Brown
92c4762a1bSJed Brown TEST*/
93