xref: /petsc/src/mat/utils/axpy.c (revision 48a46eb9bd028bec07ec0f396b1a3abb43f14558)
16f79c3a4SBarry Smith 
2af0996ceSBarry Smith #include <petsc/private/matimpl.h> /*I   "petscmat.h"  I*/
36f79c3a4SBarry Smith 
49371c9d4SSatish Balay static PetscErrorCode MatTransposeAXPY_Private(Mat Y, PetscScalar a, Mat X, MatStructure str, Mat T) {
5d54c938cSPierre Jolivet   Mat A, F;
65f80ce2aSJacob Faibussowitsch   PetscErrorCode (*f)(Mat, Mat *);
7d54c938cSPierre Jolivet 
8d54c938cSPierre Jolivet   PetscFunctionBegin;
99566063dSJacob Faibussowitsch   PetscCall(PetscObjectQueryFunction((PetscObject)T, "MatTransposeGetMat_C", &f));
10d54c938cSPierre Jolivet   if (f) {
119566063dSJacob Faibussowitsch     PetscCall(MatTransposeGetMat(T, &A));
12d54c938cSPierre Jolivet     if (T == X) {
139566063dSJacob Faibussowitsch       PetscCall(PetscInfo(NULL, "Explicitly transposing X of type MATTRANSPOSEMAT to perform MatAXPY()\n"));
149566063dSJacob Faibussowitsch       PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &F));
15d54c938cSPierre Jolivet       A = Y;
16d54c938cSPierre Jolivet     } else {
179566063dSJacob Faibussowitsch       PetscCall(PetscInfo(NULL, "Transposing X because Y of type MATTRANSPOSEMAT to perform MatAXPY()\n"));
189566063dSJacob Faibussowitsch       PetscCall(MatTranspose(X, MAT_INITIAL_MATRIX, &F));
19d54c938cSPierre Jolivet     }
20d54c938cSPierre Jolivet   } else {
219566063dSJacob Faibussowitsch     PetscCall(MatHermitianTransposeGetMat(T, &A));
22d54c938cSPierre Jolivet     if (T == X) {
239566063dSJacob Faibussowitsch       PetscCall(PetscInfo(NULL, "Explicitly Hermitian transposing X of type MATTRANSPOSEMAT to perform MatAXPY()\n"));
249566063dSJacob Faibussowitsch       PetscCall(MatHermitianTranspose(A, MAT_INITIAL_MATRIX, &F));
25d54c938cSPierre Jolivet       A = Y;
26d54c938cSPierre Jolivet     } else {
279566063dSJacob Faibussowitsch       PetscCall(PetscInfo(NULL, "Hermitian transposing X because Y of type MATTRANSPOSEMAT to perform MatAXPY()\n"));
289566063dSJacob Faibussowitsch       PetscCall(MatHermitianTranspose(X, MAT_INITIAL_MATRIX, &F));
29d54c938cSPierre Jolivet     }
30d54c938cSPierre Jolivet   }
319566063dSJacob Faibussowitsch   PetscCall(MatAXPY(A, a, F, str));
329566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&F));
33d54c938cSPierre Jolivet   PetscFunctionReturn(0);
34d54c938cSPierre Jolivet }
35d54c938cSPierre Jolivet 
3606be10caSBarry Smith /*@
3721c89e3eSBarry Smith    MatAXPY - Computes Y = a*X + Y.
386f79c3a4SBarry Smith 
393f9fe445SBarry Smith    Logically  Collective on Mat
40fee21e36SBarry Smith 
4198a79cdbSBarry Smith    Input Parameters:
42607cd303SBarry Smith +  a - the scalar multiplier
43607cd303SBarry Smith .  X - the first matrix
44607cd303SBarry Smith .  Y - the second matrix
45531d0c6aSBarry Smith -  str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN, UNKNOWN_NONZERO_PATTERN, or SUBSET_NONZERO_PATTERN (nonzeros of X is a subset of Y's)
4698a79cdbSBarry Smith 
47edf9a46cSStefano Zampini    Notes: No operation is performed when a is zero.
48edf9a46cSStefano Zampini 
492860a424SLois Curfman McInnes    Level: intermediate
502860a424SLois Curfman McInnes 
51db781477SPatrick Sanan .seealso: `MatAYPX()`
5206be10caSBarry Smith  @*/
539371c9d4SSatish Balay PetscErrorCode MatAXPY(Mat Y, PetscScalar a, Mat X, MatStructure str) {
54646531bbSStefano Zampini   PetscInt  M1, M2, N1, N2;
55c1ac3661SBarry Smith   PetscInt  m1, m2, n1, n2;
56a683ea4aSPierre Jolivet   MatType   t1, t2;
57a683ea4aSPierre Jolivet   PetscBool sametype, transpose;
586f79c3a4SBarry Smith 
593a40ed3dSBarry Smith   PetscFunctionBegin;
600700a824SBarry Smith   PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
61c5eb9154SBarry Smith   PetscValidLogicalCollectiveScalar(Y, a, 2);
62646531bbSStefano Zampini   PetscValidHeaderSpecific(X, MAT_CLASSID, 3);
63646531bbSStefano Zampini   PetscCheckSameComm(Y, 1, X, 3);
649566063dSJacob Faibussowitsch   PetscCall(MatGetSize(X, &M1, &N1));
659566063dSJacob Faibussowitsch   PetscCall(MatGetSize(Y, &M2, &N2));
669566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(X, &m1, &n1));
679566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(Y, &m2, &n2));
682c71b3e2SJacob Faibussowitsch   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);
692c71b3e2SJacob Faibussowitsch   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);
702c71b3e2SJacob Faibussowitsch   PetscCheck(Y->assembled, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix (Y)");
712c71b3e2SJacob Faibussowitsch   PetscCheck(X->assembled, PetscObjectComm((PetscObject)X), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix (X)");
72edf9a46cSStefano Zampini   if (a == (PetscScalar)0.0) PetscFunctionReturn(0);
73edf9a46cSStefano Zampini   if (Y == X) {
749566063dSJacob Faibussowitsch     PetscCall(MatScale(Y, 1.0 + a));
75edf9a46cSStefano Zampini     PetscFunctionReturn(0);
76edf9a46cSStefano Zampini   }
779566063dSJacob Faibussowitsch   PetscCall(MatGetType(X, &t1));
789566063dSJacob Faibussowitsch   PetscCall(MatGetType(Y, &t2));
799566063dSJacob Faibussowitsch   PetscCall(PetscStrcmp(t1, t2, &sametype));
809566063dSJacob Faibussowitsch   PetscCall(PetscLogEventBegin(MAT_AXPY, Y, 0, 0, 0));
8193c87e65SStefano Zampini   if (Y->ops->axpy && (sametype || X->ops->axpy == Y->ops->axpy)) {
82dbbe0bcdSBarry Smith     PetscUseTypeMethod(Y, axpy, a, X, str);
83d4bb536fSBarry Smith   } else {
849566063dSJacob Faibussowitsch     PetscCall(PetscStrcmp(t1, MATTRANSPOSEMAT, &transpose));
85a683ea4aSPierre Jolivet     if (transpose) {
869566063dSJacob Faibussowitsch       PetscCall(MatTransposeAXPY_Private(Y, a, X, str, X));
87a683ea4aSPierre Jolivet     } else {
889566063dSJacob Faibussowitsch       PetscCall(PetscStrcmp(t2, MATTRANSPOSEMAT, &transpose));
89a683ea4aSPierre Jolivet       if (transpose) {
909566063dSJacob Faibussowitsch         PetscCall(MatTransposeAXPY_Private(Y, a, X, str, Y));
91a683ea4aSPierre Jolivet       } else {
929566063dSJacob Faibussowitsch         PetscCall(MatAXPY_Basic(Y, a, X, str));
93607cd303SBarry Smith       }
94a683ea4aSPierre Jolivet     }
95a683ea4aSPierre Jolivet   }
969566063dSJacob Faibussowitsch   PetscCall(PetscLogEventEnd(MAT_AXPY, Y, 0, 0, 0));
97607cd303SBarry Smith   PetscFunctionReturn(0);
98607cd303SBarry Smith }
99607cd303SBarry Smith 
1009371c9d4SSatish Balay PetscErrorCode MatAXPY_Basic_Preallocate(Mat Y, Mat X, Mat *B) {
10127afe9e8SStefano Zampini   PetscErrorCode (*preall)(Mat, Mat, Mat *) = NULL;
102646531bbSStefano Zampini 
103646531bbSStefano Zampini   PetscFunctionBegin;
104646531bbSStefano Zampini   /* look for any available faster alternative to the general preallocator */
1059566063dSJacob Faibussowitsch   PetscCall(PetscObjectQueryFunction((PetscObject)Y, "MatAXPYGetPreallocation_C", &preall));
106646531bbSStefano Zampini   if (preall) {
1079566063dSJacob Faibussowitsch     PetscCall((*preall)(Y, X, B));
108646531bbSStefano Zampini   } else { /* Use MatPrellocator, assumes same row-col distribution */
109646531bbSStefano Zampini     Mat      preallocator;
110646531bbSStefano Zampini     PetscInt r, rstart, rend;
111646531bbSStefano Zampini     PetscInt m, n, M, N;
112646531bbSStefano Zampini 
1139566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(Y));
1149566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(X));
1159566063dSJacob Faibussowitsch     PetscCall(MatGetSize(Y, &M, &N));
1169566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSize(Y, &m, &n));
1179566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), &preallocator));
1189566063dSJacob Faibussowitsch     PetscCall(MatSetType(preallocator, MATPREALLOCATOR));
1199566063dSJacob Faibussowitsch     PetscCall(MatSetLayouts(preallocator, Y->rmap, Y->cmap));
1209566063dSJacob Faibussowitsch     PetscCall(MatSetUp(preallocator));
1219566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(preallocator, &rstart, &rend));
122646531bbSStefano Zampini     for (r = rstart; r < rend; ++r) {
123646531bbSStefano Zampini       PetscInt           ncols;
124646531bbSStefano Zampini       const PetscInt    *row;
125646531bbSStefano Zampini       const PetscScalar *vals;
126646531bbSStefano Zampini 
1279566063dSJacob Faibussowitsch       PetscCall(MatGetRow(Y, r, &ncols, &row, &vals));
1289566063dSJacob Faibussowitsch       PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
1299566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(Y, r, &ncols, &row, &vals));
1309566063dSJacob Faibussowitsch       PetscCall(MatGetRow(X, r, &ncols, &row, &vals));
1319566063dSJacob Faibussowitsch       PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
1329566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(X, r, &ncols, &row, &vals));
133646531bbSStefano Zampini     }
1349566063dSJacob Faibussowitsch     PetscCall(MatSetOption(preallocator, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
1359566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(preallocator, MAT_FINAL_ASSEMBLY));
1369566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(preallocator, MAT_FINAL_ASSEMBLY));
1379566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(Y));
1389566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(X));
139646531bbSStefano Zampini 
1409566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), B));
1419566063dSJacob Faibussowitsch     PetscCall(MatSetType(*B, ((PetscObject)Y)->type_name));
1429566063dSJacob Faibussowitsch     PetscCall(MatSetLayouts(*B, Y->rmap, Y->cmap));
1439566063dSJacob Faibussowitsch     PetscCall(MatPreallocatorPreallocate(preallocator, PETSC_FALSE, *B));
1449566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&preallocator));
145646531bbSStefano Zampini   }
146646531bbSStefano Zampini   PetscFunctionReturn(0);
147646531bbSStefano Zampini }
148646531bbSStefano Zampini 
1499371c9d4SSatish Balay PetscErrorCode MatAXPY_Basic(Mat Y, PetscScalar a, Mat X, MatStructure str) {
150be705e3aSPierre Jolivet   PetscBool isshell, isdense, isnest;
151edf9a46cSStefano Zampini 
152edf9a46cSStefano Zampini   PetscFunctionBegin;
1539566063dSJacob Faibussowitsch   PetscCall(MatIsShell(Y, &isshell));
154edf9a46cSStefano Zampini   if (isshell) { /* MatShell has special support for AXPY */
155edf9a46cSStefano Zampini     PetscErrorCode (*f)(Mat, PetscScalar, Mat, MatStructure);
156edf9a46cSStefano Zampini 
1579566063dSJacob Faibussowitsch     PetscCall(MatGetOperation(Y, MATOP_AXPY, (void (**)(void)) & f));
158edf9a46cSStefano Zampini     if (f) {
1599566063dSJacob Faibussowitsch       PetscCall((*f)(Y, a, X, str));
160edf9a46cSStefano Zampini       PetscFunctionReturn(0);
161edf9a46cSStefano Zampini     }
162edf9a46cSStefano Zampini   }
16393c87e65SStefano Zampini   /* no need to preallocate if Y is dense */
1649566063dSJacob Faibussowitsch   PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)Y, &isdense, MATSEQDENSE, MATMPIDENSE, ""));
165be705e3aSPierre Jolivet   if (isdense) {
1669566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)X, MATNEST, &isnest));
167be705e3aSPierre Jolivet     if (isnest) {
1689566063dSJacob Faibussowitsch       PetscCall(MatAXPY_Dense_Nest(Y, a, X));
169be705e3aSPierre Jolivet       PetscFunctionReturn(0);
170be705e3aSPierre Jolivet     }
171be705e3aSPierre Jolivet     if (str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN) str = SUBSET_NONZERO_PATTERN;
172be705e3aSPierre Jolivet   }
173a6ab3590SBarry Smith   if (str != DIFFERENT_NONZERO_PATTERN && str != UNKNOWN_NONZERO_PATTERN) {
174edf9a46cSStefano Zampini     PetscInt           i, start, end, j, ncols, m, n;
17538baddfdSBarry Smith     const PetscInt    *row;
176b3cc6726SBarry Smith     PetscScalar       *val;
177b3cc6726SBarry Smith     const PetscScalar *vals;
178607cd303SBarry Smith 
1799566063dSJacob Faibussowitsch     PetscCall(MatGetSize(X, &m, &n));
1809566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(X, &start, &end));
1819566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(X));
182f4df32b1SMatthew Knepley     if (a == 1.0) {
183d4bb536fSBarry Smith       for (i = start; i < end; i++) {
1849566063dSJacob Faibussowitsch         PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
1859566063dSJacob Faibussowitsch         PetscCall(MatSetValues(Y, 1, &i, ncols, row, vals, ADD_VALUES));
1869566063dSJacob Faibussowitsch         PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
187d4bb536fSBarry Smith       }
188d4bb536fSBarry Smith     } else {
18921a3365eSStefano Zampini       PetscInt vs = 100;
19021a3365eSStefano Zampini       /* realloc if needed, as this function may be used in parallel */
1919566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(vs, &val));
19206be10caSBarry Smith       for (i = start; i < end; i++) {
1939566063dSJacob Faibussowitsch         PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
19421a3365eSStefano Zampini         if (vs < ncols) {
19521a3365eSStefano Zampini           vs = PetscMin(2 * ncols, n);
1969566063dSJacob Faibussowitsch           PetscCall(PetscRealloc(vs * sizeof(*val), &val));
1976f79c3a4SBarry Smith         }
19821a3365eSStefano Zampini         for (j = 0; j < ncols; j++) val[j] = a * vals[j];
1999566063dSJacob Faibussowitsch         PetscCall(MatSetValues(Y, 1, &i, ncols, row, val, ADD_VALUES));
2009566063dSJacob Faibussowitsch         PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
2016f79c3a4SBarry Smith       }
2029566063dSJacob Faibussowitsch       PetscCall(PetscFree(val));
203d4bb536fSBarry Smith     }
2049566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(X));
2059566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
2069566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
207edf9a46cSStefano Zampini   } else {
208edf9a46cSStefano Zampini     Mat B;
209edf9a46cSStefano Zampini 
2109566063dSJacob Faibussowitsch     PetscCall(MatAXPY_Basic_Preallocate(Y, X, &B));
2119566063dSJacob Faibussowitsch     PetscCall(MatAXPY_BasicWithPreallocation(B, Y, a, X, str));
2129566063dSJacob Faibussowitsch     PetscCall(MatHeaderMerge(Y, &B));
213edf9a46cSStefano Zampini   }
2143a40ed3dSBarry Smith   PetscFunctionReturn(0);
2156f79c3a4SBarry Smith }
216052efed2SBarry Smith 
2179371c9d4SSatish Balay PetscErrorCode MatAXPY_BasicWithPreallocation(Mat B, Mat Y, PetscScalar a, Mat X, MatStructure str) {
218ec7775f6SShri Abhyankar   PetscInt           i, start, end, j, ncols, m, n;
219ec7775f6SShri Abhyankar   const PetscInt    *row;
220ec7775f6SShri Abhyankar   PetscScalar       *val;
221ec7775f6SShri Abhyankar   const PetscScalar *vals;
222ec7775f6SShri Abhyankar 
223ec7775f6SShri Abhyankar   PetscFunctionBegin;
2249566063dSJacob Faibussowitsch   PetscCall(MatGetSize(X, &m, &n));
2259566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(X, &start, &end));
2269566063dSJacob Faibussowitsch   PetscCall(MatGetRowUpperTriangular(Y));
2279566063dSJacob Faibussowitsch   PetscCall(MatGetRowUpperTriangular(X));
228ec7775f6SShri Abhyankar   if (a == 1.0) {
229ec7775f6SShri Abhyankar     for (i = start; i < end; i++) {
2309566063dSJacob Faibussowitsch       PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
2319566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
2329566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));
233ec7775f6SShri Abhyankar 
2349566063dSJacob Faibussowitsch       PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
2359566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
2369566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
237ec7775f6SShri Abhyankar     }
238ec7775f6SShri Abhyankar   } else {
23921a3365eSStefano Zampini     PetscInt vs = 100;
24021a3365eSStefano Zampini     /* realloc if needed, as this function may be used in parallel */
2419566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(vs, &val));
242ec7775f6SShri Abhyankar     for (i = start; i < end; i++) {
2439566063dSJacob Faibussowitsch       PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
2449566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
2459566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));
246ec7775f6SShri Abhyankar 
2479566063dSJacob Faibussowitsch       PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
24821a3365eSStefano Zampini       if (vs < ncols) {
24921a3365eSStefano Zampini         vs = PetscMin(2 * ncols, n);
2509566063dSJacob Faibussowitsch         PetscCall(PetscRealloc(vs * sizeof(*val), &val));
251ec7775f6SShri Abhyankar       }
25221a3365eSStefano Zampini       for (j = 0; j < ncols; j++) val[j] = a * vals[j];
2539566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, val, ADD_VALUES));
2549566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
255ec7775f6SShri Abhyankar     }
2569566063dSJacob Faibussowitsch     PetscCall(PetscFree(val));
257ec7775f6SShri Abhyankar   }
2589566063dSJacob Faibussowitsch   PetscCall(MatRestoreRowUpperTriangular(Y));
2599566063dSJacob Faibussowitsch   PetscCall(MatRestoreRowUpperTriangular(X));
2609566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2619566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
262ec7775f6SShri Abhyankar   PetscFunctionReturn(0);
263ec7775f6SShri Abhyankar }
264ec7775f6SShri Abhyankar 
265052efed2SBarry Smith /*@
26687828ca2SBarry Smith    MatShift - Computes Y =  Y + a I, where a is a PetscScalar and I is the identity matrix.
267052efed2SBarry Smith 
2683f9fe445SBarry Smith    Neighbor-wise Collective on Mat
269fee21e36SBarry Smith 
27098a79cdbSBarry Smith    Input Parameters:
27198a79cdbSBarry Smith +  Y - the matrices
27287828ca2SBarry Smith -  a - the PetscScalar
27398a79cdbSBarry Smith 
2742860a424SLois Curfman McInnes    Level: intermediate
2752860a424SLois Curfman McInnes 
27695452b02SPatrick Sanan    Notes:
27784e19e3eSJunchao Zhang     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)
27884e19e3eSJunchao Zhang 
27995452b02SPatrick Sanan     If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
2806f33a894SBarry Smith    fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
281edf9a46cSStefano Zampini    entry). No operation is performed when a is zero.
2826f33a894SBarry Smith 
2830c0fd78eSBarry Smith    To form Y = Y + diag(V) use MatDiagonalSet()
2840c0fd78eSBarry Smith 
285db781477SPatrick Sanan .seealso: `MatDiagonalSet()`, `MatScale()`, `MatDiagonalScale()`
286052efed2SBarry Smith  @*/
2879371c9d4SSatish Balay PetscErrorCode MatShift(Mat Y, PetscScalar a) {
2883a40ed3dSBarry Smith   PetscFunctionBegin;
2890700a824SBarry Smith   PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
2905f80ce2aSJacob Faibussowitsch   PetscCheck(Y->assembled, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix");
2915f80ce2aSJacob Faibussowitsch   PetscCheck(!Y->factortype, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
2924994cf47SJed Brown   MatCheckPreallocated(Y, 1);
2931b71c1a7SBarry Smith   if (a == 0.0) PetscFunctionReturn(0);
294b50b34bdSBarry Smith 
295dbbe0bcdSBarry Smith   if (Y->ops->shift) PetscUseTypeMethod(Y, shift, a);
2969566063dSJacob Faibussowitsch   else PetscCall(MatShift_Basic(Y, a));
2977d68702bSBarry Smith 
2989566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)Y));
2993a40ed3dSBarry Smith   PetscFunctionReturn(0);
300052efed2SBarry Smith }
3016d84be18SBarry Smith 
3029371c9d4SSatish Balay PetscErrorCode MatDiagonalSet_Default(Mat Y, Vec D, InsertMode is) {
30367576b8bSBarry Smith   PetscInt           i, start, end;
304fa2e578bSStefano Zampini   const PetscScalar *v;
30509f38230SBarry Smith 
30609f38230SBarry Smith   PetscFunctionBegin;
3079566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(Y, &start, &end));
3089566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(D, &v));
309*48a46eb9SPierre Jolivet   for (i = start; i < end; i++) PetscCall(MatSetValues(Y, 1, &i, 1, &i, v + i - start, is));
3109566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(D, &v));
3119566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
3129566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
31309f38230SBarry Smith   PetscFunctionReturn(0);
31409f38230SBarry Smith }
31509f38230SBarry Smith 
3166d84be18SBarry Smith /*@
317f56f2b3fSBarry Smith    MatDiagonalSet - Computes Y = Y + D, where D is a diagonal matrix
318f56f2b3fSBarry Smith    that is represented as a vector. Or Y[i,i] = D[i] if InsertMode is
319f56f2b3fSBarry Smith    INSERT_VALUES.
3206d84be18SBarry Smith 
32148e586daSJose E. Roman    Neighbor-wise Collective on Mat
32248e586daSJose E. Roman 
3236d84be18SBarry Smith    Input Parameters:
32498a79cdbSBarry Smith +  Y - the input matrix
325f56f2b3fSBarry Smith .  D - the diagonal matrix, represented as a vector
326f56f2b3fSBarry Smith -  i - INSERT_VALUES or ADD_VALUES
3276d84be18SBarry Smith 
32895452b02SPatrick Sanan    Notes:
32995452b02SPatrick Sanan     If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
3306f33a894SBarry Smith    fill the matrix so that all diagonal entries have a value (with a value of zero for those locations that would not have an
3316f33a894SBarry Smith    entry).
3326f33a894SBarry Smith 
3332860a424SLois Curfman McInnes    Level: intermediate
3342860a424SLois Curfman McInnes 
335db781477SPatrick Sanan .seealso: `MatShift()`, `MatScale()`, `MatDiagonalScale()`
3366d84be18SBarry Smith @*/
3379371c9d4SSatish Balay PetscErrorCode MatDiagonalSet(Mat Y, Vec D, InsertMode is) {
33867576b8bSBarry Smith   PetscInt matlocal, veclocal;
3396d84be18SBarry Smith 
3403a40ed3dSBarry Smith   PetscFunctionBegin;
3410700a824SBarry Smith   PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
3420700a824SBarry Smith   PetscValidHeaderSpecific(D, VEC_CLASSID, 2);
3439566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(Y, &matlocal, NULL));
3449566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(D, &veclocal));
3455f80ce2aSJacob Faibussowitsch   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);
346dbbe0bcdSBarry Smith   if (Y->ops->diagonalset) PetscUseTypeMethod(Y, diagonalset, D, is);
347dbbe0bcdSBarry Smith   else PetscCall(MatDiagonalSet_Default(Y, D, is));
3489566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)Y));
3493a40ed3dSBarry Smith   PetscFunctionReturn(0);
3506d84be18SBarry Smith }
351d4bb536fSBarry Smith 
352d4bb536fSBarry Smith /*@
35304aac2b0SHong Zhang    MatAYPX - Computes Y = a*Y + X.
354d4bb536fSBarry Smith 
3553f9fe445SBarry Smith    Logically on Mat
356fee21e36SBarry Smith 
35798a79cdbSBarry Smith    Input Parameters:
35804aac2b0SHong Zhang +  a - the PetscScalar multiplier
35904aac2b0SHong Zhang .  Y - the first matrix
36004aac2b0SHong Zhang .  X - the second matrix
361531d0c6aSBarry Smith -  str - either SAME_NONZERO_PATTERN, DIFFERENT_NONZERO_PATTERN, UNKNOWN_NONZERO_PATTERN, or SUBSET_NONZERO_PATTERN (nonzeros of X is a subset of Y's)
362d4bb536fSBarry Smith 
3632860a424SLois Curfman McInnes    Level: intermediate
3642860a424SLois Curfman McInnes 
365db781477SPatrick Sanan .seealso: `MatAXPY()`
366d4bb536fSBarry Smith  @*/
3679371c9d4SSatish Balay PetscErrorCode MatAYPX(Mat Y, PetscScalar a, Mat X, MatStructure str) {
3683a40ed3dSBarry Smith   PetscFunctionBegin;
3699566063dSJacob Faibussowitsch   PetscCall(MatScale(Y, a));
3709566063dSJacob Faibussowitsch   PetscCall(MatAXPY(Y, 1.0, X, str));
3713a40ed3dSBarry Smith   PetscFunctionReturn(0);
372d4bb536fSBarry Smith }
373b0a32e0cSBarry Smith 
374b0a32e0cSBarry Smith /*@
3750bacdadaSStefano Zampini     MatComputeOperator - Computes the explicit matrix
376b0a32e0cSBarry Smith 
377b0a32e0cSBarry Smith     Collective on Mat
378b0a32e0cSBarry Smith 
379d8d19677SJose E. Roman     Input Parameters:
380186a3e20SStefano Zampini +   inmat - the matrix
381186a3e20SStefano Zampini -   mattype - the matrix type for the explicit operator
382b0a32e0cSBarry Smith 
383b0a32e0cSBarry Smith     Output Parameter:
384a5b23f4aSJose E. Roman .   mat - the explicit  operator
385b0a32e0cSBarry Smith 
386b0a32e0cSBarry Smith     Notes:
387186a3e20SStefano Zampini     This computation is done by applying the operators to columns of the identity matrix.
388186a3e20SStefano Zampini     This routine is costly in general, and is recommended for use only with relatively small systems.
389186a3e20SStefano Zampini     Currently, this routine uses a dense matrix format if mattype == NULL.
390b0a32e0cSBarry Smith 
391b0a32e0cSBarry Smith     Level: advanced
392b0a32e0cSBarry Smith 
393b0a32e0cSBarry Smith @*/
3949371c9d4SSatish Balay PetscErrorCode MatComputeOperator(Mat inmat, MatType mattype, Mat *mat) {
395b0a32e0cSBarry Smith   PetscFunctionBegin;
3960700a824SBarry Smith   PetscValidHeaderSpecific(inmat, MAT_CLASSID, 1);
397186a3e20SStefano Zampini   PetscValidPointer(mat, 3);
3989566063dSJacob Faibussowitsch   PetscCall(MatConvert_Shell(inmat, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
39924f910e3SHong Zhang   PetscFunctionReturn(0);
40024f910e3SHong Zhang }
4014325cce7SMatthew G Knepley 
4024325cce7SMatthew G Knepley /*@
4030bacdadaSStefano Zampini     MatComputeOperatorTranspose - Computes the explicit matrix representation of
404f3b1f45cSBarry Smith         a give matrix that can apply MatMultTranspose()
405f3b1f45cSBarry Smith 
406f3b1f45cSBarry Smith     Collective on Mat
407f3b1f45cSBarry Smith 
4086b867d5aSJose E. Roman     Input Parameters:
4096b867d5aSJose E. Roman +   inmat - the matrix
4106b867d5aSJose E. Roman -   mattype - the matrix type for the explicit operator
411f3b1f45cSBarry Smith 
412f3b1f45cSBarry Smith     Output Parameter:
413a5b23f4aSJose E. Roman .   mat - the explicit  operator transposed
414f3b1f45cSBarry Smith 
415f3b1f45cSBarry Smith     Notes:
416186a3e20SStefano Zampini     This computation is done by applying the transpose of the operator to columns of the identity matrix.
417186a3e20SStefano Zampini     This routine is costly in general, and is recommended for use only with relatively small systems.
418186a3e20SStefano Zampini     Currently, this routine uses a dense matrix format if mattype == NULL.
419f3b1f45cSBarry Smith 
420f3b1f45cSBarry Smith     Level: advanced
421f3b1f45cSBarry Smith 
422f3b1f45cSBarry Smith @*/
4239371c9d4SSatish Balay PetscErrorCode MatComputeOperatorTranspose(Mat inmat, MatType mattype, Mat *mat) {
424186a3e20SStefano Zampini   Mat A;
425f3b1f45cSBarry Smith 
426f3b1f45cSBarry Smith   PetscFunctionBegin;
427f3b1f45cSBarry Smith   PetscValidHeaderSpecific(inmat, MAT_CLASSID, 1);
428186a3e20SStefano Zampini   PetscValidPointer(mat, 3);
4299566063dSJacob Faibussowitsch   PetscCall(MatCreateTranspose(inmat, &A));
4309566063dSJacob Faibussowitsch   PetscCall(MatConvert_Shell(A, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
4319566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&A));
432f3b1f45cSBarry Smith   PetscFunctionReturn(0);
433f3b1f45cSBarry Smith }
434f3b1f45cSBarry Smith 
435f3b1f45cSBarry Smith /*@
4364325cce7SMatthew G Knepley   MatChop - Set all values in the matrix less than the tolerance to zero
4374325cce7SMatthew G Knepley 
4384325cce7SMatthew G Knepley   Input Parameters:
4394325cce7SMatthew G Knepley + A   - The matrix
4404325cce7SMatthew G Knepley - tol - The zero tolerance
4414325cce7SMatthew G Knepley 
4424325cce7SMatthew G Knepley   Output Parameters:
4434325cce7SMatthew G Knepley . A - The chopped matrix
4444325cce7SMatthew G Knepley 
4454325cce7SMatthew G Knepley   Level: intermediate
4464325cce7SMatthew G Knepley 
447db781477SPatrick Sanan .seealso: `MatCreate()`, `MatZeroEntries()`
4483fc99919SSatish Balay  @*/
4499371c9d4SSatish Balay PetscErrorCode MatChop(Mat A, PetscReal tol) {
450038df967SPierre Jolivet   Mat          a;
4514325cce7SMatthew G Knepley   PetscScalar *newVals;
452cf5d3338SPierre Jolivet   PetscInt    *newCols, rStart, rEnd, numRows, maxRows, r, colMax = 0;
453038df967SPierre Jolivet   PetscBool    flg;
4544325cce7SMatthew G Knepley 
4554325cce7SMatthew G Knepley   PetscFunctionBegin;
4569566063dSJacob Faibussowitsch   PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)A, &flg, MATSEQDENSE, MATMPIDENSE, ""));
457038df967SPierre Jolivet   if (flg) {
4589566063dSJacob Faibussowitsch     PetscCall(MatDenseGetLocalMatrix(A, &a));
4599566063dSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(a, &r));
4609566063dSJacob Faibussowitsch     PetscCall(MatGetSize(a, &rStart, &rEnd));
4619566063dSJacob Faibussowitsch     PetscCall(MatDenseGetArray(a, &newVals));
462038df967SPierre Jolivet     for (; colMax < rEnd; ++colMax) {
4639371c9d4SSatish Balay       for (maxRows = 0; maxRows < rStart; ++maxRows) { newVals[maxRows + colMax * r] = PetscAbsScalar(newVals[maxRows + colMax * r]) < tol ? 0.0 : newVals[maxRows + colMax * r]; }
464038df967SPierre Jolivet     }
4659566063dSJacob Faibussowitsch     PetscCall(MatDenseRestoreArray(a, &newVals));
466038df967SPierre Jolivet   } else {
4679566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd));
4689566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(A));
4694325cce7SMatthew G Knepley     for (r = rStart; r < rEnd; ++r) {
4704325cce7SMatthew G Knepley       PetscInt ncols;
4714325cce7SMatthew G Knepley 
4729566063dSJacob Faibussowitsch       PetscCall(MatGetRow(A, r, &ncols, NULL, NULL));
4735f80ce2aSJacob Faibussowitsch       colMax = PetscMax(colMax, ncols);
4749566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(A, r, &ncols, NULL, NULL));
4754325cce7SMatthew G Knepley     }
4764325cce7SMatthew G Knepley     numRows = rEnd - rStart;
4771c2dc1cbSBarry Smith     PetscCall(MPIU_Allreduce(&numRows, &maxRows, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)A)));
4789566063dSJacob Faibussowitsch     PetscCall(PetscMalloc2(colMax, &newCols, colMax, &newVals));
4799566063dSJacob Faibussowitsch     PetscCall(MatGetOption(A, MAT_NO_OFF_PROC_ENTRIES, &flg)); /* cache user-defined value */
4809566063dSJacob Faibussowitsch     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
481cf5d3338SPierre Jolivet     /* short-circuit code in MatAssemblyBegin() and MatAssemblyEnd()             */
482cf5d3338SPierre Jolivet     /* that are potentially called many times depending on the distribution of A */
4834325cce7SMatthew G Knepley     for (r = rStart; r < rStart + maxRows; ++r) {
4844325cce7SMatthew G Knepley       const PetscScalar *vals;
4854325cce7SMatthew G Knepley       const PetscInt    *cols;
486fad45679SMatthew G. Knepley       PetscInt           ncols, newcols, c;
4874325cce7SMatthew G Knepley 
4884325cce7SMatthew G Knepley       if (r < rEnd) {
4899566063dSJacob Faibussowitsch         PetscCall(MatGetRow(A, r, &ncols, &cols, &vals));
4904325cce7SMatthew G Knepley         for (c = 0; c < ncols; ++c) {
4914325cce7SMatthew G Knepley           newCols[c] = cols[c];
4924325cce7SMatthew G Knepley           newVals[c] = PetscAbsScalar(vals[c]) < tol ? 0.0 : vals[c];
4934325cce7SMatthew G Knepley         }
494fad45679SMatthew G. Knepley         newcols = ncols;
4959566063dSJacob Faibussowitsch         PetscCall(MatRestoreRow(A, r, &ncols, &cols, &vals));
4969566063dSJacob Faibussowitsch         PetscCall(MatSetValues(A, 1, &r, newcols, newCols, newVals, INSERT_VALUES));
4974325cce7SMatthew G Knepley       }
4989566063dSJacob Faibussowitsch       PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
4999566063dSJacob Faibussowitsch       PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
5004325cce7SMatthew G Knepley     }
5019566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(A));
5029566063dSJacob Faibussowitsch     PetscCall(PetscFree2(newCols, newVals));
5039566063dSJacob Faibussowitsch     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, flg)); /* reset option to its user-defined value */
504038df967SPierre Jolivet   }
5054325cce7SMatthew G Knepley   PetscFunctionReturn(0);
5064325cce7SMatthew G Knepley }
507