Stokhos: Propagate Belos change to PseudoBlockCGIter

Signed-off-by: Christian Glusa <[email protected]>
trilinos · Jan 28, 2025 · 23198c5 · 23198c5
1 parent 0a95f12
commit 23198c5
Showing 1 changed file with 29 additions and 37 deletions.
diff --git a/packages/stokhos/src/sacado/kokkos/vector/belos/Belos_PseudoBlockCGIter_MP_Vector.hpp b/packages/stokhos/src/sacado/kokkos/vector/belos/Belos_PseudoBlockCGIter_MP_Vector.hpp
@@ -110,15 +110,14 @@ namespace Belos {
      * \note For any pointer in \c newstate which directly points to the multivectors in
      * the solver, the data is not copied.
      */
-    void initializeCG(CGIterationState<ScalarType,MV>& newstate);
+    void initializeCG(Teuchos::RCP<CGIterationStateBase<ScalarType,MV> > newstate, Teuchos::RCP<MV> R_0);
 
     /*! \brief Initialize the solver with the initial vectors from the linear problem
      *  or random data.
      */
     void initialize()
     {
-      CGIterationState<ScalarType,MV> empty;
-      initializeCG(empty);
+      initializeCG(Teuchos::null, Teuchos::null);
     }
 
     /*! \brief Get the current state of the linear solver.
@@ -128,15 +127,23 @@ namespace Belos {
      * \returns A CGIterationState object containing const pointers to the current
      * solver state.
      */
-    CGIterationState<ScalarType,MV> getState() const {
-      CGIterationState<ScalarType,MV> state;
-      state.R = R_;
-      state.P = P_;
-      state.AP = AP_;
-      state.Z = Z_;
+    Teuchos::RCP<CGIterationStateBase<ScalarType,MV> > getState() const {
+      auto state = Teuchos::rcp(new PseudoBlockCGIterationState<ScalarType,MV>());
+      state->R = R_;
+      state->P = P_;
+      state->AP = AP_;
+      state->Z = Z_;
       return state;
     }
 
+    void setState(Teuchos::RCP<CGIterationStateBase<ScalarType,MV> > state) {
+      auto s = Teuchos::rcp_dynamic_cast<PseudoBlockCGIterationState<ScalarType,MV> >(state, true);
+      R_ = s->R;
+      Z_ = s->Z;
+      P_ = s->P;
+      AP_ = s->AP;
+    }
+
     //@}
 
 
@@ -290,8 +297,8 @@ namespace Belos {
   // Initialize this iteration object
   template <class StorageType, class MV, class OP>
   void PseudoBlockCGIter<Sacado::MP::Vector<StorageType>,MV,OP>::
-  initializeCG(CGIterationState<ScalarType,MV>& newstate)
-  {
+  initializeCG(Teuchos::RCP<CGIterationStateBase<ScalarType,MV> >  newstate, Teuchos::RCP<MV> R_0) {
+
     // Check if there is any mltivector to clone from.
     Teuchos::RCP<const MV> lhsMV = lp_->getCurrLHSVec();
     Teuchos::RCP<const MV> rhsMV = lp_->getCurrRHSVec();
@@ -305,40 +312,30 @@ namespace Belos {
     int numRHS = MVT::GetNumberVecs(*tmp);
     numRHS_ = numRHS;
 
-    // Initialize the state storage
-    // If the subspace has not be initialized before or has changed sizes, generate it using the LHS or RHS from lp_.
-    if (Teuchos::is_null(R_) || MVT::GetNumberVecs(*R_)!=numRHS_) {
-      R_ = MVT::Clone( *tmp, numRHS_ );
-      Z_ = MVT::Clone( *tmp, numRHS_ );
-      P_ = MVT::Clone( *tmp, numRHS_ );
-      AP_ = MVT::Clone( *tmp, numRHS_ );
-    }
+    // Initialize the state storage if it isn't already.
+    if (!Teuchos::rcp_dynamic_cast<PseudoBlockCGIterationState<ScalarType,MV> >(newstate, true)->matches(tmp, numRHS_))
+      newstate->initialize(tmp, numRHS_);
+    setState(newstate);
 
     // Tracking information for condition number estimation
     if(numEntriesForCondEst_ > 0) {
       diag_.resize(numEntriesForCondEst_);
       offdiag_.resize(numEntriesForCondEst_-1);
     }
 
-    // NOTE:  In CGIter R_, the initial residual, is required!!!
-    //
     std::string errstr("Belos::BlockPseudoCGIter::initialize(): Specified multivectors must have a consistent length and width.");
 
-    // Create convenience variables for zero and one.
-    const ScalarType one = Teuchos::ScalarTraits<ScalarType>::one();
-    const MagnitudeType zero = Teuchos::ScalarTraits<MagnitudeType>::zero();
-
-    if (!Teuchos::is_null(newstate.R)) {
+    {
 
-      TEUCHOS_TEST_FOR_EXCEPTION( MVT::GetGlobalLength(*newstate.R) != MVT::GetGlobalLength(*R_),
+      TEUCHOS_TEST_FOR_EXCEPTION( MVT::GetGlobalLength(*R_0) != MVT::GetGlobalLength(*R_),
                           std::invalid_argument, errstr );
-      TEUCHOS_TEST_FOR_EXCEPTION( MVT::GetNumberVecs(*newstate.R) != numRHS_,
+      TEUCHOS_TEST_FOR_EXCEPTION( MVT::GetNumberVecs(*R_0) != numRHS_,
                           std::invalid_argument, errstr );
 
       // Copy basis vectors from newstate into V
-      if (newstate.R != R_) {
+      if (R_0 != R_) {
         // copy over the initial residual (unpreconditioned).
-        MVT::MvAddMv( one, *newstate.R, zero, *newstate.R, *R_ );
+        MVT::Assign( *R_0, *R_ );
       }
 
       // Compute initial direction vectors
@@ -356,14 +353,9 @@ namespace Belos {
         lp_->applyRightPrec( *R_, *Z_ );
       }
       else {
-        Z_ = R_;
+        MVT::Assign( *R_, *Z_ );
       }
-      MVT::MvAddMv( one, *Z_, zero, *Z_, *P_ );
-    }
-    else {
-
-      TEUCHOS_TEST_FOR_EXCEPTION(Teuchos::is_null(newstate.R),std::invalid_argument,
-                         "Belos::CGIter::initialize(): CGStateIterState does not have initial residual.");
+      MVT::Assign( *Z_, *P_ );
     }
 
     // The solver is initialized