VariableMatrix.hpp

Go to the documentation of this file.

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <algorithm>
#include <concepts>
#include <initializer_list>
#include <iterator>
#include <span>
#include <utility>
#include <vector>
 
#include <Eigen/Core>
 
#include "sleipnir/autodiff/Slice.hpp"
#include "sleipnir/autodiff/Variable.hpp"
#include "sleipnir/autodiff/VariableBlock.hpp"
#include "sleipnir/util/Assert.hpp"
#include "sleipnir/util/FunctionRef.hpp"
#include "sleipnir/util/SymbolExports.hpp"
#include "sleipnir/util/small_vector.hpp"
 
namespace sleipnir {
 
class SLEIPNIR_DLLEXPORT VariableMatrix {
 public:
  struct empty_t {};
  static constexpr empty_t empty{};
 
  VariableMatrix() = default;
 
  explicit VariableMatrix(int rows) : m_rows{rows}, m_cols{1} {
    m_storage.reserve(Rows());
    for (int row = 0; row < Rows(); ++row) {
      m_storage.emplace_back();
    }
  }
 
  VariableMatrix(int rows, int cols) : m_rows{rows}, m_cols{cols} {
    m_storage.reserve(Rows() * Cols());
    for (int index = 0; index < Rows() * Cols(); ++index) {
      m_storage.emplace_back();
    }
  }
 
  VariableMatrix(empty_t, int rows, int cols) : m_rows{rows}, m_cols{cols} {
    m_storage.reserve(Rows() * Cols());
    for (int index = 0; index < Rows() * Cols(); ++index) {
      m_storage.emplace_back(nullptr);
    }
  }
 
  VariableMatrix(  // NOLINT
      std::initializer_list<std::initializer_list<Variable>> list) {
    // Get row and column counts for destination matrix
    m_rows = list.size();
    m_cols = 0;
    if (list.size() > 0) {
      m_cols = list.begin()->size();
    }
 
    // Assert the first and latest column counts are the same
    for ([[maybe_unused]]
         const auto& row : list) {
      Assert(list.begin()->size() == row.size());
    }
 
    m_storage.reserve(Rows() * Cols());
    for (const auto& row : list) {
      std::copy(row.begin(), row.end(), std::back_inserter(m_storage));
    }
  }
 
  VariableMatrix(const std::vector<std::vector<double>>& list) {  // NOLINT
    // Get row and column counts for destination matrix
    m_rows = list.size();
    m_cols = 0;
    if (list.size() > 0) {
      m_cols = list.begin()->size();
    }
 
    // Assert the first and latest column counts are the same
    for ([[maybe_unused]]
         const auto& row : list) {
      Assert(list.begin()->size() == row.size());
    }
 
    m_storage.reserve(Rows() * Cols());
    for (const auto& row : list) {
      std::copy(row.begin(), row.end(), std::back_inserter(m_storage));
    }
  }
 
  VariableMatrix(const std::vector<std::vector<Variable>>& list) {  // NOLINT
    // Get row and column counts for destination matrix
    m_rows = list.size();
    m_cols = 0;
    if (list.size() > 0) {
      m_cols = list.begin()->size();
    }
 
    // Assert the first and latest column counts are the same
    for ([[maybe_unused]]
         const auto& row : list) {
      Assert(list.begin()->size() == row.size());
    }
 
    m_storage.reserve(Rows() * Cols());
    for (const auto& row : list) {
      std::copy(row.begin(), row.end(), std::back_inserter(m_storage));
    }
  }
 
  template <typename Derived>
  VariableMatrix(const Eigen::MatrixBase<Derived>& values)  // NOLINT
      : m_rows{static_cast<int>(values.rows())},
        m_cols{static_cast<int>(values.cols())} {
    m_storage.reserve(values.rows() * values.cols());
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        m_storage.emplace_back(values(row, col));
      }
    }
  }
 
  template <typename Derived>
  VariableMatrix(const Eigen::DiagonalBase<Derived>& values)  // NOLINT
      : m_rows{static_cast<int>(values.rows())},
        m_cols{static_cast<int>(values.cols())} {
    m_storage.reserve(values.rows() * values.cols());
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        if (row == col) {
          m_storage.emplace_back(values.diagonal()(row));
        } else {
          m_storage.emplace_back(0.0);
        }
      }
    }
  }
 
  template <typename Derived>
  VariableMatrix& operator=(const Eigen::MatrixBase<Derived>& values) {
    Assert(Rows() == values.rows());
    Assert(Cols() == values.cols());
 
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        (*this)(row, col) = values(row, col);
      }
    }
 
    return *this;
  }
 
  template <typename Derived>
    requires std::same_as<typename Derived::Scalar, double>
  void SetValue(const Eigen::MatrixBase<Derived>& values) {
    Assert(Rows() == values.rows());
    Assert(Cols() == values.cols());
 
    for (int row = 0; row < values.rows(); ++row) {
      for (int col = 0; col < values.cols(); ++col) {
        (*this)(row, col).SetValue(values(row, col));
      }
    }
  }
 
  VariableMatrix(const Variable& variable)  // NOLINT
      : m_rows{1}, m_cols{1} {
    m_storage.emplace_back(variable);
  }
 
  VariableMatrix(Variable&& variable) : m_rows{1}, m_cols{1} {  // NOLINT
    m_storage.emplace_back(std::move(variable));
  }
 
  VariableMatrix(const VariableBlock<VariableMatrix>& values)  // NOLINT
      : m_rows{values.Rows()}, m_cols{values.Cols()} {
    m_storage.reserve(Rows() * Cols());
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        m_storage.emplace_back(values(row, col));
      }
    }
  }
 
  VariableMatrix(const VariableBlock<const VariableMatrix>& values)  // NOLINT
      : m_rows{values.Rows()}, m_cols{values.Cols()} {
    m_storage.reserve(Rows() * Cols());
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        m_storage.emplace_back(values(row, col));
      }
    }
  }
 
  explicit VariableMatrix(std::span<const Variable> values)
      : m_rows{static_cast<int>(values.size())}, m_cols{1} {
    m_storage.reserve(Rows() * Cols());
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        m_storage.emplace_back(values[row * Cols() + col]);
      }
    }
  }
 
  VariableMatrix(std::span<const Variable> values, int rows, int cols)
      : m_rows{rows}, m_cols{cols} {
    Assert(static_cast<int>(values.size()) == Rows() * Cols());
    m_storage.reserve(Rows() * Cols());
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        m_storage.emplace_back(values[row * Cols() + col]);
      }
    }
  }
 
  Variable& operator()(int row, int col) {
    Assert(row >= 0 && row < Rows());
    Assert(col >= 0 && col < Cols());
    return m_storage[row * Cols() + col];
  }
 
  const Variable& operator()(int row, int col) const {
    Assert(row >= 0 && row < Rows());
    Assert(col >= 0 && col < Cols());
    return m_storage[row * Cols() + col];
  }
 
  Variable& operator()(int row) {
    Assert(row >= 0 && row < Rows() * Cols());
    return m_storage[row];
  }
 
  const Variable& operator()(int row) const {
    Assert(row >= 0 && row < Rows() * Cols());
    return m_storage[row];
  }
 
  VariableBlock<VariableMatrix> Block(int rowOffset, int colOffset,
                                      int blockRows, int blockCols) {
    Assert(rowOffset >= 0 && rowOffset <= Rows());
    Assert(colOffset >= 0 && colOffset <= Cols());
    Assert(blockRows >= 0 && blockRows <= Rows() - rowOffset);
    Assert(blockCols >= 0 && blockCols <= Cols() - colOffset);
    return VariableBlock{*this, rowOffset, colOffset, blockRows, blockCols};
  }
 
  const VariableBlock<const VariableMatrix> Block(int rowOffset, int colOffset,
                                                  int blockRows,
                                                  int blockCols) const {
    Assert(rowOffset >= 0 && rowOffset <= Rows());
    Assert(colOffset >= 0 && colOffset <= Cols());
    Assert(blockRows >= 0 && blockRows <= Rows() - rowOffset);
    Assert(blockCols >= 0 && blockCols <= Cols() - colOffset);
    return VariableBlock{*this, rowOffset, colOffset, blockRows, blockCols};
  }
 
  VariableBlock<VariableMatrix> operator()(Slice rowSlice, Slice colSlice) {
    int rowSliceLength = rowSlice.Adjust(Rows());
    int colSliceLength = colSlice.Adjust(Cols());
    return VariableBlock{*this, std::move(rowSlice), rowSliceLength,
                         std::move(colSlice), colSliceLength};
  }
 
  const VariableBlock<const VariableMatrix> operator()(Slice rowSlice,
                                                       Slice colSlice) const {
    int rowSliceLength = rowSlice.Adjust(Rows());
    int colSliceLength = colSlice.Adjust(Cols());
    return VariableBlock{*this, std::move(rowSlice), rowSliceLength,
                         std::move(colSlice), colSliceLength};
  }
 
  VariableBlock<VariableMatrix> operator()(Slice rowSlice, int rowSliceLength,
                                           Slice colSlice, int colSliceLength) {
    return VariableBlock{*this, std::move(rowSlice), rowSliceLength,
                         std::move(colSlice), colSliceLength};
  }
 
  const VariableBlock<const VariableMatrix> operator()(
      Slice rowSlice, int rowSliceLength, Slice colSlice,
      int colSliceLength) const {
    return VariableBlock{*this, std::move(rowSlice), rowSliceLength,
                         std::move(colSlice), colSliceLength};
  }
 
  VariableBlock<VariableMatrix> Segment(int offset, int length) {
    Assert(offset >= 0 && offset < Rows() * Cols());
    Assert(length >= 0 && length <= Rows() * Cols() - offset);
    return Block(offset, 0, length, 1);
  }
 
  const VariableBlock<const VariableMatrix> Segment(int offset,
                                                    int length) const {
    Assert(offset >= 0 && offset < Rows() * Cols());
    Assert(length >= 0 && length <= Rows() * Cols() - offset);
    return Block(offset, 0, length, 1);
  }
 
  VariableBlock<VariableMatrix> Row(int row) {
    Assert(row >= 0 && row < Rows());
    return Block(row, 0, 1, Cols());
  }
 
  const VariableBlock<const VariableMatrix> Row(int row) const {
    Assert(row >= 0 && row < Rows());
    return Block(row, 0, 1, Cols());
  }
 
  VariableBlock<VariableMatrix> Col(int col) {
    Assert(col >= 0 && col < Cols());
    return Block(0, col, Rows(), 1);
  }
 
  const VariableBlock<const VariableMatrix> Col(int col) const {
    Assert(col >= 0 && col < Cols());
    return Block(0, col, Rows(), 1);
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix
  operator*(const VariableMatrix& lhs, const VariableMatrix& rhs) {
    Assert(lhs.Cols() == rhs.Rows());
 
    VariableMatrix result{lhs.Rows(), rhs.Cols()};
 
    for (int i = 0; i < lhs.Rows(); ++i) {
      for (int j = 0; j < rhs.Cols(); ++j) {
        Variable sum;
        for (int k = 0; k < lhs.Cols(); ++k) {
          sum += lhs(i, k) * rhs(k, j);
        }
        result(i, j) = sum;
      }
    }
 
    return result;
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix operator*(const VariableMatrix& lhs,
                                                     const Variable& rhs) {
    VariableMatrix result{lhs.Rows(), lhs.Cols()};
 
    for (int row = 0; row < result.Rows(); ++row) {
      for (int col = 0; col < result.Cols(); ++col) {
        result(row, col) = lhs(row, col) * rhs;
      }
    }
 
    return result;
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix operator*(const VariableMatrix& lhs,
                                                     double rhs) {
    return lhs * Variable{rhs};
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix
  operator*(const Variable& lhs, const VariableMatrix& rhs) {
    VariableMatrix result{rhs.Rows(), rhs.Cols()};
 
    for (int row = 0; row < result.Rows(); ++row) {
      for (int col = 0; col < result.Cols(); ++col) {
        result(row, col) = rhs(row, col) * lhs;
      }
    }
 
    return result;
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix
  operator*(double lhs, const VariableMatrix& rhs) {
    return Variable{lhs} * rhs;
  }
 
  VariableMatrix& operator*=(const VariableMatrix& rhs) {
    Assert(Cols() == rhs.Rows() && Cols() == rhs.Cols());
 
    for (int i = 0; i < Rows(); ++i) {
      for (int j = 0; j < rhs.Cols(); ++j) {
        Variable sum;
        for (int k = 0; k < Cols(); ++k) {
          sum += (*this)(i, k) * rhs(k, j);
        }
        (*this)(i, j) = sum;
      }
    }
 
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix operator/(const VariableMatrix& lhs,
                                                     const Variable& rhs) {
    VariableMatrix result{lhs.Rows(), lhs.Cols()};
 
    for (int row = 0; row < result.Rows(); ++row) {
      for (int col = 0; col < result.Cols(); ++col) {
        result(row, col) = lhs(row, col) / rhs;
      }
    }
 
    return result;
  }
 
  VariableMatrix& operator/=(const Variable& rhs) {
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        (*this)(row, col) /= rhs;
      }
    }
 
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix
  operator+(const VariableMatrix& lhs, const VariableMatrix& rhs) {
    VariableMatrix result{lhs.Rows(), lhs.Cols()};
 
    for (int row = 0; row < result.Rows(); ++row) {
      for (int col = 0; col < result.Cols(); ++col) {
        result(row, col) = lhs(row, col) + rhs(row, col);
      }
    }
 
    return result;
  }
 
  VariableMatrix& operator+=(const VariableMatrix& rhs) {
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        (*this)(row, col) += rhs(row, col);
      }
    }
 
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix
  operator-(const VariableMatrix& lhs, const VariableMatrix& rhs) {
    VariableMatrix result{lhs.Rows(), lhs.Cols()};
 
    for (int row = 0; row < result.Rows(); ++row) {
      for (int col = 0; col < result.Cols(); ++col) {
        result(row, col) = lhs(row, col) - rhs(row, col);
      }
    }
 
    return result;
  }
 
  VariableMatrix& operator-=(const VariableMatrix& rhs) {
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        (*this)(row, col) -= rhs(row, col);
      }
    }
 
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT VariableMatrix
  operator-(const VariableMatrix& lhs) {
    VariableMatrix result{lhs.Rows(), lhs.Cols()};
 
    for (int row = 0; row < result.Rows(); ++row) {
      for (int col = 0; col < result.Cols(); ++col) {
        result(row, col) = -lhs(row, col);
      }
    }
 
    return result;
  }
 
  operator Variable() const {  // NOLINT
    Assert(Rows() == 1 && Cols() == 1);
    return (*this)(0, 0);
  }
 
  VariableMatrix T() const {
    VariableMatrix result{Cols(), Rows()};
 
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        result(col, row) = (*this)(row, col);
      }
    }
 
    return result;
  }
 
  int Rows() const { return m_rows; }
 
  int Cols() const { return m_cols; }
 
  double Value(int row, int col) {
    Assert(row >= 0 && row < Rows());
    Assert(col >= 0 && col < Cols());
    return m_storage[row * Cols() + col].Value();
  }
 
  double Value(int index) {
    Assert(index >= 0 && index < Rows() * Cols());
    return m_storage[index].Value();
  }
 
  Eigen::MatrixXd Value() {
    Eigen::MatrixXd result{Rows(), Cols()};
 
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        result(row, col) = Value(row, col);
      }
    }
 
    return result;
  }
 
  VariableMatrix CwiseTransform(
      function_ref<Variable(const Variable& x)> unaryOp) const {
    VariableMatrix result{Rows(), Cols()};
 
    for (int row = 0; row < Rows(); ++row) {
      for (int col = 0; col < Cols(); ++col) {
        result(row, col) = unaryOp((*this)(row, col));
      }
    }
 
    return result;
  }
 
  class iterator {
   public:
    using iterator_category = std::forward_iterator_tag;
    using value_type = Variable;
    using difference_type = std::ptrdiff_t;
    using pointer = Variable*;
    using reference = Variable&;
 
    explicit iterator(small_vector<Variable>::iterator it) : m_it{it} {}
 
    iterator& operator++() {
      ++m_it;
      return *this;
    }
    iterator operator++(int) {
      iterator retval = *this;
      ++(*this);
      return retval;
    }
    bool operator==(const iterator&) const = default;
    reference operator*() { return *m_it; }
 
   private:
    small_vector<Variable>::iterator m_it;
  };
 
  class const_iterator {
   public:
    using iterator_category = std::forward_iterator_tag;
    using value_type = Variable;
    using difference_type = std::ptrdiff_t;
    using pointer = Variable*;
    using const_reference = const Variable&;
 
    explicit const_iterator(small_vector<Variable>::const_iterator it)
        : m_it{it} {}
 
    const_iterator& operator++() {
      ++m_it;
      return *this;
    }
    const_iterator operator++(int) {
      const_iterator retval = *this;
      ++(*this);
      return retval;
    }
    bool operator==(const const_iterator&) const = default;
    const_reference operator*() const { return *m_it; }
 
   private:
    small_vector<Variable>::const_iterator m_it;
  };
 
  iterator begin() { return iterator{m_storage.begin()}; }
 
  iterator end() { return iterator{m_storage.end()}; }
 
  const_iterator begin() const { return const_iterator{m_storage.begin()}; }
 
  const_iterator end() const { return const_iterator{m_storage.end()}; }
 
  const_iterator cbegin() const { return const_iterator{m_storage.cbegin()}; }
 
  const_iterator cend() const { return const_iterator{m_storage.cend()}; }
 
  size_t size() const { return m_storage.size(); }
 
  static VariableMatrix Zero(int rows, int cols) {
    VariableMatrix result{rows, cols};
 
    for (auto& elem : result) {
      elem = 0.0;
    }
 
    return result;
  }
 
  static VariableMatrix Ones(int rows, int cols) {
    VariableMatrix result{rows, cols};
 
    for (auto& elem : result) {
      elem = 1.0;
    }
 
    return result;
  }
 
 private:
  small_vector<Variable> m_storage;
  int m_rows = 0;
  int m_cols = 0;
};
 
SLEIPNIR_DLLEXPORT inline VariableMatrix CwiseReduce(
    const VariableMatrix& lhs, const VariableMatrix& rhs,
    function_ref<Variable(const Variable& x, const Variable& y)> binaryOp) {
  Assert(lhs.Rows() == rhs.Rows());
  Assert(lhs.Rows() == rhs.Rows());
 
  VariableMatrix result{lhs.Rows(), lhs.Cols()};
 
  for (int row = 0; row < lhs.Rows(); ++row) {
    for (int col = 0; col < lhs.Cols(); ++col) {
      result(row, col) = binaryOp(lhs(row, col), rhs(row, col));
    }
  }
 
  return result;
}
 
SLEIPNIR_DLLEXPORT inline VariableMatrix Block(
    std::initializer_list<std::initializer_list<VariableMatrix>> list) {
  // Get row and column counts for destination matrix
  int rows = 0;
  int cols = -1;
  for (const auto& row : list) {
    if (row.size() > 0) {
      rows += row.begin()->Rows();
    }
 
    // Get number of columns in this row
    int latestCols = 0;
    for (const auto& elem : row) {
      // Assert the first and latest row have the same height
      Assert(row.begin()->Rows() == elem.Rows());
 
      latestCols += elem.Cols();
    }
 
    // If this is the first row, record the column count. Otherwise, assert the
    // first and latest column counts are the same.
    if (cols == -1) {
      cols = latestCols;
    } else {
      Assert(cols == latestCols);
    }
  }
 
  VariableMatrix result{rows, cols};
 
  int rowOffset = 0;
  for (const auto& row : list) {
    int colOffset = 0;
    for (const auto& elem : row) {
      result.Block(rowOffset, colOffset, elem.Rows(), elem.Cols()) = elem;
      colOffset += elem.Cols();
    }
    rowOffset += row.begin()->Rows();
  }
 
  return result;
}
 
SLEIPNIR_DLLEXPORT inline VariableMatrix Block(
    const std::vector<std::vector<VariableMatrix>>& list) {
  // Get row and column counts for destination matrix
  int rows = 0;
  int cols = -1;
  for (const auto& row : list) {
    if (row.size() > 0) {
      rows += row.begin()->Rows();
    }
 
    // Get number of columns in this row
    int latestCols = 0;
    for (const auto& elem : row) {
      // Assert the first and latest row have the same height
      Assert(row.begin()->Rows() == elem.Rows());
 
      latestCols += elem.Cols();
    }
 
    // If this is the first row, record the column count. Otherwise, assert the
    // first and latest column counts are the same.
    if (cols == -1) {
      cols = latestCols;
    } else {
      Assert(cols == latestCols);
    }
  }
 
  VariableMatrix result{rows, cols};
 
  int rowOffset = 0;
  for (const auto& row : list) {
    int colOffset = 0;
    for (const auto& elem : row) {
      result.Block(rowOffset, colOffset, elem.Rows(), elem.Cols()) = elem;
      colOffset += elem.Cols();
    }
    rowOffset += row.begin()->Rows();
  }
 
  return result;
}
 
SLEIPNIR_DLLEXPORT VariableMatrix Solve(const VariableMatrix& A,
                                        const VariableMatrix& B);
 
}  // namespace sleipnir