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) { 24 PetscCall(DMCreateMatrix(rlink->dm, &sub)); 25 } else PetscCheck(!com->FormCoupleLocations, PetscObjectComm((PetscObject)dm), PETSC_ERR_SUP, "Cannot manage off-diagonal parts yet"); 26 submats[i * n + j] = sub; 27 } 28 } 29 30 PetscCall(MatCreateNest(PetscObjectComm((PetscObject)dm), n, isg, n, isg, submats, J)); 31 32 /* Disown references */ 33 for (i = 0; i < n; i++) PetscCall(ISDestroy(&isg[i])); 34 PetscCall(PetscFree(isg)); 35 36 for (i = 0; i < n * n; i++) { 37 if (submats[i]) PetscCall(MatDestroy(&submats[i])); 38 } 39 PetscCall(PetscFree(submats)); 40 PetscFunctionReturn(PETSC_SUCCESS); 41 } 42 43 static PetscErrorCode DMCreateMatrix_Composite_AIJ(DM dm, Mat *J) 44 { 45 DM_Composite *com = (DM_Composite *)dm->data; 46 struct DMCompositeLink *next; 47 PetscInt m, *dnz, *onz, i, j, mA; 48 Mat Atmp; 49 PetscMPIInt rank; 50 PetscBool dense = PETSC_FALSE; 51 52 PetscFunctionBegin; 53 /* use global vector to determine layout needed for matrix */ 54 m = com->n; 55 56 PetscCall(MatCreate(PetscObjectComm((PetscObject)dm), J)); 57 PetscCall(MatSetSizes(*J, m, m, PETSC_DETERMINE, PETSC_DETERMINE)); 58 PetscCall(MatSetType(*J, dm->mattype)); 59 60 /* 61 Extremely inefficient but will compute entire Jacobian for testing 62 */ 63 PetscCall(PetscOptionsGetBool(((PetscObject)dm)->options, ((PetscObject)dm)->prefix, "-dmcomposite_dense_jacobian", &dense, NULL)); 64 if (dense) { 65 PetscInt rstart, rend, *indices; 66 PetscScalar *values; 67 68 mA = com->N; 69 PetscCall(MatMPIAIJSetPreallocation(*J, mA, NULL, mA - m, NULL)); 70 PetscCall(MatSeqAIJSetPreallocation(*J, mA, NULL)); 71 72 PetscCall(MatGetOwnershipRange(*J, &rstart, &rend)); 73 PetscCall(PetscMalloc2(mA, &values, mA, &indices)); 74 PetscCall(PetscArrayzero(values, mA)); 75 for (i = 0; i < mA; i++) indices[i] = i; 76 for (i = rstart; i < rend; i++) PetscCall(MatSetValues(*J, 1, &i, mA, indices, values, INSERT_VALUES)); 77 PetscCall(PetscFree2(values, indices)); 78 PetscCall(MatAssemblyBegin(*J, MAT_FINAL_ASSEMBLY)); 79 PetscCall(MatAssemblyEnd(*J, MAT_FINAL_ASSEMBLY)); 80 PetscFunctionReturn(PETSC_SUCCESS); 81 } 82 83 PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)dm), &rank)); 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, (const PetscInt **)&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, (const PetscInt **)&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