1 #include <petsc/private/matimpl.h> /*I "petscmat.h" I*/ 2 3 PETSC_INTERN PetscErrorCode MatCreateADA(Mat, Vec, Vec, Mat *); 4 5 typedef struct { 6 Mat A; 7 Vec D1; 8 Vec D2; 9 Vec W; 10 Vec W2; 11 Vec ADADiag; 12 PetscInt GotDiag; 13 } _p_TaoMatADACtx; 14 typedef _p_TaoMatADACtx *TaoMatADACtx; 15 16 static PetscErrorCode MatMult_ADA(Mat mat, Vec a, Vec y) { 17 TaoMatADACtx ctx; 18 PetscReal one = 1.0; 19 20 PetscFunctionBegin; 21 PetscCall(MatShellGetContext(mat, &ctx)); 22 PetscCall(MatMult(ctx->A, a, ctx->W)); 23 if (ctx->D1) PetscCall(VecPointwiseMult(ctx->W, ctx->D1, ctx->W)); 24 PetscCall(MatMultTranspose(ctx->A, ctx->W, y)); 25 if (ctx->D2) { 26 PetscCall(VecPointwiseMult(ctx->W2, ctx->D2, a)); 27 PetscCall(VecAXPY(y, one, ctx->W2)); 28 } 29 PetscFunctionReturn(0); 30 } 31 32 static PetscErrorCode MatMultTranspose_ADA(Mat mat, Vec a, Vec y) { 33 PetscFunctionBegin; 34 PetscCall(MatMult_ADA(mat, a, y)); 35 PetscFunctionReturn(0); 36 } 37 38 static PetscErrorCode MatDiagonalSet_ADA(Mat M, Vec D, InsertMode mode) { 39 TaoMatADACtx ctx; 40 PetscReal zero = 0.0, one = 1.0; 41 42 PetscFunctionBegin; 43 PetscCheck(mode != INSERT_VALUES, PetscObjectComm((PetscObject)M), PETSC_ERR_SUP, "Cannot insert diagonal entries of this matrix type, can only add"); 44 PetscCall(MatShellGetContext(M, &ctx)); 45 if (!ctx->D2) { 46 PetscCall(VecDuplicate(D, &ctx->D2)); 47 PetscCall(VecSet(ctx->D2, zero)); 48 } 49 PetscCall(VecAXPY(ctx->D2, one, D)); 50 PetscFunctionReturn(0); 51 } 52 53 static PetscErrorCode MatDestroy_ADA(Mat mat) { 54 TaoMatADACtx ctx; 55 56 PetscFunctionBegin; 57 PetscCall(MatShellGetContext(mat, &ctx)); 58 PetscCall(VecDestroy(&ctx->W)); 59 PetscCall(VecDestroy(&ctx->W2)); 60 PetscCall(VecDestroy(&ctx->ADADiag)); 61 PetscCall(MatDestroy(&ctx->A)); 62 PetscCall(VecDestroy(&ctx->D1)); 63 PetscCall(VecDestroy(&ctx->D2)); 64 PetscCall(PetscFree(ctx)); 65 PetscFunctionReturn(0); 66 } 67 68 static PetscErrorCode MatView_ADA(Mat mat, PetscViewer viewer) { 69 PetscFunctionBegin; 70 PetscFunctionReturn(0); 71 } 72 73 static PetscErrorCode MatShift_ADA(Mat Y, PetscReal a) { 74 TaoMatADACtx ctx; 75 76 PetscFunctionBegin; 77 PetscCall(MatShellGetContext(Y, &ctx)); 78 PetscCall(VecShift(ctx->D2, a)); 79 PetscFunctionReturn(0); 80 } 81 82 static PetscErrorCode MatDuplicate_ADA(Mat mat, MatDuplicateOption op, Mat *M) { 83 TaoMatADACtx ctx; 84 Mat A2; 85 Vec D1b = NULL, D2b; 86 87 PetscFunctionBegin; 88 PetscCall(MatShellGetContext(mat, &ctx)); 89 PetscCall(MatDuplicate(ctx->A, op, &A2)); 90 if (ctx->D1) { 91 PetscCall(VecDuplicate(ctx->D1, &D1b)); 92 PetscCall(VecCopy(ctx->D1, D1b)); 93 } 94 PetscCall(VecDuplicate(ctx->D2, &D2b)); 95 PetscCall(VecCopy(ctx->D2, D2b)); 96 PetscCall(MatCreateADA(A2, D1b, D2b, M)); 97 if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1b)); 98 PetscCall(PetscObjectDereference((PetscObject)D2b)); 99 PetscCall(PetscObjectDereference((PetscObject)A2)); 100 PetscFunctionReturn(0); 101 } 102 103 static PetscErrorCode MatEqual_ADA(Mat A, Mat B, PetscBool *flg) { 104 TaoMatADACtx ctx1, ctx2; 105 106 PetscFunctionBegin; 107 PetscCall(MatShellGetContext(A, &ctx1)); 108 PetscCall(MatShellGetContext(B, &ctx2)); 109 PetscCall(VecEqual(ctx1->D2, ctx2->D2, flg)); 110 if (*flg == PETSC_TRUE) PetscCall(VecEqual(ctx1->D1, ctx2->D1, flg)); 111 if (*flg == PETSC_TRUE) PetscCall(MatEqual(ctx1->A, ctx2->A, flg)); 112 PetscFunctionReturn(0); 113 } 114 115 static PetscErrorCode MatScale_ADA(Mat mat, PetscReal a) { 116 TaoMatADACtx ctx; 117 118 PetscFunctionBegin; 119 PetscCall(MatShellGetContext(mat, &ctx)); 120 PetscCall(VecScale(ctx->D1, a)); 121 if (ctx->D2) PetscCall(VecScale(ctx->D2, a)); 122 PetscFunctionReturn(0); 123 } 124 125 static PetscErrorCode MatTranspose_ADA(Mat mat, MatReuse reuse, Mat *B) { 126 TaoMatADACtx ctx; 127 128 PetscFunctionBegin; 129 if (reuse == MAT_REUSE_MATRIX) PetscCall(MatTransposeCheckNonzeroState_Private(mat, *B)); 130 PetscCall(MatShellGetContext(mat, &ctx)); 131 if (reuse == MAT_INITIAL_MATRIX) { 132 PetscCall(MatDuplicate(mat, MAT_COPY_VALUES, B)); 133 } else if (reuse == MAT_REUSE_MATRIX) { 134 PetscCall(MatCopy(mat, *B, SAME_NONZERO_PATTERN)); 135 } else SETERRQ(PetscObjectComm((PetscObject)mat), PETSC_ERR_SUP, "Does not support inplace transpose"); 136 PetscFunctionReturn(0); 137 } 138 139 static PetscErrorCode MatADAComputeDiagonal(Mat mat) { 140 PetscInt i, m, n, low, high; 141 PetscScalar *dtemp, *dptr; 142 TaoMatADACtx ctx; 143 144 PetscFunctionBegin; 145 PetscCall(MatShellGetContext(mat, &ctx)); 146 PetscCall(MatGetOwnershipRange(mat, &low, &high)); 147 PetscCall(MatGetSize(mat, &m, &n)); 148 149 PetscCall(PetscMalloc1(n, &dtemp)); 150 for (i = 0; i < n; i++) { 151 PetscCall(MatGetColumnVector(ctx->A, ctx->W, i)); 152 PetscCall(VecPointwiseMult(ctx->W, ctx->W, ctx->W)); 153 PetscCall(VecDotBegin(ctx->D1, ctx->W, dtemp + i)); 154 } 155 for (i = 0; i < n; i++) { PetscCall(VecDotEnd(ctx->D1, ctx->W, dtemp + i)); } 156 157 PetscCall(VecGetArray(ctx->ADADiag, &dptr)); 158 for (i = low; i < high; i++) { dptr[i - low] = dtemp[i]; } 159 PetscCall(VecRestoreArray(ctx->ADADiag, &dptr)); 160 PetscCall(PetscFree(dtemp)); 161 PetscFunctionReturn(0); 162 } 163 164 static PetscErrorCode MatGetDiagonal_ADA(Mat mat, Vec v) { 165 PetscReal one = 1.0; 166 TaoMatADACtx ctx; 167 168 PetscFunctionBegin; 169 PetscCall(MatShellGetContext(mat, &ctx)); 170 PetscCall(MatADAComputeDiagonal(mat)); 171 PetscCall(VecCopy(ctx->ADADiag, v)); 172 if (ctx->D2) PetscCall(VecAXPY(v, one, ctx->D2)); 173 PetscFunctionReturn(0); 174 } 175 176 static PetscErrorCode MatCreateSubMatrix_ADA(Mat mat, IS isrow, IS iscol, MatReuse cll, Mat *newmat) { 177 PetscInt low, high; 178 IS ISrow; 179 Vec D1, D2; 180 Mat Atemp; 181 TaoMatADACtx ctx; 182 PetscBool isequal; 183 184 PetscFunctionBegin; 185 PetscCall(ISEqual(isrow, iscol, &isequal)); 186 PetscCheck(isequal, PETSC_COMM_SELF, PETSC_ERR_SUP, "Only for identical column and row indices"); 187 PetscCall(MatShellGetContext(mat, &ctx)); 188 189 PetscCall(MatGetOwnershipRange(ctx->A, &low, &high)); 190 PetscCall(ISCreateStride(PetscObjectComm((PetscObject)mat), high - low, low, 1, &ISrow)); 191 PetscCall(MatCreateSubMatrix(ctx->A, ISrow, iscol, cll, &Atemp)); 192 PetscCall(ISDestroy(&ISrow)); 193 194 if (ctx->D1) { 195 PetscCall(VecDuplicate(ctx->D1, &D1)); 196 PetscCall(VecCopy(ctx->D1, D1)); 197 } else { 198 D1 = NULL; 199 } 200 201 if (ctx->D2) { 202 Vec D2sub; 203 204 PetscCall(VecGetSubVector(ctx->D2, isrow, &D2sub)); 205 PetscCall(VecDuplicate(D2sub, &D2)); 206 PetscCall(VecCopy(D2sub, D2)); 207 PetscCall(VecRestoreSubVector(ctx->D2, isrow, &D2sub)); 208 } else { 209 D2 = NULL; 210 } 211 212 PetscCall(MatCreateADA(Atemp, D1, D2, newmat)); 213 PetscCall(MatShellGetContext(*newmat, &ctx)); 214 PetscCall(PetscObjectDereference((PetscObject)Atemp)); 215 if (ctx->D1) PetscCall(PetscObjectDereference((PetscObject)D1)); 216 if (ctx->D2) PetscCall(PetscObjectDereference((PetscObject)D2)); 217 PetscFunctionReturn(0); 218 } 219 220 static PetscErrorCode MatCreateSubMatrices_ADA(Mat A, PetscInt n, IS *irow, IS *icol, MatReuse scall, Mat **B) { 221 PetscInt i; 222 223 PetscFunctionBegin; 224 if (scall == MAT_INITIAL_MATRIX) { PetscCall(PetscCalloc1(n + 1, B)); } 225 for (i = 0; i < n; i++) { PetscCall(MatCreateSubMatrix_ADA(A, irow[i], icol[i], scall, &(*B)[i])); } 226 PetscFunctionReturn(0); 227 } 228 229 static PetscErrorCode MatGetColumnVector_ADA(Mat mat, Vec Y, PetscInt col) { 230 PetscInt low, high; 231 PetscScalar zero = 0.0, one = 1.0; 232 233 PetscFunctionBegin; 234 PetscCall(VecSet(Y, zero)); 235 PetscCall(VecGetOwnershipRange(Y, &low, &high)); 236 if (col >= low && col < high) { PetscCall(VecSetValue(Y, col, one, INSERT_VALUES)); } 237 PetscCall(VecAssemblyBegin(Y)); 238 PetscCall(VecAssemblyEnd(Y)); 239 PetscCall(MatMult_ADA(mat, Y, Y)); 240 PetscFunctionReturn(0); 241 } 242 243 PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat, MatType newtype, Mat *NewMat) { 244 PetscMPIInt size; 245 PetscBool sametype, issame, isdense, isseqdense; 246 TaoMatADACtx ctx; 247 248 PetscFunctionBegin; 249 PetscCall(MatShellGetContext(mat, &ctx)); 250 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size)); 251 252 PetscCall(PetscObjectTypeCompare((PetscObject)mat, newtype, &sametype)); 253 PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSAME, &issame)); 254 PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATMPIDENSE, &isdense)); 255 PetscCall(PetscObjectTypeCompare((PetscObject)mat, MATSEQDENSE, &isseqdense)); 256 257 if (sametype || issame) { 258 PetscCall(MatDuplicate(mat, MAT_COPY_VALUES, NewMat)); 259 } else if (isdense) { 260 PetscInt i, j, low, high, m, n, M, N; 261 const PetscScalar *dptr; 262 Vec X; 263 264 PetscCall(VecDuplicate(ctx->D2, &X)); 265 PetscCall(MatGetSize(mat, &M, &N)); 266 PetscCall(MatGetLocalSize(mat, &m, &n)); 267 PetscCall(MatCreateDense(PetscObjectComm((PetscObject)mat), m, m, N, N, NULL, NewMat)); 268 PetscCall(MatGetOwnershipRange(*NewMat, &low, &high)); 269 for (i = 0; i < M; i++) { 270 PetscCall(MatGetColumnVector_ADA(mat, X, i)); 271 PetscCall(VecGetArrayRead(X, &dptr)); 272 for (j = 0; j < high - low; j++) { PetscCall(MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES)); } 273 PetscCall(VecRestoreArrayRead(X, &dptr)); 274 } 275 PetscCall(MatAssemblyBegin(*NewMat, MAT_FINAL_ASSEMBLY)); 276 PetscCall(MatAssemblyEnd(*NewMat, MAT_FINAL_ASSEMBLY)); 277 PetscCall(VecDestroy(&X)); 278 } else if (isseqdense && size == 1) { 279 PetscInt i, j, low, high, m, n, M, N; 280 const PetscScalar *dptr; 281 Vec X; 282 283 PetscCall(VecDuplicate(ctx->D2, &X)); 284 PetscCall(MatGetSize(mat, &M, &N)); 285 PetscCall(MatGetLocalSize(mat, &m, &n)); 286 PetscCall(MatCreateSeqDense(PetscObjectComm((PetscObject)mat), N, N, NULL, NewMat)); 287 PetscCall(MatGetOwnershipRange(*NewMat, &low, &high)); 288 for (i = 0; i < M; i++) { 289 PetscCall(MatGetColumnVector_ADA(mat, X, i)); 290 PetscCall(VecGetArrayRead(X, &dptr)); 291 for (j = 0; j < high - low; j++) { PetscCall(MatSetValue(*NewMat, low + j, i, dptr[j], INSERT_VALUES)); } 292 PetscCall(VecRestoreArrayRead(X, &dptr)); 293 } 294 PetscCall(MatAssemblyBegin(*NewMat, MAT_FINAL_ASSEMBLY)); 295 PetscCall(MatAssemblyEnd(*NewMat, MAT_FINAL_ASSEMBLY)); 296 PetscCall(VecDestroy(&X)); 297 } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "No support to convert objects to that type"); 298 PetscFunctionReturn(0); 299 } 300 301 static PetscErrorCode MatNorm_ADA(Mat mat, NormType type, PetscReal *norm) { 302 TaoMatADACtx ctx; 303 304 PetscFunctionBegin; 305 PetscCall(MatShellGetContext(mat, &ctx)); 306 if (type == NORM_FROBENIUS) { 307 *norm = 1.0; 308 } else if (type == NORM_1 || type == NORM_INFINITY) { 309 *norm = 1.0; 310 } else SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "No two norm"); 311 PetscFunctionReturn(0); 312 } 313 314 /*@C 315 MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal 316 317 Collective on matrix 318 319 Input Parameters: 320 + mat - matrix of arbitrary type 321 . d1 - A vector defining a diagonal matrix 322 - d2 - A vector defining a diagonal matrix 323 324 Output Parameters: 325 . J - New matrix whose operations are defined in terms of mat, D1, and D2. 326 327 Notes: 328 The user provides the input data and is responsible for destroying 329 this data after matrix J has been destroyed. 330 331 Level: developer 332 333 .seealso: `MatCreate()` 334 @*/ 335 PetscErrorCode MatCreateADA(Mat mat, Vec d1, Vec d2, Mat *J) { 336 MPI_Comm comm = PetscObjectComm((PetscObject)mat); 337 TaoMatADACtx ctx; 338 PetscInt nloc, n; 339 340 PetscFunctionBegin; 341 PetscCall(PetscNew(&ctx)); 342 ctx->A = mat; 343 ctx->D1 = d1; 344 ctx->D2 = d2; 345 if (d1) { 346 PetscCall(VecDuplicate(d1, &ctx->W)); 347 PetscCall(PetscObjectReference((PetscObject)d1)); 348 } else { 349 ctx->W = NULL; 350 } 351 if (d2) { 352 PetscCall(VecDuplicate(d2, &ctx->W2)); 353 PetscCall(VecDuplicate(d2, &ctx->ADADiag)); 354 PetscCall(PetscObjectReference((PetscObject)d2)); 355 } else { 356 ctx->W2 = NULL; 357 ctx->ADADiag = NULL; 358 } 359 360 ctx->GotDiag = 0; 361 PetscCall(PetscObjectReference((PetscObject)mat)); 362 363 PetscCall(VecGetLocalSize(d2, &nloc)); 364 PetscCall(VecGetSize(d2, &n)); 365 366 PetscCall(MatCreateShell(comm, nloc, nloc, n, n, ctx, J)); 367 PetscCall(MatShellSetManageScalingShifts(*J)); 368 PetscCall(MatShellSetOperation(*J, MATOP_MULT, (void (*)(void))MatMult_ADA)); 369 PetscCall(MatShellSetOperation(*J, MATOP_DESTROY, (void (*)(void))MatDestroy_ADA)); 370 PetscCall(MatShellSetOperation(*J, MATOP_VIEW, (void (*)(void))MatView_ADA)); 371 PetscCall(MatShellSetOperation(*J, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_ADA)); 372 PetscCall(MatShellSetOperation(*J, MATOP_DIAGONAL_SET, (void (*)(void))MatDiagonalSet_ADA)); 373 PetscCall(MatShellSetOperation(*J, MATOP_SHIFT, (void (*)(void))MatShift_ADA)); 374 PetscCall(MatShellSetOperation(*J, MATOP_EQUAL, (void (*)(void))MatEqual_ADA)); 375 PetscCall(MatShellSetOperation(*J, MATOP_SCALE, (void (*)(void))MatScale_ADA)); 376 PetscCall(MatShellSetOperation(*J, MATOP_TRANSPOSE, (void (*)(void))MatTranspose_ADA)); 377 PetscCall(MatShellSetOperation(*J, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiagonal_ADA)); 378 PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRICES, (void (*)(void))MatCreateSubMatrices_ADA)); 379 PetscCall(MatShellSetOperation(*J, MATOP_NORM, (void (*)(void))MatNorm_ADA)); 380 PetscCall(MatShellSetOperation(*J, MATOP_DUPLICATE, (void (*)(void))MatDuplicate_ADA)); 381 PetscCall(MatShellSetOperation(*J, MATOP_CREATE_SUBMATRIX, (void (*)(void))MatCreateSubMatrix_ADA)); 382 383 PetscCall(PetscLogObjectParent((PetscObject)(*J), (PetscObject)ctx->W)); 384 PetscCall(PetscLogObjectParent((PetscObject)mat, (PetscObject)(*J))); 385 386 PetscCall(MatSetOption(*J, MAT_SYMMETRIC, PETSC_TRUE)); 387 PetscFunctionReturn(0); 388 } 389