Teko_BlockedReordering.cpp

/*
// @HEADER
// 
// ***********************************************************************
// 
//      Teko: A package for block and physics based preconditioning
//                  Copyright 2010 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 Eric C. Cyr (eccyr@sandia.gov)
// 
// ***********************************************************************
// 
// @HEADER

*/

#include <iostream>
#include <string>
#include <vector>
#include <stack>

#include "Teko_BlockedReordering.hpp"
#include "Teko_Utilities.hpp"

#include "Teuchos_VerboseObject.hpp"
#include "Teuchos_RCP.hpp"
#include "Teuchos_StrUtils.hpp"

#include "Thyra_DefaultProductMultiVector.hpp"
#include "Thyra_DefaultProductVectorSpace.hpp"

using Teuchos::RCP;
using Teuchos::rcp;
using Teuchos::rcp_dynamic_cast;
using Teuchos::Array;

namespace Teko {

void BlockReorderManager::SetBlock(int blockIndex,int reorder)
{ 
   TEUCHOS_ASSERT(blockIndex<(int) children_.size());

   RCP<BlockReorderManager> child = rcp(new BlockReorderLeaf(reorder));

   children_[blockIndex] = child;
}

void BlockReorderManager::SetBlock(int blockIndex,const RCP<BlockReorderManager> & reorder)
{
   TEUCHOS_ASSERT(blockIndex<(int) children_.size());

   children_[blockIndex] = reorder;
}

const Teuchos::RCP<BlockReorderManager> BlockReorderManager::GetBlock(int blockIndex)
{
   TEUCHOS_ASSERT(blockIndex<(int) children_.size());

   if(children_[blockIndex]==Teuchos::null)
      children_[blockIndex] = rcp(new BlockReorderManager());

   return children_[blockIndex];
}

const Teuchos::RCP<const BlockReorderManager> BlockReorderManager::GetBlock(int blockIndex) const
{
   TEUCHOS_ASSERT(blockIndex<(int) children_.size());

   return children_[blockIndex];
}

std::string BlockReorderManager::toString() const
{
   // build the string by recursively calling each child
   std::stringstream ss;
   ss << "[";
   for(unsigned int i=0;i<children_.size();i++) {
      if(children_[i]==Teuchos::null) 
         ss << " <NULL> ";
      else
         ss << " " << children_[i]->toString() << " "; 
   }
   ss << "]";

   return ss.str();
}

int BlockReorderManager::LargestIndex() const
{
   int max = 0;
   for(unsigned int i=0;i<children_.size();i++) {
      // see if current child is larger 
      if(children_[i]!=Teuchos::null) {
         int subMax = children_[i]->LargestIndex();
         max = max > subMax ? max : subMax;
      }
   }

   return max;
} 

Teuchos::RCP<const Thyra::LinearOpBase<double> >
buildReorderedLinearOp(const BlockReorderManager & bmm,
                       const Teuchos::RCP<const Thyra::BlockedLinearOpBase<double> > & blkOp)
{
   return buildReorderedLinearOp(bmm,bmm,blkOp);
}

Teuchos::RCP<const Thyra::LinearOpBase<double> >
buildReorderedLinearOp(const BlockReorderManager & rowMgr,const BlockReorderManager & colMgr,
                       const Teuchos::RCP<const Thyra::BlockedLinearOpBase<double> > & blkOp)
{
   typedef RCP<const BlockReorderManager> BRMptr;

   int rowSz = rowMgr.GetNumBlocks();
   int colSz = colMgr.GetNumBlocks();

   if(rowSz==0 && colSz==0) {
      // both are leaf nodes
      const BlockReorderLeaf & rowLeaf = dynamic_cast<const BlockReorderLeaf &>(rowMgr);
      const BlockReorderLeaf & colLeaf = dynamic_cast<const BlockReorderLeaf &>(colMgr);

      // simply return entry in matrix
      Teko::LinearOp linOp = blkOp->getBlock(rowLeaf.GetIndex(),colLeaf.GetIndex());

      // somehow we need to set this operator up
      if(linOp==Teuchos::null) {
         linOp = Thyra::zero(blkOp->productRange()->getBlock(rowLeaf.GetIndex()),
                             blkOp->productDomain()->getBlock(colLeaf.GetIndex()));
      }

      return linOp;
   }
   else if(rowSz==0) {
      Teko::BlockedLinearOp reBlkOp = Teko::createBlockedOp();

      // operator will be rowSz by colSz
      reBlkOp->beginBlockFill(1,colSz);   

      // fill the column entries
      for(int col=0;col<colSz;col++) {
         BRMptr colPtr = colMgr.GetBlock(col);

         reBlkOp->setBlock(0,col,buildReorderedLinearOp(rowMgr,*colPtr,blkOp));
      } 

      // done building
      reBlkOp->endBlockFill();   

      return reBlkOp;
   }
   else if(colSz==0) {
      Teko::BlockedLinearOp reBlkOp = Teko::createBlockedOp();

      // operator will be rowSz by colSz
      reBlkOp->beginBlockFill(rowSz,1);   

      // fill the row entries
      for(int row=0;row<rowSz;row++) {
         BRMptr rowPtr = rowMgr.GetBlock(row);

         reBlkOp->setBlock(row,0,buildReorderedLinearOp(*rowPtr,colMgr,blkOp));
      } 

      // done building
      reBlkOp->endBlockFill();   

      return reBlkOp;
   }
   else {
      Teko::BlockedLinearOp reBlkOp = Teko::createBlockedOp();
  
      // this is the general case
      TEUCHOS_ASSERT(rowSz>0);
      TEUCHOS_ASSERT(colSz>0);

      // operator will be rowSz by colSz
      reBlkOp->beginBlockFill(rowSz,colSz);   
   
      for(int row=0;row<rowSz;row++) {
         BRMptr rowPtr = rowMgr.GetBlock(row);
   
         for(int col=0;col<colSz;col++) {
            BRMptr colPtr = colMgr.GetBlock(col);
   
            reBlkOp->setBlock(row,col,buildReorderedLinearOp(*rowPtr,*colPtr,blkOp));
         } 
      }

      // done building
      reBlkOp->endBlockFill();   

      return reBlkOp;
   }
}

Teuchos::RCP<const Thyra::VectorSpaceBase<double> >
buildReorderedVectorSpace(const BlockReorderManager & mgr,
                          const Teuchos::RCP<const Thyra::ProductVectorSpaceBase<double> > & blkSpc)
{
   typedef RCP<const BlockReorderManager> BRMptr;

   int sz = mgr.GetNumBlocks();

   if(sz==0) {
      // its a  leaf nodes
      const BlockReorderLeaf & leaf = dynamic_cast<const BlockReorderLeaf &>(mgr);

      // simply return entry in matrix
      return blkSpc->getBlock(leaf.GetIndex());
   } 
   else {
      Array<RCP<const Thyra::VectorSpaceBase<double> > > vecSpaces;

      // loop over each row
      for(int i=0;i<sz;i++) {
         BRMptr blkMgr = mgr.GetBlock(i);

         const RCP<const Thyra::VectorSpaceBase<double> > lvs = buildReorderedVectorSpace(*blkMgr,blkSpc);

         vecSpaces.push_back(lvs);
      }

      // build a vector space
      const RCP<const Thyra::DefaultProductVectorSpace<double> > vs 
            = Thyra::productVectorSpace<double>(vecSpaces);

      // build the vector
      return vs;
   }
}

Teuchos::RCP<Thyra::MultiVectorBase<double> >
buildReorderedMultiVector(const BlockReorderManager & mgr,
                          const Teuchos::RCP<Thyra::ProductMultiVectorBase<double> > & blkVec)
{
   using Teuchos::rcp_const_cast;

   // give vector const so that the right function is called
   const Teuchos::RCP<const Thyra::ProductMultiVectorBase<double> > blkVecConst
      = rcp_const_cast<const Thyra::ProductMultiVectorBase<double> >(blkVec);
   
   // its not really const, so take it away
   const Teuchos::RCP<Thyra::MultiVectorBase<double> > result
         = rcp_const_cast<Thyra::MultiVectorBase<double> >(buildReorderedMultiVector(mgr,blkVecConst));

   return result; 
}

Teuchos::RCP<const Thyra::MultiVectorBase<double> >
buildReorderedMultiVector(const BlockReorderManager & mgr,
                          const Teuchos::RCP<const Thyra::ProductMultiVectorBase<double> > & blkVec)
{
   typedef RCP<const BlockReorderManager> BRMptr;

   int sz = mgr.GetNumBlocks();

   if(sz==0) {
      // its a  leaf nodes
      const BlockReorderLeaf & leaf = dynamic_cast<const BlockReorderLeaf &>(mgr);

      // simply return entry in matrix
      return blkVec->getMultiVectorBlock(leaf.GetIndex());
   } 
   else {
      Array<RCP<const Thyra::MultiVectorBase<double> > > multiVecs;
      Array<RCP<const Thyra::VectorSpaceBase<double> > > vecSpaces;

      // loop over each row
      for(int i=0;i<sz;i++) {
         BRMptr blkMgr = mgr.GetBlock(i);

         const RCP<const Thyra::MultiVectorBase<double> > lmv = buildReorderedMultiVector(*blkMgr,blkVec);
         const RCP<const Thyra::VectorSpaceBase<double> > lvs = lmv->range();

         multiVecs.push_back(lmv);
         vecSpaces.push_back(lvs);
      }

      // build a vector space
      const RCP<const Thyra::DefaultProductVectorSpace<double> > vs 
            = Thyra::productVectorSpace<double>(vecSpaces);

      // build the vector
      return Thyra::defaultProductMultiVector<double>(vs,multiVecs);
   }
}

void buildNonconstFlatMultiVector(const BlockReorderManager & mgr,
                                  const RCP<Thyra::MultiVectorBase<double> > & blkVec,
                                  Array<RCP<Thyra::MultiVectorBase<double> > > & multivecs,
                                  Array<RCP<const Thyra::VectorSpaceBase<double> > > & vecspaces)
{
   typedef RCP<const BlockReorderManager> BRMptr;

   int sz = mgr.GetNumBlocks();

   if(sz==0) {
      // its a  leaf nodes
      const BlockReorderLeaf & leaf = dynamic_cast<const BlockReorderLeaf &>(mgr);
      int index = leaf.GetIndex();

      // simply return entry in matrix
      multivecs[index] = blkVec;
      vecspaces[index] = blkVec->range();
   } 
   else {
      const RCP<Thyra::ProductMultiVectorBase<double> > prodMV
            = rcp_dynamic_cast<Thyra::ProductMultiVectorBase<double> >(blkVec);
 
      // get flattened elements from each child
      for(int i=0;i<sz;i++) {
         const RCP<Thyra::MultiVectorBase<double> > mv = prodMV->getNonconstMultiVectorBlock(i);
         buildNonconstFlatMultiVector(*mgr.GetBlock(i),mv,multivecs,vecspaces);
      }
   }
   
}

void buildFlatMultiVector(const BlockReorderManager & mgr,
                          const RCP<const Thyra::MultiVectorBase<double> > & blkVec,
                          Array<RCP<const Thyra::MultiVectorBase<double> > > & multivecs,
                          Array<RCP<const Thyra::VectorSpaceBase<double> > > & vecspaces)
{
   typedef RCP<const BlockReorderManager> BRMptr;

   int sz = mgr.GetNumBlocks();

   if(sz==0) {
      // its a  leaf nodes
      const BlockReorderLeaf & leaf = dynamic_cast<const BlockReorderLeaf &>(mgr);
      int index = leaf.GetIndex();

      // simply return entry in matrix
      multivecs[index] = blkVec;
      vecspaces[index] = blkVec->range();
   } 
   else {
      const RCP<const Thyra::ProductMultiVectorBase<double> > prodMV
            = rcp_dynamic_cast<const Thyra::ProductMultiVectorBase<double> >(blkVec);
 
      // get flattened elements from each child
      for(int i=0;i<sz;i++) {
         RCP<const Thyra::MultiVectorBase<double> > mv = prodMV->getMultiVectorBlock(i);
         buildFlatMultiVector(*mgr.GetBlock(i),mv,multivecs,vecspaces);
      }
   }
   
}

Teuchos::RCP<Thyra::MultiVectorBase<double> >
buildFlatMultiVector(const BlockReorderManager & mgr,
                     const Teuchos::RCP<Thyra::ProductMultiVectorBase<double> > & blkVec)
{
   int numBlocks = mgr.LargestIndex()+1;
 
   Array<RCP<Thyra::MultiVectorBase<double> > > multivecs(numBlocks);
   Array<RCP<const Thyra::VectorSpaceBase<double> > > vecspaces(numBlocks);

   // flatten everything into a vector first
   buildNonconstFlatMultiVector(mgr,blkVec,multivecs,vecspaces);

   // build a vector space
   const RCP<Thyra::DefaultProductVectorSpace<double> > vs 
         = Thyra::productVectorSpace<double>(vecspaces);

   // build the vector
   return Thyra::defaultProductMultiVector<double>(vs,multivecs);
}

Teuchos::RCP<const Thyra::MultiVectorBase<double> >
buildFlatMultiVector(const BlockReorderManager & mgr,
                     const Teuchos::RCP<const Thyra::ProductMultiVectorBase<double> > & blkVec)
{
   int numBlocks = mgr.LargestIndex()+1;
 
   Array<RCP<const Thyra::MultiVectorBase<double> > > multivecs(numBlocks);
   Array<RCP<const Thyra::VectorSpaceBase<double> > > vecspaces(numBlocks);

   // flatten everything into a vector first
   buildFlatMultiVector(mgr,blkVec,multivecs,vecspaces);

   // build a vector space
   const RCP<const Thyra::DefaultProductVectorSpace<double> > vs 
         = Thyra::productVectorSpace<double>(vecspaces);

   // build the vector
   return Thyra::defaultProductMultiVector<double>(vs,multivecs);
}

// The next three functions are useful for parsing the string
// description of a BlockReorderManager.

// this function tokenizes a string, breaking out whitespace but saving the
// brackets [,] as special tokens.
static void tokenize(std::string srcInput,std::string whitespace,std::string prefer,
              std::vector<std::string> & tokens)
{
   std::string input = srcInput;
   std::vector<std::string> wsTokens;
   std::size_t endPos   = input.length()-1;
   while(endPos<input.length()) {
      std::size_t next = input.find_first_of(whitespace);


      // get the sub string
      std::string s;
      if(next!=std::string::npos) {
         s = input.substr(0,next);

         // break out the old substring
         input = input.substr(next+1,endPos);
      }
      else {
         s = input;
         input = "";
      }

      endPos   = input.length()-1;

      // add it to the WS tokens list
      if(s=="") continue;
      wsTokens.push_back(s);
   }

   for(unsigned int i=0;i<wsTokens.size();i++) {
      // get string to break up
      input = wsTokens[i];

      std::size_t endPos   = input.length()-1;
      while(endPos<input.length()) {
         std::size_t next = input.find_first_of(prefer);

         std::string s = input;
         if(next>0 && next<input.length()) {

            // get the sub string
            s = input.substr(0,next);
 
            input = input.substr(next,endPos);
         }
         else if(next==0) {
            // get the sub string
            s = input.substr(0,next+1);
 
            input = input.substr(next+1,endPos);
         }
         else input = "";

         // break out the old substring
         endPos   = input.length()-1;

         // add it to the tokens list
         tokens.push_back(s);
      }
   }
}

// this function takes a set of tokens and returns the first "block", i.e. those
// values (including) brackets that correspond to the first block
static std::vector<std::string>::const_iterator 
buildSubBlock(std::vector<std::string>::const_iterator begin,
              std::vector<std::string>::const_iterator end, std::vector<std::string> & subBlock)
{
   std::stack<std::string> matched;
   std::vector<std::string>::const_iterator itr;
   for(itr=begin;itr!=end;++itr) { 

      subBlock.push_back(*itr);

      // push/pop brackets as they are discovered
      if(*itr=="[") matched.push("["); 
      else if(*itr=="]") matched.pop(); 

      // found all matching brackets
      if(matched.empty()) 
         return itr;
   }

   TEUCHOS_ASSERT(matched.empty());

   return itr-1;
}

// This function takes a tokenized vector and converts it to a block reorder manager
static RCP<BlockReorderManager> blockedReorderFromTokens(const std::vector<std::string> & tokens)
{
   // base case
   if(tokens.size()==1)
      return rcp(new Teko::BlockReorderLeaf(Teuchos::StrUtils::atoi(tokens[0])));

   // check first and last character
   TEUCHOS_ASSERT(*(tokens.begin())=="[")
   TEUCHOS_ASSERT(*(tokens.end()-1)=="]");

   std::vector<RCP<Teko::BlockReorderManager> > vecRMgr;
   std::vector<std::string>::const_iterator itr = tokens.begin()+1; 
   while(itr!=tokens.end()-1) {
      // figure out which tokens are relevant for this block
      std::vector<std::string> subBlock;
      itr = buildSubBlock(itr,tokens.end()-1,subBlock); 

      // build the child block reorder manager
      vecRMgr.push_back(blockedReorderFromTokens(subBlock));

      // move the iterator one more
      itr++;
   }

   // build the parent reorder manager
   RCP<Teko::BlockReorderManager> rMgr = rcp(new Teko::BlockReorderManager(vecRMgr.size()));
   for(unsigned int i=0;i<vecRMgr.size();i++)
      rMgr->SetBlock(i,vecRMgr[i]);

   return rMgr;
}


Teuchos::RCP<const BlockReorderManager> blockedReorderFromString(std::string & reorder)
{
   // vector of tokens to use
   std::vector<std::string> tokens;
   
   // manager to be returned

   // build tokens vector
   tokenize(reorder," \t\n","[]",tokens);

   // parse recursively and build reorder manager
   Teuchos::RCP<BlockReorderManager> mgr = blockedReorderFromTokens(tokens);

   return mgr;
}

} // end namespace Teko