aboutsummaryrefslogtreecommitdiff
path: root/libpsi-core/src/maths/matrix.rs
diff options
context:
space:
mode:
Diffstat (limited to 'libpsi-core/src/maths/matrix.rs')
-rw-r--r--libpsi-core/src/maths/matrix.rs190
1 files changed, 190 insertions, 0 deletions
diff --git a/libpsi-core/src/maths/matrix.rs b/libpsi-core/src/maths/matrix.rs
new file mode 100644
index 0000000..975c92d
--- /dev/null
+++ b/libpsi-core/src/maths/matrix.rs
@@ -0,0 +1,190 @@
+use super::Float;
+use core::{fmt, ops};
+
+#[macro_export]
+macro_rules! matrix {
+ ( $( $( $x:expr ),* );* ) => {{
+ let mut data = Vec::new();
+ let mut rows = 0;
+ let mut cols = 0;
+
+ $(
+ let row_data = $( $x )*;
+ if cols == 0 {
+ cols = row_data.len();
+ }
+ assert_eq!(cols, row_data.len(), "All rows must have the same number of columns.");
+ data.extend(row_data);
+ rows += 1;
+ )*
+
+ $crate::Matrix::new(rows, cols, data)
+ }};
+}
+
+#[derive(Clone)]
+pub struct Matrix<T: Float> {
+ pub data: Vec<T>,
+ pub rows: usize,
+ pub cols: usize,
+}
+
+impl<T: Float> Matrix<T> {
+ pub fn new(rows: usize, cols: usize, data: Vec<T>) -> Self {
+ Matrix { data, rows, cols }
+ }
+
+ pub fn get(&self, row: usize, col: usize) -> T {
+ self.data[row * self.cols + col]
+ }
+
+ pub fn set(&mut self, row: usize, col: usize, value: T) {
+ self.data[row * self.cols + col] = value;
+ }
+
+ pub fn dot(&self, other: &Self) -> Option<Matrix<T>> {
+ if self.cols != other.rows {
+ return None;
+ }
+
+ let mut result = Matrix::new(
+ self.rows,
+ other.cols,
+ vec![T::zero(); self.rows * other.cols],
+ );
+ for i in 0..self.rows {
+ for j in 0..other.cols {
+ let mut sum = T::zero();
+ for k in 0..self.cols {
+ sum = sum + (self.get(i, k) * other.get(k, j));
+ }
+ result.set(i, j, sum);
+ }
+ }
+ Some(result)
+ }
+
+ pub fn kronecker(&self, other: &Self) -> Matrix<T> {
+ let new_rows = self.rows * other.rows;
+ let new_cols = self.cols * other.cols;
+
+ let mut result = Matrix::new(new_rows, new_cols, vec![T::zero(); new_rows * new_cols]);
+
+ for i in 0..self.rows {
+ for j in 0..self.cols {
+ let self_val = self.get(i, j);
+ for k in 0..other.rows {
+ for l in 0..other.cols {
+ let result_row = i * other.rows + k;
+ let result_col = j * other.cols + l;
+ result.set(result_row, result_col, self_val.clone() * other.get(k, l));
+ }
+ }
+ }
+ }
+
+ result
+ }
+
+ pub fn transpose(&self) -> Matrix<T> {
+ let mut result = Matrix::new(self.cols, self.rows, vec![T::zero(); self.cols * self.rows]);
+
+ for i in 0..self.rows {
+ for j in 0..self.cols {
+ let value = self.get(i, j);
+ result.set(j, i, value);
+ }
+ }
+
+ result
+ }
+
+ pub fn add_to(&self, other: &Self) -> Option<Matrix<T>> {
+ if self.rows != other.rows || self.cols != other.cols {
+ return None;
+ }
+
+ let mut result = Matrix::new(self.rows, self.cols, vec![T::zero(); self.rows * self.cols]);
+
+ for i in 0..self.rows {
+ for j in 0..self.cols {
+ let sum = self.get(i, j) + other.get(i, j);
+ result.set(i, j, sum);
+ }
+ }
+ Some(result)
+ }
+
+ pub fn subtract(&self, other: &Self) -> Option<Matrix<T>> {
+ if self.rows != other.rows || self.cols != other.cols {
+ return None;
+ }
+
+ let mut result = Matrix::new(self.rows, self.cols, vec![T::zero(); self.rows * self.cols]);
+
+ for i in 0..self.rows {
+ for j in 0..self.cols {
+ let diff = self.get(i, j) - other.get(i, j);
+ result.set(i, j, diff);
+ }
+ }
+ Some(result)
+ }
+
+ pub fn scale(&self, scalar: T) -> Matrix<T> {
+ let mut result = Matrix::new(self.rows, self.cols, vec![T::zero(); self.rows * self.cols]);
+
+ for i in 0..self.rows {
+ for j in 0..self.cols {
+ let scaled_value = self.get(i, j) * scalar;
+ result.set(i, j, scaled_value);
+ }
+ }
+ result
+ }
+}
+
+impl<T: Float> ops::Index<(usize, usize)> for Matrix<T> {
+ type Output = T;
+
+ fn index(&self, index: (usize, usize)) -> &Self::Output {
+ &self.data[index.0 * self.cols + index.1]
+ }
+}
+
+impl<T: Float> ops::IndexMut<(usize, usize)> for Matrix<T> {
+ fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
+ &mut self.data[index.0 * self.cols + index.1]
+ }
+}
+
+impl<T: Float + fmt::Debug> fmt::Debug for Matrix<T> {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ for i in 0..self.rows {
+ for j in 0..self.cols {
+ write!(f, "{:?} ", self.get(i, j))?;
+ }
+ writeln!(f)?;
+ }
+ Ok(())
+ }
+}
+
+impl<T: Float + fmt::Display> fmt::Display for Matrix<T> {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ for i in 0..self.rows {
+ write!(f, "[")?;
+ for j in 0..self.cols {
+ write!(f, "{}", self.get(i, j))?;
+ if j != self.cols - 1 {
+ write!(f, " ")?;
+ }
+ }
+ write!(f, "]")?;
+ if i != self.rows - 1 {
+ write!(f, "\n")?;
+ }
+ }
+ Ok(())
+ }
+}