PlainMatrix.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
#pragma once
 
#include <ostream>
#include <array>
#include <algorithm>
#include <cassert>
 
namespace Belle2 {
  namespace TrackFindingCDC {
 
    template <class T, int M, int N>
    class PlainMatrix {
 
      static_assert(M > 0, "Dimensions must be greater zero");
      static_assert(N > 0, "Dimensions must be greater zero");
 
      static const int S = M * N;
 
    public:
      PlainMatrix() = default;
 
      explicit PlainMatrix(std::initializer_list<T> values)
        : m_values{}
      {
        assert(m_values.size() >= values.size());
        std::copy(values.begin(), values.end(), m_values.begin());
      }
 
      explicit PlainMatrix(std::array<T, S> values)
        : m_values(values)
      {
      }
 
      static PlainMatrix<T, M, N> Zero()
      {
        PlainMatrix<T, M, N> result{{}}; // Value initialize to zero
        return result;
      }
 
      static PlainMatrix<T, M, N> Identity()
      {
        PlainMatrix<T, M, N> result{{}}; // Value initialize to zero
        for (int s = 0; s < M * N; s += M + 1)
          result.m_values[s] = 1;
        return result;
      }
 
      static PlainMatrix<T, M, N> Constant(T t)
      {
        PlainMatrix<T, M, N> result;
        result.m_values.fill(t);
        return result;
      }
 
    public:
      T* data()
      {
        return m_values.data();
      }
 
      const T* data() const
      {
        return m_values.data();
      }
 
      T& operator[](int s)
      {
        return m_values[s];
      }
 
      T operator[](int s) const
      {
        return m_values[s];
      }
 
      T& operator()(int i, int j = 0)
      {
        return m_values.operator[](j * M + i);
      }
 
      T operator()(int i, int j = 0) const
      {
        return m_values.operator[](j * M + i);
      }
 
      friend std::ostream& operator<<(std::ostream& out, const PlainMatrix<T, M, N>& rhs)
      {
        for (int i = 0; i < rhs.rows(); ++i) {
          out << rhs(i, 0) << ", ";
          for (int j = 1; j < rhs.cols(); ++j) {
            out << rhs(i, j);
          }
          out << "\n";
        }
        return out;
      }
 
      PlainMatrix<T, M, N> operator+(const PlainMatrix<T, M, N>& rhs) const
      {
        PlainMatrix<T, M, N> result;
        std::transform(begin(), end(), rhs.begin(), result.begin(), std::plus<T>());
        return result;
      }
 
      PlainMatrix<T, M, N> operator-(const PlainMatrix<T, M, N>& rhs) const
      {
        PlainMatrix<T, M, N> result;
        std::transform(begin(), end(), rhs.begin(), result.begin(), std::minus<T>());
        return result;
      }
 
      template <int O>
      PlainMatrix<T, M, O> operator*(const PlainMatrix<T, N, O>& rhs) const
      {
        PlainMatrix<T, M, O> result{{}}; // Value initialize to zero
        for (int m = 0; m < M; ++m) {
          for (int o = 0; o < O; ++o) {
            for (int n = 0; n < N; ++n) {
              result(m, o) += operator()(m, n) * rhs(n, o);
            }
          }
        }
        return result;
      }
 
      PlainMatrix<T, M, N> operator/(T rhs) const
      {
        PlainMatrix<T, M, N> result;
        std::transform(begin(), end(), result.begin(), [&rhs](const T & t) { return t / rhs; });
        return result;
      }
 
      friend PlainMatrix<T, M, N> operator/(T lhs, PlainMatrix<T, M, N>& rhs)
      {
        PlainMatrix<T, M, N> result;
        std::transform(rhs.begin(), rhs.end(), result.begin(), [&lhs](const T & t) {
          return lhs / t;
        });
        return result;
      }
 
      template <int K, int L>
      PlainMatrix<T, K, L> block(int i = 0, int j = 0) const
      {
        assert(K + i <= M && "Selected block reaches outside of the matrix");
        assert(L + j <= N && "Selected block reaches outside of the matrix");
 
        PlainMatrix<T, K, L> result;
        const int skipFront = j * M + i;
        for (int iCol = 0; iCol < L; ++iCol) {
          std::copy(begin() + skipFront + iCol * M,
                    begin() + skipFront + iCol * M + K,
                    result.data() + iCol * L);
        }
        return result;
      }
 
      template <int K>
      PlainMatrix<T, K, N> head() const
      {
        return block<K, N>(0, 0);
      }
 
      template <int K>
      PlainMatrix<T, K, N> tail() const
      {
        return block<K, N>(M - K, 0);
      }
 
      static constexpr size_t size()
      {
        return M * N;
      }
 
      static constexpr int rows()
      {
        return M;
      }
 
      static constexpr int cols()
      {
        return N;
      }
 
    private:
      T* begin()
      {
        return data();
      }
 
      T* end()
      {
        return data() + size();
      }
 
      const T* begin() const
      {
        return data();
      }
 
      const T* end() const
      {
        return data() + size();
      }
 
    private:
      std::array<T, M* N> m_values;
    };
  }
}