xref: /petsc/src/mat/utils/axpy.c (revision 013e2dc7137d1d9dc4e6b44a828a8fb0f1bd6f29)
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) {
13*013e2dc7SBarry Smith       PetscCall(PetscInfo(NULL, "Explicitly transposing X of type MATTRANSPOSEVIRTUAL to perform MatAXPY()\n"));
149566063dSJacob Faibussowitsch       PetscCall(MatTranspose(A, MAT_INITIAL_MATRIX, &F));
15d54c938cSPierre Jolivet       A = Y;
16d54c938cSPierre Jolivet     } else {
17*013e2dc7SBarry Smith       PetscCall(PetscInfo(NULL, "Transposing X because Y of type MATTRANSPOSEVIRTUAL 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) {
23*013e2dc7SBarry Smith       PetscCall(PetscInfo(NULL, "Explicitly Hermitian transposing X of type MATHERITIANTRANSPOSEVIRTUAL to perform MatAXPY()\n"));
249566063dSJacob Faibussowitsch       PetscCall(MatHermitianTranspose(A, MAT_INITIAL_MATRIX, &F));
25d54c938cSPierre Jolivet       A = Y;
26d54c938cSPierre Jolivet     } else {
27*013e2dc7SBarry Smith       PetscCall(PetscInfo(NULL, "Hermitian transposing X because Y of type MATHERITIANTRANSPOSEVIRTUAL 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 
4711a5261eSBarry Smith    Note:
4811a5261eSBarry Smith    No operation is performed when a is zero.
49edf9a46cSStefano Zampini 
502860a424SLois Curfman McInnes    Level: intermediate
512860a424SLois Curfman McInnes 
52db781477SPatrick Sanan .seealso: `MatAYPX()`
5306be10caSBarry Smith  @*/
549371c9d4SSatish Balay PetscErrorCode MatAXPY(Mat Y, PetscScalar a, Mat X, MatStructure str) {
55646531bbSStefano Zampini   PetscInt  M1, M2, N1, N2;
56c1ac3661SBarry Smith   PetscInt  m1, m2, n1, n2;
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   }
77*013e2dc7SBarry Smith   PetscCall(PetscObjectObjectTypeCompare((PetscObject)X, (PetscObject)Y, &sametype));
789566063dSJacob Faibussowitsch   PetscCall(PetscLogEventBegin(MAT_AXPY, Y, 0, 0, 0));
7993c87e65SStefano Zampini   if (Y->ops->axpy && (sametype || X->ops->axpy == Y->ops->axpy)) {
80dbbe0bcdSBarry Smith     PetscUseTypeMethod(Y, axpy, a, X, str);
81d4bb536fSBarry Smith   } else {
82*013e2dc7SBarry Smith     PetscCall(PetscObjectTypeCompareAny((PetscObject)X, &transpose, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
83a683ea4aSPierre Jolivet     if (transpose) {
849566063dSJacob Faibussowitsch       PetscCall(MatTransposeAXPY_Private(Y, a, X, str, X));
85a683ea4aSPierre Jolivet     } else {
86*013e2dc7SBarry Smith       PetscCall(PetscObjectTypeCompareAny((PetscObject)Y, &transpose, MATTRANSPOSEVIRTUAL, MATHERMITIANTRANSPOSEVIRTUAL, ""));
87a683ea4aSPierre Jolivet       if (transpose) {
889566063dSJacob Faibussowitsch         PetscCall(MatTransposeAXPY_Private(Y, a, X, str, Y));
89a683ea4aSPierre Jolivet       } else {
909566063dSJacob Faibussowitsch         PetscCall(MatAXPY_Basic(Y, a, X, str));
91607cd303SBarry Smith       }
92a683ea4aSPierre Jolivet     }
93a683ea4aSPierre Jolivet   }
949566063dSJacob Faibussowitsch   PetscCall(PetscLogEventEnd(MAT_AXPY, Y, 0, 0, 0));
95607cd303SBarry Smith   PetscFunctionReturn(0);
96607cd303SBarry Smith }
97607cd303SBarry Smith 
989371c9d4SSatish Balay PetscErrorCode MatAXPY_Basic_Preallocate(Mat Y, Mat X, Mat *B) {
9927afe9e8SStefano Zampini   PetscErrorCode (*preall)(Mat, Mat, Mat *) = NULL;
100646531bbSStefano Zampini 
101646531bbSStefano Zampini   PetscFunctionBegin;
102646531bbSStefano Zampini   /* look for any available faster alternative to the general preallocator */
1039566063dSJacob Faibussowitsch   PetscCall(PetscObjectQueryFunction((PetscObject)Y, "MatAXPYGetPreallocation_C", &preall));
104646531bbSStefano Zampini   if (preall) {
1059566063dSJacob Faibussowitsch     PetscCall((*preall)(Y, X, B));
106646531bbSStefano Zampini   } else { /* Use MatPrellocator, assumes same row-col distribution */
107646531bbSStefano Zampini     Mat      preallocator;
108646531bbSStefano Zampini     PetscInt r, rstart, rend;
109646531bbSStefano Zampini     PetscInt m, n, M, N;
110646531bbSStefano Zampini 
1119566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(Y));
1129566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(X));
1139566063dSJacob Faibussowitsch     PetscCall(MatGetSize(Y, &M, &N));
1149566063dSJacob Faibussowitsch     PetscCall(MatGetLocalSize(Y, &m, &n));
1159566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), &preallocator));
1169566063dSJacob Faibussowitsch     PetscCall(MatSetType(preallocator, MATPREALLOCATOR));
1179566063dSJacob Faibussowitsch     PetscCall(MatSetLayouts(preallocator, Y->rmap, Y->cmap));
1189566063dSJacob Faibussowitsch     PetscCall(MatSetUp(preallocator));
1199566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(preallocator, &rstart, &rend));
120646531bbSStefano Zampini     for (r = rstart; r < rend; ++r) {
121646531bbSStefano Zampini       PetscInt           ncols;
122646531bbSStefano Zampini       const PetscInt    *row;
123646531bbSStefano Zampini       const PetscScalar *vals;
124646531bbSStefano Zampini 
1259566063dSJacob Faibussowitsch       PetscCall(MatGetRow(Y, r, &ncols, &row, &vals));
1269566063dSJacob Faibussowitsch       PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
1279566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(Y, r, &ncols, &row, &vals));
1289566063dSJacob Faibussowitsch       PetscCall(MatGetRow(X, r, &ncols, &row, &vals));
1299566063dSJacob Faibussowitsch       PetscCall(MatSetValues(preallocator, 1, &r, ncols, row, vals, INSERT_VALUES));
1309566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(X, r, &ncols, &row, &vals));
131646531bbSStefano Zampini     }
1329566063dSJacob Faibussowitsch     PetscCall(MatSetOption(preallocator, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
1339566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(preallocator, MAT_FINAL_ASSEMBLY));
1349566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(preallocator, MAT_FINAL_ASSEMBLY));
1359566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(Y));
1369566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(X));
137646531bbSStefano Zampini 
1389566063dSJacob Faibussowitsch     PetscCall(MatCreate(PetscObjectComm((PetscObject)Y), B));
1399566063dSJacob Faibussowitsch     PetscCall(MatSetType(*B, ((PetscObject)Y)->type_name));
1409566063dSJacob Faibussowitsch     PetscCall(MatSetLayouts(*B, Y->rmap, Y->cmap));
1419566063dSJacob Faibussowitsch     PetscCall(MatPreallocatorPreallocate(preallocator, PETSC_FALSE, *B));
1429566063dSJacob Faibussowitsch     PetscCall(MatDestroy(&preallocator));
143646531bbSStefano Zampini   }
144646531bbSStefano Zampini   PetscFunctionReturn(0);
145646531bbSStefano Zampini }
146646531bbSStefano Zampini 
1479371c9d4SSatish Balay PetscErrorCode MatAXPY_Basic(Mat Y, PetscScalar a, Mat X, MatStructure str) {
148be705e3aSPierre Jolivet   PetscBool isshell, isdense, isnest;
149edf9a46cSStefano Zampini 
150edf9a46cSStefano Zampini   PetscFunctionBegin;
1519566063dSJacob Faibussowitsch   PetscCall(MatIsShell(Y, &isshell));
152edf9a46cSStefano Zampini   if (isshell) { /* MatShell has special support for AXPY */
153edf9a46cSStefano Zampini     PetscErrorCode (*f)(Mat, PetscScalar, Mat, MatStructure);
154edf9a46cSStefano Zampini 
1559566063dSJacob Faibussowitsch     PetscCall(MatGetOperation(Y, MATOP_AXPY, (void (**)(void)) & f));
156edf9a46cSStefano Zampini     if (f) {
1579566063dSJacob Faibussowitsch       PetscCall((*f)(Y, a, X, str));
158edf9a46cSStefano Zampini       PetscFunctionReturn(0);
159edf9a46cSStefano Zampini     }
160edf9a46cSStefano Zampini   }
16193c87e65SStefano Zampini   /* no need to preallocate if Y is dense */
1629566063dSJacob Faibussowitsch   PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)Y, &isdense, MATSEQDENSE, MATMPIDENSE, ""));
163be705e3aSPierre Jolivet   if (isdense) {
1649566063dSJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare((PetscObject)X, MATNEST, &isnest));
165be705e3aSPierre Jolivet     if (isnest) {
1669566063dSJacob Faibussowitsch       PetscCall(MatAXPY_Dense_Nest(Y, a, X));
167be705e3aSPierre Jolivet       PetscFunctionReturn(0);
168be705e3aSPierre Jolivet     }
169be705e3aSPierre Jolivet     if (str == DIFFERENT_NONZERO_PATTERN || str == UNKNOWN_NONZERO_PATTERN) str = SUBSET_NONZERO_PATTERN;
170be705e3aSPierre Jolivet   }
171a6ab3590SBarry Smith   if (str != DIFFERENT_NONZERO_PATTERN && str != UNKNOWN_NONZERO_PATTERN) {
172edf9a46cSStefano Zampini     PetscInt           i, start, end, j, ncols, m, n;
17338baddfdSBarry Smith     const PetscInt    *row;
174b3cc6726SBarry Smith     PetscScalar       *val;
175b3cc6726SBarry Smith     const PetscScalar *vals;
176607cd303SBarry Smith 
1779566063dSJacob Faibussowitsch     PetscCall(MatGetSize(X, &m, &n));
1789566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(X, &start, &end));
1799566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(X));
180f4df32b1SMatthew Knepley     if (a == 1.0) {
181d4bb536fSBarry Smith       for (i = start; i < end; i++) {
1829566063dSJacob Faibussowitsch         PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
1839566063dSJacob Faibussowitsch         PetscCall(MatSetValues(Y, 1, &i, ncols, row, vals, ADD_VALUES));
1849566063dSJacob Faibussowitsch         PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
185d4bb536fSBarry Smith       }
186d4bb536fSBarry Smith     } else {
18721a3365eSStefano Zampini       PetscInt vs = 100;
18821a3365eSStefano Zampini       /* realloc if needed, as this function may be used in parallel */
1899566063dSJacob Faibussowitsch       PetscCall(PetscMalloc1(vs, &val));
19006be10caSBarry Smith       for (i = start; i < end; i++) {
1919566063dSJacob Faibussowitsch         PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
19221a3365eSStefano Zampini         if (vs < ncols) {
19321a3365eSStefano Zampini           vs = PetscMin(2 * ncols, n);
1949566063dSJacob Faibussowitsch           PetscCall(PetscRealloc(vs * sizeof(*val), &val));
1956f79c3a4SBarry Smith         }
19621a3365eSStefano Zampini         for (j = 0; j < ncols; j++) val[j] = a * vals[j];
1979566063dSJacob Faibussowitsch         PetscCall(MatSetValues(Y, 1, &i, ncols, row, val, ADD_VALUES));
1989566063dSJacob Faibussowitsch         PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
1996f79c3a4SBarry Smith       }
2009566063dSJacob Faibussowitsch       PetscCall(PetscFree(val));
201d4bb536fSBarry Smith     }
2029566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(X));
2039566063dSJacob Faibussowitsch     PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
2049566063dSJacob Faibussowitsch     PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
205edf9a46cSStefano Zampini   } else {
206edf9a46cSStefano Zampini     Mat B;
207edf9a46cSStefano Zampini 
2089566063dSJacob Faibussowitsch     PetscCall(MatAXPY_Basic_Preallocate(Y, X, &B));
2099566063dSJacob Faibussowitsch     PetscCall(MatAXPY_BasicWithPreallocation(B, Y, a, X, str));
2109566063dSJacob Faibussowitsch     PetscCall(MatHeaderMerge(Y, &B));
211edf9a46cSStefano Zampini   }
2123a40ed3dSBarry Smith   PetscFunctionReturn(0);
2136f79c3a4SBarry Smith }
214052efed2SBarry Smith 
2159371c9d4SSatish Balay PetscErrorCode MatAXPY_BasicWithPreallocation(Mat B, Mat Y, PetscScalar a, Mat X, MatStructure str) {
216ec7775f6SShri Abhyankar   PetscInt           i, start, end, j, ncols, m, n;
217ec7775f6SShri Abhyankar   const PetscInt    *row;
218ec7775f6SShri Abhyankar   PetscScalar       *val;
219ec7775f6SShri Abhyankar   const PetscScalar *vals;
220ec7775f6SShri Abhyankar 
221ec7775f6SShri Abhyankar   PetscFunctionBegin;
2229566063dSJacob Faibussowitsch   PetscCall(MatGetSize(X, &m, &n));
2239566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(X, &start, &end));
2249566063dSJacob Faibussowitsch   PetscCall(MatGetRowUpperTriangular(Y));
2259566063dSJacob Faibussowitsch   PetscCall(MatGetRowUpperTriangular(X));
226ec7775f6SShri Abhyankar   if (a == 1.0) {
227ec7775f6SShri Abhyankar     for (i = start; i < end; i++) {
2289566063dSJacob Faibussowitsch       PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
2299566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
2309566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));
231ec7775f6SShri Abhyankar 
2329566063dSJacob Faibussowitsch       PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
2339566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
2349566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
235ec7775f6SShri Abhyankar     }
236ec7775f6SShri Abhyankar   } else {
23721a3365eSStefano Zampini     PetscInt vs = 100;
23821a3365eSStefano Zampini     /* realloc if needed, as this function may be used in parallel */
2399566063dSJacob Faibussowitsch     PetscCall(PetscMalloc1(vs, &val));
240ec7775f6SShri Abhyankar     for (i = start; i < end; i++) {
2419566063dSJacob Faibussowitsch       PetscCall(MatGetRow(Y, i, &ncols, &row, &vals));
2429566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, vals, ADD_VALUES));
2439566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(Y, i, &ncols, &row, &vals));
244ec7775f6SShri Abhyankar 
2459566063dSJacob Faibussowitsch       PetscCall(MatGetRow(X, i, &ncols, &row, &vals));
24621a3365eSStefano Zampini       if (vs < ncols) {
24721a3365eSStefano Zampini         vs = PetscMin(2 * ncols, n);
2489566063dSJacob Faibussowitsch         PetscCall(PetscRealloc(vs * sizeof(*val), &val));
249ec7775f6SShri Abhyankar       }
25021a3365eSStefano Zampini       for (j = 0; j < ncols; j++) val[j] = a * vals[j];
2519566063dSJacob Faibussowitsch       PetscCall(MatSetValues(B, 1, &i, ncols, row, val, ADD_VALUES));
2529566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(X, i, &ncols, &row, &vals));
253ec7775f6SShri Abhyankar     }
2549566063dSJacob Faibussowitsch     PetscCall(PetscFree(val));
255ec7775f6SShri Abhyankar   }
2569566063dSJacob Faibussowitsch   PetscCall(MatRestoreRowUpperTriangular(Y));
2579566063dSJacob Faibussowitsch   PetscCall(MatRestoreRowUpperTriangular(X));
2589566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY));
2599566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY));
260ec7775f6SShri Abhyankar   PetscFunctionReturn(0);
261ec7775f6SShri Abhyankar }
262ec7775f6SShri Abhyankar 
263052efed2SBarry Smith /*@
26487828ca2SBarry Smith    MatShift - Computes Y =  Y + a I, where a is a PetscScalar and I is the identity matrix.
265052efed2SBarry Smith 
2663f9fe445SBarry Smith    Neighbor-wise Collective on Mat
267fee21e36SBarry Smith 
26898a79cdbSBarry Smith    Input Parameters:
26998a79cdbSBarry Smith +  Y - the matrices
27087828ca2SBarry Smith -  a - the PetscScalar
27198a79cdbSBarry Smith 
2722860a424SLois Curfman McInnes    Level: intermediate
2732860a424SLois Curfman McInnes 
27495452b02SPatrick Sanan    Notes:
27584e19e3eSJunchao 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)
27684e19e3eSJunchao Zhang 
27795452b02SPatrick Sanan     If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
2786f33a894SBarry 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
279edf9a46cSStefano Zampini    entry). No operation is performed when a is zero.
2806f33a894SBarry Smith 
2810c0fd78eSBarry Smith    To form Y = Y + diag(V) use MatDiagonalSet()
2820c0fd78eSBarry Smith 
283db781477SPatrick Sanan .seealso: `MatDiagonalSet()`, `MatScale()`, `MatDiagonalScale()`
284052efed2SBarry Smith  @*/
2859371c9d4SSatish Balay PetscErrorCode MatShift(Mat Y, PetscScalar a) {
2863a40ed3dSBarry Smith   PetscFunctionBegin;
2870700a824SBarry Smith   PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
2885f80ce2aSJacob Faibussowitsch   PetscCheck(Y->assembled, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for unassembled matrix");
2895f80ce2aSJacob Faibussowitsch   PetscCheck(!Y->factortype, PetscObjectComm((PetscObject)Y), PETSC_ERR_ARG_WRONGSTATE, "Not for factored matrix");
2904994cf47SJed Brown   MatCheckPreallocated(Y, 1);
2911b71c1a7SBarry Smith   if (a == 0.0) PetscFunctionReturn(0);
292b50b34bdSBarry Smith 
293dbbe0bcdSBarry Smith   if (Y->ops->shift) PetscUseTypeMethod(Y, shift, a);
2949566063dSJacob Faibussowitsch   else PetscCall(MatShift_Basic(Y, a));
2957d68702bSBarry Smith 
2969566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)Y));
2973a40ed3dSBarry Smith   PetscFunctionReturn(0);
298052efed2SBarry Smith }
2996d84be18SBarry Smith 
3009371c9d4SSatish Balay PetscErrorCode MatDiagonalSet_Default(Mat Y, Vec D, InsertMode is) {
30167576b8bSBarry Smith   PetscInt           i, start, end;
302fa2e578bSStefano Zampini   const PetscScalar *v;
30309f38230SBarry Smith 
30409f38230SBarry Smith   PetscFunctionBegin;
3059566063dSJacob Faibussowitsch   PetscCall(MatGetOwnershipRange(Y, &start, &end));
3069566063dSJacob Faibussowitsch   PetscCall(VecGetArrayRead(D, &v));
30748a46eb9SPierre Jolivet   for (i = start; i < end; i++) PetscCall(MatSetValues(Y, 1, &i, 1, &i, v + i - start, is));
3089566063dSJacob Faibussowitsch   PetscCall(VecRestoreArrayRead(D, &v));
3099566063dSJacob Faibussowitsch   PetscCall(MatAssemblyBegin(Y, MAT_FINAL_ASSEMBLY));
3109566063dSJacob Faibussowitsch   PetscCall(MatAssemblyEnd(Y, MAT_FINAL_ASSEMBLY));
31109f38230SBarry Smith   PetscFunctionReturn(0);
31209f38230SBarry Smith }
31309f38230SBarry Smith 
3146d84be18SBarry Smith /*@
315f56f2b3fSBarry Smith    MatDiagonalSet - Computes Y = Y + D, where D is a diagonal matrix
316f56f2b3fSBarry Smith    that is represented as a vector. Or Y[i,i] = D[i] if InsertMode is
317f56f2b3fSBarry Smith    INSERT_VALUES.
3186d84be18SBarry Smith 
31948e586daSJose E. Roman    Neighbor-wise Collective on Mat
32048e586daSJose E. Roman 
3216d84be18SBarry Smith    Input Parameters:
32298a79cdbSBarry Smith +  Y - the input matrix
323f56f2b3fSBarry Smith .  D - the diagonal matrix, represented as a vector
324f56f2b3fSBarry Smith -  i - INSERT_VALUES or ADD_VALUES
3256d84be18SBarry Smith 
32611a5261eSBarry Smith    Note:
32795452b02SPatrick Sanan     If the matrix Y is missing some diagonal entries this routine can be very slow. To make it fast one should initially
3286f33a894SBarry 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
3296f33a894SBarry Smith    entry).
3306f33a894SBarry Smith 
3312860a424SLois Curfman McInnes    Level: intermediate
3322860a424SLois Curfman McInnes 
333db781477SPatrick Sanan .seealso: `MatShift()`, `MatScale()`, `MatDiagonalScale()`
3346d84be18SBarry Smith @*/
3359371c9d4SSatish Balay PetscErrorCode MatDiagonalSet(Mat Y, Vec D, InsertMode is) {
33667576b8bSBarry Smith   PetscInt matlocal, veclocal;
3376d84be18SBarry Smith 
3383a40ed3dSBarry Smith   PetscFunctionBegin;
3390700a824SBarry Smith   PetscValidHeaderSpecific(Y, MAT_CLASSID, 1);
3400700a824SBarry Smith   PetscValidHeaderSpecific(D, VEC_CLASSID, 2);
3419566063dSJacob Faibussowitsch   PetscCall(MatGetLocalSize(Y, &matlocal, NULL));
3429566063dSJacob Faibussowitsch   PetscCall(VecGetLocalSize(D, &veclocal));
3435f80ce2aSJacob 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);
344dbbe0bcdSBarry Smith   if (Y->ops->diagonalset) PetscUseTypeMethod(Y, diagonalset, D, is);
345dbbe0bcdSBarry Smith   else PetscCall(MatDiagonalSet_Default(Y, D, is));
3469566063dSJacob Faibussowitsch   PetscCall(PetscObjectStateIncrease((PetscObject)Y));
3473a40ed3dSBarry Smith   PetscFunctionReturn(0);
3486d84be18SBarry Smith }
349d4bb536fSBarry Smith 
350d4bb536fSBarry Smith /*@
35104aac2b0SHong Zhang    MatAYPX - Computes Y = a*Y + X.
352d4bb536fSBarry Smith 
3533f9fe445SBarry Smith    Logically on Mat
354fee21e36SBarry Smith 
35598a79cdbSBarry Smith    Input Parameters:
35604aac2b0SHong Zhang +  a - the PetscScalar multiplier
35704aac2b0SHong Zhang .  Y - the first matrix
35804aac2b0SHong Zhang .  X - the second matrix
359531d0c6aSBarry 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)
360d4bb536fSBarry Smith 
3612860a424SLois Curfman McInnes    Level: intermediate
3622860a424SLois Curfman McInnes 
363db781477SPatrick Sanan .seealso: `MatAXPY()`
364d4bb536fSBarry Smith  @*/
3659371c9d4SSatish Balay PetscErrorCode MatAYPX(Mat Y, PetscScalar a, Mat X, MatStructure str) {
3663a40ed3dSBarry Smith   PetscFunctionBegin;
3679566063dSJacob Faibussowitsch   PetscCall(MatScale(Y, a));
3689566063dSJacob Faibussowitsch   PetscCall(MatAXPY(Y, 1.0, X, str));
3693a40ed3dSBarry Smith   PetscFunctionReturn(0);
370d4bb536fSBarry Smith }
371b0a32e0cSBarry Smith 
372b0a32e0cSBarry Smith /*@
3730bacdadaSStefano Zampini     MatComputeOperator - Computes the explicit matrix
374b0a32e0cSBarry Smith 
375b0a32e0cSBarry Smith     Collective on Mat
376b0a32e0cSBarry Smith 
377d8d19677SJose E. Roman     Input Parameters:
378186a3e20SStefano Zampini +   inmat - the matrix
379186a3e20SStefano Zampini -   mattype - the matrix type for the explicit operator
380b0a32e0cSBarry Smith 
381b0a32e0cSBarry Smith     Output Parameter:
382a5b23f4aSJose E. Roman .   mat - the explicit  operator
383b0a32e0cSBarry Smith 
38411a5261eSBarry Smith     Note:
385186a3e20SStefano Zampini     This computation is done by applying the operators to columns of the identity matrix.
386186a3e20SStefano Zampini     This routine is costly in general, and is recommended for use only with relatively small systems.
387186a3e20SStefano Zampini     Currently, this routine uses a dense matrix format if mattype == NULL.
388b0a32e0cSBarry Smith 
389b0a32e0cSBarry Smith     Level: advanced
390b0a32e0cSBarry Smith 
391b0a32e0cSBarry Smith @*/
3929371c9d4SSatish Balay PetscErrorCode MatComputeOperator(Mat inmat, MatType mattype, Mat *mat) {
393b0a32e0cSBarry Smith   PetscFunctionBegin;
3940700a824SBarry Smith   PetscValidHeaderSpecific(inmat, MAT_CLASSID, 1);
395186a3e20SStefano Zampini   PetscValidPointer(mat, 3);
3969566063dSJacob Faibussowitsch   PetscCall(MatConvert_Shell(inmat, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
39724f910e3SHong Zhang   PetscFunctionReturn(0);
39824f910e3SHong Zhang }
3994325cce7SMatthew G Knepley 
4004325cce7SMatthew G Knepley /*@
4010bacdadaSStefano Zampini     MatComputeOperatorTranspose - Computes the explicit matrix representation of
402f3b1f45cSBarry Smith         a give matrix that can apply MatMultTranspose()
403f3b1f45cSBarry Smith 
404f3b1f45cSBarry Smith     Collective on Mat
405f3b1f45cSBarry Smith 
4066b867d5aSJose E. Roman     Input Parameters:
4076b867d5aSJose E. Roman +   inmat - the matrix
4086b867d5aSJose E. Roman -   mattype - the matrix type for the explicit operator
409f3b1f45cSBarry Smith 
410f3b1f45cSBarry Smith     Output Parameter:
411a5b23f4aSJose E. Roman .   mat - the explicit  operator transposed
412f3b1f45cSBarry Smith 
41311a5261eSBarry Smith     Note:
414186a3e20SStefano Zampini     This computation is done by applying the transpose of the operator to columns of the identity matrix.
415186a3e20SStefano Zampini     This routine is costly in general, and is recommended for use only with relatively small systems.
416186a3e20SStefano Zampini     Currently, this routine uses a dense matrix format if mattype == NULL.
417f3b1f45cSBarry Smith 
418f3b1f45cSBarry Smith     Level: advanced
419f3b1f45cSBarry Smith 
420f3b1f45cSBarry Smith @*/
4219371c9d4SSatish Balay PetscErrorCode MatComputeOperatorTranspose(Mat inmat, MatType mattype, Mat *mat) {
422186a3e20SStefano Zampini   Mat A;
423f3b1f45cSBarry Smith 
424f3b1f45cSBarry Smith   PetscFunctionBegin;
425f3b1f45cSBarry Smith   PetscValidHeaderSpecific(inmat, MAT_CLASSID, 1);
426186a3e20SStefano Zampini   PetscValidPointer(mat, 3);
4279566063dSJacob Faibussowitsch   PetscCall(MatCreateTranspose(inmat, &A));
4289566063dSJacob Faibussowitsch   PetscCall(MatConvert_Shell(A, mattype ? mattype : MATDENSE, MAT_INITIAL_MATRIX, mat));
4299566063dSJacob Faibussowitsch   PetscCall(MatDestroy(&A));
430f3b1f45cSBarry Smith   PetscFunctionReturn(0);
431f3b1f45cSBarry Smith }
432f3b1f45cSBarry Smith 
433f3b1f45cSBarry Smith /*@
4344325cce7SMatthew G Knepley   MatChop - Set all values in the matrix less than the tolerance to zero
4354325cce7SMatthew G Knepley 
4364325cce7SMatthew G Knepley   Input Parameters:
4374325cce7SMatthew G Knepley + A   - The matrix
4384325cce7SMatthew G Knepley - tol - The zero tolerance
4394325cce7SMatthew G Knepley 
4404325cce7SMatthew G Knepley   Output Parameters:
4414325cce7SMatthew G Knepley . A - The chopped matrix
4424325cce7SMatthew G Knepley 
4434325cce7SMatthew G Knepley   Level: intermediate
4444325cce7SMatthew G Knepley 
445db781477SPatrick Sanan .seealso: `MatCreate()`, `MatZeroEntries()`
4463fc99919SSatish Balay  @*/
4479371c9d4SSatish Balay PetscErrorCode MatChop(Mat A, PetscReal tol) {
448038df967SPierre Jolivet   Mat          a;
4494325cce7SMatthew G Knepley   PetscScalar *newVals;
450cf5d3338SPierre Jolivet   PetscInt    *newCols, rStart, rEnd, numRows, maxRows, r, colMax = 0;
451038df967SPierre Jolivet   PetscBool    flg;
4524325cce7SMatthew G Knepley 
4534325cce7SMatthew G Knepley   PetscFunctionBegin;
4549566063dSJacob Faibussowitsch   PetscCall(PetscObjectBaseTypeCompareAny((PetscObject)A, &flg, MATSEQDENSE, MATMPIDENSE, ""));
455038df967SPierre Jolivet   if (flg) {
4569566063dSJacob Faibussowitsch     PetscCall(MatDenseGetLocalMatrix(A, &a));
4579566063dSJacob Faibussowitsch     PetscCall(MatDenseGetLDA(a, &r));
4589566063dSJacob Faibussowitsch     PetscCall(MatGetSize(a, &rStart, &rEnd));
4599566063dSJacob Faibussowitsch     PetscCall(MatDenseGetArray(a, &newVals));
460038df967SPierre Jolivet     for (; colMax < rEnd; ++colMax) {
461ad540459SPierre Jolivet       for (maxRows = 0; maxRows < rStart; ++maxRows) newVals[maxRows + colMax * r] = PetscAbsScalar(newVals[maxRows + colMax * r]) < tol ? 0.0 : newVals[maxRows + colMax * r];
462038df967SPierre Jolivet     }
4639566063dSJacob Faibussowitsch     PetscCall(MatDenseRestoreArray(a, &newVals));
464038df967SPierre Jolivet   } else {
4659566063dSJacob Faibussowitsch     PetscCall(MatGetOwnershipRange(A, &rStart, &rEnd));
4669566063dSJacob Faibussowitsch     PetscCall(MatGetRowUpperTriangular(A));
4674325cce7SMatthew G Knepley     for (r = rStart; r < rEnd; ++r) {
4684325cce7SMatthew G Knepley       PetscInt ncols;
4694325cce7SMatthew G Knepley 
4709566063dSJacob Faibussowitsch       PetscCall(MatGetRow(A, r, &ncols, NULL, NULL));
4715f80ce2aSJacob Faibussowitsch       colMax = PetscMax(colMax, ncols);
4729566063dSJacob Faibussowitsch       PetscCall(MatRestoreRow(A, r, &ncols, NULL, NULL));
4734325cce7SMatthew G Knepley     }
4744325cce7SMatthew G Knepley     numRows = rEnd - rStart;
4751c2dc1cbSBarry Smith     PetscCall(MPIU_Allreduce(&numRows, &maxRows, 1, MPIU_INT, MPI_MAX, PetscObjectComm((PetscObject)A)));
4769566063dSJacob Faibussowitsch     PetscCall(PetscMalloc2(colMax, &newCols, colMax, &newVals));
4779566063dSJacob Faibussowitsch     PetscCall(MatGetOption(A, MAT_NO_OFF_PROC_ENTRIES, &flg)); /* cache user-defined value */
4789566063dSJacob Faibussowitsch     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, PETSC_TRUE));
479cf5d3338SPierre Jolivet     /* short-circuit code in MatAssemblyBegin() and MatAssemblyEnd()             */
480cf5d3338SPierre Jolivet     /* that are potentially called many times depending on the distribution of A */
4814325cce7SMatthew G Knepley     for (r = rStart; r < rStart + maxRows; ++r) {
4824325cce7SMatthew G Knepley       const PetscScalar *vals;
4834325cce7SMatthew G Knepley       const PetscInt    *cols;
484fad45679SMatthew G. Knepley       PetscInt           ncols, newcols, c;
4854325cce7SMatthew G Knepley 
4864325cce7SMatthew G Knepley       if (r < rEnd) {
4879566063dSJacob Faibussowitsch         PetscCall(MatGetRow(A, r, &ncols, &cols, &vals));
4884325cce7SMatthew G Knepley         for (c = 0; c < ncols; ++c) {
4894325cce7SMatthew G Knepley           newCols[c] = cols[c];
4904325cce7SMatthew G Knepley           newVals[c] = PetscAbsScalar(vals[c]) < tol ? 0.0 : vals[c];
4914325cce7SMatthew G Knepley         }
492fad45679SMatthew G. Knepley         newcols = ncols;
4939566063dSJacob Faibussowitsch         PetscCall(MatRestoreRow(A, r, &ncols, &cols, &vals));
4949566063dSJacob Faibussowitsch         PetscCall(MatSetValues(A, 1, &r, newcols, newCols, newVals, INSERT_VALUES));
4954325cce7SMatthew G Knepley       }
4969566063dSJacob Faibussowitsch       PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
4979566063dSJacob Faibussowitsch       PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
4984325cce7SMatthew G Knepley     }
4999566063dSJacob Faibussowitsch     PetscCall(MatRestoreRowUpperTriangular(A));
5009566063dSJacob Faibussowitsch     PetscCall(PetscFree2(newCols, newVals));
5019566063dSJacob Faibussowitsch     PetscCall(MatSetOption(A, MAT_NO_OFF_PROC_ENTRIES, flg)); /* reset option to its user-defined value */
502038df967SPierre Jolivet   }
5034325cce7SMatthew G Knepley   PetscFunctionReturn(0);
5044325cce7SMatthew G Knepley }
505