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

Actual source code: lmvmpc.c

      1: /*
      2:    This provides a thin wrapper around LMVM matrices in order to use their MatLMVMSolve
      3:    methods as preconditioner applications in KSP solves.
      4: */

      6: #include <petsc/private/pcimpl.h>
      7: #include <petsc/private/matimpl.h>

      9: typedef struct {
     10:   Vec              xwork, ywork;
     11:   IS               inactive;
     12:   Mat              B;
     13:   Vec              X;
     14:   PetscObjectState Xstate;
     15:   PetscBool        setfromoptionscalled;
     16: } PC_LMVM;

     18: /*@
     19:   PCLMVMSetUpdateVec - Set the vector to be used as solution update for the internal LMVM matrix.

     21:   Input Parameters:
     22: + pc - The preconditioner
     23: - X  - Solution vector

     25:   Level: intermediate

     27:   Notes:
     28:   This is only needed if you want the preconditioner to automatically update the internal matrix.
     29:   It is called in some `SNES` implementations to update the preconditioner.
     30:   The right-hand side of the linear system is used as function vector.

     32: .seealso: `MatLMVMUpdate()`, `PCLMVMSetMatLMVM()`
     33: @*/
     34: PetscErrorCode PCLMVMSetUpdateVec(PC pc, Vec X)
     35: {
     36:   PC_LMVM  *ctx;
     37:   PetscBool same;

     39:   PetscFunctionBegin;
     42:   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCLMVM, &same));
     43:   if (!same) PetscFunctionReturn(PETSC_SUCCESS);
     44:   ctx = (PC_LMVM *)pc->data;
     45:   PetscCall(PetscObjectReference((PetscObject)X));
     46:   PetscCall(VecDestroy(&ctx->X));
     47:   ctx->X      = X;
     48:   ctx->Xstate = -1;
     49:   PetscFunctionReturn(PETSC_SUCCESS);
     50: }

     52: /*@
     53:   PCLMVMSetMatLMVM - Replaces the `MATLMVM` matrix inside the preconditioner with the one provided by the user.

     55:   Input Parameters:
     56: + pc - An `PCLMVM` preconditioner
     57: - B  - An `MATLMVM` type matrix

     59:   Level: intermediate

     61: .seealso: [](ch_ksp), `PCLMVMGetMatLMVM()`
     62: @*/
     63: PetscErrorCode PCLMVMSetMatLMVM(PC pc, Mat B)
     64: {
     65:   PC_LMVM  *ctx;
     66:   PetscBool same;

     68:   PetscFunctionBegin;
     71:   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCLMVM, &same));
     72:   if (!same) PetscFunctionReturn(PETSC_SUCCESS);
     73:   PetscCall(PetscObjectBaseTypeCompare((PetscObject)B, MATLMVM, &same));
     74:   PetscCheck(same, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Matrix must be an MATLMVM.");
     75:   ctx = (PC_LMVM *)pc->data;
     76:   PetscCall(PetscObjectReference((PetscObject)B));
     77:   PetscCall(MatDestroy(&ctx->B));
     78:   ctx->B = B;
     79:   PetscFunctionReturn(PETSC_SUCCESS);
     80: }

     82: /*@
     83:   PCLMVMGetMatLMVM - Returns a pointer to the underlying `MATLMVM` matrix.

     85:   Input Parameter:
     86: . pc - An `PCLMVM` preconditioner

     88:   Output Parameter:
     89: . B - `MATLMVM` matrix used by the preconditioner

     91:   Level: intermediate

     93: .seealso: [](ch_ksp), `PCLMVMSetMatLMVM()`
     94: @*/
     95: PetscErrorCode PCLMVMGetMatLMVM(PC pc, Mat *B)
     96: {
     97:   PC_LMVM  *ctx;
     98:   PetscBool same;

    100:   PetscFunctionBegin;
    102:   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCLMVM, &same));
    103:   PetscCheck(same, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "PC must be a PCLMVM type.");
    104:   ctx = (PC_LMVM *)pc->data;
    105:   if (!ctx->B) {
    106:     Mat J;

    108:     if (pc->useAmat) J = pc->mat;
    109:     else J = pc->pmat;
    110:     PetscCall(PetscObjectBaseTypeCompare((PetscObject)J, MATLMVM, &same));
    111:     if (same) *B = J;
    112:     else {
    113:       const char *prefix;

    115:       PetscCall(PCGetOptionsPrefix(pc, &prefix));
    116:       PetscCall(MatCreate(PetscObjectComm((PetscObject)pc), &ctx->B));
    117:       PetscCall(MatSetOptionsPrefix(ctx->B, prefix));
    118:       PetscCall(MatAppendOptionsPrefix(ctx->B, "pc_lmvm_"));
    119:       PetscCall(MatSetType(ctx->B, MATLMVMBFGS));
    120:       PetscCall(PetscObjectIncrementTabLevel((PetscObject)ctx->B, (PetscObject)pc, 1));
    121:       *B = ctx->B;
    122:     }
    123:   } else *B = ctx->B;
    124:   PetscFunctionReturn(PETSC_SUCCESS);
    125: }

    127: /*@
    128:   PCLMVMSetIS - Sets the index sets that reduce the `PC` application.

    130:   Input Parameters:
    131: + pc       - An `PCLMVM` preconditioner
    132: - inactive - Index set defining the variables removed from the problem

    134:   Level: intermediate

    136:   Developer Notes:
    137:   Need to explain the purpose of this `IS`

    139: .seealso: [](ch_ksp), `PCLMVMClearIS()`
    140: @*/
    141: PetscErrorCode PCLMVMSetIS(PC pc, IS inactive)
    142: {
    143:   PC_LMVM  *ctx;
    144:   PetscBool same;

    146:   PetscFunctionBegin;
    149:   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCLMVM, &same));
    150:   if (!same) PetscFunctionReturn(PETSC_SUCCESS);
    151:   ctx = (PC_LMVM *)pc->data;
    152:   PetscCall(PCLMVMClearIS(pc));
    153:   PetscCall(PetscObjectReference((PetscObject)inactive));
    154:   ctx->inactive = inactive;
    155:   PetscFunctionReturn(PETSC_SUCCESS);
    156: }

    158: /*@
    159:   PCLMVMClearIS - Removes the inactive variable index set from a `PCLMVM`

    161:   Input Parameter:
    162: . pc - An `PCLMVM` preconditioner

    164:   Level: intermediate

    166: .seealso: [](ch_ksp), `PCLMVMSetIS()`
    167: @*/
    168: PetscErrorCode PCLMVMClearIS(PC pc)
    169: {
    170:   PC_LMVM  *ctx;
    171:   PetscBool same;

    173:   PetscFunctionBegin;
    175:   PetscCall(PetscObjectTypeCompare((PetscObject)pc, PCLMVM, &same));
    176:   if (!same) PetscFunctionReturn(PETSC_SUCCESS);
    177:   ctx = (PC_LMVM *)pc->data;
    178:   PetscCall(ISDestroy(&ctx->inactive));
    179:   PetscFunctionReturn(PETSC_SUCCESS);
    180: }

    182: static PetscErrorCode PCApply_LMVM(PC pc, Vec x, Vec y)
    183: {
    184:   PC_LMVM *ctx = (PC_LMVM *)pc->data;
    185:   Vec      xsub, ysub, Bx = x, By = y;
    186:   Mat      B = ctx->B ? ctx->B : (pc->useAmat ? pc->mat : pc->pmat);

    188:   PetscFunctionBegin;
    189:   if (ctx->inactive) {
    190:     if (!ctx->xwork) PetscCall(MatCreateVecs(B, &ctx->xwork, &ctx->ywork));
    191:     PetscCall(VecZeroEntries(ctx->xwork));
    192:     PetscCall(VecGetSubVector(ctx->xwork, ctx->inactive, &xsub));
    193:     PetscCall(VecCopy(x, xsub));
    194:     PetscCall(VecRestoreSubVector(ctx->xwork, ctx->inactive, &xsub));
    195:     Bx = ctx->xwork;
    196:     By = ctx->ywork;
    197:   }
    198:   PetscCall(MatSolve(B, Bx, By));
    199:   if (ctx->inactive) {
    200:     PetscCall(VecGetSubVector(ctx->ywork, ctx->inactive, &ysub));
    201:     PetscCall(VecCopy(ysub, y));
    202:     PetscCall(VecRestoreSubVector(ctx->ywork, ctx->inactive, &ysub));
    203:   }
    204:   PetscFunctionReturn(PETSC_SUCCESS);
    205: }

    207: static PetscErrorCode PCReset_LMVM(PC pc)
    208: {
    209:   PC_LMVM *ctx = (PC_LMVM *)pc->data;

    211:   PetscFunctionBegin;
    212:   PetscCall(VecDestroy(&ctx->xwork));
    213:   PetscCall(VecDestroy(&ctx->ywork));
    214:   PetscFunctionReturn(PETSC_SUCCESS);
    215: }

    217: static PetscErrorCode PCSetUp_LMVM(PC pc)
    218: {
    219:   PetscInt  n, N;
    220:   PetscBool allocated;
    221:   Mat       B;
    222:   PC_LMVM  *ctx = (PC_LMVM *)pc->data;

    224:   PetscFunctionBegin;
    225:   PetscCall(PCLMVMGetMatLMVM(pc, &B));
    226:   PetscCall(MatLMVMIsAllocated(B, &allocated));
    227:   if (!allocated) {
    228:     Vec t1, t2;

    230:     PetscCall(MatCreateVecs(pc->mat, &t1, &t2));
    231:     PetscCall(VecGetLocalSize(t1, &n));
    232:     PetscCall(VecGetSize(t1, &N));
    233:     PetscCall(MatSetSizes(B, n, n, N, N));
    234:     PetscCall(MatLMVMAllocate(B, t1, t2));
    235:     PetscCall(VecDestroy(&t1));
    236:     PetscCall(VecDestroy(&t2));
    237:   }
    238:   /* Only call SetFromOptions if we internally handle the LMVM matrix */
    239:   if (B == ctx->B && ctx->setfromoptionscalled) PetscCall(MatSetFromOptions(ctx->B));
    240:   ctx->setfromoptionscalled = PETSC_FALSE;
    241:   PetscFunctionReturn(PETSC_SUCCESS);
    242: }

    244: static PetscErrorCode PCView_LMVM(PC pc, PetscViewer viewer)
    245: {
    246:   PC_LMVM  *ctx = (PC_LMVM *)pc->data;
    247:   PetscBool iascii;

    249:   PetscFunctionBegin;
    250:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
    251:   if (iascii && ctx->B && ctx->B->assembled) {
    252:     PetscCall(PetscViewerPushFormat(viewer, PETSC_VIEWER_ASCII_INFO));
    253:     PetscCall(MatView(ctx->B, viewer));
    254:     PetscCall(PetscViewerPopFormat(viewer));
    255:   }
    256:   PetscFunctionReturn(PETSC_SUCCESS);
    257: }

    259: static PetscErrorCode PCSetFromOptions_LMVM(PC pc, PetscOptionItems PetscOptionsObject)
    260: {
    261:   PC_LMVM *ctx = (PC_LMVM *)pc->data;

    263:   PetscFunctionBegin;
    264:   /* defer SetFromOptions calls to PCSetUp_LMVM */
    265:   ctx->setfromoptionscalled = PETSC_TRUE;
    266:   PetscFunctionReturn(PETSC_SUCCESS);
    267: }

    269: static PetscErrorCode PCPreSolve_LMVM(PC pc, KSP ksp, Vec F, Vec X)
    270: {
    271:   PC_LMVM *ctx = (PC_LMVM *)pc->data;

    273:   PetscFunctionBegin;
    274:   if (ctx->X && ctx->B) { /* Perform update only if requested. Otherwise we assume the user, e.g. TAO, has already taken care of it */
    275:     PetscObjectState Xstate;

    277:     PetscCall(PetscObjectStateGet((PetscObject)ctx->X, &Xstate));
    278:     if (ctx->Xstate != Xstate) PetscCall(MatLMVMUpdate(ctx->B, ctx->X, F));
    279:     ctx->Xstate = Xstate;
    280:   }
    281:   PetscFunctionReturn(PETSC_SUCCESS);
    282: }

    284: static PetscErrorCode PCDestroy_LMVM(PC pc)
    285: {
    286:   PC_LMVM *ctx = (PC_LMVM *)pc->data;

    288:   PetscFunctionBegin;
    289:   PetscCall(ISDestroy(&ctx->inactive));
    290:   PetscCall(VecDestroy(&ctx->xwork));
    291:   PetscCall(VecDestroy(&ctx->ywork));
    292:   PetscCall(VecDestroy(&ctx->X));
    293:   PetscCall(MatDestroy(&ctx->B));
    294:   PetscCall(PetscFree(pc->data));
    295:   PetscFunctionReturn(PETSC_SUCCESS);
    296: }

    298: /*MC
    299:    PCLMVM - A preconditioner constructed from a `MATLMVM` matrix.
    300:             If the preconditioner matrix is not of type `MATLMVM`, an internal matrix is used.
    301:             Options for the internal `MATLMVM` matrix can be accessed with the -pc_lmvm_ prefix.
    302:             Alternatively, the user can pass a suitable matrix with `PCLMVMSetMatLMVM()`.

    304:    Level: intermediate

    306: .seealso: [](ch_ksp), `PCCreate()`, `PCSetType()`, `PCLMVMSetUpdateVec()`, `PCLMVMSetMatLMVM()`, `PCLMVMGetMatLMVM()`
    307: M*/
    308: PETSC_EXTERN PetscErrorCode PCCreate_LMVM(PC pc)
    309: {
    310:   PC_LMVM *ctx;

    312:   PetscFunctionBegin;
    313:   PetscCall(PetscNew(&ctx));
    314:   pc->data = (void *)ctx;

    316:   pc->ops->reset               = PCReset_LMVM;
    317:   pc->ops->setup               = PCSetUp_LMVM;
    318:   pc->ops->destroy             = PCDestroy_LMVM;
    319:   pc->ops->view                = PCView_LMVM;
    320:   pc->ops->apply               = PCApply_LMVM;
    321:   pc->ops->setfromoptions      = PCSetFromOptions_LMVM;
    322:   pc->ops->applysymmetricleft  = NULL;
    323:   pc->ops->applysymmetricright = NULL;
    324:   pc->ops->applytranspose      = NULL;
    325:   pc->ops->applyrichardson     = NULL;
    326:   pc->ops->presolve            = PCPreSolve_LMVM;
    327:   pc->ops->postsolve           = NULL;
    328:   PetscFunctionReturn(PETSC_SUCCESS);
    329: }