C++ MatGetDiagonal函数代码示例

// updated
PetscErrorCode BSSCR_MatStokesMVBlockDefaultBuildScaling( MatStokesBlockScaling BA, Mat A, Vec b, Vec x, PetscTruth sym )
{
	Mat K,G,D,C;
	Vec rG;
	PetscScalar rg2, rg, ra;  
	PetscInt N;
	Vec rA, rC;
	Vec L1,L2, R1,R2;
	Mat S;
	
	VecNestGetSubVec( BA->Lz, 0, &L1 );
	VecNestGetSubVec( BA->Lz, 1, &L2 );
	
	VecNestGetSubVec( BA->Rz, 0, &R1 );
	VecNestGetSubVec( BA->Rz, 1, &R2 );
	
	rA = L1;
	rC = L2;
	
	MatNestGetSubMat( A, 0,0, &K );
	MatNestGetSubMat( A, 0,1, &G );
	MatNestGetSubMat( A, 1,0, &D );
	MatNestGetSubMat( A, 1,1, &C );
	
	VecDuplicate( rA, &rG );
	
	
	/* Get magnitude of K */  
	//px_MatGetAbsRowSum( K, rA );
        //MatGetRowMax( K, rA, PETSC_NULL );
	MatGetDiagonal( K, rA);

	VecSqrt( rA );  
	VecReciprocal( rA );
	
	/* VecDot( rA,rA, &ra ); */
	/* VecGetSize( rA, &N ); */
	/* ra = PetscSqrtScalar( ra/N ); */
	
	
	/* Get magnitude of G */
	//px_MatGetAbsRowSum( G, rG );
	//MatGetRowMax( G, rG, PETSC_NULL );

	Mat A21_cpy;
	Mat Shat;
	Vec diag; /* same as rA*rA */
	
	
	MatGetVecs( K, &diag, PETSC_NULL );
	MatGetDiagonal( K, diag );
	VecReciprocal( diag );

	if( sym )
#if( PETSC_VERSION_MAJOR <= 2 )
	    MatTranspose( G, &A21_cpy );
#else
	    MatTranspose( G, MAT_INITIAL_MATRIX, &A21_cpy );
#endif
	else {

int main(int argc,char **args)
{
  Mat            A;
  Vec            x;
  PetscErrorCode ierr;
  PetscViewer    fd;              /* viewer */
  char           file[PETSC_MAX_PATH_LEN]; /* input file name */
  PetscReal      norm;
  PetscBool      flg;

  PetscInitialize(&argc,&args,(char*)0,help);

  /* Determine file from which we read the matrix A */
  ierr = PetscOptionsGetString(NULL,"-f",file,PETSC_MAX_PATH_LEN,&flg);CHKERRQ(ierr);
  if (!flg) SETERRQ(PETSC_COMM_WORLD,1,"Must indicate binary file with the -f option");

  /* Load matrix A */
  ierr = PetscViewerBinaryOpen(PETSC_COMM_WORLD,file,FILE_MODE_READ,&fd);CHKERRQ(ierr);
  ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
  ierr = MatLoad(A,fd);CHKERRQ(ierr);
  ierr = PetscViewerDestroy(&fd);CHKERRQ(ierr);
  ierr = MatCreateVecs(A,&x,NULL);CHKERRQ(ierr);
  ierr = MatGetDiagonal(A,x);CHKERRQ(ierr);
  ierr = VecScale(x,-1.0);CHKERRQ(ierr);
  ierr = MatDiagonalSet(A,x,ADD_VALUES);CHKERRQ(ierr);
  ierr = MatGetDiagonal(A,x);CHKERRQ(ierr);
  ierr = VecNorm(x,NORM_2,&norm);CHKERRQ(ierr);
  ierr = PetscPrintf(PETSC_COMM_WORLD,"Norm %g\n",(double)norm);CHKERRQ(ierr);

  /* Free data structures */
  ierr = MatDestroy(&A);CHKERRQ(ierr);
  ierr = VecDestroy(&x);CHKERRQ(ierr);
  ierr = PetscFinalize();
  return 0;
}

PetscErrorCode bsscr_DGMiGtD( Mat *_K2, Mat K, Mat G, Mat M){
    Mat K2;
    Vec diag;
    Vec Mdiag;
    Mat MinvGt;
    Mat Gtrans;
    PetscErrorCode ierr;

    PetscFunctionBegin;
    MatGetVecs( K, &diag, PETSC_NULL );
    MatGetDiagonal( K, diag );
    VecSqrt(diag);
    //VecReciprocal(diag);/* trying something different here */
    MatGetVecs( M, &Mdiag, PETSC_NULL );
    MatGetDiagonal( M, Mdiag );
    VecReciprocal(Mdiag);
    #if( PETSC_VERSION_MAJOR <= 2 )
    ierr=MatTranspose(G, &Gtrans);CHKERRQ(ierr);
    #else
    ierr=MatTranspose(G, MAT_INITIAL_MATRIX,&Gtrans);CHKERRQ(ierr);
    #endif
    ierr=MatConvert(Gtrans, MATSAME, MAT_INITIAL_MATRIX, &MinvGt);CHKERRQ(ierr);/* copy Gtrans -> MinvGt */
    MatDiagonalScale(MinvGt, Mdiag, PETSC_NULL);/* Minv*Gtrans */
    /* MAT_INITIAL_MATRIX -> creates K2 matrix : PETSC_DEFAULT for fill ratio: run with -info to find what it should be*/
    ierr=MatMatMult( G, MinvGt, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &K2);CHKERRQ(ierr);/* K2 = G*Minv*Gtrans */
    MatDiagonalScale(K2, diag, diag );/* K2 = D*K2*D  = D*G*Minv*Gtrans*D */
    Stg_MatDestroy(&Gtrans);
    Stg_MatDestroy(&MinvGt);
    Stg_VecDestroy(&diag);
    Stg_VecDestroy(&Mdiag);
    *_K2=K2;
    PetscFunctionReturn(0);
}

PetscErrorCode StokesSetupApproxSchur(Stokes *s)
{
  Vec            diag;
  PetscErrorCode ierr;

  PetscFunctionBeginUser;
  /* Schur complement approximation: myS = A11 - A10 diag(A00)^(-1) A01 */
  /* note: A11 is zero */
  /* note: in real life this matrix would be build directly, */
  /* i.e. without MatMatMult */

  /* inverse of diagonal of A00 */
  ierr = VecCreate(PETSC_COMM_WORLD,&diag);CHKERRQ(ierr);
  ierr = VecSetSizes(diag,PETSC_DECIDE,2*s->nx*s->ny);CHKERRQ(ierr);
  ierr = VecSetType(diag,VECMPI);CHKERRQ(ierr);
  ierr = MatGetDiagonal(s->subA[0],diag);
  ierr = VecReciprocal(diag);

  /* compute: - A10 diag(A00)^(-1) A01 */
  ierr = MatDiagonalScale(s->subA[1],diag,NULL); /* (*warning* overwrites subA[1]) */
  ierr = MatMatMult(s->subA[2],s->subA[1],MAT_INITIAL_MATRIX,PETSC_DEFAULT,&s->myS);CHKERRQ(ierr);
  ierr = MatScale(s->myS,-1.0);CHKERRQ(ierr);

  /* restore A10 */
  ierr = MatGetDiagonal(s->subA[0],diag);
  ierr = MatDiagonalScale(s->subA[1],diag,NULL);
  ierr = VecDestroy(&diag);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode PCSetUp_Eisenstat(PC pc)
{
  PetscErrorCode ierr;
  PetscInt       M,N,m,n;
  PC_Eisenstat   *eis = (PC_Eisenstat*)pc->data;

  PetscFunctionBegin;
  if (!pc->setupcalled) {
    ierr = MatGetSize(pc->mat,&M,&N);CHKERRQ(ierr);
    ierr = MatGetLocalSize(pc->mat,&m,&n);CHKERRQ(ierr);
    ierr = MatCreate(((PetscObject)pc)->comm,&eis->shell);CHKERRQ(ierr);
    ierr = MatSetSizes(eis->shell,m,n,M,N);CHKERRQ(ierr);
    ierr = MatSetType(eis->shell,MATSHELL);CHKERRQ(ierr);
    ierr = MatSetUp(eis->shell);CHKERRQ(ierr);
    ierr = MatShellSetContext(eis->shell,(void*)pc);CHKERRQ(ierr);
    ierr = PetscLogObjectParent(pc,eis->shell);CHKERRQ(ierr);
    ierr = MatShellSetOperation(eis->shell,MATOP_MULT,(void(*)(void))PCMult_Eisenstat);CHKERRQ(ierr);
  }
  if (!eis->usediag) PetscFunctionReturn(0);
  if (!pc->setupcalled) {
    ierr = MatGetVecs(pc->pmat,&eis->diag,0);CHKERRQ(ierr);
    ierr = PetscLogObjectParent(pc,eis->diag);CHKERRQ(ierr);
  }
  ierr = MatGetDiagonal(pc->pmat,eis->diag);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode PCSetUp_LSC(PC pc)
{
  PC_LSC         *lsc = (PC_LSC*)pc->data;
  Mat            L,Lp,B,C;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = PCLSCAllocate_Private(pc);CHKERRQ(ierr);
  ierr = PetscObjectQuery((PetscObject)pc->mat,"LSC_L",(PetscObject*)&L);CHKERRQ(ierr);
  if (!L) {ierr = PetscObjectQuery((PetscObject)pc->pmat,"LSC_L",(PetscObject*)&L);CHKERRQ(ierr);}
  ierr = PetscObjectQuery((PetscObject)pc->pmat,"LSC_Lp",(PetscObject*)&Lp);CHKERRQ(ierr);
  if (!Lp) {ierr = PetscObjectQuery((PetscObject)pc->mat,"LSC_Lp",(PetscObject*)&Lp);CHKERRQ(ierr);}
  if (!L) {
    ierr = MatSchurComplementGetSubMatrices(pc->mat,NULL,NULL,&B,&C,NULL);CHKERRQ(ierr);
    if (!lsc->L) {
      ierr = MatMatMult(C,B,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&lsc->L);CHKERRQ(ierr);
    } else {
      ierr = MatMatMult(C,B,MAT_REUSE_MATRIX,PETSC_DEFAULT,&lsc->L);CHKERRQ(ierr);
    }
    Lp = L = lsc->L;
  }
  if (lsc->scale) {
    Mat Ap;
    ierr = MatSchurComplementGetSubMatrices(pc->mat,NULL,&Ap,NULL,NULL,NULL);CHKERRQ(ierr);
    ierr = MatGetDiagonal(Ap,lsc->scale);CHKERRQ(ierr); /* Should be the mass matrix, but we don't have plumbing for that yet */
    ierr = VecReciprocal(lsc->scale);CHKERRQ(ierr);
  }
  ierr = KSPSetOperators(lsc->kspL,L,Lp);CHKERRQ(ierr);
  ierr = KSPSetFromOptions(lsc->kspL);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

void FETI_Operations::set_jacobi_precond_vector()
{
	homemade_assert_msg(m_bC_RR_MatrixSet,"Preconditioner matrix not set yet!");

	// Create and set the vector
	VecCreate(m_comm.get(),&m_coupling_jacobi_precond_vec);
	VecSetSizes(m_coupling_jacobi_precond_vec,m_C_RR_M_local,m_C_RR_M);
	VecSetFromOptions(m_coupling_jacobi_precond_vec);

	// Get the diagonal
	MatGetDiagonal(m_C_RR,m_coupling_jacobi_precond_vec);

	// Calculate the reciprocal
	VecReciprocal(m_coupling_jacobi_precond_vec);

	// Export it
	write_PETSC_vector(m_coupling_jacobi_precond_vec,m_scratch_folder_path + "/precond_Jacobi_vector.petscvec",m_comm.rank(),m_comm.get());

// Print MatLab debugging output? Variable defined at "carl_headers.h"
#ifdef PRINT_MATLAB_DEBUG
	write_PETSC_vector_MATLAB(m_coupling_jacobi_precond_vec,m_scratch_folder_path + "/precond_Jacobi_vector.m",m_comm.get());
#endif

	// Set flag
	m_bCreatedPrecondJacobiVec = true;
}

void PetscSparseStorage::getDiagonal( OoqpVector& vec_in )
{
  PetscVector & vec = dynamic_cast<PetscVector &>(vec_in);

  int ierr;
  ierr = MatGetDiagonal(M, vec.pv); assert( ierr  == 0);
}

void
NRSolver :: applyConstraintsToLoadIncrement(int nite, const SparseMtrx &k, FloatArray &R,
                                            referenceLoadInputModeType rlm, TimeStep *tStep)
{
    double factor = engngModel->giveDomain(1)->giveFunction(prescribedDisplacementTF)->evaluateAtTime( tStep->giveTargetTime() );
    if ( ( rlm == rlm_total ) && ( !tStep->isTheFirstStep() ) ) {
        //factor -= engngModel->giveDomain(1)->giveFunction(prescribedDisplacementTF)->
        // at(tStep->givePreviousStep()->giveTime()) ;
        factor -= engngModel->giveDomain(1)->giveFunction(prescribedDisplacementTF)->
        evaluateAtTime( tStep->giveTargetTime() - tStep->giveTimeIncrement() );
    }

    if ( nite == 0 ) {
#if 0
 #ifdef __PETSC_MODULE
        if ( solverType == ST_Petsc ) {
            //Natural2LocalOrdering* n2lpm = engngModel->giveParallelContext(1)->giveN2Lmap();
            //IntArray* map = n2lpm->giveN2Lmap();
            for ( i = 1; i <= prescribedEqs.giveSize(); i++ ) {
                eq = prescribedEqs.at(i);
                R.at(eq) = prescribedDofsValues.at(i) * factor; // local eq
            }

            return;
        }

 #endif
#else
 #ifdef __PETSC_MODULE
        const PetscSparseMtrx *lhs = dynamic_cast< const PetscSparseMtrx * >(&k);
        if ( lhs ) {
            Vec diag;
            PetscScalar *ptr;
            lhs->createVecGlobal(& diag);
            MatGetDiagonal(* ( const_cast< PetscSparseMtrx * >(lhs)->giveMtrx() ), diag);
            VecGetArray(diag, & ptr);

            for ( int i = 1; i <= numberOfPrescribedDofs; i++ ) {
                int eq = prescribedEqs.at(i) - 1;
                R.at(eq + 1) = ptr [ eq ] * prescribedDofsValues.at(i) * factor;
            }

            VecRestoreArray(diag, & ptr);
            VecDestroy(& diag);
            return;
        }
 #endif
#endif
        for ( int i = 1; i <= numberOfPrescribedDofs; i++ ) {
            int eq = prescribedEqs.at(i);
            R.at(eq) = k.at(eq, eq) * prescribedDofsValues.at(i) * factor;
        }
    } else {
        for ( int i = 1; i <= numberOfPrescribedDofs; i++ ) {
            R.at( prescribedEqs.at(i) ) = 0.0;
        }
    }
}

PetscErrorCode MatGetDiagonal_SMF(Mat mat,Vec v)
{
  PetscErrorCode   ierr;
  MatSubMatFreeCtx ctx;

  PetscFunctionBegin;
  ierr = MatShellGetContext(mat,(void **)&ctx);CHKERRQ(ierr);
  ierr = MatGetDiagonal(ctx->A,v);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

static PetscErrorCode MatGetDiagonal_User(Mat A,Vec X)
{
  User           user;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  ierr = MatShellGetContext(A,&user);CHKERRQ(ierr);
  ierr = MatGetDiagonal(user->B,X);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

PetscErrorCode PCGAMGOptProl_Classical_Jacobi(PC pc,const Mat A,Mat *P)
{

  PetscErrorCode    ierr;
  PetscInt          f,s,n,cf,cs,i,idx;
  PetscInt          *coarserows;
  PetscInt          ncols;
  const PetscInt    *pcols;
  const PetscScalar *pvals;
  Mat               Pnew;
  Vec               diag;
  PC_MG             *mg          = (PC_MG*)pc->data;
  PC_GAMG           *pc_gamg     = (PC_GAMG*)mg->innerctx;
  PC_GAMG_Classical *cls         = (PC_GAMG_Classical*)pc_gamg->subctx;

  PetscFunctionBegin;
  if (cls->nsmooths == 0) {
    ierr = PCGAMGTruncateProlongator_Private(pc,P);CHKERRQ(ierr);
    PetscFunctionReturn(0);
  }
  ierr = MatGetOwnershipRange(*P,&s,&f);CHKERRQ(ierr);
  n = f-s;
  ierr = MatGetOwnershipRangeColumn(*P,&cs,&cf);CHKERRQ(ierr);
  ierr = PetscMalloc(sizeof(PetscInt)*n,&coarserows);CHKERRQ(ierr);
  /* identify the rows corresponding to coarse unknowns */
  idx = 0;
  for (i=s;i<f;i++) {
    ierr = MatGetRow(*P,i,&ncols,&pcols,&pvals);CHKERRQ(ierr);
    /* assume, for now, that it's a coarse unknown if it has a single unit entry */
    if (ncols == 1) {
      if (pvals[0] == 1.) {
        coarserows[idx] = i;
        idx++;
      }
    }
    ierr = MatRestoreRow(*P,i,&ncols,&pcols,&pvals);CHKERRQ(ierr);
  }
  ierr = MatGetVecs(A,&diag,0);CHKERRQ(ierr);
  ierr = MatGetDiagonal(A,diag);CHKERRQ(ierr);
  ierr = VecReciprocal(diag);CHKERRQ(ierr);
  for (i=0;i<cls->nsmooths;i++) {
    ierr = MatMatMult(A,*P,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&Pnew);CHKERRQ(ierr);
    ierr = MatZeroRows(Pnew,idx,coarserows,0.,NULL,NULL);CHKERRQ(ierr);
    ierr = MatDiagonalScale(Pnew,diag,0);CHKERRQ(ierr);
    ierr = MatAYPX(Pnew,-1.0,*P,DIFFERENT_NONZERO_PATTERN);CHKERRQ(ierr);
    ierr = MatDestroy(P);CHKERRQ(ierr);
    *P  = Pnew;
    Pnew = NULL;
  }
  ierr = VecDestroy(&diag);CHKERRQ(ierr);
  ierr = PetscFree(coarserows);CHKERRQ(ierr);
  ierr = PCGAMGTruncateProlongator_Private(pc,P);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}

/*
   SampleShellPCSetUp - This routine sets up a user-defined
   preconditioner context.

   Input Parameters:
.  pc    - preconditioner object
.  pmat  - preconditioner matrix
.  x     - vector

   Output Parameter:
.  shell - fully set up user-defined preconditioner context

   Notes:
   In this example, we define the shell preconditioner to be Jacobi's
   method.  Thus, here we create a work vector for storing the reciprocal
   of the diagonal of the preconditioner matrix; this vector is then
   used within the routine SampleShellPCApply().
*/
PetscErrorCode SampleShellPCSetUp(PC pc,Mat pmat,Vec x)
{
  SampleShellPC  *shell;
  Vec            diag;
  PetscErrorCode ierr;

  ierr = PCShellGetContext(pc,(void**)&shell);CHKERRQ(ierr);
  ierr = VecDuplicate(x,&diag);CHKERRQ(ierr);
  ierr = MatGetDiagonal(pmat,diag);CHKERRQ(ierr);
  ierr = VecReciprocal(diag);CHKERRQ(ierr);

  shell->diag = diag;
  return 0;
}

/*@C
  TaoMatGetSubMat - Gets a submatrix using the IS

  Input Parameters:
+ M - the full matrix (n x n)
. is - the index set for the submatrix (both row and column index sets need to be the same)
. v1 - work vector of dimension n, needed for TAO_SUBSET_MASK option
- subset_type - the method TAO is using for subsetting (TAO_SUBSET_SUBVEC, TAO_SUBSET_MASK,
  TAO_SUBSET_MATRIXFREE)

  Output Parameters:
. Msub - the submatrix
@*/
PetscErrorCode TaoMatGetSubMat(Mat M, IS is, Vec v1, TaoSubsetType subset_type, Mat *Msub)
{
  PetscErrorCode ierr;
  IS             iscomp;
  PetscBool      flg = PETSC_FALSE;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(M,MAT_CLASSID,1);
  PetscValidHeaderSpecific(is,IS_CLASSID,2);
  ierr = MatDestroy(Msub);CHKERRQ(ierr);
  switch (subset_type) {
  case TAO_SUBSET_SUBVEC:
    ierr = MatGetSubMatrix(M, is, is, MAT_INITIAL_MATRIX, Msub);CHKERRQ(ierr);
    break;

  case TAO_SUBSET_MASK:
    /* Get Reduced Hessian
     Msub[i,j] = M[i,j] if i,j in Free_Local or i==j
     Msub[i,j] = 0      if i!=j and i or j not in Free_Local
     */
    ierr = PetscOptionsBegin(PetscObjectComm((PetscObject)M),NULL,NULL,NULL);CHKERRQ(ierr);
    ierr = PetscOptionsBool("-different_submatrix","use separate hessian matrix when computing submatrices","TaoSubsetType",flg,&flg,NULL);CHKERRQ(ierr);
    ierr = PetscOptionsEnd();CHKERRQ(ierr);
    if (flg) {
      ierr = MatDuplicate(M, MAT_COPY_VALUES, Msub);CHKERRQ(ierr);
    } else {
      /* Act on hessian directly (default) */
      ierr = PetscObjectReference((PetscObject)M);CHKERRQ(ierr);
      *Msub = M;
    }
    /* Save the diagonal to temporary vector */
    ierr = MatGetDiagonal(*Msub,v1);CHKERRQ(ierr);

    /* Zero out rows and columns */
    ierr = ISComplementVec(is,v1,&iscomp);CHKERRQ(ierr);

    /* Use v1 instead of 0 here because of PETSc bug */
    ierr = MatZeroRowsColumnsIS(*Msub,iscomp,1.0,v1,v1);CHKERRQ(ierr);

    ierr = ISDestroy(&iscomp);CHKERRQ(ierr);
    break;
  case TAO_SUBSET_MATRIXFREE:
    ierr = ISComplementVec(is,v1,&iscomp);CHKERRQ(ierr);
    ierr = MatCreateSubMatrixFree(M,iscomp,iscomp,Msub);CHKERRQ(ierr);
    ierr = ISDestroy(&iscomp);CHKERRQ(ierr);
    break;
  }
  PetscFunctionReturn(0);
}

PetscErrorCode MatGetDiagonal_IS(Mat A, Vec v)
{
  Mat_IS         *is = (Mat_IS*)A->data;
  PetscErrorCode ierr;

  PetscFunctionBegin;
  /* get diagonal of the local matrix */
  ierr = MatGetDiagonal(is->A,is->y);CHKERRQ(ierr);

  /* scatter diagonal back into global vector */
  ierr = VecSet(v,0);CHKERRQ(ierr);
  ierr = VecScatterBegin(is->rctx,is->y,v,ADD_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
  ierr = VecScatterEnd(is->rctx,is->y,v,ADD_VALUES,SCATTER_REVERSE);CHKERRQ(ierr);
  PetscFunctionReturn(0);
}