retile.h

Go to the documentation of this file.

/*
 *  This file is a part of TiledArray.
 *  Copyright (C) 2020  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/>.
 *
 *  retile.h
 *  May 11, 2020
 */
 
#ifndef TILEDARRAY_RETILE_H
#define TILEDARRAY_RETILE_H
 
#include "TiledArray/util/annotation.h"
#include "TiledArray/special/diagonal_array.h"
 
 
 
namespace TiledArray {
 
template <typename TileType, typename PolicyType>
auto retile(const DistArray<TileType, PolicyType>& tensor,
            const TiledRange& new_trange) {
  // Make sure ranks match
  auto rank = new_trange.rank();
  auto tensor_rank = tensor.trange().rank();
  assert((rank == tensor_rank) && "TiledRanges are of different ranks");
 
  // Makes the annotations for the contraction step
  auto annotations =
      [&](std::size_t target_dim) -> std::tuple<std::string, std::string> {
    std::ostringstream final, switcher;
    switcher << "i" << target_dim << ",iX";
    if (target_dim == 0) {
      final << "iX";
    } else {
      final << "i0";
    }
    for (unsigned int d = 1; d < rank; ++d) {
      if (d == target_dim) {
        final << ",iX";
      } else {
        final << ",i" << d;
      }
    }
    return {final.str(), switcher.str()};
  };
 
  // Check the different dimensions and contract when needed
  using tensor_type = DistArray<TileType, PolicyType>;
  auto start = detail::dummy_annotation(rank);
  tensor_type output_tensor;
  for (auto i = 0; i < rank; ++i) {
    if (i == 0) { output_tensor(start) = tensor(start); }
    if (new_trange.dim(i) != tensor.trange().dim(i)) {
      // Make identity for contraction
      TiledRange retiler{tensor.trange().dim(i), new_trange.dim(i)};
      auto identity = diagonal_array<tensor_type>(tensor.world(), retiler);
 
      // Make indices for contraction
      auto [finish, change] = annotations(i);
 
      // Retile
      output_tensor(finish) = output_tensor(start) * identity(change);
    }
  }
 
  return output_tensor;
}
 
} // namespace TiledArray
 
 
#endif  // TILEDARRAY_RETILE_H