Date: 2025-03-29T05:58:34+00:00

Actual source code: lrc.c

      1: #include <petsc/private/matimpl.h>

      3: PETSC_EXTERN PetscErrorCode VecGetRootType_Private(Vec, VecType *);

      5: typedef struct {
      6:   Mat A;            /* sparse matrix */
      7:   Mat U, V;         /* dense tall-skinny matrices */
      8:   Vec c;            /* sequential vector containing the diagonal of C */
      9:   Vec work1, work2; /* sequential vectors that hold partial products */
     10:   Vec xl, yl;       /* auxiliary sequential vectors for matmult operation */
     11: } Mat_LRC;

     13: static PetscErrorCode MatMult_LRC_kernel(Mat N, Vec x, Vec y, PetscBool transpose)
     14: {
     15:   Mat_LRC    *Na = (Mat_LRC *)N->data;
     16:   PetscMPIInt size;
     17:   Mat         U, V;

     19:   PetscFunctionBegin;
     20:   U = transpose ? Na->V : Na->U;
     21:   V = transpose ? Na->U : Na->V;
     22:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)N), &size));
     23:   if (size == 1) {
     24:     PetscCall(MatMultHermitianTranspose(V, x, Na->work1));
     25:     if (Na->c) PetscCall(VecPointwiseMult(Na->work1, Na->c, Na->work1));
     26:     if (Na->A) {
     27:       if (transpose) {
     28:         PetscCall(MatMultTranspose(Na->A, x, y));
     29:       } else {
     30:         PetscCall(MatMult(Na->A, x, y));
     31:       }
     32:       PetscCall(MatMultAdd(U, Na->work1, y, y));
     33:     } else {
     34:       PetscCall(MatMult(U, Na->work1, y));
     35:     }
     36:   } else {
     37:     Mat                Uloc, Vloc;
     38:     Vec                yl, xl;
     39:     const PetscScalar *w1;
     40:     PetscScalar       *w2;
     41:     PetscInt           nwork;

     43:     xl = transpose ? Na->yl : Na->xl;
     44:     yl = transpose ? Na->xl : Na->yl;
     45:     PetscCall(VecGetLocalVector(y, yl));
     46:     PetscCall(MatDenseGetLocalMatrix(U, &Uloc));
     47:     PetscCall(MatDenseGetLocalMatrix(V, &Vloc));

     49:     /* multiply the local part of V with the local part of x */
     50:     PetscCall(VecGetLocalVectorRead(x, xl));
     51:     PetscCall(MatMultHermitianTranspose(Vloc, xl, Na->work1));
     52:     PetscCall(VecRestoreLocalVectorRead(x, xl));

     54:     /* form the sum of all the local multiplies: this is work2 = V'*x =
     55:        sum_{all processors} work1 */
     56:     PetscCall(VecGetArrayRead(Na->work1, &w1));
     57:     PetscCall(VecGetArrayWrite(Na->work2, &w2));
     58:     PetscCall(VecGetLocalSize(Na->work1, &nwork));
     59:     PetscCallMPI(MPIU_Allreduce(w1, w2, nwork, MPIU_SCALAR, MPIU_SUM, PetscObjectComm((PetscObject)N)));
     60:     PetscCall(VecRestoreArrayRead(Na->work1, &w1));
     61:     PetscCall(VecRestoreArrayWrite(Na->work2, &w2));

     63:     if (Na->c) { /* work2 = C*work2 */
     64:       PetscCall(VecPointwiseMult(Na->work2, Na->c, Na->work2));
     65:     }

     67:     if (Na->A) {
     68:       /* form y = A*x or A^t*x */
     69:       if (transpose) {
     70:         PetscCall(MatMultTranspose(Na->A, x, y));
     71:       } else {
     72:         PetscCall(MatMult(Na->A, x, y));
     73:       }
     74:       /* multiply-add y = y + U*work2 */
     75:       PetscCall(MatMultAdd(Uloc, Na->work2, yl, yl));
     76:     } else {
     77:       /* multiply y = U*work2 */
     78:       PetscCall(MatMult(Uloc, Na->work2, yl));
     79:     }

     81:     PetscCall(VecRestoreLocalVector(y, yl));
     82:   }
     83:   PetscFunctionReturn(PETSC_SUCCESS);
     84: }

     86: static PetscErrorCode MatMult_LRC(Mat N, Vec x, Vec y)
     87: {
     88:   PetscFunctionBegin;
     89:   PetscCall(MatMult_LRC_kernel(N, x, y, PETSC_FALSE));
     90:   PetscFunctionReturn(PETSC_SUCCESS);
     91: }

     93: static PetscErrorCode MatMultTranspose_LRC(Mat N, Vec x, Vec y)
     94: {
     95:   PetscFunctionBegin;
     96:   PetscCall(MatMult_LRC_kernel(N, x, y, PETSC_TRUE));
     97:   PetscFunctionReturn(PETSC_SUCCESS);
     98: }

    100: static PetscErrorCode MatDestroy_LRC(Mat N)
    101: {
    102:   Mat_LRC *Na = (Mat_LRC *)N->data;

    104:   PetscFunctionBegin;
    105:   PetscCall(MatDestroy(&Na->A));
    106:   PetscCall(MatDestroy(&Na->U));
    107:   PetscCall(MatDestroy(&Na->V));
    108:   PetscCall(VecDestroy(&Na->c));
    109:   PetscCall(VecDestroy(&Na->work1));
    110:   PetscCall(VecDestroy(&Na->work2));
    111:   PetscCall(VecDestroy(&Na->xl));
    112:   PetscCall(VecDestroy(&Na->yl));
    113:   PetscCall(PetscFree(N->data));
    114:   PetscCall(PetscObjectComposeFunction((PetscObject)N, "MatLRCGetMats_C", NULL));
    115:   PetscCall(PetscObjectComposeFunction((PetscObject)N, "MatLRCSetMats_C", NULL));
    116:   PetscFunctionReturn(PETSC_SUCCESS);
    117: }

    119: static PetscErrorCode MatLRCGetMats_LRC(Mat N, Mat *A, Mat *U, Vec *c, Mat *V)
    120: {
    121:   Mat_LRC *Na = (Mat_LRC *)N->data;

    123:   PetscFunctionBegin;
    124:   if (A) *A = Na->A;
    125:   if (U) *U = Na->U;
    126:   if (c) *c = Na->c;
    127:   if (V) *V = Na->V;
    128:   PetscFunctionReturn(PETSC_SUCCESS);
    129: }

    131: static PetscErrorCode MatLRCSetMats_LRC(Mat N, Mat A, Mat U, Vec c, Mat V)
    132: {
    133:   Mat_LRC *Na = (Mat_LRC *)N->data;

    135:   PetscFunctionBegin;
    136:   PetscCall(PetscObjectReference((PetscObject)A));
    137:   PetscCall(PetscObjectReference((PetscObject)U));
    138:   PetscCall(PetscObjectReference((PetscObject)V));
    139:   PetscCall(PetscObjectReference((PetscObject)c));
    140:   PetscCall(MatDestroy(&Na->A));
    141:   PetscCall(MatDestroy(&Na->U));
    142:   PetscCall(MatDestroy(&Na->V));
    143:   PetscCall(VecDestroy(&Na->c));
    144:   Na->A = A;
    145:   Na->U = U;
    146:   Na->c = c;
    147:   Na->V = V;
    148:   PetscFunctionReturn(PETSC_SUCCESS);
    149: }

    151: /*@
    152:   MatLRCGetMats - Returns the constituents of an LRC matrix

    154:   Not collective

    156:   Input Parameter:
    157: . N - matrix of type `MATLRC`

    159:   Output Parameters:
    160: + A - the (sparse) matrix
    161: . U - first dense rectangular (tall and skinny) matrix
    162: . c - a sequential vector containing the diagonal of C
    163: - V - second dense rectangular (tall and skinny) matrix

    165:   Level: intermediate

    167:   Notes:
    168:   The returned matrices should not be destroyed by the caller.

    170:   `U`, `c`, `V` may be `NULL` if not needed

    172: .seealso: [](ch_matrices), `MatLRCSetMats()`, `Mat`, `MATLRC`, `MatCreateLRC()`
    173: @*/
    174: PetscErrorCode MatLRCGetMats(Mat N, Mat *A, Mat *U, Vec *c, Mat *V)
    175: {
    176:   PetscFunctionBegin;
    177:   PetscUseMethod(N, "MatLRCGetMats_C", (Mat, Mat *, Mat *, Vec *, Mat *), (N, A, U, c, V));
    178:   PetscFunctionReturn(PETSC_SUCCESS);
    179: }

    181: /*@
    182:   MatLRCSetMats - Sets the constituents of an LRC matrix

    184:   Logically collective

    186:   Input Parameters:
    187: + N - matrix of type `MATLRC`
    188: . A - the (sparse) matrix
    189: . U - first dense rectangular (tall and skinny) matrix
    190: . c - a sequential vector containing the diagonal of C
    191: - V - second dense rectangular (tall and skinny) matrix

    193:   Level: intermediate

    195:   Note:
    196:   If `V` is `NULL`, then it is assumed to be identical to `U`.

    198: .seealso: [](ch_matrices), `MatLRCGetMats()`, `Mat`, `MATLRC`, `MatCreateLRC()`
    199: @*/
    200: PetscErrorCode MatLRCSetMats(Mat N, Mat A, Mat U, Vec c, Mat V)
    201: {
    202:   PetscInt  k, k1, m, n, m1, n1;
    203:   PetscBool match;

    205:   PetscFunctionBegin;
    209:   if (V) {
    211:     PetscCheckSameComm(U, 3, V, 5);
    212:   }
    213:   if (A) PetscCheckSameComm(A, 2, U, 3);
    214:   if (!V) V = U;
    215:   PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)U, &match, MATSEQDENSE, MATMPIDENSE, ""));
    216:   PetscCheck(match, PetscObjectComm((PetscObject)U), PETSC_ERR_SUP, "Matrix U must be of type dense, found %s", ((PetscObject)U)->type_name);
    217:   PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)V, &match, MATSEQDENSE, MATMPIDENSE, ""));
    218:   PetscCheck(match, PetscObjectComm((PetscObject)U), PETSC_ERR_SUP, "Matrix V must be of type dense, found %s", ((PetscObject)V)->type_name);
    219:   PetscCall(PetscStrcmp(U->defaultvectype, V->defaultvectype, &match));
    220:   PetscCheck(match, PetscObjectComm((PetscObject)U), PETSC_ERR_ARG_WRONG, "Matrix U and V must have the same VecType %s != %s", U->defaultvectype, V->defaultvectype);
    221:   if (A) {
    222:     PetscCall(PetscStrcmp(A->defaultvectype, U->defaultvectype, &match));
    223:     PetscCheck(match, PetscObjectComm((PetscObject)U), PETSC_ERR_ARG_WRONG, "Matrix A and U must have the same VecType %s != %s", A->defaultvectype, U->defaultvectype);
    224:   }
    225:   PetscCall(MatGetSize(U, NULL, &k));
    226:   PetscCall(MatGetSize(V, NULL, &k1));
    227:   PetscCheck(k == k1, PetscObjectComm((PetscObject)U), PETSC_ERR_ARG_INCOMP, "U and V have different number of columns (%" PetscInt_FMT " vs %" PetscInt_FMT ")", k, k1);
    228:   PetscCall(MatGetLocalSize(U, &m, NULL));
    229:   PetscCall(MatGetLocalSize(V, &n, NULL));
    230:   if (A) {
    231:     PetscCall(MatGetLocalSize(A, &m1, &n1));
    232:     PetscCheck(m == m1, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Local dimensions of U %" PetscInt_FMT " and A %" PetscInt_FMT " do not match", m, m1);
    233:     PetscCheck(n == n1, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Local dimensions of V %" PetscInt_FMT " and A %" PetscInt_FMT " do not match", n, n1);
    234:   }
    235:   if (c) {
    236:     PetscMPIInt size, csize;

    238:     PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)U), &size));
    239:     PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)c), &csize));
    240:     PetscCall(VecGetSize(c, &k1));
    241:     PetscCheck(k == k1, PetscObjectComm((PetscObject)c), PETSC_ERR_ARG_INCOMP, "The length of c %" PetscInt_FMT " does not match the number of columns of U and V (%" PetscInt_FMT ")", k1, k);
    242:     PetscCheck(csize == 1 || csize == size, PetscObjectComm((PetscObject)c), PETSC_ERR_ARG_INCOMP, "U and c must have the same communicator size %d != %d", size, csize);
    243:   }
    244:   PetscCall(MatSetSizes(N, m, n, PETSC_DECIDE, PETSC_DECIDE));

    246:   PetscUseMethod(N, "MatLRCSetMats_C", (Mat, Mat, Mat, Vec, Mat), (N, A, U, c, V));
    247:   PetscFunctionReturn(PETSC_SUCCESS);
    248: }

    250: static PetscErrorCode MatSetUp_LRC(Mat N)
    251: {
    252:   Mat_LRC    *Na = (Mat_LRC *)N->data;
    253:   Mat         A  = Na->A;
    254:   Mat         U  = Na->U;
    255:   Mat         V  = Na->V;
    256:   Vec         c  = Na->c;
    257:   Mat         Uloc;
    258:   PetscMPIInt size, csize = 0;
    259:   PetscBool   sym = (PetscBool)(U == V), dummy;

    261:   PetscFunctionBegin;
    262:   PetscCall(MatSetVecType(N, U->defaultvectype));
    263:   // Flag matrix as symmetric if A is symmetric and U == V
    264:   if (A && sym) PetscCall(MatIsSymmetricKnown(A, &dummy, &sym));
    265:   PetscCall(MatSetOption(N, MAT_SYMMETRIC, sym));
    266:   PetscCall(MatDenseGetLocalMatrix(Na->U, &Uloc));
    267:   PetscCall(MatCreateVecs(Uloc, &Na->work1, NULL));

    269:   PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)U), &size));
    270:   if (c) PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)c), &csize));
    271:   if (size != 1) {
    272:     Mat Vloc;

    274:     if (Na->c && csize != 1) { /* scatter parallel vector to sequential */
    275:       VecScatter sct;

    277:       PetscCall(VecScatterCreateToAll(Na->c, &sct, &c));
    278:       PetscCall(VecScatterBegin(sct, Na->c, c, INSERT_VALUES, SCATTER_FORWARD));
    279:       PetscCall(VecScatterEnd(sct, Na->c, c, INSERT_VALUES, SCATTER_FORWARD));
    280:       PetscCall(VecScatterDestroy(&sct));
    281:       PetscCall(VecDestroy(&Na->c));
    282:       Na->c = c;
    283:     }
    284:     PetscCall(MatDenseGetLocalMatrix(Na->V, &Vloc));
    285:     PetscCall(VecDuplicate(Na->work1, &Na->work2));
    286:     PetscCall(MatCreateVecs(Vloc, NULL, &Na->xl));
    287:     PetscCall(MatCreateVecs(Uloc, NULL, &Na->yl));
    288:   }
    289:   // Internally create a scaling vector if roottypes do not match
    290:   if (Na->c) {
    291:     VecType   rt1, rt2;
    292:     PetscBool match;

    294:     PetscCall(VecGetRootType_Private(Na->work1, &rt1));
    295:     PetscCall(VecGetRootType_Private(Na->c, &rt2));
    296:     PetscCall(PetscStrcmp(rt1, rt2, &match));
    297:     if (!match) {
    298:       PetscCall(VecDuplicate(Na->c, &c));
    299:       PetscCall(VecCopy(Na->c, c));
    300:       PetscCall(VecDestroy(&Na->c));
    301:       Na->c = c;
    302:     }
    303:   }
    304:   N->assembled    = PETSC_TRUE;
    305:   N->preallocated = PETSC_TRUE;
    306:   PetscFunctionReturn(PETSC_SUCCESS);
    307: }

    309: PETSC_EXTERN PetscErrorCode MatCreate_LRC(Mat N)
    310: {
    311:   Mat_LRC *Na;

    313:   PetscFunctionBegin;
    314:   PetscCall(PetscObjectChangeTypeName((PetscObject)N, MATLRC));
    315:   PetscCall(PetscNew(&Na));
    316:   N->data               = (void *)Na;
    317:   N->ops->destroy       = MatDestroy_LRC;
    318:   N->ops->setup         = MatSetUp_LRC;
    319:   N->ops->mult          = MatMult_LRC;
    320:   N->ops->multtranspose = MatMultTranspose_LRC;

    322:   PetscCall(PetscObjectComposeFunction((PetscObject)N, "MatLRCGetMats_C", MatLRCGetMats_LRC));
    323:   PetscCall(PetscObjectComposeFunction((PetscObject)N, "MatLRCSetMats_C", MatLRCSetMats_LRC));
    324:   PetscFunctionReturn(PETSC_SUCCESS);
    325: }

    327: /*MC
    328:   MATLRC -  "lrc" - a matrix object that behaves like A + U*C*V'

    330:   Note:
    331:    The matrix A + U*C*V' is not formed! Rather the matrix  object performs the matrix-vector product `MatMult()`, by first multiplying by
    332:    A and then adding the other term.

    334:   Level: advanced

    336: .seealso: [](ch_matrices), `Mat`, `MatCreateLRC()`, `MatMult()`, `MatLRCGetMats()`, `MatLRCSetMats()`
    337: M*/

    339: /*@
    340:   MatCreateLRC - Creates a new matrix object that behaves like A + U*C*V' of type `MATLRC`

    342:   Collective

    344:   Input Parameters:
    345: + A - the (sparse) matrix (can be `NULL`)
    346: . U - dense rectangular (tall and skinny) matrix
    347: . V - dense rectangular (tall and skinny) matrix
    348: - c - a vector containing the diagonal of C (can be `NULL`)

    350:   Output Parameter:
    351: . N - the matrix that represents A + U*C*V'

    353:   Level: intermediate

    355:   Notes:
    356:   The matrix A + U*C*V' is not formed! Rather the new matrix
    357:   object performs the matrix-vector product `MatMult()`, by first multiplying by
    358:   A and then adding the other term.

    360:   `C` is a diagonal matrix (represented as a vector) of order k,
    361:   where k is the number of columns of both `U` and `V`.

    363:   If `A` is `NULL` then the new object behaves like a low-rank matrix U*C*V'.

    365:   Use `V`=`U` (or `V`=`NULL`) for a symmetric low-rank correction, A + U*C*U'.

    367:   If `c` is `NULL` then the low-rank correction is just U*V'.
    368:   If a sequential `c` vector is used for a parallel matrix,
    369:   PETSc assumes that the values of the vector are consistently set across processors.

    371: .seealso: [](ch_matrices), `Mat`, `MATLRC`, `MatLRCGetMats()`
    372: @*/
    373: PetscErrorCode MatCreateLRC(Mat A, Mat U, Vec c, Mat V, Mat *N)
    374: {
    375:   PetscFunctionBegin;
    376:   PetscCall(MatCreate(PetscObjectComm((PetscObject)U), N));
    377:   PetscCall(MatSetType(*N, MATLRC));
    378:   PetscCall(MatLRCSetMats(*N, A, U, c, V));
    379:   PetscCall(MatSetUp(*N));
    380:   PetscFunctionReturn(PETSC_SUCCESS);
    381: }