Adding a vector to matrix diagonal (#4087)

src/shogun/mathematics/linalg/LinalgBackendBase.h

@@ -123,6 +123,23 @@ namespace shogun
 		DEFINE_FOR_ALL_PTYPE(BACKEND_GENERIC_ADD_COL_VEC, SGMatrix)
 #undef BACKEND_GENERIC_ADD_COL_VEC
 
+/**
+ * Wrapper method of add diagonal vector A.diagonal = alpha * A.diagonal +
+ * beta * b.
+ *
+ * @see linalg::add_diag
+ */
+#define BACKEND_GENERIC_ADD_DIAG(Type, Container)                              \
+	virtual void add_diag(                                                     \
+	    SGMatrix<Type>& A, const SGVector<Type>& b, Type alpha, Type beta)     \
+	    const                                                                  \
+	{                                                                          \
+		SG_SNOTIMPLEMENTED;                                                    \
+		return;                                                                \
+	}
+		DEFINE_FOR_ALL_PTYPE(BACKEND_GENERIC_ADD_DIAG, SGMatrix);
+#undef BACKEND_GENERIC_ADD_DIAG
+
 /**
  * Wrapper method of add vector to each column of matrix.
  *

src/shogun/mathematics/linalg/LinalgBackendEigen.h

@@ -65,6 +65,14 @@ namespace shogun
 		DEFINE_FOR_NUMERIC_PTYPE(BACKEND_GENERIC_ADD_COL_VEC, SGMatrix)
 #undef BACKEND_GENERIC_ADD_COL_VEC
 
+/** Implementation of @see LinalgBackendBase::add_diag */
+#define BACKEND_GENERIC_ADD_DIAG(Type, Container)                              \
+	virtual void add_diag(                                                     \
+	    SGMatrix<Type>& A, const SGVector<Type>& b, Type alpha, Type beta)     \
+	    const;
+		DEFINE_FOR_NUMERIC_PTYPE(BACKEND_GENERIC_ADD_DIAG, SGMatrix)
+#undef BACKEND_GENERIC_ADD_DIAG
+
 /** Implementation of @see LinalgBackendBase::add_vector */
 #define BACKEND_GENERIC_ADD(Type, Container)                                   \
 	virtual void add_vector(                                                   \
@@ -391,6 +399,11 @@ namespace shogun
 		    const SGMatrix<T>& A, index_t i, const SGVector<T>& b,
 		    SGVector<T>& result, T alpha, T beta) const;
 
+		/** Eigen3 add diagonal vector method */
+		template <typename T>
+		void add_diag_impl(
+		    SGMatrix<T>& A, const SGVector<T>& b, T alpha, T beta) const;
+
 		/** Eigen3 add vector to each column of matrix method */
 		template <typename T>
 		void add_vector_impl(

src/shogun/mathematics/linalg/LinalgNamespace.h

@@ -445,6 +445,32 @@ namespace shogun
 			    ->add_col_vec(A, i, b, result, alpha, beta);
 		}
 
+		/**
+		 * Performs the operation A.diagonal = alpha * A.diagonal + beta * b.
+		 * The matrix is not required to be square.
+		 *
+		 * @param A The matrix
+		 * @param b The vector
+		 * @param alpha Constant to be multiplied by the main diagonal of the
+		 * matrix
+		 * @param beta Constant to be multiplied by the vector
+		 */
+		template <typename T>
+		void
+		add_diag(SGMatrix<T>& A, const SGVector<T>& b, T alpha = 1, T beta = 1)
+		{
+			auto diagonal_len = CMath::min(A.num_cols, A.num_rows);
+			REQUIRE(
+			    diagonal_len == b.vlen, "Length of main diagonal of matrix A "
+			                            "(%d) doesn't match length of vector b "
+			                            "(%d).\n",
+			    diagonal_len, b.vlen);
+			REQUIRE(
+			    diagonal_len > 0 && b.vlen > 0, "Matrix / vector can't be "
+			                                    "empty.\n");
+			infer_backend(A, SGMatrix<T>(b))->add_diag(A, b, alpha, beta);
+		}
+
 		/**
 		 * Performs the operation result = alpha * A.col(i) + beta * b,
 		 * for each column of A.

src/shogun/mathematics/linalg/backend/eigen/BasicOps.cpp

@@ -57,6 +57,16 @@ DEFINE_FOR_NUMERIC_PTYPE(BACKEND_GENERIC_ADD_COL_VEC, SGVector)
 DEFINE_FOR_NUMERIC_PTYPE(BACKEND_GENERIC_ADD_COL_VEC, SGMatrix)
 #undef BACKEND_GENERIC_ADD_COL_VEC
 
+#define BACKEND_GENERIC_ADD_DIAG(Type, Container)                              \
+	void LinalgBackendEigen::add_diag(                                         \
+	    SGMatrix<Type>& A, const SGVector<Type>& b, Type alpha, Type beta)     \
+	    const                                                                  \
+	{                                                                          \
+		add_diag_impl(A, b, alpha, beta);                                      \
+	}
+DEFINE_FOR_NUMERIC_PTYPE(BACKEND_GENERIC_ADD_DIAG, SGMatrix)
+#undef BACKEND_GENERIC_ADD_DIAG
+
 #define BACKEND_GENERIC_ADD_VECTOR(Type, Container)                            \
 	void LinalgBackendEigen::add_vector(                                       \
 	    const SGMatrix<Type>& A, const SGVector<Type>& b,                      \
@@ -179,6 +189,16 @@ void LinalgBackendEigen::add_col_vec_impl(
 	result_eig.col(i) = alpha * A_eig.col(i) + beta * b_eig;
 }
 
+template <typename T>
+void LinalgBackendEigen::add_diag_impl(
+    SGMatrix<T>& A, const SGVector<T>& b, T alpha, T beta) const
+{
+	typename SGMatrix<T>::EigenMatrixXtMap A_eig = A;
+	typename SGVector<T>::EigenVectorXtMap b_eig = b;
+
+	A_eig.diagonal() = alpha * A_eig.diagonal() + beta * b_eig;
+}
+
 template <typename T>
 void LinalgBackendEigen::add_col_vec_impl(
     const SGMatrix<T>& A, index_t i, const SGVector<T>& b, SGVector<T>& result,

tests/unit/mathematics/linalg/operations/Eigen3_operations_unittest.cc

@@ -1,10 +1,11 @@
 #include <gtest/gtest.h>
 
+#include <shogun/base/range.h>
+#include <shogun/lib/ShogunException.h>
 #include <shogun/lib/config.h>
 #include <shogun/mathematics/Math.h>
 #include <shogun/mathematics/linalg/LinalgNamespace.h>
 #include <shogun/mathematics/linalg/LinalgSpecialPurposes.h>
-#include <shogun/lib/ShogunException.h>
 
 using namespace shogun;
 using namespace linalg;
@@ -189,6 +190,41 @@ TEST(LinalgBackendEigen, SGMatrix_add_col_vec_in_place)
 		}
 }
 
+TEST(LinalgBackendEigen, add_diag)
+{
+	SGMatrix<float64_t> A1(2, 3);
+	SGVector<float64_t> b1(2);
+
+	A1(0, 0) = 1;
+	A1(0, 1) = 2;
+	A1(0, 2) = 3;
+	A1(1, 0) = 4;
+	A1(1, 1) = 5;
+	A1(1, 2) = 6;
+
+	b1[0] = 1;
+	b1[1] = 2;
+
+	add_diag(A1, b1, 0.5, 2.0);
+
+	EXPECT_NEAR(A1(0, 0), 2.5, 1e-15);
+	EXPECT_NEAR(A1(0, 1), 2, 1e-15);
+	EXPECT_NEAR(A1(0, 2), 3, 1e-15);
+	EXPECT_NEAR(A1(1, 0), 4, 1e-15);
+	EXPECT_NEAR(A1(1, 1), 6.5, 1e-15);
+	EXPECT_NEAR(A1(1, 2), 6, 1e-15);
+
+	// test error cases
+	SGMatrix<float64_t> A2(2, 2);
+	SGVector<float64_t> b2(3);
+	SGMatrix<float64_t> A3;
+	SGVector<float64_t> b3;
+	EXPECT_THROW(add_diag(A2, b2), ShogunException);
+	EXPECT_THROW(add_diag(A2, b3), ShogunException);
+	EXPECT_THROW(add_diag(A3, b2), ShogunException);
+	EXPECT_THROW(add_diag(A3, b3), ShogunException);
+}
+
 TEST(LinalgBackendEigen, add_vector)
 {
 	const float64_t alpha = 0.7;