From 05b7f2d56095ba5adb7c99f6c120666ab4eed775 Mon Sep 17 00:00:00 2001
From: lisyarus <lisyarus@gmail.com>
Date: Thu, 28 Dec 2023 17:53:42 +0300
Subject: [PATCH] Add cholesky decomposition implementation & tests

---
 libs/geom/include/psemek/geom/cholesky.hpp | 66 ++++++++++++++++++++++
 libs/geom/tests/matrix.cpp                 | 54 +++++++++++++++++-
 2 files changed, 118 insertions(+), 2 deletions(-)
 create mode 100644 libs/geom/include/psemek/geom/cholesky.hpp

diff --git a/libs/geom/include/psemek/geom/cholesky.hpp b/libs/geom/include/psemek/geom/cholesky.hpp
new file mode 100644
index 00000000..015a792e
--- /dev/null
+++ b/libs/geom/include/psemek/geom/cholesky.hpp
@@ -0,0 +1,66 @@
+#pragma once
+
+#include <psemek/geom/matrix.hpp>
+
+namespace psemek::geom
+{
+
+	template <typename T, std::size_t N>
+	struct cholesky_ldl_result
+	{
+		matrix<T, N, N> l; // lower unit-triangular
+		vector<T, N> d;
+	};
+
+	// returns lower-triangular l such that m = l * l^T
+	// m must be symmetric positive-definite
+	template <typename T, std::size_t N>
+	matrix<T, N, N> cholesky(matrix<T, N, N> const & m)
+	{
+		matrix<T, N, N> l = matrix<T, N, N>::zero();
+
+		for (std::size_t i = 0; i < N; ++i)
+		{
+			for (std::size_t j = 0; j <= i; ++j)
+			{
+				T sum = T{};
+				for (std::size_t k = 0; k < j; ++k)
+					sum += l[i][k] * l[j][k];
+
+				if (i == j)
+					l[i][j] = std::sqrt(std::max(T{}, sum));
+				else
+					l[i][j] = (m[i][j] - sum) / l[j][j];
+			}
+		}
+
+		return l;
+	}
+
+	// returns lower unit-triangular l and a diagonal d such that m = l * d * l^T
+	// m must be symmetric
+	template <typename T, std::size_t N>
+	cholesky_ldl_result<T, N> cholesky_ldl(matrix<T, N, N> const & m)
+	{
+		matrix<T, N, N> l = matrix<T, N, N>::identity();
+		vector<T, N> d = vector<T, N>::zero();
+
+		for (std::size_t i = 0; i < N; ++i)
+		{
+			for (std::size_t j = 0; j <= i; ++j)
+			{
+				T sum = T{};
+				for (std::size_t k = 0; k < j; ++k)
+					sum += l[i][k] * l[j][k] * d[k];
+
+				if (i == j)
+					d[i] = m[i][i] - sum;
+				else
+					l[i][j] = (m[i][j] - sum) / d[j];
+			}
+		}
+
+		return {l, d};
+	}
+
+}
diff --git a/libs/geom/tests/matrix.cpp b/libs/geom/tests/matrix.cpp
index 16fd5626..c9f4b2fd 100644
--- a/libs/geom/tests/matrix.cpp
+++ b/libs/geom/tests/matrix.cpp
@@ -3,6 +3,7 @@
 #include <psemek/geom/matrix.hpp>
 #include <psemek/geom/gauss.hpp>
 #include <psemek/geom/qr.hpp>
+#include <psemek/geom/cholesky.hpp>
 #include <psemek/random/generator.hpp>
 #include <psemek/random/normal.hpp>
 
@@ -241,8 +242,6 @@ test_case(geom_matrix_qr__eig_symmetric)
 			for (std::size_t j = i + 1; j < m.columns(); ++j)
 				expect_small(result.r[j][i], 1e-3);
 
-		std::string test;
-
 		for (std::size_t i = 0; i < m.columns(); ++i)
 		{
 			auto v = column(result.q, i);
@@ -284,3 +283,54 @@ test_case(geom_matrix_qr__eig_general)
 		}
 	}
 }
+
+test_case(geom_matrix_cholesky)
+{
+	random::generator rng;
+	random::normal_distribution<float> normal;
+
+	for (int iteration = 0; iteration < 64; ++iteration)
+	{
+		matrix<float, 8, 8> m;
+		for (std::size_t i = 0; i < m.rows(); ++i)
+			for (std::size_t j = 0; j < m.columns(); ++j)
+				m[i][j] = normal(rng);
+
+		m = transpose(m) * m;
+
+		auto l = cholesky(m);
+
+		expect_small(frobenius_norm(m - l * transpose(l)), 1e-6);
+
+		for (std::size_t i = 0; i < m.rows(); ++i)
+			for (std::size_t j = i + 1; j < m.columns(); ++j)
+				expect_small(l[i][j], 1e-6);
+	}
+}
+
+test_case(geom_matrix_cholesky__ldl)
+{
+	random::generator rng;
+	random::normal_distribution<float> normal;
+
+	for (int iteration = 0; iteration < 64; ++iteration)
+	{
+		matrix<float, 8, 8> m;
+		for (std::size_t i = 0; i < m.rows(); ++i)
+			for (std::size_t j = 0; j < m.columns(); ++j)
+				m[i][j] = normal(rng);
+
+		m = (m + transpose(m)) * 0.5f;
+
+		auto ldl = cholesky_ldl(m);
+
+		expect_small(frobenius_norm(m - ldl.l * m.diagonal(ldl.d) * transpose(ldl.l)), 1e-4);
+
+		for (std::size_t i = 0; i < m.rows(); ++i)
+			expect_close(ldl.l[i][i], 1.f, 1e-6);
+
+		for (std::size_t i = 0; i < m.rows(); ++i)
+			for (std::size_t j = i + 1; j < m.columns(); ++j)
+				expect_small(ldl.l[i][j], 1e-6);
+	}
+}