Date: 2021-03-03T22:26:20+00:00

Actual source code: kaij.c
petsc-3.14.5 2021-03-03
Report Typos and Errors

      2: /*
      3:   Defines the basic matrix operations for the KAIJ  matrix storage format.
      4:   This format is used to evaluate matrices of the form:

      6:     [I \otimes S + A \otimes T]

      8:   where
      9:     S is a dense (p \times q) matrix
     10:     T is a dense (p \times q) matrix
     11:     A is an AIJ  (n \times n) matrix
     12:     I is the identity matrix

     14:   The resulting matrix is (np \times nq)

     16:   We provide:
     17:      MatMult()
     18:      MatMultAdd()
     19:      MatInvertBlockDiagonal()
     20:   and
     21:      MatCreateKAIJ(Mat,PetscInt,PetscInt,const PetscScalar[],const PetscScalar[],Mat*)

     23:   This single directory handles both the sequential and parallel codes
     24: */

     26: #include <../src/mat/impls/kaij/kaij.h>
     27: #include <../src/mat/utils/freespace.h>
     28: #include <petsc/private/vecimpl.h>

     30: /*@C
     31:    MatKAIJGetAIJ - Get the AIJ matrix describing the blockwise action of the KAIJ matrix

     33:    Not Collective, but if the KAIJ matrix is parallel, the AIJ matrix is also parallel

     35:    Input Parameter:
     36: .  A - the KAIJ matrix

     38:    Output Parameter:
     39: .  B - the AIJ matrix

     41:    Level: advanced

     43:    Notes: The reference count on the AIJ matrix is not increased so you should not destroy it.

     45: .seealso: MatCreateKAIJ()
     46: @*/
     47: PetscErrorCode  MatKAIJGetAIJ(Mat A,Mat *B)
     48: {
     50:   PetscBool      ismpikaij,isseqkaij;

     53:   PetscObjectTypeCompare((PetscObject)A,MATMPIKAIJ,&ismpikaij);
     54:   PetscObjectTypeCompare((PetscObject)A,MATSEQKAIJ,&isseqkaij);
     55:   if (ismpikaij) {
     56:     Mat_MPIKAIJ *b = (Mat_MPIKAIJ*)A->data;

     58:     *B = b->A;
     59:   } else if (isseqkaij) {
     60:     Mat_SeqKAIJ *b = (Mat_SeqKAIJ*)A->data;

     62:     *B = b->AIJ;
     63:   } else SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_ARG_WRONG,"Matrix passed in is not of type KAIJ");
     64:   return(0);
     65: }

     67: /*@C
     68:    MatKAIJGetS - Get the S matrix describing the shift action of the KAIJ matrix

     70:    Not Collective; the entire S is stored and returned independently on all processes.

     72:    Input Parameter:
     73: .  A - the KAIJ matrix

     75:    Output Parameters:
     76: +  m - the number of rows in S
     77: .  n - the number of columns in S
     78: -  S - the S matrix, in form of a scalar array in column-major format

     80:    Note: All output parameters are optional (pass NULL or PETSC_IGNORE if not desired)

     82:    Level: advanced

     84: .seealso: MatCreateKAIJ(), MatGetBlockSizes()
     85: @*/
     86: PetscErrorCode MatKAIJGetS(Mat A,PetscInt *m,PetscInt *n,PetscScalar **S)
     87: {
     88:   Mat_SeqKAIJ *b = (Mat_SeqKAIJ*)A->data;
     90:   if (m) *m = b->p;
     91:   if (n) *n = b->q;
     92:   if (S) *S = b->S;
     93:   return(0);
     94: }

     96: /*@C
     97:    MatKAIJGetSRead - Get a read-only pointer to the S matrix describing the shift action of the KAIJ matrix

     99:    Not Collective; the entire S is stored and returned independently on all processes.

    101:    Input Parameter:
    102: .  A - the KAIJ matrix

    104:    Output Parameters:
    105: +  m - the number of rows in S
    106: .  n - the number of columns in S
    107: -  S - the S matrix, in form of a scalar array in column-major format

    109:    Note: All output parameters are optional (pass NULL or PETSC_IGNORE if not desired)

    111:    Level: advanced

    113: .seealso: MatCreateKAIJ(), MatGetBlockSizes()
    114: @*/
    115: PetscErrorCode MatKAIJGetSRead(Mat A,PetscInt *m,PetscInt *n,const PetscScalar **S)
    116: {
    117:   Mat_SeqKAIJ *b = (Mat_SeqKAIJ*)A->data;
    119:   if (m) *m = b->p;
    120:   if (n) *n = b->q;
    121:   if (S) *S = b->S;
    122:   return(0);
    123: }

    125: /*@C
    126:   MatKAIJRestoreS - Restore array obtained with MatKAIJGetS()

    128:   Not collective

    130:   Input Parameter:
    131: . A - the KAIJ matrix

    133:   Output Parameter:
    134: . S - location of pointer to array obtained with MatKAIJGetS()

    136:   Note: This routine zeros the array pointer to prevent accidental reuse after it has been restored.
    137:   If NULL is passed, it will not attempt to zero the array pointer.

    139:   Level: advanced
    140: .seealso: MatKAIJGetS(), MatKAIJGetSRead(), MatKAIJRestoreSRead()
    141: @*/
    142: PetscErrorCode MatKAIJRestoreS(Mat A,PetscScalar **S)
    143: {

    147:   if (S) *S = NULL;
    148:   PetscObjectStateIncrease((PetscObject)A);
    149:   return(0);
    150: }

    152: /*@C
    153:   MatKAIJRestoreSRead - Restore array obtained with MatKAIJGetSRead()

    155:   Not collective

    157:   Input Parameter:
    158: . A - the KAIJ matrix

    160:   Output Parameter:
    161: . S - location of pointer to array obtained with MatKAIJGetS()

    163:   Note: This routine zeros the array pointer to prevent accidental reuse after it has been restored.
    164:   If NULL is passed, it will not attempt to zero the array pointer.

    166:   Level: advanced
    167: .seealso: MatKAIJGetS(), MatKAIJGetSRead(), MatKAIJRestoreSRead()
    168: @*/
    169: PetscErrorCode MatKAIJRestoreSRead(Mat A,const PetscScalar **S)
    170: {
    172:   if (S) *S = NULL;
    173:   return(0);
    174: }

    176: /*@C
    177:    MatKAIJGetT - Get the transformation matrix T associated with the KAIJ matrix

    179:    Not Collective; the entire T is stored and returned independently on all processes

    181:    Input Parameter:
    182: .  A - the KAIJ matrix

    184:    Output Parameter:
    185: +  m - the number of rows in T
    186: .  n - the number of columns in T
    187: -  T - the T matrix, in form of a scalar array in column-major format

    189:    Note: All output parameters are optional (pass NULL or PETSC_IGNORE if not desired)

    191:    Level: advanced

    193: .seealso: MatCreateKAIJ(), MatGetBlockSizes()
    194: @*/
    195: PetscErrorCode MatKAIJGetT(Mat A,PetscInt *m,PetscInt *n,PetscScalar **T)
    196: {
    197:   Mat_SeqKAIJ *b = (Mat_SeqKAIJ*)A->data;
    199:   if (m) *m = b->p;
    200:   if (n) *n = b->q;
    201:   if (T) *T = b->T;
    202:   return(0);
    203: }

    205: /*@C
    206:    MatKAIJGetTRead - Get a read-only pointer to the transformation matrix T associated with the KAIJ matrix

    208:    Not Collective; the entire T is stored and returned independently on all processes

    210:    Input Parameter:
    211: .  A - the KAIJ matrix

    213:    Output Parameter:
    214: +  m - the number of rows in T
    215: .  n - the number of columns in T
    216: -  T - the T matrix, in form of a scalar array in column-major format

    218:    Note: All output parameters are optional (pass NULL or PETSC_IGNORE if not desired)

    220:    Level: advanced

    222: .seealso: MatCreateKAIJ(), MatGetBlockSizes()
    223: @*/
    224: PetscErrorCode MatKAIJGetTRead(Mat A,PetscInt *m,PetscInt *n,const PetscScalar **T)
    225: {
    226:   Mat_SeqKAIJ *b = (Mat_SeqKAIJ*)A->data;
    228:   if (m) *m = b->p;
    229:   if (n) *n = b->q;
    230:   if (T) *T = b->T;
    231:   return(0);
    232: }

    234: /*@C
    235:   MatKAIJRestoreT - Restore array obtained with MatKAIJGetT()

    237:   Not collective

    239:   Input Parameter:
    240: . A - the KAIJ matrix

    242:   Output Parameter:
    243: . T - location of pointer to array obtained with MatKAIJGetS()

    245:   Note: This routine zeros the array pointer to prevent accidental reuse after it has been restored.
    246:   If NULL is passed, it will not attempt to zero the array pointer.

    248:   Level: advanced
    249: .seealso: MatKAIJGetT(), MatKAIJGetTRead(), MatKAIJRestoreTRead()
    250: @*/
    251: PetscErrorCode MatKAIJRestoreT(Mat A,PetscScalar **T)
    252: {

    256:   if (T) *T = NULL;
    257:   PetscObjectStateIncrease((PetscObject)A);
    258:   return(0);
    259: }

    261: /*@C
    262:   MatKAIJRestoreTRead - Restore array obtained with MatKAIJGetTRead()

    264:   Not collective

    266:   Input Parameter:
    267: . A - the KAIJ matrix

    269:   Output Parameter:
    270: . T - location of pointer to array obtained with MatKAIJGetS()

    272:   Note: This routine zeros the array pointer to prevent accidental reuse after it has been restored.
    273:   If NULL is passed, it will not attempt to zero the array pointer.

    275:   Level: advanced
    276: .seealso: MatKAIJGetT(), MatKAIJGetTRead(), MatKAIJRestoreTRead()
    277: @*/
    278: PetscErrorCode MatKAIJRestoreTRead(Mat A,const PetscScalar **T)
    279: {
    281:   if (T) *T = NULL;
    282:   return(0);
    283: }

    285: /*@
    286:    MatKAIJSetAIJ - Set the AIJ matrix describing the blockwise action of the KAIJ matrix

    288:    Logically Collective; if the AIJ matrix is parallel, the KAIJ matrix is also parallel

    290:    Input Parameters:
    291: +  A - the KAIJ matrix
    292: -  B - the AIJ matrix

    294:    Notes:
    295:    This function increases the reference count on the AIJ matrix, so the user is free to destroy the matrix if it is not needed.
    296:    Changes to the entries of the AIJ matrix will immediately affect the KAIJ matrix.

    298:    Level: advanced

    300: .seealso: MatKAIJGetAIJ(), MatKAIJSetS(), MatKAIJSetT()
    301: @*/
    302: PetscErrorCode MatKAIJSetAIJ(Mat A,Mat B)
    303: {
    305:   PetscMPIInt    size;

    308:   MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);
    309:   if (size == 1) {
    310:     Mat_SeqKAIJ *a = (Mat_SeqKAIJ*)A->data;
    311:     a->AIJ = B;
    312:   } else {
    313:     Mat_MPIKAIJ *a = (Mat_MPIKAIJ*)A->data;
    314:     a->A = B;
    315:   }
    316:   PetscObjectReference((PetscObject)B);
    317:   return(0);
    318: }

    320: /*@C
    321:    MatKAIJSetS - Set the S matrix describing the shift action of the KAIJ matrix

    323:    Logically Collective; the entire S is stored independently on all processes.

    325:    Input Parameters:
    326: +  A - the KAIJ matrix
    327: .  p - the number of rows in S
    328: .  q - the number of columns in S
    329: -  S - the S matrix, in form of a scalar array in column-major format

    331:    Notes: The dimensions p and q must match those of the transformation matrix T associated with the KAIJ matrix.
    332:    The S matrix is copied, so the user can destroy this array.

    334:    Level: Advanced

    336: .seealso: MatKAIJGetS(), MatKAIJSetT(), MatKAIJSetAIJ()
    337: @*/
    338: PetscErrorCode MatKAIJSetS(Mat A,PetscInt p,PetscInt q,const PetscScalar S[])
    339: {
    341:   Mat_SeqKAIJ    *a = (Mat_SeqKAIJ*)A->data;

    344:   PetscFree(a->S);
    345:   if (S) {
    346:     PetscMalloc1(p*q*sizeof(PetscScalar),&a->S);
    347:     PetscMemcpy(a->S,S,p*q*sizeof(PetscScalar));
    348:   } else  a->S = NULL;

    350:   a->p = p;
    351:   a->q = q;
    352:   return(0);
    353: }

    355: /*@C
    356:    MatKAIJGetScaledIdentity - Check if both S and T are scaled identities.

    358:    Logically Collective.

    360:    Input Parameter:
    361: .  A - the KAIJ matrix

    363:   Output Parameter:
    364: .  identity - the Boolean value

    366:    Level: Advanced

    368: .seealso: MatKAIJGetS(), MatKAIJGetT()
    369: @*/
    370: PetscErrorCode MatKAIJGetScaledIdentity(Mat A,PetscBool* identity)
    371: {
    372:   Mat_SeqKAIJ *a = (Mat_SeqKAIJ*)A->data;
    373:   PetscInt    i,j;

    376:   if (a->p != a->q) {
    377:     *identity = PETSC_FALSE;
    378:     return(0);
    379:   } else *identity = PETSC_TRUE;
    380:   if (!a->isTI || a->S) {
    381:     for (i=0; i<a->p && *identity; i++) {
    382:       for (j=0; j<a->p && *identity; j++) {
    383:         if (i != j) {
    384:           if (a->S && PetscAbsScalar(a->S[i+j*a->p]) > PETSC_SMALL) *identity = PETSC_FALSE;
    385:           if (a->T && PetscAbsScalar(a->T[i+j*a->p]) > PETSC_SMALL) *identity = PETSC_FALSE;
    386:         } else {
    387:           if (a->S && PetscAbsScalar(a->S[i*(a->p+1)]-a->S[0]) > PETSC_SMALL) *identity = PETSC_FALSE;
    388:           if (a->T && PetscAbsScalar(a->T[i*(a->p+1)]-a->T[0]) > PETSC_SMALL) *identity = PETSC_FALSE;
    389:         }
    390:       }
    391:     }
    392:   }
    393:   return(0);
    394: }

    396: /*@C
    397:    MatKAIJSetT - Set the transformation matrix T associated with the KAIJ matrix

    399:    Logically Collective; the entire T is stored independently on all processes.

    401:    Input Parameters:
    402: +  A - the KAIJ matrix
    403: .  p - the number of rows in S
    404: .  q - the number of columns in S
    405: -  T - the T matrix, in form of a scalar array in column-major format

    407:    Notes: The dimensions p and q must match those of the shift matrix S associated with the KAIJ matrix.
    408:    The T matrix is copied, so the user can destroy this array.

    410:    Level: Advanced

    412: .seealso: MatKAIJGetT(), MatKAIJSetS(), MatKAIJSetAIJ()
    413: @*/
    414: PetscErrorCode MatKAIJSetT(Mat A,PetscInt p,PetscInt q,const PetscScalar T[])
    415: {
    417:   PetscInt       i,j;
    418:   Mat_SeqKAIJ    *a = (Mat_SeqKAIJ*)A->data;
    419:   PetscBool      isTI = PETSC_FALSE;

    422:   /* check if T is an identity matrix */
    423:   if (T && (p == q)) {
    424:     isTI = PETSC_TRUE;
    425:     for (i=0; i<p; i++) {
    426:       for (j=0; j<q; j++) {
    427:         if (i == j) {
    428:           /* diagonal term must be 1 */
    429:           if (T[i+j*p] != 1.0) isTI = PETSC_FALSE;
    430:         } else {
    431:           /* off-diagonal term must be 0 */
    432:           if (T[i+j*p] != 0.0) isTI = PETSC_FALSE;
    433:         }
    434:       }
    435:     }
    436:   }
    437:   a->isTI = isTI;

    439:   PetscFree(a->T);
    440:   if (T && (!isTI)) {
    441:     PetscMalloc1(p*q*sizeof(PetscScalar),&a->T);
    442:     PetscMemcpy(a->T,T,p*q*sizeof(PetscScalar));
    443:   } else a->T = NULL;

    445:   a->p = p;
    446:   a->q = q;
    447:   return(0);
    448: }

    450: PetscErrorCode MatDestroy_SeqKAIJ(Mat A)
    451: {
    453:   Mat_SeqKAIJ    *b = (Mat_SeqKAIJ*)A->data;

    456:   MatDestroy(&b->AIJ);
    457:   PetscFree(b->S);
    458:   PetscFree(b->T);
    459:   PetscFree(b->ibdiag);
    460:   PetscFree5(b->sor.w,b->sor.y,b->sor.work,b->sor.t,b->sor.arr);
    461:   PetscFree(A->data);
    462:   return(0);
    463: }

    465: PetscErrorCode MatSetUp_KAIJ(Mat A)
    466: {
    468:   PetscInt       n;
    469:   PetscMPIInt    size;
    470:   Mat_SeqKAIJ    *seqkaij = (Mat_SeqKAIJ*)A->data;

    473:   MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);
    474:   if (size == 1) {
    475:     MatSetSizes(A,seqkaij->p*seqkaij->AIJ->rmap->n,seqkaij->q*seqkaij->AIJ->cmap->n,seqkaij->p*seqkaij->AIJ->rmap->N,seqkaij->q*seqkaij->AIJ->cmap->N);
    476:     PetscLayoutSetBlockSize(A->rmap,seqkaij->p);
    477:     PetscLayoutSetBlockSize(A->cmap,seqkaij->q);
    478:     PetscLayoutSetUp(A->rmap);
    479:     PetscLayoutSetUp(A->cmap);
    480:   } else {
    481:     Mat_MPIKAIJ *a;
    482:     Mat_MPIAIJ  *mpiaij;
    483:     IS          from,to;
    484:     Vec         gvec;
    485:     PetscScalar *T;
    486:     PetscInt    i,j;

    488:     a = (Mat_MPIKAIJ*)A->data;
    489:     mpiaij = (Mat_MPIAIJ*)a->A->data;
    490:     MatSetSizes(A,a->p*a->A->rmap->n,a->q*a->A->cmap->n,a->p*a->A->rmap->N,a->q*a->A->cmap->N);
    491:     PetscLayoutSetBlockSize(A->rmap,seqkaij->p);
    492:     PetscLayoutSetBlockSize(A->cmap,seqkaij->q);
    493:     PetscLayoutSetUp(A->rmap);
    494:     PetscLayoutSetUp(A->cmap);

    496:     if (a->isTI) {
    497:       /* If the transformation matrix associated with the parallel matrix A is the identity matrix, then a->T will be NULL.
    498:        * In this case, if we pass a->T directly to the MatCreateKAIJ() calls to create the sequential submatrices, the routine will
    499:        * not be able to tell that transformation matrix should be set to the identity; thus we create a temporary identity matrix
    500:        * to pass in. */
    501:       PetscMalloc1(a->p*a->q*sizeof(PetscScalar),&T);
    502:       for (i=0; i<a->p; i++) {
    503:         for (j=0; j<a->q; j++) {
    504:           if (i==j) T[i+j*a->p] = 1.0;
    505:           else      T[i+j*a->p] = 0.0;
    506:         }
    507:       }
    508:     } else T = a->T;
    509:     MatCreateKAIJ(mpiaij->A,a->p,a->q,a->S,T,&a->AIJ);
    510:     MatCreateKAIJ(mpiaij->B,a->p,a->q,NULL,T,&a->OAIJ);
    511:     if (a->isTI) {
    512:       PetscFree(T);
    513:     }

    515:     VecGetSize(mpiaij->lvec,&n);
    516:     VecCreate(PETSC_COMM_SELF,&a->w);
    517:     VecSetSizes(a->w,n*a->q,n*a->q);
    518:     VecSetBlockSize(a->w,a->q);
    519:     VecSetType(a->w,VECSEQ);

    521:     /* create two temporary Index sets for build scatter gather */
    522:     ISCreateBlock(PetscObjectComm((PetscObject)a->A),a->q,n,mpiaij->garray,PETSC_COPY_VALUES,&from);
    523:     ISCreateStride(PETSC_COMM_SELF,n*a->q,0,1,&to);

    525:     /* create temporary global vector to generate scatter context */
    526:     VecCreateMPIWithArray(PetscObjectComm((PetscObject)a->A),a->q,a->q*a->A->cmap->n,a->q*a->A->cmap->N,NULL,&gvec);

    528:     /* generate the scatter context */
    529:     VecScatterCreate(gvec,from,a->w,to,&a->ctx);

    531:     ISDestroy(&from);
    532:     ISDestroy(&to);
    533:     VecDestroy(&gvec);
    534:   }

    536:   A->assembled = PETSC_TRUE;
    537:   return(0);
    538: }

    540: PetscErrorCode MatView_KAIJ(Mat A,PetscViewer viewer)
    541: {
    542:   PetscViewerFormat format;
    543:   Mat_SeqKAIJ       *a = (Mat_SeqKAIJ*)A->data;
    544:   Mat               B;
    545:   PetscInt          i;
    546:   PetscErrorCode    ierr;
    547:   PetscBool         ismpikaij;

    550:   PetscObjectTypeCompare((PetscObject)A,MATMPIKAIJ,&ismpikaij);
    551:   PetscViewerGetFormat(viewer,&format);
    552:   if (format == PETSC_VIEWER_ASCII_INFO || format == PETSC_VIEWER_ASCII_INFO_DETAIL || format == PETSC_VIEWER_ASCII_IMPL) {
    553:     PetscViewerASCIIPrintf(viewer,"S and T have %D rows and %D columns\n",a->p,a->q);

    555:     /* Print appropriate details for S. */
    556:     if (!a->S) {
    557:       PetscViewerASCIIPrintf(viewer,"S is NULL\n");
    558:     } else if (format == PETSC_VIEWER_ASCII_IMPL) {
    559:       PetscViewerASCIIPrintf(viewer,"Entries of S are ");
    560:       for (i=0; i<(a->p * a->q); i++) {
    561: #if defined(PETSC_USE_COMPLEX)
    562:         PetscViewerASCIIPrintf(viewer,"%18.16e %18.16e ",(double)PetscRealPart(a->S[i]),(double)PetscImaginaryPart(a->S[i]));
    563: #else
    564:         PetscViewerASCIIPrintf(viewer,"%18.16e ",(double)PetscRealPart(a->S[i]));
    565: #endif
    566:       }
    567:       PetscViewerASCIIPrintf(viewer,"\n");
    568:     }

    570:     /* Print appropriate details for T. */
    571:     if (a->isTI) {
    572:       PetscViewerASCIIPrintf(viewer,"T is the identity matrix\n");
    573:     } else if (!a->T) {
    574:       PetscViewerASCIIPrintf(viewer,"T is NULL\n");
    575:     } else if (format == PETSC_VIEWER_ASCII_IMPL) {
    576:       PetscViewerASCIIPrintf(viewer,"Entries of T are ");
    577:       for (i=0; i<(a->p * a->q); i++) {
    578: #if defined(PETSC_USE_COMPLEX)
    579:         PetscViewerASCIIPrintf(viewer,"%18.16e %18.16e ",(double)PetscRealPart(a->T[i]),(double)PetscImaginaryPart(a->T[i]));
    580: #else
    581:         PetscViewerASCIIPrintf(viewer,"%18.16e ",(double)PetscRealPart(a->T[i]));
    582: #endif
    583:       }
    584:       PetscViewerASCIIPrintf(viewer,"\n");
    585:     }

    587:     /* Now print details for the AIJ matrix, using the AIJ viewer. */
    588:     PetscViewerASCIIPrintf(viewer,"Now viewing the associated AIJ matrix:\n");
    589:     if (ismpikaij) {
    590:       Mat_MPIKAIJ *b = (Mat_MPIKAIJ*)A->data;
    591:       MatView(b->A,viewer);
    592:     } else {
    593:       MatView(a->AIJ,viewer);
    594:     }

    596:   } else {
    597:     /* For all other matrix viewer output formats, simply convert to an AIJ matrix and call MatView() on that. */
    598:     if (ismpikaij) {
    599:       MatConvert(A,MATMPIAIJ,MAT_INITIAL_MATRIX,&B);
    600:     } else {
    601:       MatConvert(A,MATSEQAIJ,MAT_INITIAL_MATRIX,&B);
    602:     }
    603:     MatView(B,viewer);
    604:     MatDestroy(&B);
    605:   }
    606:   return(0);
    607: }

    609: PetscErrorCode MatDestroy_MPIKAIJ(Mat A)
    610: {
    612:   Mat_MPIKAIJ    *b = (Mat_MPIKAIJ*)A->data;

    615:   MatDestroy(&b->AIJ);
    616:   MatDestroy(&b->OAIJ);
    617:   MatDestroy(&b->A);
    618:   VecScatterDestroy(&b->ctx);
    619:   VecDestroy(&b->w);
    620:   PetscFree(b->S);
    621:   PetscFree(b->T);
    622:   PetscFree(b->ibdiag);
    623:   PetscFree(A->data);
    624:   return(0);
    625: }

    627: /* --------------------------------------------------------------------------------------*/

    629: /* zz = yy + Axx */
    630: PetscErrorCode MatMultAdd_SeqKAIJ(Mat A,Vec xx,Vec yy,Vec zz)
    631: {
    632:   Mat_SeqKAIJ       *b = (Mat_SeqKAIJ*)A->data;
    633:   Mat_SeqAIJ        *a = (Mat_SeqAIJ*)b->AIJ->data;
    634:   const PetscScalar *s = b->S, *t = b->T;
    635:   const PetscScalar *x,*v,*bx;
    636:   PetscScalar       *y,*sums;
    637:   PetscErrorCode    ierr;
    638:   const PetscInt    m = b->AIJ->rmap->n,*idx,*ii;
    639:   PetscInt          n,i,jrow,j,l,p=b->p,q=b->q,k;

    642:   if (!yy) {
    643:     VecSet(zz,0.0);
    644:   } else {
    645:     VecCopy(yy,zz);
    646:   }
    647:   if ((!s) && (!t) && (!b->isTI)) return(0);

    649:   VecGetArrayRead(xx,&x);
    650:   VecGetArray(zz,&y);
    651:   idx  = a->j;
    652:   v    = a->a;
    653:   ii   = a->i;

    655:   if (b->isTI) {
    656:     for (i=0; i<m; i++) {
    657:       jrow = ii[i];
    658:       n    = ii[i+1] - jrow;
    659:       sums = y + p*i;
    660:       for (j=0; j<n; j++) {
    661:         for (k=0; k<p; k++) {
    662:           sums[k] += v[jrow+j]*x[q*idx[jrow+j]+k];
    663:         }
    664:       }
    665:     }
    666:     PetscLogFlops(3.0*(a->nz)*p);
    667:   } else if (t) {
    668:     for (i=0; i<m; i++) {
    669:       jrow = ii[i];
    670:       n    = ii[i+1] - jrow;
    671:       sums = y + p*i;
    672:       for (j=0; j<n; j++) {
    673:         for (k=0; k<p; k++) {
    674:           for (l=0; l<q; l++) {
    675:             sums[k] += v[jrow+j]*t[k+l*p]*x[q*idx[jrow+j]+l];
    676:           }
    677:         }
    678:       }
    679:     }
    680:     /* The flop count below assumes that v[jrow+j] is hoisted out (which an optimizing compiler is likely to do),
    681:      * and also that T part is hoisted outside this loop (in exchange for temporary storage) as (A \otimes I) (I \otimes T),
    682:      * so that this multiply doesn't have to be redone for each matrix entry, but just once per column. The latter
    683:      * transformation is much less likely to be applied, but we nonetheless count the minimum flops required. */
    684:     PetscLogFlops((2.0*p*q-p)*m+2.0*p*a->nz);
    685:   }
    686:   if (s) {
    687:     for (i=0; i<m; i++) {
    688:       sums = y + p*i;
    689:       bx   = x + q*i;
    690:       if (i < b->AIJ->cmap->n) {
    691:         for (j=0; j<q; j++) {
    692:           for (k=0; k<p; k++) {
    693:             sums[k] += s[k+j*p]*bx[j];
    694:           }
    695:         }
    696:       }
    697:     }
    698:     PetscLogFlops(2.0*m*p*q);
    699:   }

    701:   VecRestoreArrayRead(xx,&x);
    702:   VecRestoreArray(zz,&y);
    703:   return(0);
    704: }

    706: PetscErrorCode MatMult_SeqKAIJ(Mat A,Vec xx,Vec yy)
    707: {
    710:   MatMultAdd_SeqKAIJ(A,xx,PETSC_NULL,yy);
    711:   return(0);
    712: }

    714: #include <petsc/private/kernels/blockinvert.h>

    716: PetscErrorCode MatInvertBlockDiagonal_SeqKAIJ(Mat A,const PetscScalar **values)
    717: {
    718:   Mat_SeqKAIJ       *b  = (Mat_SeqKAIJ*)A->data;
    719:   Mat_SeqAIJ        *a  = (Mat_SeqAIJ*)b->AIJ->data;
    720:   const PetscScalar *S  = b->S;
    721:   const PetscScalar *T  = b->T;
    722:   const PetscScalar *v  = a->a;
    723:   const PetscInt     p  = b->p, q = b->q, m = b->AIJ->rmap->n, *idx = a->j, *ii = a->i;
    724:   PetscErrorCode    ierr;
    725:   PetscInt          i,j,*v_pivots,dof,dof2;
    726:   PetscScalar       *diag,aval,*v_work;

    729:   if (p != q) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_SUP,"MATKAIJ: Block size must be square to calculate inverse.");
    730:   if ((!S) && (!T) && (!b->isTI)) SETERRQ(PetscObjectComm((PetscObject)A),PETSC_ERR_SUP,"MATKAIJ: Cannot invert a zero matrix.");

    732:   dof  = p;
    733:   dof2 = dof*dof;

    735:   if (b->ibdiagvalid) {
    736:     if (values) *values = b->ibdiag;
    737:     return(0);
    738:   }
    739:   if (!b->ibdiag) {
    740:     PetscMalloc1(dof2*m*sizeof(PetscScalar),&b->ibdiag);
    741:     PetscLogObjectMemory((PetscObject)A,dof2*m*sizeof(PetscScalar));
    742:   }
    743:   if (values) *values = b->ibdiag;
    744:   diag = b->ibdiag;

    746:   PetscMalloc2(dof,&v_work,dof,&v_pivots);
    747:   for (i=0; i<m; i++) {
    748:     if (S) {
    749:       PetscMemcpy(diag,S,dof2*sizeof(PetscScalar));
    750:     } else {
    751:       PetscMemzero(diag,dof2*sizeof(PetscScalar));
    752:     }
    753:     if (b->isTI) {
    754:       aval = 0;
    755:       for (j=ii[i]; j<ii[i+1]; j++) if (idx[j] == i) aval = v[j];
    756:       for (j=0; j<dof; j++) diag[j+dof*j] += aval;
    757:     } else if (T) {
    758:       aval = 0;
    759:       for (j=ii[i]; j<ii[i+1]; j++) if (idx[j] == i) aval = v[j];
    760:       for (j=0; j<dof2; j++) diag[j] += aval*T[j];
    761:     }
    762:     PetscKernel_A_gets_inverse_A(dof,diag,v_pivots,v_work,PETSC_FALSE,NULL);
    763:     diag += dof2;
    764:   }
    765:   PetscFree2(v_work,v_pivots);

    767:   b->ibdiagvalid = PETSC_TRUE;
    768:   return(0);
    769: }

    771: static PetscErrorCode MatGetDiagonalBlock_MPIKAIJ(Mat A,Mat *B)
    772: {
    773:   Mat_MPIKAIJ *kaij = (Mat_MPIKAIJ*) A->data;

    776:   *B = kaij->AIJ;
    777:   return(0);
    778: }

    780: PetscErrorCode MatSOR_SeqKAIJ(Mat A,Vec bb,PetscReal omega,MatSORType flag,PetscReal fshift,PetscInt its,PetscInt lits,Vec xx)
    781: {
    782:   PetscErrorCode    ierr;
    783:   Mat_SeqKAIJ       *kaij = (Mat_SeqKAIJ*) A->data;
    784:   Mat_SeqAIJ        *a = (Mat_SeqAIJ*)kaij->AIJ->data;
    785:   const PetscScalar *aa = a->a, *T = kaij->T, *v;
    786:   const PetscInt    m  = kaij->AIJ->rmap->n, *ai=a->i, *aj=a->j, p = kaij->p, q = kaij->q, *diag, *vi;
    787:   const PetscScalar *b, *xb, *idiag;
    788:   PetscScalar       *x, *work, *workt, *w, *y, *arr, *t, *arrt;
    789:   PetscInt          i, j, k, i2, bs, bs2, nz;

    792:   its = its*lits;
    793:   if (flag & SOR_EISENSTAT) SETERRQ (PETSC_COMM_SELF,PETSC_ERR_SUP,"No support yet for Eisenstat");
    794:   if (its <= 0)             SETERRQ2(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Relaxation requires global its %D and local its %D both positive",its,lits);
    795:   if (fshift)               SETERRQ (PETSC_COMM_SELF,PETSC_ERR_SUP,"No support for diagonal shift");
    796:   if ((flag & SOR_APPLY_UPPER) || (flag & SOR_APPLY_LOWER)) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No support for applying upper or lower triangular parts");
    797:   if (p != q) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"MatSOR for KAIJ: No support for non-square dense blocks");
    798:   else        {bs = p; bs2 = bs*bs; }

    800:   if (!m) return(0);

    802:   if (!kaij->ibdiagvalid) { MatInvertBlockDiagonal_SeqKAIJ(A,NULL); }
    803:   idiag = kaij->ibdiag;
    804:   diag  = a->diag;

    806:   if (!kaij->sor.setup) {
    807:     PetscMalloc5(bs,&kaij->sor.w,bs,&kaij->sor.y,m*bs,&kaij->sor.work,m*bs,&kaij->sor.t,m*bs2,&kaij->sor.arr);
    808:     kaij->sor.setup = PETSC_TRUE;
    809:   }
    810:   y     = kaij->sor.y;
    811:   w     = kaij->sor.w;
    812:   work  = kaij->sor.work;
    813:   t     = kaij->sor.t;
    814:   arr   = kaij->sor.arr;

    816:   VecGetArray(xx,&x);
    817:   VecGetArrayRead(bb,&b);

    819:   if (flag & SOR_ZERO_INITIAL_GUESS) {
    820:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
    821:       PetscKernel_w_gets_Ar_times_v(bs,bs,b,idiag,x);                            /* x[0:bs] <- D^{-1} b[0:bs] */
    822:        PetscMemcpy(t,b,bs*sizeof(PetscScalar));
    823:       i2     =  bs;
    824:       idiag  += bs2;
    825:       for (i=1; i<m; i++) {
    826:         v  = aa + ai[i];
    827:         vi = aj + ai[i];
    828:         nz = diag[i] - ai[i];

    830:         if (T) {                /* b - T (Arow * x) */
    831:           PetscMemzero(w,bs*sizeof(PetscScalar));
    832:           for (j=0; j<nz; j++) {
    833:             for (k=0; k<bs; k++) w[k] -= v[j] * x[vi[j]*bs+k];
    834:           }
    835:           PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs,w,T,&t[i2]);
    836:           for (k=0; k<bs; k++) t[i2+k] += b[i2+k];
    837:         } else if (kaij->isTI) {
    838:           PetscMemcpy(t+i2,b+i2,bs*sizeof(PetscScalar));
    839:           for (j=0; j<nz; j++) {
    840:             for (k=0; k<bs; k++) t[i2+k] -= v[j] * x[vi[j]*bs+k];
    841:           }
    842:         } else {
    843:           PetscMemcpy(t+i2,b+i2,bs*sizeof(PetscScalar));
    844:         }

    846:         PetscKernel_w_gets_Ar_times_v(bs,bs,t+i2,idiag,y);
    847:         for (j=0; j<bs; j++) x[i2+j] = omega * y[j];

    849:         idiag += bs2;
    850:         i2    += bs;
    851:       }
    852:       /* for logging purposes assume number of nonzero in lower half is 1/2 of total */
    853:       PetscLogFlops(1.0*bs2*a->nz);
    854:       xb = t;
    855:     } else xb = b;
    856:     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
    857:       idiag = kaij->ibdiag+bs2*(m-1);
    858:       i2    = bs * (m-1);
    859:       PetscMemcpy(w,xb+i2,bs*sizeof(PetscScalar));
    860:       PetscKernel_w_gets_Ar_times_v(bs,bs,w,idiag,x+i2);
    861:       i2    -= bs;
    862:       idiag -= bs2;
    863:       for (i=m-2; i>=0; i--) {
    864:         v  = aa + diag[i] + 1 ;
    865:         vi = aj + diag[i] + 1;
    866:         nz = ai[i+1] - diag[i] - 1;

    868:         if (T) {                /* FIXME: This branch untested */
    869:           PetscMemcpy(w,xb+i2,bs*sizeof(PetscScalar));
    870:           /* copy all rows of x that are needed into contiguous space */
    871:           workt = work;
    872:           for (j=0; j<nz; j++) {
    873:             PetscMemcpy(workt,x + bs*(*vi++),bs*sizeof(PetscScalar));
    874:             workt += bs;
    875:           }
    876:           arrt = arr;
    877:           for (j=0; j<nz; j++) {
    878:             PetscMemcpy(arrt,T,bs2*sizeof(PetscScalar));
    879:             for (k=0; k<bs2; k++) arrt[k] *= v[j];
    880:             arrt += bs2;
    881:           }
    882:           PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    883:         } else if (kaij->isTI) {
    884:           PetscMemcpy(w,t+i2,bs*sizeof(PetscScalar));
    885:           for (j=0; j<nz; j++) {
    886:             for (k=0; k<bs; k++) w[k] -= v[j] * x[vi[j]*bs+k];
    887:           }
    888:         }

    890:         PetscKernel_w_gets_Ar_times_v(bs,bs,w,idiag,y); /* RHS incorrect for omega != 1.0 */
    891:         for (j=0; j<bs; j++) x[i2+j] = (1.0-omega) * x[i2+j] + omega * y[j];

    893:         idiag -= bs2;
    894:         i2    -= bs;
    895:       }
    896:       PetscLogFlops(1.0*bs2*(a->nz));
    897:     }
    898:     its--;
    899:   }
    900:   while (its--) {               /* FIXME: This branch not updated */
    901:     if (flag & SOR_FORWARD_SWEEP || flag & SOR_LOCAL_FORWARD_SWEEP) {
    902:       i2     =  0;
    903:       idiag  = kaij->ibdiag;
    904:       for (i=0; i<m; i++) {
    905:         PetscMemcpy(w,b+i2,bs*sizeof(PetscScalar));

    907:         v  = aa + ai[i];
    908:         vi = aj + ai[i];
    909:         nz = diag[i] - ai[i];
    910:         workt = work;
    911:         for (j=0; j<nz; j++) {
    912:           PetscMemcpy(workt,x + bs*(*vi++),bs*sizeof(PetscScalar));
    913:           workt += bs;
    914:         }
    915:         arrt = arr;
    916:         if (T) {
    917:           for (j=0; j<nz; j++) {
    918:             PetscMemcpy(arrt,T,bs2*sizeof(PetscScalar));
    919:             for (k=0; k<bs2; k++) arrt[k] *= v[j];
    920:             arrt += bs2;
    921:           }
    922:           PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    923:         } else if (kaij->isTI) {
    924:           for (j=0; j<nz; j++) {
    925:             PetscMemzero(arrt,bs2*sizeof(PetscScalar));
    926:             for (k=0; k<bs; k++) arrt[k+bs*k] = v[j];
    927:             arrt += bs2;
    928:           }
    929:           PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    930:         }
    931:         PetscMemcpy(t+i2,w,bs*sizeof(PetscScalar));

    933:         v  = aa + diag[i] + 1;
    934:         vi = aj + diag[i] + 1;
    935:         nz = ai[i+1] - diag[i] - 1;
    936:         workt = work;
    937:         for (j=0; j<nz; j++) {
    938:           PetscMemcpy(workt,x + bs*(*vi++),bs*sizeof(PetscScalar));
    939:           workt += bs;
    940:         }
    941:         arrt = arr;
    942:         if (T) {
    943:           for (j=0; j<nz; j++) {
    944:             PetscMemcpy(arrt,T,bs2*sizeof(PetscScalar));
    945:             for (k=0; k<bs2; k++) arrt[k] *= v[j];
    946:             arrt += bs2;
    947:           }
    948:           PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    949:         } else if (kaij->isTI) {
    950:           for (j=0; j<nz; j++) {
    951:             PetscMemzero(arrt,bs2*sizeof(PetscScalar));
    952:             for (k=0; k<bs; k++) arrt[k+bs*k] = v[j];
    953:             arrt += bs2;
    954:           }
    955:           PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    956:         }

    958:         PetscKernel_w_gets_Ar_times_v(bs,bs,w,idiag,y);
    959:         for (j=0; j<bs; j++) *(x+i2+j) = (1.0-omega) * *(x+i2+j) + omega * *(y+j);

    961:         idiag += bs2;
    962:         i2    += bs;
    963:       }
    964:       xb = t;
    965:     }
    966:     else xb = b;
    967:     if (flag & SOR_BACKWARD_SWEEP || flag & SOR_LOCAL_BACKWARD_SWEEP) {
    968:       idiag = kaij->ibdiag+bs2*(m-1);
    969:       i2    = bs * (m-1);
    970:       if (xb == b) {
    971:         for (i=m-1; i>=0; i--) {
    972:           PetscMemcpy(w,b+i2,bs*sizeof(PetscScalar));

    974:           v  = aa + ai[i];
    975:           vi = aj + ai[i];
    976:           nz = diag[i] - ai[i];
    977:           workt = work;
    978:           for (j=0; j<nz; j++) {
    979:             PetscMemcpy(workt,x + bs*(*vi++),bs*sizeof(PetscScalar));
    980:             workt += bs;
    981:           }
    982:           arrt = arr;
    983:           if (T) {
    984:             for (j=0; j<nz; j++) {
    985:               PetscMemcpy(arrt,T,bs2*sizeof(PetscScalar));
    986:               for (k=0; k<bs2; k++) arrt[k] *= v[j];
    987:               arrt += bs2;
    988:             }
    989:             PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    990:           } else if (kaij->isTI) {
    991:             for (j=0; j<nz; j++) {
    992:               PetscMemzero(arrt,bs2*sizeof(PetscScalar));
    993:               for (k=0; k<bs; k++) arrt[k+bs*k] = v[j];
    994:               arrt += bs2;
    995:             }
    996:             PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    997:           }

    999:           v  = aa + diag[i] + 1;
    1000:           vi = aj + diag[i] + 1;
    1001:           nz = ai[i+1] - diag[i] - 1;
    1002:           workt = work;
    1003:           for (j=0; j<nz; j++) {
    1004:             PetscMemcpy(workt,x + bs*(*vi++),bs*sizeof(PetscScalar));
    1005:             workt += bs;
    1006:           }
    1007:           arrt = arr;
    1008:           if (T) {
    1009:             for (j=0; j<nz; j++) {
    1010:               PetscMemcpy(arrt,T,bs2*sizeof(PetscScalar));
    1011:               for (k=0; k<bs2; k++) arrt[k] *= v[j];
    1012:               arrt += bs2;
    1013:             }
    1014:             PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    1015:           } else if (kaij->isTI) {
    1016:             for (j=0; j<nz; j++) {
    1017:               PetscMemzero(arrt,bs2*sizeof(PetscScalar));
    1018:               for (k=0; k<bs; k++) arrt[k+bs*k] = v[j];
    1019:               arrt += bs2;
    1020:             }
    1021:             PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    1022:           }

    1024:           PetscKernel_w_gets_Ar_times_v(bs,bs,w,idiag,y);
    1025:           for (j=0; j<bs; j++) *(x+i2+j) = (1.0-omega) * *(x+i2+j) + omega * *(y+j);
    1026:         }
    1027:       } else {
    1028:         for (i=m-1; i>=0; i--) {
    1029:           PetscMemcpy(w,xb+i2,bs*sizeof(PetscScalar));
    1030:           v  = aa + diag[i] + 1;
    1031:           vi = aj + diag[i] + 1;
    1032:           nz = ai[i+1] - diag[i] - 1;
    1033:           workt = work;
    1034:           for (j=0; j<nz; j++) {
    1035:             PetscMemcpy(workt,x + bs*(*vi++),bs*sizeof(PetscScalar));
    1036:             workt += bs;
    1037:           }
    1038:           arrt = arr;
    1039:           if (T) {
    1040:             for (j=0; j<nz; j++) {
    1041:               PetscMemcpy(arrt,T,bs2*sizeof(PetscScalar));
    1042:               for (k=0; k<bs2; k++) arrt[k] *= v[j];
    1043:               arrt += bs2;
    1044:             }
    1045:             PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    1046:           } else if (kaij->isTI) {
    1047:             for (j=0; j<nz; j++) {
    1048:               PetscMemzero(arrt,bs2*sizeof(PetscScalar));
    1049:               for (k=0; k<bs; k++) arrt[k+bs*k] = v[j];
    1050:               arrt += bs2;
    1051:             }
    1052:             PetscKernel_w_gets_w_minus_Ar_times_v(bs,bs*nz,w,arr,work);
    1053:           }
    1054:           PetscKernel_w_gets_Ar_times_v(bs,bs,w,idiag,y);
    1055:           for (j=0; j<bs; j++) *(x+i2+j) = (1.0-omega) * *(x+i2+j) + omega * *(y+j);
    1056:         }
    1057:       }
    1058:       PetscLogFlops(1.0*bs2*(a->nz));
    1059:     }
    1060:   }

    1062:   VecRestoreArray(xx,&x);
    1063:   VecRestoreArrayRead(bb,&b);
    1064:   return(0);
    1065: }

    1067: /*===================================================================================*/

    1069: PetscErrorCode MatMultAdd_MPIKAIJ(Mat A,Vec xx,Vec yy,Vec zz)
    1070: {
    1071:   Mat_MPIKAIJ    *b = (Mat_MPIKAIJ*)A->data;

    1075:   if (!yy) {
    1076:     VecSet(zz,0.0);
    1077:   } else {
    1078:     VecCopy(yy,zz);
    1079:   }
    1080:   /* start the scatter */
    1081:   VecScatterBegin(b->ctx,xx,b->w,INSERT_VALUES,SCATTER_FORWARD);
    1082:   (*b->AIJ->ops->multadd)(b->AIJ,xx,zz,zz);
    1083:   VecScatterEnd(b->ctx,xx,b->w,INSERT_VALUES,SCATTER_FORWARD);
    1084:   (*b->OAIJ->ops->multadd)(b->OAIJ,b->w,zz,zz);
    1085:   return(0);
    1086: }

    1088: PetscErrorCode MatMult_MPIKAIJ(Mat A,Vec xx,Vec yy)
    1089: {
    1092:   MatMultAdd_MPIKAIJ(A,xx,PETSC_NULL,yy);
    1093:   return(0);
    1094: }

    1096: PetscErrorCode MatInvertBlockDiagonal_MPIKAIJ(Mat A,const PetscScalar **values)
    1097: {
    1098:   Mat_MPIKAIJ     *b = (Mat_MPIKAIJ*)A->data;
    1099:   PetscErrorCode  ierr;

    1102:   (*b->AIJ->ops->invertblockdiagonal)(b->AIJ,values);
    1103:   return(0);
    1104: }

    1106: /* ----------------------------------------------------------------*/

    1108: PetscErrorCode MatGetRow_SeqKAIJ(Mat A,PetscInt row,PetscInt *ncols,PetscInt **cols,PetscScalar **values)
    1109: {
    1110:   Mat_SeqKAIJ     *b   = (Mat_SeqKAIJ*) A->data;
    1111:   PetscErrorCode  diag = PETSC_FALSE;
    1112:   PetscErrorCode  ierr;
    1113:   PetscInt        nzaij,nz,*colsaij,*idx,i,j,p=b->p,q=b->q,r=row/p,s=row%p,c;
    1114:   PetscScalar     *vaij,*v,*S=b->S,*T=b->T;

    1117:   if (b->getrowactive) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Already active");
    1118:   b->getrowactive = PETSC_TRUE;
    1119:   if (row < 0 || row >= A->rmap->n) SETERRQ1(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Row %D out of range",row);

    1121:   if ((!S) && (!T) && (!b->isTI)) {
    1122:     if (ncols)    *ncols  = 0;
    1123:     if (cols)     *cols   = NULL;
    1124:     if (values)   *values = NULL;
    1125:     return(0);
    1126:   }

    1128:   if (T || b->isTI) {
    1129:     MatGetRow_SeqAIJ(b->AIJ,r,&nzaij,&colsaij,&vaij);
    1130:     c     = nzaij;
    1131:     for (i=0; i<nzaij; i++) {
    1132:       /* check if this row contains a diagonal entry */
    1133:       if (colsaij[i] == r) {
    1134:         diag = PETSC_TRUE;
    1135:         c = i;
    1136:       }
    1137:     }
    1138:   } else nzaij = c = 0;

    1140:   /* calculate size of row */
    1141:   nz = 0;
    1142:   if (S)            nz += q;
    1143:   if (T || b->isTI) nz += (diag && S ? (nzaij-1)*q : nzaij*q);

    1145:   if (cols || values) {
    1146:     PetscMalloc2(nz,&idx,nz,&v);
    1147:     for (i=0; i<q; i++) {
    1148:       /* We need to initialize the v[i] to zero to handle the case in which T is NULL (not the identity matrix). */
    1149:       v[i] = 0.0;
    1150:     }
    1151:     if (b->isTI) {
    1152:       for (i=0; i<nzaij; i++) {
    1153:         for (j=0; j<q; j++) {
    1154:           idx[i*q+j] = colsaij[i]*q+j;
    1155:           v[i*q+j]   = (j==s ? vaij[i] : 0);
    1156:         }
    1157:       }
    1158:     } else if (T) {
    1159:       for (i=0; i<nzaij; i++) {
    1160:         for (j=0; j<q; j++) {
    1161:           idx[i*q+j] = colsaij[i]*q+j;
    1162:           v[i*q+j]   = vaij[i]*T[s+j*p];
    1163:         }
    1164:       }
    1165:     }
    1166:     if (S) {
    1167:       for (j=0; j<q; j++) {
    1168:         idx[c*q+j] = r*q+j;
    1169:         v[c*q+j]  += S[s+j*p];
    1170:       }
    1171:     }
    1172:   }

    1174:   if (ncols)    *ncols  = nz;
    1175:   if (cols)     *cols   = idx;
    1176:   if (values)   *values = v;
    1177:   return(0);
    1178: }

    1180: PetscErrorCode MatRestoreRow_SeqKAIJ(Mat A,PetscInt row,PetscInt *nz,PetscInt **idx,PetscScalar **v)
    1181: {
    1184:   PetscFree2(*idx,*v);
    1185:   ((Mat_SeqKAIJ*)A->data)->getrowactive = PETSC_FALSE;
    1186:   return(0);
    1187: }

    1189: PetscErrorCode MatGetRow_MPIKAIJ(Mat A,PetscInt row,PetscInt *ncols,PetscInt **cols,PetscScalar **values)
    1190: {
    1191:   Mat_MPIKAIJ     *b      = (Mat_MPIKAIJ*) A->data;
    1192:   Mat             MatAIJ  = ((Mat_SeqKAIJ*)b->AIJ->data)->AIJ;
    1193:   Mat             MatOAIJ = ((Mat_SeqKAIJ*)b->OAIJ->data)->AIJ;
    1194:   Mat             AIJ     = b->A;
    1195:   PetscBool       diag    = PETSC_FALSE;
    1196:   PetscErrorCode  ierr;
    1197:   const PetscInt  rstart=A->rmap->rstart,rend=A->rmap->rend,p=b->p,q=b->q,*garray;
    1198:   PetscInt        nz,*idx,ncolsaij = 0,ncolsoaij = 0,*colsaij,*colsoaij,r,s,c,i,j,lrow;
    1199:   PetscScalar     *v,*vals,*ovals,*S=b->S,*T=b->T;

    1202:   if (b->getrowactive) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Already active");
    1203:   b->getrowactive = PETSC_TRUE;
    1204:   if (row < rstart || row >= rend) SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Only local rows");
    1205:   lrow = row - rstart;

    1207:   if ((!S) && (!T) && (!b->isTI)) {
    1208:     if (ncols)    *ncols  = 0;
    1209:     if (cols)     *cols   = NULL;
    1210:     if (values)   *values = NULL;
    1211:     return(0);
    1212:   }

    1214:   r = lrow/p;
    1215:   s = lrow%p;

    1217:   if (T || b->isTI) {
    1218:     MatMPIAIJGetSeqAIJ(AIJ,NULL,NULL,&garray);
    1219:     MatGetRow_SeqAIJ(MatAIJ,lrow/p,&ncolsaij,&colsaij,&vals);
    1220:     MatGetRow_SeqAIJ(MatOAIJ,lrow/p,&ncolsoaij,&colsoaij,&ovals);

    1222:     c     = ncolsaij + ncolsoaij;
    1223:     for (i=0; i<ncolsaij; i++) {
    1224:       /* check if this row contains a diagonal entry */
    1225:       if (colsaij[i] == r) {
    1226:         diag = PETSC_TRUE;
    1227:         c = i;
    1228:       }
    1229:     }
    1230:   } else c = 0;

    1232:   /* calculate size of row */
    1233:   nz = 0;
    1234:   if (S)            nz += q;
    1235:   if (T || b->isTI) nz += (diag && S ? (ncolsaij+ncolsoaij-1)*q : (ncolsaij+ncolsoaij)*q);

    1237:   if (cols || values) {
    1238:     PetscMalloc2(nz,&idx,nz,&v);
    1239:     for (i=0; i<q; i++) {
    1240:       /* We need to initialize the v[i] to zero to handle the case in which T is NULL (not the identity matrix). */
    1241:       v[i] = 0.0;
    1242:     }
    1243:     if (b->isTI) {
    1244:       for (i=0; i<ncolsaij; i++) {
    1245:         for (j=0; j<q; j++) {
    1246:           idx[i*q+j] = (colsaij[i]+rstart/p)*q+j;
    1247:           v[i*q+j]   = (j==s ? vals[i] : 0.0);
    1248:         }
    1249:       }
    1250:       for (i=0; i<ncolsoaij; i++) {
    1251:         for (j=0; j<q; j++) {
    1252:           idx[(i+ncolsaij)*q+j] = garray[colsoaij[i]]*q+j;
    1253:           v[(i+ncolsaij)*q+j]   = (j==s ? ovals[i]: 0.0);
    1254:         }
    1255:       }
    1256:     } else if (T) {
    1257:       for (i=0; i<ncolsaij; i++) {
    1258:         for (j=0; j<q; j++) {
    1259:           idx[i*q+j] = (colsaij[i]+rstart/p)*q+j;
    1260:           v[i*q+j]   = vals[i]*T[s+j*p];
    1261:         }
    1262:       }
    1263:       for (i=0; i<ncolsoaij; i++) {
    1264:         for (j=0; j<q; j++) {
    1265:           idx[(i+ncolsaij)*q+j] = garray[colsoaij[i]]*q+j;
    1266:           v[(i+ncolsaij)*q+j]   = ovals[i]*T[s+j*p];
    1267:         }
    1268:       }
    1269:     }
    1270:     if (S) {
    1271:       for (j=0; j<q; j++) {
    1272:         idx[c*q+j] = (r+rstart/p)*q+j;
    1273:         v[c*q+j]  += S[s+j*p];
    1274:       }
    1275:     }
    1276:   }

    1278:   if (ncols)  *ncols  = nz;
    1279:   if (cols)   *cols   = idx;
    1280:   if (values) *values = v;
    1281:   return(0);
    1282: }

    1284: PetscErrorCode MatRestoreRow_MPIKAIJ(Mat A,PetscInt row,PetscInt *nz,PetscInt **idx,PetscScalar **v)
    1285: {
    1288:   PetscFree2(*idx,*v);
    1289:   ((Mat_SeqKAIJ*)A->data)->getrowactive = PETSC_FALSE;
    1290:   return(0);
    1291: }

    1293: PetscErrorCode  MatCreateSubMatrix_KAIJ(Mat mat,IS isrow,IS iscol,MatReuse cll,Mat *newmat)
    1294: {
    1296:   Mat            A;

    1299:   MatConvert(mat,MATAIJ,MAT_INITIAL_MATRIX,&A);
    1300:   MatCreateSubMatrix(A,isrow,iscol,cll,newmat);
    1301:   MatDestroy(&A);
    1302:   return(0);
    1303: }

    1305: /* ---------------------------------------------------------------------------------- */
    1306: /*@C
    1307:   MatCreateKAIJ - Creates a matrix type to be used for matrices of the following form:

    1309:     [I \otimes S + A \otimes T]

    1311:   where
    1312:     S is a dense (p \times q) matrix
    1313:     T is a dense (p \times q) matrix
    1314:     A is an AIJ  (n \times n) matrix
    1315:     I is the identity matrix
    1316:   The resulting matrix is (np \times nq)

    1318:   S and T are always stored independently on all processes as PetscScalar arrays in column-major format.

    1320:   Collective

    1322:   Input Parameters:
    1323: + A - the AIJ matrix
    1324: . p - number of rows in S and T
    1325: . q - number of columns in S and T
    1326: . S - the S matrix (can be PETSC_NULL), stored as a PetscScalar array (column-major)
    1327: - T - the T matrix (can be PETSC_NULL), stored as a PetscScalar array (column-major)

    1329:   Output Parameter:
    1330: . kaij - the new KAIJ matrix

    1332:   Notes:
    1333:   This function increases the reference count on the AIJ matrix, so the user is free to destroy the matrix if it is not needed.
    1334:   Changes to the entries of the AIJ matrix will immediately affect the KAIJ matrix.

    1336:   Level: advanced

    1338: .seealso: MatKAIJSetAIJ(), MatKAIJSetS(), MatKAIJSetT(), MatKAIJGetAIJ(), MatKAIJGetS(), MatKAIJGetT(), MATKAIJ
    1339: @*/
    1340: PetscErrorCode  MatCreateKAIJ(Mat A,PetscInt p,PetscInt q,const PetscScalar S[],const PetscScalar T[],Mat *kaij)
    1341: {
    1343:   PetscMPIInt    size;

    1346:   MatCreate(PetscObjectComm((PetscObject)A),kaij);
    1347:   MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);
    1348:   if (size == 1) {
    1349:     MatSetType(*kaij,MATSEQKAIJ);
    1350:   } else {
    1351:     MatSetType(*kaij,MATMPIKAIJ);
    1352:   }
    1353:   MatKAIJSetAIJ(*kaij,A);
    1354:   MatKAIJSetS(*kaij,p,q,S);
    1355:   MatKAIJSetT(*kaij,p,q,T);
    1356:   MatSetUp(*kaij);
    1357:   return(0);
    1358: }

    1360: /*MC
    1361:   MATKAIJ - MATKAIJ = "kaij" - A matrix type to be used to evaluate matrices of form
    1362:     [I \otimes S + A \otimes T],
    1363:   where
    1364:     S is a dense (p \times q) matrix,
    1365:     T is a dense (p \times q) matrix,
    1366:     A is an AIJ  (n \times n) matrix,
    1367:     and I is the identity matrix.
    1368:   The resulting matrix is (np \times nq).

    1370:   S and T are always stored independently on all processes as PetscScalar arrays in column-major format.

    1372:   Notes:
    1373:   A linear system with multiple right-hand sides, AX = B, can be expressed in the KAIJ-friendly form of (A \otimes I) x = b,
    1374:   where x and b are column vectors containing the row-major representations of X and B.

    1376:   Level: advanced

    1378: .seealso: MatKAIJSetAIJ(), MatKAIJSetS(), MatKAIJSetT(), MatKAIJGetAIJ(), MatKAIJGetS(), MatKAIJGetT(), MatCreateKAIJ()
    1379: M*/

    1381: PETSC_EXTERN PetscErrorCode MatCreate_KAIJ(Mat A)
    1382: {
    1384:   Mat_MPIKAIJ    *b;
    1385:   PetscMPIInt    size;

    1388:   PetscNewLog(A,&b);
    1389:   A->data  = (void*)b;

    1391:   PetscMemzero(A->ops,sizeof(struct _MatOps));

    1393:   A->ops->setup = MatSetUp_KAIJ;

    1395:   b->w    = NULL;
    1396:   MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);
    1397:   if (size == 1) {
    1398:     PetscObjectChangeTypeName((PetscObject)A,MATSEQKAIJ);
    1399:     A->ops->setup               = MatSetUp_KAIJ;
    1400:     A->ops->destroy             = MatDestroy_SeqKAIJ;
    1401:     A->ops->view                = MatView_KAIJ;
    1402:     A->ops->mult                = MatMult_SeqKAIJ;
    1403:     A->ops->multadd             = MatMultAdd_SeqKAIJ;
    1404:     A->ops->invertblockdiagonal = MatInvertBlockDiagonal_SeqKAIJ;
    1405:     A->ops->getrow              = MatGetRow_SeqKAIJ;
    1406:     A->ops->restorerow          = MatRestoreRow_SeqKAIJ;
    1407:     A->ops->sor                 = MatSOR_SeqKAIJ;
    1408:   } else {
    1409:     PetscObjectChangeTypeName((PetscObject)A,MATMPIKAIJ);
    1410:     A->ops->setup               = MatSetUp_KAIJ;
    1411:     A->ops->destroy             = MatDestroy_MPIKAIJ;
    1412:     A->ops->view                = MatView_KAIJ;
    1413:     A->ops->mult                = MatMult_MPIKAIJ;
    1414:     A->ops->multadd             = MatMultAdd_MPIKAIJ;
    1415:     A->ops->invertblockdiagonal = MatInvertBlockDiagonal_MPIKAIJ;
    1416:     A->ops->getrow              = MatGetRow_MPIKAIJ;
    1417:     A->ops->restorerow          = MatRestoreRow_MPIKAIJ;
    1418:     PetscObjectComposeFunction((PetscObject)A,"MatGetDiagonalBlock_C",MatGetDiagonalBlock_MPIKAIJ);
    1419:   }
    1420:   A->ops->createsubmatrix = MatCreateSubMatrix_KAIJ;
    1421:   return(0);
    1422: }