1 #include <../src/dm/impls/composite/packimpl.h> /*I "petscdmcomposite.h" I*/ 2 3 static PetscErrorCode DMCreateMatrix_Composite_Nest(DM dm, Mat *J) 4 { 5 const DM_Composite *com = (DM_Composite *)dm->data; 6 const struct DMCompositeLink *rlink, *clink; 7 IS *isg; 8 Mat *submats; 9 PetscInt i, j, n; 10 11 PetscFunctionBegin; 12 n = com->nDM; /* Total number of entries */ 13 14 /* Explicit index sets are not required for MatCreateNest, but getting them here allows MatNest to do consistency 15 * checking and allows ISEqual to compare by identity instead of by contents. */ 16 PetscCall(DMCompositeGetGlobalISs(dm, &isg)); 17 18 /* Get submatrices */ 19 PetscCall(PetscMalloc1(n * n, &submats)); 20 for (i = 0, rlink = com->next; rlink; i++, rlink = rlink->next) { 21 for (j = 0, clink = com->next; clink; j++, clink = clink->next) { 22 Mat sub = NULL; 23 if (i == j) PetscCall(DMCreateMatrix(rlink->dm, &sub)); 24 else PetscCheck(!com->FormCoupleLocations, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot manage off-diagonal parts yet"); 25 submats[i * n + j] = sub; 26 } 27 } 28 29 PetscCall(MatCreateNest(PetscObjectComm((PetscObject)dm), n, isg, n, isg, submats, J)); 30 31 /* Disown references */ 32 for (i = 0; i < n; i++) PetscCall(ISDestroy(&isg[i])); 33 PetscCall(PetscFree(isg)); 34 35 for (i = 0; i < n * n; i++) { 36 if (submats[i]) PetscCall(MatDestroy(&submats[i])); 37 } 38 PetscCall(PetscFree(submats)); 39 PetscFunctionReturn(PETSC_SUCCESS); 40 } 41 42 static PetscErrorCode DMCreateMatrix_Composite_AIJ(DM dm, Mat *J) 43 { 44 DM_Composite *com = (DM_Composite *)dm->data; 45 struct DMCompositeLink *next; 46 PetscInt m, *dnz, *onz, i, j, mA; 47 Mat Atmp; 48 PetscMPIInt rank; 49 PetscBool dense = PETSC_FALSE; 50 51 PetscFunctionBegin; 52 /* use global vector to determine layout needed for matrix */ 53 m = com->n; 54 55 PetscCall(MatCreate(PetscObjectComm((PetscObject)dm), J)); 56 PetscCall(MatSetSizes(*J, m, m, PETSC_DETERMINE, PETSC_DETERMINE)); 57 PetscCall(MatSetType(*J, dm->mattype)); 58 59 /* 60 Extremely inefficient but will compute entire Jacobian for testing 61 */ 62 PetscCall(PetscOptionsGetBool(((PetscObject)dm)->options, ((PetscObject)dm)->prefix, "-dmcomposite_dense_jacobian", &dense, NULL)); 63 if (dense) { 64 PetscInt rstart, rend, *indices; 65 PetscScalar *values; 66 67 mA = com->N; 68 PetscCall(MatMPIAIJSetPreallocation(*J, mA, NULL, mA - m, NULL)); 69 PetscCall(MatSeqAIJSetPreallocation(*J, mA, NULL)); 70 71 PetscCall(MatGetOwnershipRange(*J, &rstart, &rend)); 72 PetscCall(PetscMalloc2(mA, &values, mA, &indices)); 73 PetscCall(PetscArrayzero(values, mA)); 74 for (i = 0; i < mA; i++) indices[i] = i; 75 for (i = rstart; i < rend; i++) PetscCall(MatSetValues(*J, 1, &i, mA, indices, values, INSERT_VALUES)); 76 PetscCall(PetscFree2(values, indices)); 77 PetscCall(MatAssemblyBegin(*J, MAT_FINAL_ASSEMBLY)); 78 PetscCall(MatAssemblyEnd(*J, MAT_FINAL_ASSEMBLY)); 79 PetscFunctionReturn(PETSC_SUCCESS); 80 } 81 82 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank)); 83 if (dm->prealloc_skip) PetscFunctionReturn(PETSC_SUCCESS); 84 MatPreallocateBegin(PetscObjectComm((PetscObject)dm), m, m, dnz, onz); 85 /* loop over packed objects, handling one at a time */ 86 next = com->next; 87 while (next) { 88 PetscInt nc, rstart, *ccols, maxnc; 89 const PetscInt *cols, *rstarts; 90 PetscMPIInt proc; 91 92 PetscCall(DMCreateMatrix(next->dm, &Atmp)); 93 PetscCall(MatGetOwnershipRange(Atmp, &rstart, NULL)); 94 PetscCall(MatGetOwnershipRanges(Atmp, &rstarts)); 95 PetscCall(MatGetLocalSize(Atmp, &mA, NULL)); 96 97 maxnc = 0; 98 for (i = 0; i < mA; i++) { 99 PetscCall(MatGetRow(Atmp, rstart + i, &nc, NULL, NULL)); 100 maxnc = PetscMax(nc, maxnc); 101 PetscCall(MatRestoreRow(Atmp, rstart + i, &nc, NULL, NULL)); 102 } 103 PetscCall(PetscMalloc1(maxnc, &ccols)); 104 for (i = 0; i < mA; i++) { 105 PetscCall(MatGetRow(Atmp, rstart + i, &nc, &cols, NULL)); 106 /* remap the columns taking into how much they are shifted on each process */ 107 for (j = 0; j < nc; j++) { 108 proc = 0; 109 while (cols[j] >= rstarts[proc + 1]) proc++; 110 ccols[j] = cols[j] + next->grstarts[proc] - rstarts[proc]; 111 } 112 PetscCall(MatPreallocateSet(com->rstart + next->rstart + i, nc, ccols, dnz, onz)); 113 PetscCall(MatRestoreRow(Atmp, rstart + i, &nc, &cols, NULL)); 114 } 115 PetscCall(PetscFree(ccols)); 116 PetscCall(MatDestroy(&Atmp)); 117 next = next->next; 118 } 119 if (com->FormCoupleLocations) PetscCall((*com->FormCoupleLocations)(dm, NULL, dnz, onz, __rstart, __nrows, __start, __end)); 120 PetscCall(MatMPIAIJSetPreallocation(*J, 0, dnz, 0, onz)); 121 PetscCall(MatSeqAIJSetPreallocation(*J, 0, dnz)); 122 MatPreallocateEnd(dnz, onz); 123 124 if (dm->prealloc_only) PetscFunctionReturn(PETSC_SUCCESS); 125 126 next = com->next; 127 while (next) { 128 PetscInt nc, rstart, row, maxnc, *ccols; 129 const PetscInt *cols, *rstarts; 130 const PetscScalar *values; 131 PetscMPIInt proc; 132 133 PetscCall(DMCreateMatrix(next->dm, &Atmp)); 134 PetscCall(MatGetOwnershipRange(Atmp, &rstart, NULL)); 135 PetscCall(MatGetOwnershipRanges(Atmp, &rstarts)); 136 PetscCall(MatGetLocalSize(Atmp, &mA, NULL)); 137 maxnc = 0; 138 for (i = 0; i < mA; i++) { 139 PetscCall(MatGetRow(Atmp, rstart + i, &nc, NULL, NULL)); 140 maxnc = PetscMax(nc, maxnc); 141 PetscCall(MatRestoreRow(Atmp, rstart + i, &nc, NULL, NULL)); 142 } 143 PetscCall(PetscMalloc1(maxnc, &ccols)); 144 for (i = 0; i < mA; i++) { 145 PetscCall(MatGetRow(Atmp, rstart + i, &nc, &cols, &values)); 146 for (j = 0; j < nc; j++) { 147 proc = 0; 148 while (cols[j] >= rstarts[proc + 1]) proc++; 149 ccols[j] = cols[j] + next->grstarts[proc] - rstarts[proc]; 150 } 151 row = com->rstart + next->rstart + i; 152 PetscCall(MatSetValues(*J, 1, &row, nc, ccols, values, INSERT_VALUES)); 153 PetscCall(MatRestoreRow(Atmp, rstart + i, &nc, &cols, &values)); 154 } 155 PetscCall(PetscFree(ccols)); 156 PetscCall(MatDestroy(&Atmp)); 157 next = next->next; 158 } 159 if (com->FormCoupleLocations) { 160 PetscInt __rstart; 161 PetscCall(MatGetOwnershipRange(*J, &__rstart, NULL)); 162 PetscCall((*com->FormCoupleLocations)(dm, *J, NULL, NULL, __rstart, 0, 0, 0)); 163 } 164 PetscCall(MatAssemblyBegin(*J, MAT_FINAL_ASSEMBLY)); 165 PetscCall(MatAssemblyEnd(*J, MAT_FINAL_ASSEMBLY)); 166 PetscFunctionReturn(PETSC_SUCCESS); 167 } 168 169 PetscErrorCode DMCreateMatrix_Composite(DM dm, Mat *J) 170 { 171 PetscBool usenest; 172 ISLocalToGlobalMapping ltogmap; 173 174 PetscFunctionBegin; 175 PetscCall(DMSetFromOptions(dm)); 176 PetscCall(DMSetUp(dm)); 177 PetscCall(PetscStrcmp(dm->mattype, MATNEST, &usenest)); 178 if (usenest) PetscCall(DMCreateMatrix_Composite_Nest(dm, J)); 179 else PetscCall(DMCreateMatrix_Composite_AIJ(dm, J)); 180 181 PetscCall(DMGetLocalToGlobalMapping(dm, <ogmap)); 182 PetscCall(MatSetLocalToGlobalMapping(*J, ltogmap, ltogmap)); 183 PetscCall(MatSetDM(*J, dm)); 184 PetscFunctionReturn(PETSC_SUCCESS); 185 } 186