1 2 const char help[] = "Test MatCreateLocalRef()\n\n"; 3 4 #include <petscmat.h> 5 6 static PetscErrorCode GetLocalRef(Mat A,IS isrow,IS iscol,Mat *B) 7 { 8 IS istmp; 9 10 PetscFunctionBegin; 11 CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Extracting LocalRef with isrow:\n")); 12 CHKERRQ(ISOnComm(isrow,PETSC_COMM_WORLD,PETSC_COPY_VALUES,&istmp)); 13 CHKERRQ(ISView(istmp,PETSC_VIEWER_STDOUT_WORLD)); 14 CHKERRQ(ISDestroy(&istmp)); 15 CHKERRQ(PetscPrintf(PETSC_COMM_WORLD,"Extracting LocalRef with iscol (only rank=0 shown):\n")); 16 CHKERRQ(ISOnComm(iscol,PETSC_COMM_WORLD,PETSC_COPY_VALUES,&istmp)); 17 CHKERRQ(ISView(istmp,PETSC_VIEWER_STDOUT_WORLD)); 18 CHKERRQ(ISDestroy(&istmp)); 19 20 CHKERRQ(MatCreateLocalRef(A,isrow,iscol,B)); 21 PetscFunctionReturn(0); 22 } 23 24 int main(int argc,char *argv[]) 25 { 26 PetscErrorCode ierr; 27 MPI_Comm comm; 28 Mat J,B; 29 PetscInt i,j,k,l,rstart,rend,m,n,top_bs,row_bs,col_bs,nlocblocks,*idx,nrowblocks,ncolblocks,*ridx,*cidx,*icol,*irow; 30 PetscScalar *vals; 31 ISLocalToGlobalMapping brmap; 32 IS is0,is1; 33 PetscBool diag,blocked; 34 35 ierr = PetscInitialize(&argc,&argv,0,help);if (ierr) return ierr; 36 comm = PETSC_COMM_WORLD; 37 38 ierr = PetscOptionsBegin(comm,NULL,"LocalRef Test Options",NULL);CHKERRQ(ierr); 39 { 40 top_bs = 2; row_bs = 2; col_bs = 2; diag = PETSC_FALSE; blocked = PETSC_FALSE; 41 CHKERRQ(PetscOptionsInt("-top_bs","Block size of top-level matrix",0,top_bs,&top_bs,NULL)); 42 CHKERRQ(PetscOptionsInt("-row_bs","Block size of row map",0,row_bs,&row_bs,NULL)); 43 CHKERRQ(PetscOptionsInt("-col_bs","Block size of col map",0,col_bs,&col_bs,NULL)); 44 CHKERRQ(PetscOptionsBool("-diag","Extract a diagonal black",0,diag,&diag,NULL)); 45 CHKERRQ(PetscOptionsBool("-blocked","Use block insertion",0,blocked,&blocked,NULL)); 46 } 47 ierr = PetscOptionsEnd();CHKERRQ(ierr); 48 49 CHKERRQ(MatCreate(comm,&J)); 50 CHKERRQ(MatSetSizes(J,6,6,PETSC_DETERMINE,PETSC_DETERMINE)); 51 CHKERRQ(MatSetBlockSize(J,top_bs)); 52 CHKERRQ(MatSetFromOptions(J)); 53 CHKERRQ(MatSeqBAIJSetPreallocation(J,top_bs,PETSC_DECIDE,0)); 54 CHKERRQ(MatMPIBAIJSetPreallocation(J,top_bs,PETSC_DECIDE,0,PETSC_DECIDE,0)); 55 CHKERRQ(MatSetUp(J)); 56 CHKERRQ(MatGetSize(J,&m,&n)); 57 CHKERRQ(MatGetOwnershipRange(J,&rstart,&rend)); 58 59 nlocblocks = (rend-rstart)/top_bs + 2; 60 CHKERRQ(PetscMalloc1(nlocblocks,&idx)); 61 for (i=0; i<nlocblocks; i++) { 62 idx[i] = (rstart/top_bs + i - 1 + m/top_bs) % (m/top_bs); 63 } 64 CHKERRQ(ISLocalToGlobalMappingCreate(comm,top_bs,nlocblocks,idx,PETSC_OWN_POINTER,&brmap)); 65 CHKERRQ(PetscPrintf(comm,"Block ISLocalToGlobalMapping:\n")); 66 CHKERRQ(ISLocalToGlobalMappingView(brmap,PETSC_VIEWER_STDOUT_WORLD)); 67 68 CHKERRQ(MatSetLocalToGlobalMapping(J,brmap,brmap)); 69 CHKERRQ(ISLocalToGlobalMappingDestroy(&brmap)); 70 71 /* Create index sets for local submatrix */ 72 nrowblocks = (rend-rstart)/row_bs; 73 ncolblocks = (rend-rstart)/col_bs; 74 CHKERRQ(PetscMalloc2(nrowblocks,&ridx,ncolblocks,&cidx)); 75 for (i=0; i<nrowblocks; i++) ridx[i] = i + ((i > nrowblocks/2) ^ !rstart); 76 for (i=0; i<ncolblocks; i++) cidx[i] = i + ((i < ncolblocks/2) ^ !!rstart); 77 CHKERRQ(ISCreateBlock(PETSC_COMM_SELF,row_bs,nrowblocks,ridx,PETSC_COPY_VALUES,&is0)); 78 CHKERRQ(ISCreateBlock(PETSC_COMM_SELF,col_bs,ncolblocks,cidx,PETSC_COPY_VALUES,&is1)); 79 CHKERRQ(PetscFree2(ridx,cidx)); 80 81 if (diag) { 82 CHKERRQ(ISDestroy(&is1)); 83 CHKERRQ(PetscObjectReference((PetscObject)is0)); 84 is1 = is0; 85 ncolblocks = nrowblocks; 86 } 87 88 CHKERRQ(GetLocalRef(J,is0,is1,&B)); 89 90 CHKERRQ(MatZeroEntries(J)); 91 92 CHKERRQ(PetscMalloc3(row_bs,&irow,col_bs,&icol,row_bs*col_bs,&vals)); 93 for (i=0; i<nrowblocks; i++) { 94 for (j=0; j<ncolblocks; j++) { 95 for (k=0; k<row_bs; k++) { 96 for (l=0; l<col_bs; l++) { 97 vals[k*col_bs+l] = i * 100000 + j * 1000 + k * 10 + l; 98 } 99 irow[k] = i*row_bs+k; 100 } 101 for (l=0; l<col_bs; l++) icol[l] = j*col_bs+l; 102 if (blocked) { 103 CHKERRQ(MatSetValuesBlockedLocal(B,1,&i,1,&j,vals,ADD_VALUES)); 104 } else { 105 CHKERRQ(MatSetValuesLocal(B,row_bs,irow,col_bs,icol,vals,ADD_VALUES)); 106 } 107 } 108 } 109 CHKERRQ(PetscFree3(irow,icol,vals)); 110 111 CHKERRQ(MatAssemblyBegin(J,MAT_FINAL_ASSEMBLY)); 112 CHKERRQ(MatAssemblyEnd(J,MAT_FINAL_ASSEMBLY)); 113 114 CHKERRQ(MatView(J,PETSC_VIEWER_STDOUT_WORLD)); 115 116 CHKERRQ(ISDestroy(&is0)); 117 CHKERRQ(ISDestroy(&is1)); 118 CHKERRQ(MatDestroy(&B)); 119 CHKERRQ(MatDestroy(&J)); 120 ierr = PetscFinalize(); 121 return ierr; 122 } 123 124 /*TEST 125 126 test: 127 nsize: 2 128 filter: grep -v "type: mpi" 129 args: -blocked {{0 1}} -mat_type {{aij baij}} 130 131 TEST*/ 132