Teko_TpetraHelpers.cpp

// @HEADER
// *****************************************************************************
//      Teko: A package for block and physics based preconditioning
//
// Copyright 2010 NTESS and the Teko contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
 
#include "Teko_TpetraHelpers.hpp"
#include "Teko_ConfigDefs.hpp"
 
// Thyra Includes
#include "Thyra_BlockedLinearOpBase.hpp"
#include "Thyra_DefaultMultipliedLinearOp.hpp"
#include "Thyra_DefaultDiagonalLinearOp.hpp"
#include "Thyra_DefaultZeroLinearOp.hpp"
#include "Thyra_DefaultBlockedLinearOp.hpp"
 
#include "Thyra_SpmdVectorBase.hpp"
#include "Thyra_SpmdVectorSpaceBase.hpp"
#include "Thyra_ScalarProdVectorSpaceBase.hpp"
 
// Teko includes
#include "Teko_Utilities.hpp"
 
// Tpetra
#include "Thyra_TpetraLinearOp.hpp"
#include "Thyra_TpetraMultiVector.hpp"
#include "Tpetra_CrsMatrix.hpp"
#include "Tpetra_Vector.hpp"
#include "Thyra_TpetraThyraWrappers.hpp"
#include "TpetraExt_MatrixMatrix.hpp"
#include "Tpetra_RowMatrixTransposer.hpp"
 
using Teuchos::null;
using Teuchos::RCP;
using Teuchos::rcp;
using Teuchos::rcp_dynamic_cast;
using Teuchos::rcpFromRef;
 
namespace Teko {
namespace TpetraHelpers {
 
const Teuchos::RCP<const Thyra::LinearOpBase<ST> > thyraDiagOp(
    const RCP<const Tpetra::Vector<ST, LO, GO, NT> >& tv, const Tpetra::Map<LO, GO, NT>& map,
    const std::string& lbl) {
  const RCP<const Thyra::VectorBase<ST> > thyraVec  // need a Thyra::VectorBase object
      = Thyra::createConstVector<ST, LO, GO, NT>(
          tv, Thyra::createVectorSpace<ST, LO, GO, NT>(rcpFromRef(map)));
  Teuchos::RCP<Thyra::LinearOpBase<ST> > op =
      Teuchos::rcp(new Thyra::DefaultDiagonalLinearOp<ST>(thyraVec));
  op->setObjectLabel(lbl);
  return op;
}
 
const Teuchos::RCP<Thyra::LinearOpBase<ST> > thyraDiagOp(
    const RCP<Tpetra::Vector<ST, LO, GO, NT> >& tv, const Tpetra::Map<LO, GO, NT>& map,
    const std::string& lbl) {
  const RCP<Thyra::VectorBase<ST> > thyraVec  // need a Thyra::VectorBase object
      = Thyra::createVector<ST, LO, GO, NT>(
          tv, Thyra::createVectorSpace<ST, LO, GO, NT>(rcpFromRef(map)));
  Teuchos::RCP<Thyra::LinearOpBase<ST> > op =
      Teuchos::rcp(new Thyra::DefaultDiagonalLinearOp<ST>(thyraVec));
  op->setObjectLabel(lbl);
  return op;
}
 
void fillDefaultSpmdMultiVector(Teuchos::RCP<Thyra::TpetraMultiVector<ST, LO, GO, NT> >& spmdMV,
                                Teuchos::RCP<Tpetra::MultiVector<ST, LO, GO, NT> >& tpetraMV) {
  // first get desired range and domain
  // const RCP<const Thyra::SpmdVectorSpaceBase<ST> > range  = spmdMV->spmdSpace();
  const RCP<Thyra::TpetraVectorSpace<ST, LO, GO, NT> > range =
      Thyra::tpetraVectorSpace<ST, LO, GO, NT>(tpetraMV->getMap());
  const RCP<const Thyra::ScalarProdVectorSpaceBase<ST> > domain =
      rcp_dynamic_cast<const Thyra::ScalarProdVectorSpaceBase<ST> >(spmdMV->domain());
 
  TEUCHOS_ASSERT((size_t)domain->dim() == tpetraMV->getNumVectors());
 
  // New local view of raw data
  if (!tpetraMV->isConstantStride())
    TEUCHOS_TEST_FOR_EXCEPT(true);  // ToDo: Implement views of non-contiguous mult-vectors!
 
  // Build the MultiVector
  spmdMV->initialize(range, domain, tpetraMV);
 
  // make sure the Tpetra::MultiVector doesn't disappear prematurely
  Teuchos::set_extra_data<RCP<Tpetra::MultiVector<ST, LO, GO, NT> > >(
      tpetraMV, "Tpetra::MultiVector", Teuchos::outArg(spmdMV));
}
 
void identityRowIndices(const Tpetra::Map<LO, GO, NT>& rowMap,
                        const Tpetra::CrsMatrix<ST, LO, GO, NT>& mat, std::vector<GO>& outIndices) {
  // loop over elements owned by this processor
  for (size_t i = 0; i < rowMap.getLocalNumElements(); i++) {
    bool rowIsIdentity = true;
    GO rowGID          = rowMap.getGlobalElement(i);
 
    size_t numEntries = mat.getNumEntriesInGlobalRow(i);
    auto indices = typename Tpetra::CrsMatrix<ST, LO, GO, NT>::nonconst_global_inds_host_view_type(
        Kokkos::ViewAllocateWithoutInitializing("rowIndices"), numEntries);
    auto values = typename Tpetra::CrsMatrix<ST, LO, GO, NT>::nonconst_values_host_view_type(
        Kokkos::ViewAllocateWithoutInitializing("rowIndices"), numEntries);
 
    mat.getGlobalRowCopy(rowGID, indices, values, numEntries);
 
    // loop over the columns of this row
    for (size_t j = 0; j < numEntries; j++) {
      GO colGID = indices(j);
 
      // look at row entries
      if (colGID == rowGID)
        rowIsIdentity &= values(j) == 1.0;
      else
        rowIsIdentity &= values(j) == 0.0;
 
      // not a dirchlet row...quit
      if (not rowIsIdentity) break;
    }
 
    // save a row that is dirchlet
    if (rowIsIdentity) outIndices.push_back(rowGID);
  }
}
 
void zeroMultiVectorRowIndices(Tpetra::MultiVector<ST, LO, GO, NT>& mv,
                               const std::vector<GO>& zeroIndices) {
  LO colCnt = mv.getNumVectors();
  std::vector<GO>::const_iterator itr;
 
  // loop over the indices to zero
  for (itr = zeroIndices.begin(); itr != zeroIndices.end(); ++itr) {
    // loop over columns
    for (int j = 0; j < colCnt; j++) mv.replaceGlobalValue(*itr, j, 0.0);
  }
}
 
ZeroedOperator::ZeroedOperator(const std::vector<GO>& zeroIndices,
                               const Teuchos::RCP<const Tpetra::Operator<ST, LO, GO, NT> >& op)
    : zeroIndices_(zeroIndices), tpetraOp_(op) {}
 
void ZeroedOperator::apply(const Tpetra::MultiVector<ST, LO, GO, NT>& X,
                           Tpetra::MultiVector<ST, LO, GO, NT>& Y, Teuchos::ETransp mode, ST alpha,
                           ST beta) const {
  /*
     Tpetra::MultiVector temp(X);
     zeroMultiVectorRowIndices(temp,zeroIndices_);
     int result = TpetraOp_->Apply(temp,Y);
  */
 
  tpetraOp_->apply(X, Y, mode, alpha, beta);
 
  // zero a few of the rows
  zeroMultiVectorRowIndices(Y, zeroIndices_);
}
 
bool isTpetraLinearOp(const LinearOp& op) {
  // See if the operator is a TpetraLinearOp
  RCP<const Thyra::TpetraLinearOp<ST, LO, GO, NT> > tOp =
      rcp_dynamic_cast<const Thyra::TpetraLinearOp<ST, LO, GO, NT> >(op);
  if (!tOp.is_null()) return true;
 
  // See if the operator is a wrapped TpetraLinearOp
  ST scalar               = 0.0;
  Thyra::EOpTransp transp = Thyra::NOTRANS;
  RCP<const Thyra::LinearOpBase<ST> > wrapped_op;
  Thyra::unwrap(op, &scalar, &transp, &wrapped_op);
  tOp = rcp_dynamic_cast<const Thyra::TpetraLinearOp<ST, LO, GO, NT> >(wrapped_op);
  if (!tOp.is_null()) return true;
 
  return false;
}
 
RCP<const Tpetra::CrsMatrix<ST, LO, GO, NT> > getTpetraCrsMatrix(const LinearOp& op, ST* scalar,
                                                                 bool* transp) {
  // If the operator is a TpetraLinearOp
  RCP<const Thyra::TpetraLinearOp<ST, LO, GO, NT> > tOp =
      rcp_dynamic_cast<const Thyra::TpetraLinearOp<ST, LO, GO, NT> >(op);
  if (!tOp.is_null()) {
    RCP<const Tpetra::CrsMatrix<ST, LO, GO, NT> > matrix =
        rcp_dynamic_cast<const Tpetra::CrsMatrix<ST, LO, GO, NT> >(tOp->getConstTpetraOperator(),
                                                                   true);
    *scalar = 1.0;
    *transp = false;
    return matrix;
  }
 
  // If the operator is a wrapped TpetraLinearOp
  RCP<const Thyra::LinearOpBase<ST> > wrapped_op;
  Thyra::EOpTransp eTransp = Thyra::NOTRANS;
  Thyra::unwrap(op, scalar, &eTransp, &wrapped_op);
  tOp = rcp_dynamic_cast<const Thyra::TpetraLinearOp<ST, LO, GO, NT> >(wrapped_op, true);
  if (!tOp.is_null()) {
    RCP<const Tpetra::CrsMatrix<ST, LO, GO, NT> > matrix =
        rcp_dynamic_cast<const Tpetra::CrsMatrix<ST, LO, GO, NT> >(tOp->getConstTpetraOperator(),
                                                                   true);
    *transp = true;
    if (eTransp == Thyra::NOTRANS) *transp = false;
    return matrix;
  }
 
  return Teuchos::null;
}
 
RCP<Tpetra::CrsMatrix<ST, LO, GO, NT> > materializeTpetraCrsMatrix(const LinearOp& op) {
  ST scalar          = 0.0;
  bool transp        = false;
  auto tCrsOp        = getTpetraCrsMatrix(op, &scalar, &transp);
  auto explicitCrsOp = transp
                           ? Tpetra::RowMatrixTransposer<ST, LO, GO, NT>(tCrsOp).createTranspose()
                           : rcp(new Tpetra::CrsMatrix<ST, LO, GO, NT>(*tCrsOp, Teuchos::Copy));
  if (scalar != Teuchos::ScalarTraits<ST>::one()) explicitCrsOp->scale(scalar);
  return explicitCrsOp;
}
 
ModifiableLinearOp materializeTpetraLinearOp(const LinearOp& op, const ModifiableLinearOp& destOp) {
  auto explicitCrsOp = materializeTpetraCrsMatrix(op);
 
  RCP<Thyra::LinearOpBase<ST> > explicitOp;
  if (destOp != Teuchos::null)
    explicitOp = destOp;
  else
    explicitOp = rcp(new Thyra::TpetraLinearOp<ST, LO, GO, NT>());
 
  auto tExplicitOp = rcp_dynamic_cast<Thyra::TpetraLinearOp<ST, LO, GO, NT> >(explicitOp, true);
  tExplicitOp->initialize(Thyra::tpetraVectorSpace<ST, LO, GO, NT>(explicitCrsOp->getRangeMap()),
                          Thyra::tpetraVectorSpace<ST, LO, GO, NT>(explicitCrsOp->getDomainMap()),
                          explicitCrsOp);
  return tExplicitOp;
}
 
}  // end namespace TpetraHelpers
}  // end namespace Teko