Domi_MDMap.hpp

// @HEADER
// ***********************************************************************
//
//     Domi: Multi-dimensional Distributed Linear Algebra Services
//                 Copyright (2014) Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia
// Corporation, the U.S. Government retains certain rights in this
// software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact William F. Spotz (wfspotz@sandia.gov)
//
// ***********************************************************************
// @HEADER

#ifndef DOMI_MDMAP_HPP
#define DOMI_MDMAP_HPP

// System includes
#include <algorithm>
#include <limits>
#include <sstream>

// Domi includes
#include "Domi_ConfigDefs.hpp"
#include "Domi_Utils.hpp"
#include "Domi_getValidParameters.hpp"
#include "Domi_Slice.hpp"
#include "Domi_MDComm.hpp"
#include "Domi_MDArray.hpp"

// Teuchos includes
#include "Teuchos_Comm.hpp"
#include "Teuchos_DefaultComm.hpp"
#include "Teuchos_CommHelpers.hpp"
#include "Teuchos_Tuple.hpp"

// Kokkos includes
#include "Kokkos_Core.hpp"

#ifdef HAVE_EPETRA
#include "Epetra_Map.h"
#endif

#ifdef HAVE_TPETRA
#include "Tpetra_Map.hpp"
#include "Tpetra_Util.hpp"
typedef Tpetra::Details::DefaultTypes::local_ordinal_type  TpetraLOType;
typedef Tpetra::Details::DefaultTypes::global_ordinal_type TpetraGOType;
typedef Tpetra::Details::DefaultTypes::node_type           TpetraNodeType;
#endif

namespace Domi
{

class MDMap
{
public:


  MDMap(const Teuchos::RCP< const MDComm > mdComm,
        const Teuchos::ArrayView< const dim_type > & dimensions,
        const Teuchos::ArrayView< const int > & commPad =
          Teuchos::ArrayView< const int >(),
        const Teuchos::ArrayView< const int > & bndryPad =
          Teuchos::ArrayView< const int >(),
        const Teuchos::ArrayView< const int > & replicatedBoundary =
          Teuchos::ArrayView< const int >(),
        const Layout layout = DEFAULT_ORDER);

  MDMap(Teuchos::ParameterList & plist);

  MDMap(const Teuchos::RCP< const Teuchos::Comm< int > > teuchosComm,
        Teuchos::ParameterList & plist);

  MDMap(const Teuchos::RCP< const MDComm > mdComm,
        Teuchos::ParameterList & plist);

  MDMap(const Teuchos::RCP< const MDComm > mdComm,
        const Teuchos::ArrayView< Slice > & myGlobalBounds,
        const Teuchos::ArrayView< padding_type > & padding =
          Teuchos::ArrayView< padding_type >(),
        const Teuchos::ArrayView< const int > & replicatedBoundary =
          Teuchos::ArrayView< const int >(),
        const Layout layout = DEFAULT_ORDER);

  MDMap(const MDMap & source);

  MDMap(const MDMap & parent,
        int axis,
        dim_type index);

  MDMap(const MDMap & parent,
        int axis,
        const Slice & slice,
        int bndryPad = 0);

  MDMap(const MDMap & parent,
        const Teuchos::ArrayView< Slice > & slices,
        const Teuchos::ArrayView< int > & bndryPad =
          Teuchos::ArrayView< int >());

  ~MDMap();



  MDMap & operator=(const MDMap & source);



  inline Teuchos::RCP< const Teuchos::Comm< int > > getTeuchosComm() const;

  inline Teuchos::RCP< const MDComm > getMDComm() const;

  inline bool onSubcommunicator() const;

  inline int numDims() const;

  inline Teuchos::Array< int > getCommDims() const;

  inline int getCommDim(int axis) const;

  inline bool isPeriodic(int axis) const;

 inline  int getCommIndex(int axis) const;

  inline int getLowerNeighbor(int axis) const;

  inline int getUpperNeighbor(int axis) const;



  inline Teuchos::Array< dim_type > getGlobalDims() const;

  dim_type getGlobalDim(int axis,
                        bool withBndryPad=false) const;

  Slice getGlobalBounds(int axis,
                        bool withBndryPad=false) const;

  Slice getGlobalRankBounds(int axis,
                            bool withBndryPad=false) const;

  Teuchos::Array< dim_type > getLocalDims() const;

  dim_type getLocalDim(int axis,
                       bool withPad=false) const;

  Teuchos::ArrayView< const Slice > getLocalBounds() const;

  Slice getLocalBounds(int axis,
                       bool withPad=false) const;

  Slice getLocalInteriorBounds(int axis) const;



  bool hasPadding() const;

  int getLowerPadSize(int axis) const;

  int getUpperPadSize(int axis) const;

  int getCommPadSize(int axis) const;

  int getLowerBndryPad(int axis) const;

  int getUpperBndryPad(int axis) const;

  Teuchos::Array< int > getBndryPadSizes() const;

  int getBndryPadSize(int axis) const;

  /* \brief Return whether given local index is in the padding
   *
   * \param index [in] An array of indexes ((i) for 1D, (i,j) for 2D,
   *        (i,j,k) for 3D, etc)
   */
  bool isPad(const Teuchos::ArrayView< dim_type > & index) const;

  /* \brief Return whether given local index is in the communication
   *        padding
   *
   * \param index [in] An array of indexes ((i) for 1D, (i,j) for 2D,
   *        (i,j,k) for 3D, etc)
   */
  bool isCommPad(const Teuchos::ArrayView< dim_type > & index) const;

  /* \brief Return whether given local index is in the boundary
   *        padding
   *
   * \param index [in] An array of indexes ((i) for 1D, (i,j) for 2D,
   *        (i,j,k) for 3D, etc)
   */
  bool isBndryPad(const Teuchos::ArrayView< dim_type > & index) const;

  bool isReplicatedBoundary(int axis) const;

  Layout getLayout() const;



  Teuchos::ArrayView< Teuchos::RCP< const Domi::MDMap > >
  getAxisMaps() const;

  Teuchos::RCP< const Domi::MDMap > getAxisMap(int axis) const;

  Teuchos::RCP< const MDMap >
  getAugmentedMDMap(const dim_type leadingDim,
                    const dim_type trailingDim=0) const;

#ifdef HAVE_EPETRA

  Teuchos::RCP< const Epetra_Map >
  getEpetraMap(bool withCommPad=true) const;

  Teuchos::RCP< const Epetra_Map >
  getEpetraAxisMap(int axis,
                   bool withCommPad=true) const;

#endif

#ifdef HAVE_TPETRA

  template< class LocalOrdinal  = TpetraLOType,
            class GlobalOrdinal = TpetraGOType,
            class Node          = TpetraNodeType>
  Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > >
  getTpetraMap(bool withCommPad=true) const;

  template< class LocalOrdinal  = TpetraLOType,
            class GlobalOrdinal = TpetraGOType,
            class Node          = TpetraNodeType >
  Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > >
  getTpetraAxisMap(int axis,
                   bool withCommPad=true) const;

#endif



  Teuchos::Array< dim_type >
  getGlobalIndex(size_type globalID) const;

  Teuchos::Array< dim_type >
  getLocalIndex(size_type localID) const;

  size_type getGlobalID(size_type localID) const;

  size_type
  getGlobalID(const Teuchos::ArrayView< dim_type > & globalIndex) const;

  size_type getLocalID(size_type globalID) const;

  size_type
  getLocalID(const Teuchos::ArrayView< dim_type > & localIndex) const;



  bool isCompatible(const MDMap & mdMap) const;

  bool isSameAs(const MDMap & mdMap,
                const int verbose = 0) const;

  bool isContiguous() const;


private:

  // A private method for computing the bounds and local dimensions,
  // after the global dimensions, communication and boundary padding
  // have been properly assigned.
  void computeBounds();

  // The underlying multi-dimensional communicator.
  Teuchos::RCP< const MDComm > _mdComm;

  // The size of the global dimensions along each axis.  This includes
  // the values of the boundary padding.
  Teuchos::Array< dim_type > _globalDims;

  // Store the start and stop indexes of the global problem along each
  // axis.  This includes the values of the boundary padding.
  Teuchos::Array< Slice > _globalBounds;

  // A double array of Slices that stores the starting and stopping
  // indexes for each axis processor along each axis.  This data
  // structure allows any processor to know the map bounds on any
  // other processor.  These bounds do NOT include the boundary or
  // communication padding.  The outer array will have a size equal to
  // the number of dimensions.  Each inner array will have a size
  // equal to the number of axis processors along that axis.
  Teuchos::Array< Teuchos::Array< Slice > > _globalRankBounds;

  // The global stride between adjacent IDs.  This quantity is
  // "virtual" as it does not describe actual memory layout, but it is
  // useful for index conversion algorithms.
  Teuchos::Array< size_type > _globalStrides;

  // The minumum global ID of the global data structure, including
  // boundary padding.  This will only be non-zero on a sub-map.
  size_type _globalMin;

  // The maximum global ID of the global data structure, including
  // boundary padding.
  size_type _globalMax;

  // The size of the local dimensions along each axis.  This includes
  // the values of the padding.
  Teuchos::Array< dim_type > _localDims;

  // The local loop bounds along each axis, stored as an array of
  // Slices.  These bounds DO include the padding.  By definition, the
  // start() attribute of these Slices will all be zero by definition.
  // This may seem wasteful (we could store an array of dim_type
  // rather than an array of Slice), but this storage is convenient
  // when it comes to the getLocalBounds() method.
  Teuchos::Array< Slice > _localBounds;

  // The local stride between adjacent elements in memory.
  Teuchos::Array< size_type > _localStrides;

  // The minimum local ID of the local data structure, including
  // padding.
  size_type _localMin;

  // The maximum local ID of the local data structure, including
  // padding.
  size_type _localMax;

  // The communication padding that was specified at construction, one
  // value along each axis.
  Teuchos::Array< int > _commPadSizes;

  // The actual padding stored on this processor along each axis.  The
  // padding can be either communication padding or boundary padding
  // based on the processor position on the boundary of a domain.
  Teuchos::Array< padding_type > _pad;

  // The size of the boundary padding along each axis.
  Teuchos::Array< int > _bndryPadSizes;

  // The actual boundary padding stored along each axis.  For a full
  // communicator, the lower and upper boundary padding are both the
  // same as the corresponding value in _bndryPadSizes.  However, the
  // introduction of sub-maps creates the possibility of different
  // upper and lower boundary padding values.  If an axis is periodic,
  // then these values will be set to the communication padding.
  Teuchos::Array< padding_type > _bndryPad;

  // An array of flags, one for each axis, that specifies whether the
  // endpoints of a periodic boundary are replicated (1 or true) or
  // unique (0 or false). These flags only have meaning for axes that
  // are periodic. The arrray type is int, beacause Array< bool > is
  // specialized and sometimes difficult to work with.
  Teuchos::Array< int > _replicatedBoundary;

  // The storage order
  Layout _layout;

  // An array of axis maps that correspond to the full MDMap.  An axis
  // map describes the map along a single axis. This member is mutable
  // because it is logically const but does not get constructed until
  // requested by the user.
  mutable Teuchos::Array< Teuchos::RCP< const MDMap > > _axisMaps;

#ifdef HAVE_EPETRA
  // An RCP pointing to an Epetra_Map that is equivalent to this
  // MDMap, including communication padding.  It is mutable because we
  // do not construct it until it asked for by a get method that is
  // const.
  mutable Teuchos::RCP< const Epetra_Map > _epetraMap;

  // An RCP pointing to an Epetra_Map that is equivalent to this
  // MDMap, excluding communication padding.  The returned Epetra_Map
  // thus indicates what IDs are owned by this processor.  It is
  // mutable because we do not construct it until it asked for by a
  // get method that is const.
  mutable Teuchos::RCP< const Epetra_Map > _epetraOwnMap;

  // An ArrayRCP that stores Epetra_Maps that represent the
  // distribution of indexes along each axis, including communication
  // padding.  It is mutable because we do not construct it until it
  // asked for by a get method that is const.
  mutable Teuchos::Array< Teuchos::RCP< const Epetra_Map > > _epetraAxisMaps;

  // An ArrayRCP that stores Epetra_Maps that represent the
  // distribution of indexes along each axis, excluding communication
  // padding.  It is mutable because we do not construct it until it
  // asked for by a get method that is const.
  mutable Teuchos::Array< Teuchos::RCP< const Epetra_Map > > _epetraAxisOwnMaps;
#endif

};

// Inline implementations //


Teuchos::RCP< const Teuchos::Comm< int > >
MDMap::getTeuchosComm() const
{
  return _mdComm->getTeuchosComm();
}


Teuchos::RCP< const MDComm >
MDMap::getMDComm() const
{
  return _mdComm;
}


bool
MDMap::onSubcommunicator() const
{
  return _mdComm->onSubcommunicator();
}


int
MDMap::numDims() const
{
  return _mdComm->numDims();
}


Teuchos::Array< int >
MDMap::getCommDims() const
{
  return _mdComm->getCommDims();
}


int
MDMap::getCommDim(int axis) const
{
  return _mdComm->getCommDim(axis);
}


bool
MDMap::isPeriodic(int axis) const
{
  return _mdComm->isPeriodic(axis);
}


int
MDMap::getCommIndex(int axis) const
{
  return _mdComm->getCommIndex(axis);
}


int
MDMap::getLowerNeighbor(int axis) const
{
  return _mdComm->getLowerNeighbor(axis);
}


int
MDMap::getUpperNeighbor(int axis) const
{
  return _mdComm->getUpperNeighbor(axis);
}


Teuchos::Array< dim_type >
MDMap::
getGlobalDims() const
{
  return _globalDims;
}


// Templated method implementations //


#ifdef HAVE_TPETRA

template< class LocalOrdinal,
          class GlobalOrdinal,
          class Node >
Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > >
MDMap::getTpetraMap(bool withCommPad) const
{
  if (withCommPad)
  {
    // Allocate the elementsMDArray and the index array
    int num_dims = numDims();
    Teuchos::Array<dim_type> localDims(num_dims);
    for (int axis = 0; axis < num_dims; ++axis)
      localDims[axis] = _localDims[axis];
    MDArray< GlobalOrdinal > elementMDArray(localDims);
    Teuchos::Array< LocalOrdinal > index(num_dims);

    // Iterate over the local MDArray and assign global IDs
    for (typename MDArray< GlobalOrdinal >::iterator it = elementMDArray.begin();
         it != elementMDArray.end(); ++it)
    {
      GlobalOrdinal globalID = 0;
      for (int axis = 0; axis < num_dims; ++axis)
      {
        int axisRank        = getCommIndex(axis);
        GlobalOrdinal start = _globalRankBounds[axis][axisRank].start() -
                              _pad[axis][0];
        globalID += (start + it.index(axis)) * _globalStrides[axis];
      }
      *it = globalID;
    }

    // Return the Tpetra::Map
    const Teuchos::Array< GlobalOrdinal > & myElements =
      elementMDArray.array();
    Teuchos::RCP< const Teuchos::Comm< int > > teuchosComm =
      _mdComm->getTeuchosComm();
    return
      Teuchos::rcp(new Tpetra::Map< LocalOrdinal,
                                    GlobalOrdinal,
                                    Node >(Teuchos::OrdinalTraits< Tpetra::global_size_t >::invalid(),
                                           myElements(),
                                           0,
                                           teuchosComm));
  }
  else
  {
    // Allocate the elementMDArray MDArray and the index array
    int num_dims = numDims();
    Teuchos::Array< LocalOrdinal > index(num_dims);
    Teuchos::Array< dim_type >     myDims(num_dims);
    for (int axis = 0; axis < num_dims; ++axis)
    {
      myDims[axis] =
        _localDims[axis] - _pad[axis][0] - _pad[axis][1];
      int axisRank = getCommIndex(axis);
      if (axisRank == 0)
        myDims[axis] += _bndryPad[axis][0];
      if (axisRank == getCommDim(axis)-1)
        myDims[axis] += _bndryPad[axis][1];
    }
    MDArray< GlobalOrdinal > elementMDArray(myDims());

    // Iterate over the local MDArray and assign global IDs
    for (typename MDArray< GlobalOrdinal >::iterator it = elementMDArray.begin();
         it != elementMDArray.end(); ++it)
    {
      GlobalOrdinal globalID = 0;
      for (int axis = 0; axis < num_dims; ++axis)
      {
        int axisRank        = getCommIndex(axis);
        GlobalOrdinal start = _globalRankBounds[axis][axisRank].start();
        if (axisRank == 0)
          start -= _bndryPad[axis][0];
        if (axisRank == getCommDim(axis)-1)
          start += _bndryPad[axis][1];
        globalID += (start + it.index(axis)) * _globalStrides[axis];
      }
    }

    // Return the Tpetra::Map
    const Teuchos::Array< GlobalOrdinal> & myElements =
      elementMDArray.array();
    Teuchos::RCP< const Teuchos::Comm< int > > teuchosComm =
      _mdComm->getTeuchosComm();
    return
      Teuchos::rcp(new Tpetra::Map< LocalOrdinal,
                                    GlobalOrdinal,
                                    Node >(Teuchos::OrdinalTraits< Tpetra::global_size_t>::invalid(),
                                           myElements(),
                                           0,
                                           teuchosComm));
  }
}
#endif


#ifdef HAVE_TPETRA

template< class LocalOrdinal,
          class GlobalOrdinal,
          class Node >
Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > >
MDMap::
getTpetraAxisMap(int axis,
                 bool withCommPad) const
{
#ifdef HAVE_DOMI_ARRAY_BOUNDSCHECK
  TEUCHOS_TEST_FOR_EXCEPTION(
    ((axis < 0) || (axis >= numDims())),
    RangeError,
    "invalid axis index = " << axis << " (number of dimensions = " <<
    numDims() << ")");
#endif
  Teuchos::RCP< const Teuchos::Comm< int > > teuchosComm =
    _mdComm->getTeuchosComm();
  Teuchos::Array< GlobalOrdinal > elements(getLocalDim(axis,withCommPad));
  GlobalOrdinal start = getGlobalRankBounds(axis,true).start();
  if (withCommPad && (getCommIndex(axis) != 0)) start -= _pad[axis][0];
  for (LocalOrdinal i = 0; i < elements.size(); ++i)
    elements[i] = i + start;
  return Teuchos::rcp(new Tpetra::Map< LocalOrdinal,
                                       GlobalOrdinal,
                                       Node >(Teuchos::OrdinalTraits< Tpetra::global_size_t>::invalid(),
                                              elements,
                                              0,
                                              teuchosComm));
}
#endif


}

#endif