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 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