1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

/*
    Appellation: linalg <mod>
    Contrib: FL03 <jo3mccain@icloud.com>
*/
//! Implementations for linear algebra operations.
//!
//!
use crate::linalg::{Inverse, Matmul};
use crate::prelude::{Scalar, ShapeError, TensorError, TensorExpr, TensorResult};
use crate::tensor::{self, TensorBase};
use acme::prelude::UnaryOp;
use num::traits::{Num, NumAssign};

fn inverse_impl<T>(tensor: &TensorBase<T>) -> TensorResult<TensorBase<T>>
where
    T: Copy + NumAssign + PartialOrd,
{
    let op = TensorExpr::unary(tensor.clone(), UnaryOp::Inv);
    let n = tensor.nrows();

    if !tensor.is_square() {
        return Err(ShapeError::NotSquare.into()); // Matrix must be square for inversion
    }

    let eye = TensorBase::eye(n);

    // Construct an augmented matrix by concatenating the original matrix with an identity matrix
    let mut aug = TensorBase::zeros((n, 2 * n));
    // aug.slice_mut(s![.., ..cols]).assign(matrix);
    for i in 0..n {
        for j in 0..n {
            aug[[i, j]] = tensor[[i, j]];
        }
        for j in n..(2 * n) {
            aug[[i, j]] = eye[[i, j - n]];
        }
    }

    // Perform Gaussian elimination to reduce the left half to the identity matrix
    for i in 0..n {
        let pivot = aug[[i, i]];

        if pivot == T::zero() {
            return Err(TensorError::Singular); // Matrix is singular
        }

        for j in 0..(2 * n) {
            aug[[i, j]] = aug[[i, j]] / pivot;
        }

        for j in 0..n {
            if i != j {
                let am = aug.clone();
                let factor = aug[[j, i]];
                for k in 0..(2 * n) {
                    aug[[j, k]] -= factor * am[[i, k]];
                }
            }
        }
    }

    // Extract the inverted matrix from the augmented matrix
    let mut inv = tensor.zeros_like().with_op(op.into());
    for i in 0..n {
        for j in 0..n {
            inv[[i, j]] = aug[[i, j + n]];
        }
    }

    Ok(inv.to_owned())
}

impl<T> TensorBase<T>
where
    T: Copy,
{
    /// Creates a new tensor containing only the diagonal elements of the original tensor.
    pub fn diag(&self) -> Self {
        let n = self.nrows();
        Self::from_shape_iter(self.shape().diag(), (0..n).map(|i| self[vec![i; n]]))
    }
    /// Find the inverse of the tensor
    ///
    /// # Errors
    ///
    /// Returns an error if the matrix is not square or if the matrix is singular.
    ///
    pub fn inv(&self) -> TensorResult<Self>
    where
        T: NumAssign + PartialOrd,
    {
        inverse_impl(self)
    }
    /// Compute the trace of the matrix.
    /// The trace of a matrix is the sum of the diagonal elements.
    pub fn trace(&self) -> TensorResult<T>
    where
        T: Num,
    {
        if !self.is_square() {
            return Err(ShapeError::NotSquare.into());
        }
        let n = self.nrows();
        let trace = (0..n).fold(T::zero(), |acc, i| acc + self[[i, i]]);
        Ok(trace)
    }
}

impl<T> Inverse for TensorBase<T>
where
    T: Copy + Num + NumAssign + PartialOrd,
{
    type Output = TensorResult<Self>;

    fn inv(&self) -> Self::Output {
        inverse_impl(self)
    }
}

impl<T> Matmul<TensorBase<T>> for TensorBase<T>
where
    T: Scalar,
{
    type Output = Self;

    fn matmul(&self, other: &Self) -> Self {
        let sc = |m: usize, n: usize| m * self.ncols() + n;
        let oc = |m: usize, n: usize| m * other.ncols() + n;

        let shape = self.shape().matmul_shape(&other.shape()).unwrap();
        let mut result = vec![T::zero(); shape.size()];

        for i in 0..self.nrows() {
            for j in 0..other.ncols() {
                for k in 0..self.ncols() {
                    result[oc(i, j)] += self.data[sc(i, k)] * other.data[oc(k, j)];
                }
            }
        }
        let op = TensorExpr::matmul(self.clone(), other.clone());
        tensor::from_vec_with_op(false, op, shape, result)
    }
}