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