mat/utils/axpy.c

#include <petsc/private/matimpl.h> /*I   "petscmat.h"  I*/

static PetscErrorCode MatTransposeAXPY_Private(Mat Y, PetscScalar a, Mat X, MatStructure str, Mat T) {
  Mat A, F;
  PetscErrorCode (*f)(Mat, Mat *);

  PetscFunctionBegin;
  PetscCall(PetscObjectQueryFunction((PetscObject)T, "MatTransposeGetMat_C", &f));
  if (f) {
    PetscCall(MatTransposeGetMat(T, &A));
    if (T == X) {
      PetscCall(PetscInfo(NULL, "Explicitly transposing X of type MATTRANSPOSE to perform MatAXPY()\n"));
      PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &F));
      A = Y;
    } else {
      PetscCall(PetscInfo(NULL, "Transposing X because Y of type MATTRANSPOSE to perform MatAXPY()\n"));
      PetscCall(MatTranspose(X, MAT_INITIAL_MATRIX, &F));
    }
  } else {
    PetscCall(MatHermitianTransposeGetMat(T, &A));
    if (T == X) {
      PetscCall(PetscInfo(NULL, "Explicitly Hermitian transposing X of type MATTRANSPOSE to perform MatAXPY()\n"));
      PetscCall(MatHermitianTranspose(A, MAT_INITIAL_MATRIX, &F));
      A = Y;
    } else {
      PetscCall(PetscInfo(NULL, "Hermitian transposing X because Y of type MATTRANSPOSE to perform MatAXPY()\n"));
      PetscCall(MatHermitianTranspose(X, MAT_INITIAL_MATRIX, &F));
    }
  }
  PetscCall(MatAXPY(A, a, F, str));
  PetscCall(MatDestroy(&F));
  PetscFunctionReturn(0);
}

/*@
   MatAXPY - Computes Y = a*X + Y.

   Logically  Collective on Mat

   Input Parameters:
+  a - the scalar multiplier
.  X - the first matrix
.  Y - the second matrix
-  str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN, UNKNOWN_NONZERO_PATTERN, or SUBSET_NONZERO_PATTERN (nonzeros of X is a subset of Y's)

   Note:
   No operation is performed when a is zero.

   Level: intermediate

.seealso: `MatAYPX()`
 @*/
PetscErrorCode MatAXPY(Mat Y, PetscScalar a, Mat X, MatStructure str) {
  PetscInt  M1, M2, N1, N2;
  PetscInt  m1, m2, n1, n2;
  MatType   t1, t2;
  PetscBool sametype, transpose;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
  PetscValidLogicalCollectiveScalar(Y, a, 2);
  PetscValidHeaderSpecific(X, MAT_CLASSID, 3);
  PetscCheckSameComm(Y, 1, X, 3);
  PetscCall(MatGetSize(X, &M1, &N1));
  PetscCall(MatGetSize(Y, &M2, &N2));
  PetscCall(MatGetLocalSize(X, &m1, &n1));
  PetscCall(MatGetLocalSize(Y, &m2, &n2));
  PetscCheck(M1 == M2 && N1 == N2, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_SIZ, "Non conforming matrix add: global sizes X %" PetscInt_FMT " x %" PetscInt_FMT ", Y %" PetscInt_FMT " x %" PetscInt_FMT, M1, N1, M2, N2);
  PetscCheck(m1 == m2 && n1 == n2, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "Non conforming matrix add: local sizes X %" PetscInt_FMT " x %" PetscInt_FMT ", Y %" PetscInt_FMT " x %" PetscInt_FMT, m1, n1, m2, n2);
  PetscCheck(Y->assembled, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix (Y)");
  PetscCheck(X->assembled, PetscObjectComm((PetscObject)X), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix (X)");
  if (a == (PetscScalar)0.0) PetscFunctionReturn(0);
  if (Y == X) {
    PetscCall(MatScale(Y, 1.0 + a));
    PetscFunctionReturn(0);
  }
  PetscCall(MatGetType(X, &t1));
  PetscCall(MatGetType(Y, &t2));
  PetscCall(PetscStrcmp(t1, t2, &sametype));
  PetscCall(PetscLogEventBegin(MAT_AXPY, Y, 0, 0, 0));
  if (Y->ops->axpy && (sametype || X->ops->axpy == Y->ops->axpy)) {
    PetscUseTypeMethod(Y, axpy, a, X, str);
  } else {
    PetscCall(PetscStrcmp(t1, MATTRANSPOSE, &transpose));
    if (transpose) {
      PetscCall(MatTransposeAXPY_Private(Y, a, X, str, X));
    } else {
      PetscCall(PetscStrcmp(t2, MATTRANSPOSE, &transpose));
      if (transpose) {
        PetscCall(MatTransposeAXPY_Private(Y, a, X, str, Y));
      } else {
        PetscCall(MatAXPY_Basic(Y, a, X, str));
      }
    }
  }
  PetscCall(PetscLogEventEnd(MAT_AXPY, Y, 0, 0, 0));
  PetscFunctionReturn(0);
}

PetscErrorCode MatAXPY_Basic_Preallocate(Mat Y, Mat X, Mat *B) {
  PetscErrorCode (*preall)(Mat, Mat, Mat *) = NULL;

  PetscFunctionBegin;
  /* look for any available faster alternative to the general preallocator */
  PetscCall(PetscObjectQueryFunction((PetscObject)Y, "MatAXPYGetPreallocation_C", &preall));
  if (preall) {
    PetscCall((*preall)(Y, X, B));
  } else { /* Use MatPrellocator, assumes same row-col distribution */
    Mat      preallocator;
    PetscInt r, rstart, rend;
    PetscInt m, n, M, N;

    PetscCall(MatGetRowUpperTriangular(Y));
    PetscCall(MatGetRowUpperTriangular(X));
    PetscCall(MatGetSize(Y, &M, &N));
    PetscCall(MatGetLocalSize(Y, &m, &n));
    PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), &preallocator));
    PetscCall(MatSetType(preallocator, MATPREALLOCATOR));
    PetscCall(MatSetLayouts(preallocator, Y->rmap, Y->cmap));
    PetscCall(MatSetUp(preallocator));
    PetscCall(MatGetOwnershipRange(preallocator, &rstart, &rend));
    for (r = rstart; r < rend; ++r) {
      PetscInt           ncols;
      const PetscInt    *row;
      const PetscScalar *vals;

      PetscCall(MatGetRow(Y, r, &ncols, &row, &vals));
      PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
      PetscCall(MatRestoreRow(Y, r, &ncols, &row, &vals));
      PetscCall(MatGetRow(X, r, &ncols, &row, &vals));
      PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
      PetscCall(MatRestoreRow(X, r, &ncols, &row, &vals));
    }
    PetscCall(MatSetOption(preallocator, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    PetscCall(MatAssemblyBegin(preallocator, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(preallocator, MAT_FINAL_ASSEMBLY));
    PetscCall(MatRestoreRowUpperTriangular(Y));
    PetscCall(MatRestoreRowUpperTriangular(X));

    PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), B));
    PetscCall(MatSetType(*B, ((PetscObject)Y)->type_name));
    PetscCall(MatSetLayouts(*B, Y->rmap, Y->cmap));
    PetscCall(MatPreallocatorPreallocate(preallocator, PETSC_FALSE, *B));
    PetscCall(MatDestroy(&preallocator));
  }
  PetscFunctionReturn(0);
}

PetscErrorCode MatAXPY_Basic(Mat Y, PetscScalar a, Mat X, MatStructure str) {
  PetscBool isshell, isdense, isnest;

  PetscFunctionBegin;
  PetscCall(MatIsShell(Y, &isshell));
  if (isshell) { /* MatShell has special support for AXPY */
    PetscErrorCode (*f)(Mat, PetscScalar, Mat, MatStructure);

    PetscCall(MatGetOperation(Y, MATOP_AXPY, (void (**)(void)) & f));
    if (f) {
      PetscCall((*f)(Y, a, X, str));
      PetscFunctionReturn(0);
    }
  }
  /* no need to preallocate if Y is dense */
  PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)Y, &isdense, MATSEQDENSE, MATMPIDENSE, ""));
  if (isdense) {
    PetscCall(PetscObjectTypeCompare((PetscObject)X, MATNEST, &isnest));
    if (isnest) {
      PetscCall(MatAXPY_Dense_Nest(Y, a, X));
      PetscFunctionReturn(0);
    }
    if (str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN) str = SUBSET_NONZERO_PATTERN;
  }
  if (str != DIFFERENT_NONZERO_PATTERN && str != UNKNOWN_NONZERO_PATTERN) {
    PetscInt           i, start, end, j, ncols, m, n;
    const PetscInt    *row;
    PetscScalar       *val;
    const PetscScalar *vals;

    PetscCall(MatGetSize(X, &m, &n));
    PetscCall(MatGetOwnershipRange(X, &start, &end));
    PetscCall(MatGetRowUpperTriangular(X));
    if (a == 1.0) {
      for (i = start; i < end; i++) {
        PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
        PetscCall(MatSetValues(Y, 1, &i, ncols, row, vals, ADD_VALUES));
        PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
      }
    } else {
      PetscInt vs = 100;
      /* realloc if needed, as this function may be used in parallel */
      PetscCall(PetscMalloc1(vs, &val));
      for (i = start; i < end; i++) {
        PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
        if (vs < ncols) {
          vs = PetscMin(2 * ncols, n);
          PetscCall(PetscRealloc(vs * sizeof(*val), &val));
        }
        for (j = 0; j < ncols; j++) val[j] = a * vals[j];
        PetscCall(MatSetValues(Y, 1, &i, ncols, row, val, ADD_VALUES));
        PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
      }
      PetscCall(PetscFree(val));
    }
    PetscCall(MatRestoreRowUpperTriangular(X));
    PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
    PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
  } else {
    Mat B;

    PetscCall(MatAXPY_Basic_Preallocate(Y, X, &B));
    PetscCall(MatAXPY_BasicWithPreallocation(B, Y, a, X, str));
    PetscCall(MatHeaderMerge(Y, &B));
  }
  PetscFunctionReturn(0);
}

PetscErrorCode MatAXPY_BasicWithPreallocation(Mat B, Mat Y, PetscScalar a, Mat X, MatStructure str) {
  PetscInt           i, start, end, j, ncols, m, n;
  const PetscInt    *row;
  PetscScalar       *val;
  const PetscScalar *vals;

  PetscFunctionBegin;
  PetscCall(MatGetSize(X, &m, &n));
  PetscCall(MatGetOwnershipRange(X, &start, &end));
  PetscCall(MatGetRowUpperTriangular(Y));
  PetscCall(MatGetRowUpperTriangular(X));
  if (a == 1.0) {
    for (i = start; i < end; i++) {
      PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
      PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
      PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));

      PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
      PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
      PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
    }
  } else {
    PetscInt vs = 100;
    /* realloc if needed, as this function may be used in parallel */
    PetscCall(PetscMalloc1(vs, &val));
    for (i = start; i < end; i++) {
      PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
      PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
      PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));

      PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
      if (vs < ncols) {
        vs = PetscMin(2 * ncols, n);
        PetscCall(PetscRealloc(vs * sizeof(*val), &val));
      }
      for (j = 0; j < ncols; j++) val[j] = a * vals[j];
      PetscCall(MatSetValues(B, 1, &i, ncols, row, val, ADD_VALUES));
      PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
    }
    PetscCall(PetscFree(val));
  }
  PetscCall(MatRestoreRowUpperTriangular(Y));
  PetscCall(MatRestoreRowUpperTriangular(X));
  PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
  PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
  PetscFunctionReturn(0);
}

/*@
   MatShift - Computes Y =  Y + a I, where a is a PetscScalar and I is the identity matrix.

   Neighbor-wise Collective on Mat

   Input Parameters:
+  Y - the matrices
-  a - the PetscScalar

   Level: intermediate

   Notes:
    If Y is a rectangular matrix, the shift is done on the main diagonal Y_{ii} of the matrix (https://en.wikipedia.org/wiki/Main_diagonal)

    If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
   fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
   entry). No operation is performed when a is zero.

   To form Y = Y + diag(V) use MatDiagonalSet()

.seealso: `MatDiagonalSet()`, `MatScale()`, `MatDiagonalScale()`
 @*/
PetscErrorCode MatShift(Mat Y, PetscScalar a) {
  PetscFunctionBegin;
  PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
  PetscCheck(Y->assembled, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix");
  PetscCheck(!Y->factortype, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
  MatCheckPreallocated(Y, 1);
  if (a == 0.0) PetscFunctionReturn(0);

  if (Y->ops->shift) PetscUseTypeMethod(Y, shift, a);
  else PetscCall(MatShift_Basic(Y, a));

  PetscCall(PetscObjectStateIncrease((PetscObject)Y));
  PetscFunctionReturn(0);
}

PetscErrorCode MatDiagonalSet_Default(Mat Y, Vec D, InsertMode is) {
  PetscInt           i, start, end;
  const PetscScalar *v;

  PetscFunctionBegin;
  PetscCall(MatGetOwnershipRange(Y, &start, &end));
  PetscCall(VecGetArrayRead(D, &v));
  for (i = start; i < end; i++) PetscCall(MatSetValues(Y, 1, &i, 1, &i, v + i - start, is));
  PetscCall(VecRestoreArrayRead(D, &v));
  PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
  PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
  PetscFunctionReturn(0);
}

/*@
   MatDiagonalSet - Computes Y = Y + D, where D is a diagonal matrix
   that is represented as a vector. Or Y[i,i] = D[i] if InsertMode is
   INSERT_VALUES.

   Neighbor-wise Collective on Mat

   Input Parameters:
+  Y - the input matrix
.  D - the diagonal matrix, represented as a vector
-  i - INSERT_VALUES or ADD_VALUES

   Note:
    If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
   fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
   entry).

   Level: intermediate

.seealso: `MatShift()`, `MatScale()`, `MatDiagonalScale()`
@*/
PetscErrorCode MatDiagonalSet(Mat Y, Vec D, InsertMode is) {
  PetscInt matlocal, veclocal;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
  PetscValidHeaderSpecific(D, VEC_CLASSID, 2);
  PetscCall(MatGetLocalSize(Y, &matlocal, NULL));
  PetscCall(VecGetLocalSize(D, &veclocal));
  PetscCheck(matlocal == veclocal, PETSC_COMM_SELF, PETSC_ERR_ARG_INCOMP, "Number local rows of matrix %" PetscInt_FMT " does not match that of vector for diagonal %" PetscInt_FMT, matlocal, veclocal);
  if (Y->ops->diagonalset) PetscUseTypeMethod(Y, diagonalset, D, is);
  else PetscCall(MatDiagonalSet_Default(Y, D, is));
  PetscCall(PetscObjectStateIncrease((PetscObject)Y));
  PetscFunctionReturn(0);
}

/*@
   MatAYPX - Computes Y = a*Y + X.

   Logically on Mat

   Input Parameters:
+  a - the PetscScalar multiplier
.  Y - the first matrix
.  X - the second matrix
-  str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN, UNKNOWN_NONZERO_PATTERN, or SUBSET_NONZERO_PATTERN (nonzeros of X is a subset of Y's)

   Level: intermediate

.seealso: `MatAXPY()`
 @*/
PetscErrorCode MatAYPX(Mat Y, PetscScalar a, Mat X, MatStructure str) {
  PetscFunctionBegin;
  PetscCall(MatScale(Y, a));
  PetscCall(MatAXPY(Y, 1.0, X, str));
  PetscFunctionReturn(0);
}

/*@
    MatComputeOperator - Computes the explicit matrix

    Collective on Mat

    Input Parameters:
+   inmat - the matrix
-   mattype - the matrix type for the explicit operator

    Output Parameter:
.   mat - the explicit  operator

    Note:
    This computation is done by applying the operators to columns of the identity matrix.
    This routine is costly in general, and is recommended for use only with relatively small systems.
    Currently, this routine uses a dense matrix format if mattype == NULL.

    Level: advanced

@*/
PetscErrorCode MatComputeOperator(Mat inmat, MatType mattype, Mat *mat) {
  PetscFunctionBegin;
  PetscValidHeaderSpecific(inmat, MAT_CLASSID, 1);
  PetscValidPointer(mat, 3);
  PetscCall(MatConvert_Shell(inmat, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
  PetscFunctionReturn(0);
}

/*@
    MatComputeOperatorTranspose - Computes the explicit matrix representation of
        a give matrix that can apply MatMultTranspose()

    Collective on Mat

    Input Parameters:
+   inmat - the matrix
-   mattype - the matrix type for the explicit operator

    Output Parameter:
.   mat - the explicit  operator transposed

    Note:
    This computation is done by applying the transpose of the operator to columns of the identity matrix.
    This routine is costly in general, and is recommended for use only with relatively small systems.
    Currently, this routine uses a dense matrix format if mattype == NULL.

    Level: advanced

@*/
PetscErrorCode MatComputeOperatorTranspose(Mat inmat, MatType mattype, Mat *mat) {
  Mat A;

  PetscFunctionBegin;
  PetscValidHeaderSpecific(inmat, MAT_CLASSID, 1);
  PetscValidPointer(mat, 3);
  PetscCall(MatCreateTranspose(inmat, &A));
  PetscCall(MatConvert_Shell(A, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
  PetscCall(MatDestroy(&A));
  PetscFunctionReturn(0);
}

/*@
  MatChop - Set all values in the matrix less than the tolerance to zero

  Input Parameters:
+ A   - The matrix
- tol - The zero tolerance

  Output Parameters:
. A - The chopped matrix

  Level: intermediate

.seealso: `MatCreate()`, `MatZeroEntries()`
 @*/
PetscErrorCode MatChop(Mat A, PetscReal tol) {
  Mat          a;
  PetscScalar *newVals;
  PetscInt    *newCols, rStart, rEnd, numRows, maxRows, r, colMax = 0;
  PetscBool    flg;

  PetscFunctionBegin;
  PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)A, &flg, MATSEQDENSE, MATMPIDENSE, ""));
  if (flg) {
    PetscCall(MatDenseGetLocalMatrix(A, &a));
    PetscCall(MatDenseGetLDA(a, &r));
    PetscCall(MatGetSize(a, &rStart, &rEnd));
    PetscCall(MatDenseGetArray(a, &newVals));
    for (; colMax < rEnd; ++colMax) {
      for (maxRows = 0; maxRows < rStart; ++maxRows) newVals[maxRows + colMax * r] = PetscAbsScalar(newVals[maxRows + colMax * r]) < tol ? 0.0 : newVals[maxRows + colMax * r];
    }
    PetscCall(MatDenseRestoreArray(a, &newVals));
  } else {
    PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd));
    PetscCall(MatGetRowUpperTriangular(A));
    for (r = rStart; r < rEnd; ++r) {
      PetscInt ncols;

      PetscCall(MatGetRow(A, r, &ncols, NULL, NULL));
      colMax = PetscMax(colMax, ncols);
      PetscCall(MatRestoreRow(A, r, &ncols, NULL, NULL));
    }
    numRows = rEnd - rStart;
    PetscCall(MPIU_Allreduce(&numRows, &maxRows, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)A)));
    PetscCall(PetscMalloc2(colMax, &newCols, colMax, &newVals));
    PetscCall(MatGetOption(A, MAT_NO_OFF_PROC_ENTRIES, &flg)); /* cache user-defined value */
    PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
    /* short-circuit code in MatAssemblyBegin() and MatAssemblyEnd()             */
    /* that are potentially called many times depending on the distribution of A */
    for (r = rStart; r < rStart + maxRows; ++r) {
      const PetscScalar *vals;
      const PetscInt    *cols;
      PetscInt           ncols, newcols, c;

      if (r < rEnd) {
        PetscCall(MatGetRow(A, r, &ncols, &cols, &vals));
        for (c = 0; c < ncols; ++c) {
          newCols[c] = cols[c];
          newVals[c] = PetscAbsScalar(vals[c]) < tol ? 0.0 : vals[c];
        }
        newcols = ncols;
        PetscCall(MatRestoreRow(A, r, &ncols, &cols, &vals));
        PetscCall(MatSetValues(A, 1, &r, newcols, newCols, newVals, INSERT_VALUES));
      }
      PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
      PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
    }
    PetscCall(MatRestoreRowUpperTriangular(A));
    PetscCall(PetscFree2(newCols, newVals));
    PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, flg)); /* reset option to its user-defined value */
  }
  PetscFunctionReturn(0);
}