eigen.h

Go to the documentation of this file.

/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2013  Virginia Tech
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 *  Justus Calvin
 *  Department of Chemistry, Virginia Tech
 *
 *  eigen.h
 *  May 02, 2015
 *
 */

#ifndef TILEDARRAY_CONVERSIONS_EIGEN_H__INCLUDED
#define TILEDARRAY_CONVERSIONS_EIGEN_H__INCLUDED

#include <cstdint>
#include <tiledarray_fwd.h>
#include <TiledArray/tensor.h>
#include <TiledArray/error.h>
#include <TiledArray/math/eigen.h>
#include <TiledArray/madness.h>
#include <TiledArray/pmap/replicated_pmap.h>
#include "TiledArray/dist_array.h"

namespace TiledArray {

  // Convenience typedefs
  typedef Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> EigenMatrixXd;
  typedef Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> EigenMatrixXf;
  typedef Eigen::Matrix<std::complex<double>, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> EigenMatrixXcd;
  typedef Eigen::Matrix<std::complex<float>, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> EigenMatrixXcf;
  typedef Eigen::Matrix<int, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> EigenMatrixXi;
  typedef Eigen::Matrix<long, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor> EigenMatrixXl;
  typedef Eigen::Matrix<double, Eigen::Dynamic, 1> EigenVectorXd;
  typedef Eigen::Matrix<float, Eigen::Dynamic, 1> EigenVectorXf;
  typedef Eigen::Matrix<std::complex<double>, 1, Eigen::Dynamic> EigenVectorXcd;
  typedef Eigen::Matrix<std::complex<float>, 1, Eigen::Dynamic> EigenVectorXcf;
  typedef Eigen::Matrix<int, Eigen::Dynamic, 1> EigenVectorXi;
  typedef Eigen::Matrix<long, Eigen::Dynamic, 1> EigenVectorXl;



  template <typename T, typename A>
  inline Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>, Eigen::AutoAlign>
  eigen_map(const Tensor<T, A>& tensor, const std::size_t m, const std::size_t n) {
    TA_ASSERT((m * n) == tensor.size());

    return Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic,
        Eigen::RowMajor>, Eigen::AutoAlign>(tensor.data(), m, n);
  }


  template <typename T, typename A>
  inline Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>, Eigen::AutoAlign>
  eigen_map(Tensor<T, A>& tensor, const std::size_t m, const std::size_t n) {
    TA_ASSERT((m * n) == tensor.size());

    return Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic,
        Eigen::RowMajor>, Eigen::AutoAlign>(tensor.data(), m, n);
  }


  template <typename T, typename A>
  inline Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>, Eigen::AutoAlign>
  eigen_map(const Tensor<T, A>& tensor, const std::size_t n) {
    TA_ASSERT(n == tensor.size());

    return Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, 1>,
        Eigen::AutoAlign>(tensor.data(), n);
  }


  template <typename T, typename A>
  inline Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 1>, Eigen::AutoAlign>
  eigen_map(Tensor<T, A>& tensor, const std::size_t n) {
    TA_ASSERT(n == tensor.size());

    return Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, 1>,
        Eigen::AutoAlign>(tensor.data(), n);
  }


  template <typename T, typename A>
  inline Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>, Eigen::AutoAlign>
  eigen_map(const Tensor<T, A>& tensor) {
    TA_ASSERT((tensor.range().rank() == 2u) || (tensor.range().rank() == 1u));
    const auto* MADNESS_RESTRICT const tensor_extent = tensor.range().extent_data();
    return eigen_map(tensor, tensor_extent[0],
            (tensor.range().rank() == 2u ? tensor_extent[1] : 1ul));
  }


  template <typename T, typename A>
  inline Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>, Eigen::AutoAlign>
  eigen_map(Tensor<T, A>& tensor) {
    TA_ASSERT((tensor.range().rank() == 2u) || (tensor.range().rank() == 1u));
    const auto* MADNESS_RESTRICT const tensor_extent = tensor.range().extent_data();
    return eigen_map(tensor, tensor_extent[0],
            (tensor.range().rank() == 2u ? tensor_extent[1] : 1ul));
  }


  template <typename T, typename A, typename Derived>
  inline void eigen_submatrix_to_tensor(const Eigen::MatrixBase<Derived>& matrix, Tensor<T, A>& tensor) {
    typedef typename Tensor<T, A>::size_type size_type;
    TA_ASSERT((tensor.range().rank() == 2u) || (tensor.range().rank() == 1u));

    // Get pointers to the tensor range data
    const auto* MADNESS_RESTRICT const tensor_lower = tensor.range().lobound_data();
    const auto* MADNESS_RESTRICT const tensor_upper = tensor.range().upbound_data();
    const auto* MADNESS_RESTRICT const tensor_extent = tensor.range().extent_data();

    if(tensor.range().rank() == 2u) {
      // Get tensor range data
      const std::size_t tensor_lower_0 = tensor_lower[0];
      const std::size_t tensor_lower_1 = tensor_lower[1];
      const std::size_t tensor_upper_0 = tensor_upper[0];
      const std::size_t tensor_upper_1 = tensor_upper[1];
      const std::size_t tensor_extent_0 = tensor_extent[0];
      const std::size_t tensor_extent_1 = tensor_extent[1];

      TA_ASSERT(tensor_upper_0 <= size_type(matrix.rows()));
      TA_ASSERT(tensor_upper_1 <= size_type(matrix.cols()));

      // Copy matrix
      eigen_map(tensor, tensor_extent_0, tensor_extent_1) =
          matrix.block(tensor_lower_0, tensor_lower_1,
          tensor_extent_0, tensor_extent_1);
    } else {
      // Get tensor range data
      const std::size_t tensor_lower_0 = tensor_lower[0];
      const std::size_t tensor_upper_0 = tensor_upper[0];
      const std::size_t tensor_extent_0 = tensor_extent[0];

      // Check that matrix is a vector.
      TA_ASSERT((matrix.rows() == 1) || (matrix.cols() == 1));

      if(matrix.rows() == 1) {
        TA_ASSERT(tensor_upper_0 <= size_type(matrix.cols()));

        // Copy the row vector to tensor
        eigen_map(tensor, 1, tensor_extent_0) =
            matrix.block(0, tensor_lower_0, 1, tensor_extent_0);
      } else {
        TA_ASSERT(tensor_upper_0 <= size_type(matrix.rows()));

        // Copy the column vector to tensor
        eigen_map(tensor, tensor_extent_0, 1) =
            matrix.block(tensor_lower_0, 0, tensor_extent_0, 1);
      }
    }
  }


  template <typename T, typename A, typename Derived>
  inline void tensor_to_eigen_submatrix(const Tensor<T, A>& tensor, Eigen::MatrixBase<Derived>& matrix) {
    typedef typename Tensor<T, A>::size_type size_type;
    TA_ASSERT((tensor.range().rank() == 2u) || (tensor.range().rank() == 1u));

    // Get pointers to the tensor range data
    const auto* MADNESS_RESTRICT const tensor_lower = tensor.range().lobound_data();
    const auto* MADNESS_RESTRICT const tensor_upper = tensor.range().upbound_data();
    const auto* MADNESS_RESTRICT const tensor_extent = tensor.range().extent_data();

    if(tensor.range().rank() == 2) {
      // Get tensor range data
      const std::size_t tensor_lower_0 = tensor_lower[0];
      const std::size_t tensor_lower_1 = tensor_lower[1];
      const std::size_t tensor_upper_0 = tensor_upper[0];
      const std::size_t tensor_upper_1 = tensor_upper[1];
      const std::size_t tensor_extent_0 = tensor_extent[0];
      const std::size_t tensor_extent_1 = tensor_extent[1];

      TA_ASSERT(tensor_upper_0 <= size_type(matrix.rows()));
      TA_ASSERT(tensor_upper_1 <= size_type(matrix.cols()));

      // Copy tensor into matrix
      matrix.block(tensor_lower_0, tensor_lower_1,
          tensor_extent_0, tensor_extent_1) =
              eigen_map(tensor, tensor_extent_0, tensor_extent_1);
    } else {
      // Get tensor range data
      const std::size_t tensor_lower_0 = tensor_lower[0];
      const std::size_t tensor_upper_0 = tensor_upper[0];
      const std::size_t tensor_extent_0 = tensor_extent[0];

      TA_ASSERT((matrix.rows() == 1) || (matrix.cols() == 1));

      if(matrix.rows() == 1) {
        TA_ASSERT(tensor_upper_0 <= size_type(matrix.cols()));

        // Copy tensor into row vector
        matrix.block(0, tensor_lower_0, 1, tensor_extent_0) =
            eigen_map(tensor, 1, tensor_extent_0);
      } else {
        TA_ASSERT(tensor_upper_0 <= size_type(matrix.rows()));

        // Copy tensor into column vector
        matrix.block(tensor_lower_0, 0, tensor_extent_0, 1) =
            eigen_map(tensor, tensor_extent_0, 1);
      }
    }
  }

  namespace detail {


    template <typename A, typename Derived>
    void counted_eigen_submatrix_to_tensor(const Eigen::MatrixBase<Derived>* matrix,
        A* array, const typename A::size_type i, madness::AtomicInt* counter)
    {
      typename A::value_type tensor(array->trange().make_tile_range(i));
      eigen_submatrix_to_tensor(*matrix, tensor);
      array->set(i, tensor);
      (*counter)++;
    }


    template <typename Derived, typename T>
    void counted_tensor_to_eigen_submatrix(const T& tensor,
        Eigen::MatrixBase<Derived>* matrix, madness::AtomicInt* counter)
    {
      tensor_to_eigen_submatrix(tensor, *matrix);
      (*counter)++;
    }

  } // namespace detail


  template <typename A, typename Derived>
  A eigen_to_array(World& world, const typename A::trange_type& trange,
      const Eigen::MatrixBase<Derived>& matrix, bool replicated = false)
  {
    typedef typename A::size_type size_type;
    // Check that trange matches the dimensions of other
    if((matrix.cols() > 1) && (matrix.rows() > 1)) {
      TA_USER_ASSERT(trange.tiles_range().rank() == 2,
          "TiledArray::eigen_to_array(): The number of dimensions in trange is not equal to that of the Eigen matrix.");
      TA_USER_ASSERT(trange.elements_range().extent(0) == size_type(matrix.rows()),
          "TiledArray::eigen_to_array(): The number of rows in trange is not equal to the number of rows in the Eigen matrix.");
      TA_USER_ASSERT(trange.elements_range().extent(1) == size_type(matrix.cols()),
          "TiledArray::eigen_to_array(): The number of columns in trange is not equal to the number of columns in the Eigen matrix.");
    } else {
      TA_USER_ASSERT(trange.tiles_range().rank() == 1,
          "TiledArray::eigen_to_array(): The number of dimensions in trange must match that of the Eigen matrix.");
      TA_USER_ASSERT(trange.elements_range().extent(0) == size_type(matrix.size()),
          "TiledArray::eigen_to_array(): The size of trange must be equal to the matrix size.");
    }

    // Check that this is not a distributed computing environment
    if(! replicated)
      TA_USER_ASSERT(world.size() == 1,
          "An array cannot be assigned with an Eigen::Matrix when the number of World ranks is greater than 1.");

    // Create a new tensor
    A array = (replicated && (world.size() > 1)
                   ? A(world, trange,
                       std::static_pointer_cast<typename A::pmap_interface>(
                           std::make_shared<detail::ReplicatedPmap>(
                               world, trange.tiles_range().volume())))
                   : A(world, trange));

    // Spawn tasks to copy Eigen to an Array
    madness::AtomicInt counter;
    counter = 0;
    std::int64_t n = 0;
    for(std::size_t i = 0; i < array.size(); ++i) {
      world.taskq.add(& detail::counted_eigen_submatrix_to_tensor<A, Derived>,
          &matrix, &array, i, &counter);
      ++n;
    }

    // Wait until the write tasks are complete
    array.world().await([&counter,n] () { return counter == n; });

    return array;
  }


  template <typename Tile, typename Policy,
            unsigned int EigenStorageOrder = Eigen::ColMajor>
  Eigen::Matrix<typename Tile::value_type, Eigen::Dynamic, Eigen::Dynamic,
                EigenStorageOrder>
  array_to_eigen(const DistArray<Tile, Policy>& array) {
    typedef Eigen::Matrix<typename Tile::value_type, Eigen::Dynamic,
                          Eigen::Dynamic, EigenStorageOrder>
        EigenMatrix;

    const auto rank = array.trange().tiles_range().rank();

    // Check that the array will fit in a matrix or vector
    TA_USER_ASSERT((rank == 2u) || (rank == 1u),
        "TiledArray::array_to_eigen(): The array dimensions must be equal to 1 or 2.");

    // Check that this is not a distributed computing environment or that the
    // array is replicated
    if(! array.pmap()->is_replicated())
      TA_USER_ASSERT(array.world().size() == 1,
          "TiledArray::array_to_eigen(): Array cannot be assigned with an Eigen::Matrix when the number of World ranks is greater than 1.");

    // Construct the Eigen matrix
    const auto* MADNESS_RESTRICT const array_extent = array.trange().elements_range().extent_data();
    // if array is sparse must initialize to zero
    EigenMatrix matrix = EigenMatrix::Zero(array_extent[0], (rank == 2 ? array_extent[1] : 1));

    // Spawn tasks to copy array tiles to the Eigen matrix
    madness::AtomicInt counter;
    counter = 0;
    int n = 0;
    for(std::size_t i = 0; i < array.size(); ++i) {
      if(! array.is_zero(i)) {
        array.world().taskq.add(
            & detail::counted_tensor_to_eigen_submatrix<EigenMatrix,
            typename DistArray<Tile, Policy>::value_type>,
            array.find(i), &matrix, &counter);
        ++n;
      }
    }

    // Wait until the above tasks are complete. Tasks will be processed by this
    // thread while waiting.
    array.world().await([&counter,n] () { return counter == n; });

    return matrix;
  }


  template <typename A>
  inline A row_major_buffer_to_array(World& world, const typename A::trange_type& trange,
      const typename A::value_type::value_type* buffer, const std::size_t m,
      const std::size_t n, const bool replicated = false)
  {
    TA_USER_ASSERT(trange.elements_range().extent(0) == m,
        "TiledArray::eigen_to_array(): The number of rows in trange is not equal to m.");
    TA_USER_ASSERT(trange.elements_range().extent(1) == n,
        "TiledArray::eigen_to_array(): The number of columns in trange is not equal to n.");

    typedef Eigen::Matrix<typename A::value_type::value_type, Eigen::Dynamic,
        Eigen::Dynamic, Eigen::RowMajor> matrix_type;
    return eigen_to_array(world, trange, Eigen::Map<const matrix_type,
        Eigen::AutoAlign>(buffer, m, n), replicated);
  }


  template <typename A>
  inline A column_major_buffer_to_array(World& world, const typename A::trange_type& trange,
      const typename A::value_type::value_type* buffer, const std::size_t m,
      const std::size_t n, const bool replicated = false)
  {
    TA_USER_ASSERT(trange.elements_range().extent(0) == m,
        "TiledArray::eigen_to_array(): The number of rows in trange is not equal to m.");
    TA_USER_ASSERT(trange.elements_range().extent(1) == n,
        "TiledArray::eigen_to_array(): The number of columns in trange is not equal to n.");

    typedef Eigen::Matrix<typename A::value_type::value_type, Eigen::Dynamic,
        Eigen::Dynamic, Eigen::ColMajor> matrix_type;
    return eigen_to_array(world, trange, Eigen::Map<const matrix_type,
        Eigen::AutoAlign>(buffer, m, n), replicated);
  }

} // namespace TiledArray

#endif // TILEDARRAY_CONVERSIONS_EIGEN_H__INCLUDED