From 68e665748a6068c180edb1ad1e73d54cd71dc3be Mon Sep 17 00:00:00 2001
From: Cory <cory@balaton.dev>
Date: Sat, 16 Sep 2023 18:01:26 +0200
Subject: [PATCH] Implement functions to create our desired matrix

---
 include/matrix.hpp | 16 ++++++++++++++++
 src/matrix.cpp     | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)
 create mode 100644 include/matrix.hpp
 create mode 100644 src/matrix.cpp

diff --git a/include/matrix.hpp b/include/matrix.hpp
new file mode 100644
index 0000000..3ed9290
--- /dev/null
+++ b/include/matrix.hpp
@@ -0,0 +1,16 @@
+#ifndef __MATRIX__
+#define __MATRIX__
+
+#include <armadillo>
+
+
+arma::mat create_tridiagonal(
+    const arma::vec& a, 
+    const arma::vec& d, 
+    const arma::vec& e);
+
+arma::mat create_tridiagonal(int n, double a, double d, double e);
+
+arma::mat create_symmetric_tridiagonal(int n, double a, double d);
+
+#endif
diff --git a/src/matrix.cpp b/src/matrix.cpp
new file mode 100644
index 0000000..1fa3536
--- /dev/null
+++ b/src/matrix.cpp
@@ -0,0 +1,38 @@
+#include "matrix.hpp"
+
+arma::mat create_tridiagonal(
+    const arma::vec& a, 
+    const arma::vec& d, 
+    const arma::vec& e)
+{
+    int n = d.n_elem;
+    arma::mat A = arma::mat(n, n).fill(0.);
+
+    A(0,0) = d(0);
+    A(0,1) = e(0);
+
+    for (int i = 1; i < n-1; i++) {
+        A(i, i-1) = a(i-1);
+        A(i, i) = d(i);
+        A(i, i+1) = e(i);
+    }
+    
+    A(n-1,n-1) = d(n-1);
+    A(n-1, n-2) = a(n-1);
+
+    return A;
+}
+
+arma::mat create_tridiagonal(int n, double a, double d, double e)
+{
+    arma::vec a_vec = arma::vec(n-1, arma::fill::ones) * a;
+    arma::vec d_vec = arma::vec(n-1, arma::fill::ones) * d;
+    arma::vec e_vec = arma::vec(n-1, arma::fill::ones) * e;
+
+    return create_tridiagonal(a_vec, d_vec, e_vec);
+}
+
+arma::mat create_symmetric_tridiagonal(int n, double a, double d)
+{
+    return create_tridiagonal(n, a, d, a);
+}