From b612467593a6881f76ef4d86957704b173c93005 Mon Sep 17 00:00:00 2001
From: bresch <brescianimathieu@gmail.com>
Date: Thu, 27 Apr 2023 09:48:13 +0200
Subject: [PATCH] square matrix: add partial trace computation

This is useful when we need the sum of some variances in a large
covariance matrix
---
 src/lib/matrix/matrix/SquareMatrix.hpp   | 14 ++++++++++++--
 src/lib/matrix/test/MatrixSquareTest.cpp |  1 +
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/lib/matrix/matrix/SquareMatrix.hpp b/src/lib/matrix/matrix/SquareMatrix.hpp
index 7d72710e17..3c47fb0230 100644
--- a/src/lib/matrix/matrix/SquareMatrix.hpp
+++ b/src/lib/matrix/matrix/SquareMatrix.hpp
@@ -126,18 +126,28 @@ public:
 		return res;
 	}
 
-	Type trace() const
+	template <size_t Width>
+	Type trace(size_t first) const
 	{
+		static_assert(Width <= M, "Width bigger than matrix");
+		assert(first + Width <= M);
+
 		Type res = 0;
 		const SquareMatrix<Type, M> &self = *this;
 
-		for (size_t i = 0; i < M; i++) {
+		for (size_t i = first; i < (first + Width); i++) {
 			res += self(i, i);
 		}
 
 		return res;
 	}
 
+	Type trace() const
+	{
+		const SquareMatrix<Type, M> &self = *this;
+		return self.trace<M>(0);
+	}
+
 	// zero all offdiagonal elements and keep corresponding diagonal elements
 	template <size_t Width>
 	void uncorrelateCovariance(size_t first)
diff --git a/src/lib/matrix/test/MatrixSquareTest.cpp b/src/lib/matrix/test/MatrixSquareTest.cpp
index 58ceab940f..da9e7df76c 100644
--- a/src/lib/matrix/test/MatrixSquareTest.cpp
+++ b/src/lib/matrix/test/MatrixSquareTest.cpp
@@ -47,6 +47,7 @@ TEST(MatrixSquareTest, Square)
 
 	EXPECT_EQ(A.diag(), diag_check);
 	EXPECT_FLOAT_EQ(A.trace(), 16);
+	EXPECT_FLOAT_EQ(A.trace<2>(1), 15);
 
 	float data_check[9] = {
 		1.01158503f,  0.02190432f,  0.03238144f,