Matrix.hpp

 
//TODO:优化矩阵乘法
#pragma once
 
#include <cmath>
#include <cstdint>
#include <array>
#include <algorithm>
 
namespace roboctrl::utils
{
template <int _rows, int _cols, typename T = float>
    requires(std::is_arithmetic_v<T> && _rows > 0 && _cols > 0) //编译期检查是否为算数类型
class Matrix
{
 
public:
    constexpr Matrix() = default;
 
    constexpr explicit Matrix(T val)
    {
        for (int i = 0; i < _rows; i++)
            std::fill(data[i].begin(), data[i].end(), val);
    }
 
    constexpr explicit Matrix(const T *data) : Matrix()
    {
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                this->data[i][j] = data[i * _cols + j];
    }
 
    constexpr Matrix(const Matrix &mat) : Matrix()
    {
        this->data = mat.data;
    }
 
    constexpr Matrix(Matrix &&mat) noexcept : Matrix()
    {
        this->data = std::move(mat.data);
    }
 
    template <typename Ty>
        requires(std::is_arithmetic_v<Ty>)
    constexpr Matrix(const Matrix<_rows, _cols, Ty> &mat) : Matrix()
    {
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                data[i][j] = static_cast<T>(mat.data[i][j]);
    }
 
    ~Matrix() = default;
 
    [[nodiscard]] constexpr uint32_t rows() const
    {
        return _rows;
    }
 
    [[nodiscard]] constexpr uint32_t cols() const
    {
        return _cols;
    }
 
    constexpr std::array<T, _cols> &operator[](int row)
    {
        return data[row];
    }
 
    constexpr const std::array<T, _cols> &operator[](int row) const
    {
        return data[row];
    }
 
    constexpr Matrix &operator=(const Matrix &mat)
    {
        data = mat.data;
        return *this;
    }
 
    constexpr Matrix &operator=(Matrix &&mat) noexcept
    {
        data = std::move(mat.data);
        return *this;
    }
 
    constexpr Matrix &operator+=(const Matrix &mat)
    {
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                data[i][j] += mat.data[i][j];
 
        return *this;
    }
 
    constexpr Matrix &operator-=(const Matrix &mat)
    {
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                data[i][j] -= mat.data[i][j];
 
        return *this;
    }
 
    template <typename Ty>
        requires(std::is_arithmetic_v<Ty>)
    constexpr Matrix &operator*=(Ty val)
    {
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                data[i][j] *= val;
 
        return *this;
    }
 
    template <typename Ty>
        requires(std::is_arithmetic_v<Ty>)
    constexpr Matrix &operator/=(Ty val)
    {
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                data[i][j] /= val;
 
        return *this;
    }
 
    constexpr Matrix operator+(const Matrix &mat) const
    {
        Matrix res;
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                res.data[i][j] = data[i][j] + mat.data[i][j];
 
        return res;
    }
 
    constexpr Matrix operator-(const Matrix &mat) const
    {
        Matrix res;
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                res.data[i][j] = data[i][j] - mat.data[i][j];
 
        return res;
    }
 
    template <typename Ty>
        requires(std::is_arithmetic_v<Ty>)
    constexpr Matrix operator*(Ty val) const
    {
        Matrix res;
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                res.data[i][j] = data[i][j] * val;
 
        return res;
    }
 
    template <typename Ty>
        requires(std::is_arithmetic_v<Ty>)
    constexpr friend Matrix operator*(Ty val, const Matrix &mat)
    {
        return mat * val;
    }
 
    template <typename Ty>
        requires(std::is_arithmetic_v<Ty>)
    Matrix operator/(Ty val) const
    {
        Matrix res;
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                res.data[i][j] = data[i][j] / val;
 
        return res;
    }
 
    template <int cols2>
    constexpr friend Matrix<_rows, cols2, T> operator*(const Matrix<_rows, _cols, T> &mat1,
                                                       const Matrix<_cols, cols2, T> &mat2)
    {
        Matrix<_rows, cols2, T> res;
 
        for (int i = 0; i < _rows; i++)
            for (int k = 0; k < _cols; k++)
                for (int j = 0; j < cols2; j++)
                    res[i][j] += mat1[i][k] * mat2[k][j];
 
        return res;
    }
 
    constexpr bool operator==(const Matrix &mat) const
    {
        return data == mat.data;
    }
 
    // Submatrix
    template <int rows, int cols> constexpr Matrix<rows, cols, T> block(int start_row, int start_col) const
    {
        Matrix<rows, cols, T> res;
        for (int i = 0; i < rows; i++)
            std::copy(data[i + start_row].begin() + start_col, data[i + start_row].begin() + start_col + cols,
                      res.data[i].begin());
        return res;
    }
 
    constexpr Matrix<1, _cols, T> row(int row) const
    {
        return block<1, _cols, T>(row, 0);
    }
 
    constexpr Matrix<_rows, 1, T> col(int col) const
    {
        return block<_rows, 1, T>(0, col);
    }
 
    constexpr Matrix<_cols, _rows> trans() const
    {
        Matrix<_cols, _rows> res;
        for (int i = 0; i < _rows; i++)
            for (int j = 0; j < _cols; j++)
                res[j][i] = data[i][j];
 
        return res;
    }
 
    constexpr Matrix clone() const
    {
        Matrix res;
        res.data = data;
        return res;
    }
 
    constexpr T trace() const
    {
        T res = 0;
        for (int i = 0; i < fmin(_rows, _cols); i++)
            res += data[i][i];
        return res;
    }
 
    constexpr Matrix inv() const
    {
        Matrix res = Matrix::eye();
        static_assert(_cols == _rows, "Matrix must be square");
 
        // 创建一个临时矩阵用于 LU 分解
        Matrix temp = this->clone();
 
        // 高斯-约当消元法
        for (int k = 0; k < _cols; k++)
        {
            // 检查主元是否为 0
            if (temp[k][k] == 0.0f)
            {
                return zeros();
            }
 
            // 归一化当前行
            T pivot = temp[k][k];
            for (int j = 0; j < _rows; j++)
            {
                temp[k][j] /= pivot;
                res[k][j] /= pivot;
            }
 
            // 消去其他行的当前列
            for (int i = 0; i < _rows; i++)
            {
                if (i != k)
                {
                    T factor = temp[i][k];
                    for (int j = 0; j < _rows; j++)
                    {
                        temp[i][j] -= factor * temp[k][j];
                        res[i][j] -= factor * res[k][j];
                    }
                }
            }
        }
 
        return res;
    }
 
    /*==============================================================*/
    // Static function
    static constexpr Matrix zeros()
    {
        return Matrix(static_cast<T>(0));
    }
 
    static constexpr Matrix ones()
    {
        return Matrix(static_cast<T>(1));
    }
 
    static constexpr Matrix eye()
    {
        Matrix mat = Matrix::zeros();
 
        for (int i = 0; i < std::min(_rows, _cols); i++)
            mat[i][i] = 1;
 
        return mat;
    }
 
    static constexpr Matrix diag(Matrix<_rows, 1> vec)
    {
        Matrix res = Matrix::zeros();
        for (int i = 0; i < std::min(_rows, _cols); i++)
        {
            res[i][i] = vec[i][0];
        }
        return res;
    }
 
private:
    std::array<std::array<T, _cols>, _rows> data;
 
};
 
template <int r, int c> using Matrixf = Matrix<r, c, float>;
} // namespace Power