Amesos2_Solver_UQ_PCE.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//                           Stokhos Package
//
// Copyright 2009 NTESS and the Stokhos contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef AMESOS2_SOLVER_UQ_PCE_HPP
#define AMESOS2_SOLVER_UQ_PCE_HPP

#include "Amesos2_Solver.hpp"
#include "Amesos2_Factory.hpp"
#include "Stokhos_Sacado_Kokkos_UQ_PCE.hpp"
#include "Stokhos_Tpetra_Utilities_UQ_PCE.hpp"

namespace Amesos2 {

  template <class S, class LO, class GO, class NO>
  typename Sacado::UQ::PCE<S>::cijk_type
  get_pce_cijk(
    const Teuchos::RCP<const Tpetra::CrsMatrix<Sacado::UQ::PCE<S>, LO, GO, NO > >& A = Teuchos::null,
    const Teuchos::RCP<Tpetra::MultiVector<Sacado::UQ::PCE<S>, LO, GO, NO > >& X = Teuchos::null,
    const Teuchos::RCP<const Tpetra::MultiVector<Sacado::UQ::PCE<S>, LO, GO, NO > >& B = Teuchos::null)
  {
    if (A != Teuchos::null) {
      return Kokkos::cijk(A->getLocalValuesDevice(Tpetra::Access::ReadOnly));
    }
    else if (X != Teuchos::null) {
      return Kokkos::cijk(X->getLocalViewDevice(Tpetra::Access::ReadOnly));
    }
    else if (B != Teuchos::null) {
      return Kokkos::cijk(B->getLocalViewDevice(Tpetra::Access::ReadOnly));
    }
    return typename Sacado::UQ::PCE<S>::cijk_type();
  }

  template <class Storage, class LocalOrdinal, class GlobalOrdinal,
            class Node, template<class,class> class ConcreteSolver>
  class PCESolverAdapter :
    public Solver< Tpetra::CrsMatrix<Sacado::UQ::PCE<Storage>,
                                     LocalOrdinal,
                                     GlobalOrdinal,
                                     Node >,
                   Tpetra::MultiVector<Sacado::UQ::PCE<Storage>,
                                       LocalOrdinal,
                                       GlobalOrdinal,
                                       Node >
                   >
  {
  public:

    typedef Sacado::UQ::PCE<Storage> Scalar;
    typedef Tpetra::CrsMatrix<Scalar,LocalOrdinal,GlobalOrdinal,Node> Matrix;
    typedef Tpetra::MultiVector<Scalar,LocalOrdinal,GlobalOrdinal,Node> Vector;

    typedef typename Scalar::value_type BaseScalar;
    typedef Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> Map;
    typedef Tpetra::CrsGraph<LocalOrdinal,GlobalOrdinal,Node> FlatGraph;
    typedef Tpetra::CrsMatrix<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatMatrix;
    typedef Tpetra::MultiVector<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatVector;
    typedef ConcreteSolver<FlatMatrix,FlatVector> FlatConcreteSolver;
    typedef Solver<FlatMatrix,FlatVector> FlatSolver;

    typedef Solver<Matrix,Vector> solver_type;
    typedef typename solver_type::type type;
    typedef typename Scalar::cijk_type cijk_type;

    PCESolverAdapter(
      const Teuchos::RCP<const Matrix>& A_,
      const Teuchos::RCP<Vector>& X_,
      const Teuchos::RCP<const Vector>& B_) :
      A(A_), X(X_), B(B_) {
      cijk = get_pce_cijk(A, X, B);
      const size_t pce_size =
        Kokkos::dimension_scalar(A->getLocalMatrixDevice().values);
      flat_graph =
        Stokhos::create_flat_pce_graph(*(A->getCrsGraph()),
                                       cijk,
                                       flat_X_map,
                                       flat_B_map,
                                       cijk_graph,
                                       pce_size);
      if (A != Teuchos::null)
        flat_A = Stokhos::create_flat_matrix(*A, flat_graph, cijk_graph, cijk);
      if (X != Teuchos::null)
        flat_X = Stokhos::create_flat_vector_view(*X, flat_X_map);
      if (B != Teuchos::null)
        flat_B = Stokhos::create_flat_vector_view(*B, flat_B_map);
      flat_solver =
        create_solver_with_supported_type<ConcreteSolver,FlatMatrix,FlatVector>::apply(flat_A, flat_X, flat_B);
    }



    virtual type& preOrdering( void ) {
      flat_solver->preOrdering();
      return *this;
    }


    virtual type& symbolicFactorization( void ) {
      flat_solver->symbolicFactorization();
      return *this;
    }


    virtual type& numericFactorization( void ) {
      flat_solver->numericFactorization();
      return *this;
    }


    virtual void solve( void ) {
      flat_solver->solve();
    }


    virtual void solve(const Teuchos::Ptr<Vector>       XX,
                       const Teuchos::Ptr<const Vector> BB) const {
      flat_solver->solve(
        Stokhos::create_flat_vector_view(*XX, flat_X_map).get(),
        Stokhos::create_flat_vector_view(*BB, flat_B_map).get() );
    }


    virtual void solve(Vector* XX, const Vector* BB) const {
      flat_solver->solve(
        Stokhos::create_flat_vector_view(*XX, flat_X_map).get(),
        Stokhos::create_flat_vector_view(*BB, flat_B_map).get() );
    }



    virtual type& setParameters(
      const Teuchos::RCP<Teuchos::ParameterList> & parameterList ) {
      flat_solver->setParameters(parameterList);
      return *this;
    }


    virtual Teuchos::RCP<const Teuchos::ParameterList>
    getValidParameters( void ) const {
      return flat_solver->getValidParameters();
    }



    virtual void setA( const Teuchos::RCP<const Matrix> a,
                       EPhase keep_phase = CLEAN ) {
      A = a;

      // Rebuild flat matrix/graph
      cijk = get_pce_cijk(A);
      if (keep_phase <= CLEAN) {
        flat_X_map = Teuchos::null;
        flat_B_map = Teuchos::null;
        flat_graph = Teuchos::null;
        flat_graph =
          Stokhos::create_flat_pce_graph(*(A->getCrsGraph()),
                                         cijk,
                                         flat_X_map,
                                         flat_B_map,
                                         cijk_graph,
                                         Kokkos::dimension_scalar(A->getLocalMatrixDevice().values));
      }
      if (keep_phase <= SYMBFACT) // should this by NUMFACT???
        flat_A = Stokhos::create_flat_matrix(*a, flat_graph, cijk_graph, cijk);

      flat_solver->setA(flat_A, keep_phase);
    }

    virtual void setA( const Matrix* a, EPhase keep_phase = CLEAN ) {
      this->setA(Teuchos::rcp(a,false), keep_phase);
    }


    virtual bool matrixShapeOK( void ) {
      return flat_solver->matrixShapeOK();
    }


    virtual void setX( const Teuchos::RCP<Vector> x ) {
      X = x;
      if (x != Teuchos::null)
        flat_X = Stokhos::create_flat_vector_view(*x, flat_X_map);
      else
        flat_X = Teuchos::null;
      flat_solver->setX(flat_X);
    }


    virtual void setX( Vector* x ) {
      if (x != 0) {
        X = Teuchos::rcp(x, false);
        flat_X = Stokhos::create_flat_vector_view(*x, flat_X_map);
      }
      else {
        X = Teuchos::null;
        flat_X = Teuchos::null;
      }
      flat_solver->setX(flat_X);
    }


    virtual const Teuchos::RCP<Vector> getX( void ) {
      return X;
    }


    virtual Vector* getXRaw( void ) {
      return X.get();
    }


    virtual void setB( const Teuchos::RCP<const Vector> b ) {
      B = b;
      if (b != Teuchos::null)
        flat_B = Stokhos::create_flat_vector_view(*b, flat_B_map);
      else
        flat_B = Teuchos::null;
      flat_solver->setB(flat_B);
    }


    virtual void setB( const Vector* b ) {
      if (b != 0) {
        B = Teuchos::rcp(b, false);
        flat_B = Stokhos::create_flat_vector_view(*b, flat_B_map);
      }
      else {
        B = Teuchos::null;
        flat_B = Teuchos::null;
      }
      flat_solver->setB(flat_B);
    }


    virtual const Teuchos::RCP<const Vector> getB( void ) {
      return B;
    }


    virtual const Vector* getBRaw( void ) {
      return B.get();
    }


    virtual Teuchos::RCP<const Teuchos::Comm<int> > getComm( void ) const {
      return flat_solver->getComm();
    }


    virtual Status& getStatus() const {
      return flat_solver->getStatus();
    }


    virtual std::string name( void ) const {
      return flat_solver->name();
    }



    virtual std::string description( void ) const {
      return flat_solver->description();
    }


    virtual void describe( Teuchos::FancyOStream &out,
                           const Teuchos::EVerbosityLevel verbLevel=Teuchos::Describable::verbLevel_default ) const {
      flat_solver->describe(out, verbLevel);
    }



    virtual void printTiming( Teuchos::FancyOStream &out,
                              const Teuchos::EVerbosityLevel verbLevel = Teuchos::Describable::verbLevel_default ) const{
      flat_solver->printTiming(out, verbLevel);
    }


    virtual void getTiming( Teuchos::ParameterList& timingParameterList ) const{
      flat_solver->getTiming(timingParameterList);
    }


  protected:

    Teuchos::RCP<const Matrix> A;
    Teuchos::RCP<Vector> X;
    Teuchos::RCP<const Vector> B;
    Teuchos::RCP<const Map> flat_X_map, flat_B_map;
    Teuchos::RCP<const FlatGraph> flat_graph, cijk_graph;
    Teuchos::RCP<const FlatMatrix> flat_A;
    Teuchos::RCP<FlatVector> flat_X;
    Teuchos::RCP<const FlatVector> flat_B;
    Teuchos::RCP<FlatSolver> flat_solver;
    cijk_type cijk;

  };

  // Specialization of create_solver_with_supported_type for
  // Sacado::UQ::PCE where we create PCESolverAdapter wrapping
  // each solver
  template < template <class,class> class ConcreteSolver,
             class ST, class LO, class GO, class NO >
  struct create_solver_with_supported_type<
    ConcreteSolver,
    Tpetra::CrsMatrix<Sacado::UQ::PCE<ST>,LO,GO,NO >,
    Tpetra::MultiVector<Sacado::UQ::PCE<ST>,LO,GO,NO > > {
    typedef Sacado::UQ::PCE<ST> SC;
    typedef Tpetra::CrsMatrix<SC,LO,GO,NO> Matrix;
    typedef Tpetra::MultiVector<SC,LO,GO,NO> Vector;
    static Teuchos::RCP<Solver<Matrix,Vector> >
    apply(Teuchos::RCP<const Matrix> A,
          Teuchos::RCP<Vector>       X,
          Teuchos::RCP<const Vector> B )
    {
      ctassert<
        std::is_same_v<
          typename MatrixTraits<Matrix>::scalar_t,
          typename MultiVecAdapter<Vector>::scalar_t
        >
      > same_scalar_assertion;
      (void)same_scalar_assertion; // This stops the compiler from warning about unused declared variables

      // If our assertion did not fail, then create and return a new solver
      return Teuchos::rcp( new PCESolverAdapter<ST,LO,GO,NO,ConcreteSolver>(A, X, B) );
    }
  };

  // Specialization for solver_supports_scalar for Sacado::UQ::PCE<Storage>
  // value == true if and only if
  // solver_supprts_scalar<ConcreteSolver,Storage::value_type> == true
  template <template <class,class> class ConcreteSolver,
            typename Storage>
  struct solver_supports_scalar<ConcreteSolver, Sacado::UQ::PCE<Storage> > {
    typedef Sacado::UQ::PCE<Storage> Scalar;
    typedef typename Scalar::value_type BaseScalar;
    typedef typename solver_traits<ConcreteSolver>::supported_scalars supported_scalars;
    static const bool value =
      std::conditional_t<std::is_same_v<supported_scalars, Meta::nil_t>,
                         std::true_type,
                         Meta::type_list_contains<supported_scalars,
                                                  BaseScalar> >::value;
  };


} // namespace Amesos2

#endif // AMESOS2_SOLVER_UQ_PCE_HPP