Rythmos_AdjointModelEvaluator.hpp

//@HEADER
// ***********************************************************************
//
//                           Rythmos Package
//                 Copyright (2006) Sandia Corporation
//
// Under terms of Contract DE-AC04-94AL85000, there is a non-exclusive
// license for use of this work by or on behalf of the U.S. Government.
//
// This library is free software; you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as
// published by the Free Software Foundation; either version 2.1 of the
// License, or (at your option) any later version.
//
// This library is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// Lesser General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License along with this library; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
// USA
// Questions? Contact Todd S. Coffey (tscoffe@sandia.gov)
//
// ***********************************************************************
//@HEADER

#ifndef RYTHMOS_ADJOINT_MODEL_EVALUATOR_HPP
#define RYTHMOS_ADJOINT_MODEL_EVALUATOR_HPP


#include "Rythmos_IntegratorBase.hpp"
#include "Thyra_ModelEvaluator.hpp" // Interface
#include "Thyra_StateFuncModelEvaluatorBase.hpp" // Implementation
#include "Thyra_DefaultScaledAdjointLinearOp.hpp"
#include "Thyra_DefaultAdjointLinearOpWithSolve.hpp"
#include "Teuchos_implicit_cast.hpp"
#include "Teuchos_Assert.hpp"


namespace Rythmos {


template<class Scalar>
class AdjointModelEvaluator
  : virtual public Thyra::StateFuncModelEvaluatorBase<Scalar>
{
public:


  AdjointModelEvaluator();

  void setFwdStateModel(
    const RCP<const Thyra::ModelEvaluator<Scalar> > &fwdStateModel );

  void setFwdTimeRange( const TimeRange<Scalar> &fwdTimeRange );

  void setFwdStateSolutionBuffer(
    const RCP<const InterpolationBufferBase<Scalar> > &fwdStateSolutionBuffer );
  


  RCP<const Thyra::VectorSpaceBase<Scalar> > get_x_space() const;
  RCP<const Thyra::VectorSpaceBase<Scalar> > get_f_space() const;
  Thyra::ModelEvaluatorBase::InArgs<Scalar> getNominalValues() const;
  RCP<Thyra::LinearOpWithSolveBase<Scalar> > create_W() const;
  RCP<Thyra::LinearOpBase<Scalar> > create_W_op() const;
  Thyra::ModelEvaluatorBase::InArgs<Scalar> createInArgs() const;


private:


  Thyra::ModelEvaluatorBase::OutArgs<Scalar> createOutArgsImpl() const;
  void evalModelImpl(
    const Thyra::ModelEvaluatorBase::InArgs<Scalar> &inArgs_bar,
    const Thyra::ModelEvaluatorBase::OutArgs<Scalar> &outArgs_bar
    ) const;


private:

  // /////////////////////////
  // Private data members

  RCP<const Thyra::ModelEvaluator<Scalar> > fwdStateModel_;
  TimeRange<Scalar> fwdTimeRange_;
  RCP<const InterpolationBufferBase<Scalar> > fwdStateSolutionBuffer_;

  mutable Thyra::ModelEvaluatorBase::InArgs<Scalar> prototypeInArgs_bar_;
  mutable Thyra::ModelEvaluatorBase::OutArgs<Scalar> prototypeOutArgs_bar_;
  mutable Thyra::ModelEvaluatorBase::InArgs<Scalar> adjointNominalValues_;
  mutable RCP<Thyra::LinearOpBase<Scalar> > my_W_bar_adj_op_;
  mutable RCP<Thyra::LinearOpBase<Scalar> > my_d_f_d_x_dot_op_;

  // /////////////////////////
  // Private member functions

  // Just-in-time initialization function
  void initialize() const;

};


template<class Scalar>
RCP<AdjointModelEvaluator<Scalar> >
adjointModelEvaluator(
  const RCP<const Thyra::ModelEvaluator<Scalar> > &fwdStateModel,
  const TimeRange<Scalar> &fwdTimeRange
  )
{
  RCP<AdjointModelEvaluator<Scalar> >
    adjointModel = Teuchos::rcp(new AdjointModelEvaluator<Scalar>);
  adjointModel->setFwdStateModel(fwdStateModel);
  adjointModel->setFwdTimeRange(fwdTimeRange);
  return adjointModel;
}


// /////////////////////////////////
// Implementations


// Constructors/Intializers/Accessors


template<class Scalar>
AdjointModelEvaluator<Scalar>::AdjointModelEvaluator()
{}


template<class Scalar>
void AdjointModelEvaluator<Scalar>::setFwdStateModel(
  const RCP<const Thyra::ModelEvaluator<Scalar> > &fwdStateModel )
{
  TEST_FOR_EXCEPT(is_null(fwdStateModel));
  fwdStateModel_ = fwdStateModel;
}


template<class Scalar>
void AdjointModelEvaluator<Scalar>::setFwdTimeRange(
  const TimeRange<Scalar> &fwdTimeRange )
{
  fwdTimeRange_ = fwdTimeRange;
}


template<class Scalar>
void AdjointModelEvaluator<Scalar>::setFwdStateSolutionBuffer(
  const RCP<const InterpolationBufferBase<Scalar> > &fwdStateSolutionBuffer )
{
  TEST_FOR_EXCEPT(is_null(fwdStateSolutionBuffer));
  fwdStateSolutionBuffer_ = fwdStateSolutionBuffer;
}


// Public functions overridden from ModelEvaulator


template<class Scalar>
RCP<const Thyra::VectorSpaceBase<Scalar> >
AdjointModelEvaluator<Scalar>::get_x_space() const
{
  initialize();
  return fwdStateModel_->get_f_space();
}


template<class Scalar>
RCP<const Thyra::VectorSpaceBase<Scalar> >
AdjointModelEvaluator<Scalar>::get_f_space() const
{
  initialize();
  return fwdStateModel_->get_x_space();
}


template<class Scalar>
Thyra::ModelEvaluatorBase::InArgs<Scalar>
AdjointModelEvaluator<Scalar>::getNominalValues() const
{
  initialize();
  return adjointNominalValues_;
}


template<class Scalar>
RCP<Thyra::LinearOpWithSolveBase<Scalar> >
AdjointModelEvaluator<Scalar>::create_W() const
{
  initialize();
  return Thyra::nonconstAdjointLows<Scalar>(fwdStateModel_->create_W());
}


template<class Scalar>
RCP<Thyra::LinearOpBase<Scalar> >
AdjointModelEvaluator<Scalar>::create_W_op() const
{
  initialize();
  return Thyra::nonconstAdjoint<Scalar>(fwdStateModel_->create_W_op());
}


template<class Scalar>
Thyra::ModelEvaluatorBase::InArgs<Scalar>
AdjointModelEvaluator<Scalar>::createInArgs() const
{
  initialize();
  return prototypeInArgs_bar_;
}


// Private functions overridden from ModelEvaulatorDefaultBase


template<class Scalar>
Thyra::ModelEvaluatorBase::OutArgs<Scalar>
AdjointModelEvaluator<Scalar>::createOutArgsImpl() const
{
  initialize();
  return prototypeOutArgs_bar_;
}


template<class Scalar>
void AdjointModelEvaluator<Scalar>::evalModelImpl(
  const Thyra::ModelEvaluatorBase::InArgs<Scalar> &inArgs_bar,
  const Thyra::ModelEvaluatorBase::OutArgs<Scalar> &outArgs_bar
  ) const
{

  using Teuchos::rcp_dynamic_cast;
  using Teuchos::describe;
  typedef Teuchos::ScalarTraits<Scalar> ST;
  typedef Thyra::ModelEvaluatorBase MEB;
  typedef Thyra::DefaultScaledAdjointLinearOp<Scalar> DSALO;
  typedef Thyra::DefaultAdjointLinearOpWithSolve<Scalar> DALOWS;
  typedef Teuchos::VerboseObjectTempState<Thyra::ModelEvaluatorBase> VOTSME;

  //
  // A) Header stuff
  //

  THYRA_MODEL_EVALUATOR_DECORATOR_EVAL_MODEL_GEN_BEGIN(
    "AdjointModelEvaluator", inArgs_bar, outArgs_bar, Teuchos::null );

  initialize();

  VOTSME fwdStateModel_outputTempState(fwdStateModel_,out,verbLevel);

  //const bool trace = includesVerbLevel(verbLevel, Teuchos::VERB_LOW);
  const bool dumpAll = includesVerbLevel(localVerbLevel, Teuchos::VERB_EXTREME);

  //
  // B) Unpack the input and output arguments to see what we have to compute
  //

  // B.1) InArgs

  const RCP<const Thyra::VectorBase<Scalar> >
    lambda_rev_dot = inArgs_bar.get_x_dot().assert_not_null(), // x_bar_dot
    lambda = inArgs_bar.get_x().assert_not_null(); // x_bar
  const Scalar alpha_bar = inArgs_bar.get_alpha();
  const Scalar beta_bar = inArgs_bar.get_beta();
  // const Scalar t_bar = inArgs_bar.get_t(); // Don't need yet!

  if (dumpAll) {
    *out << "\nlambda_rev_dot = " << describe(*lambda_rev_dot, Teuchos::VERB_EXTREME);
    *out << "\nlambda = " << describe(*lambda, Teuchos::VERB_EXTREME);
    *out << "\nalpha_bar = " << alpha_bar << "\n";
    *out << "\nbeta_bar = " << beta_bar << "\n";
  }

  // B.2) OutArgs

  const RCP<Thyra::VectorBase<Scalar> > f_bar = outArgs_bar.get_f();

  RCP<DALOWS> W_bar;
  if (outArgs_bar.supports(MEB::OUT_ARG_W))
    W_bar = rcp_dynamic_cast<DALOWS>(outArgs_bar.get_W(), true);

  RCP<DSALO> W_bar_op;
  if (outArgs_bar.supports(MEB::OUT_ARG_W_op))
    W_bar_op = rcp_dynamic_cast<DSALO>(outArgs_bar.get_W_op(), true);
  
  //
  // C) Evaluate the needed quantities from the underlying forward Model
  //

  MEB::InArgs<Scalar> fwdInArgs = fwdStateModel_->createInArgs();

  // C.1) Set the required input arguments

  if (!is_null(fwdStateSolutionBuffer_)) {
    TEST_FOR_EXCEPT_MSG(true, "ToDo: Implement getting the x and x_dot from IB");
  }
  else {
    // If we don't have an IB object to get the state from, we will assume
    // that the problem is linear and, therefore, we can pass in any old value
    // of x, x_dot, and t and get the W_bar_adj object that we need
    fwdInArgs = fwdStateModel_->getNominalValues();
    // 2008/05/14: rabartl: ToDo: Implement real variable dependancy
    // communication support to make sure that this is okay!
  }

  // C.2) Evaluate W_bar_adj if needed

  RCP<Thyra::LinearOpWithSolveBase<Scalar> > W_bar_adj;
  RCP<Thyra::LinearOpBase<Scalar> > W_bar_adj_op;
  {

    MEB::OutArgs<Scalar> fwdOutArgs = fwdStateModel_->createOutArgs();
    
    // Get or create W_bar_adj or W_bar_adj_op if needed
    if (!is_null(W_bar)) {
      // If we have W_bar, the W_bar_adj was already created in
      // this->create_W()
      W_bar_adj = W_bar->getNonconstOp();
      W_bar_adj_op = W_bar_adj;
    }
    else if (!is_null(W_bar_op)) {
      TEST_FOR_EXCEPT_MSG(true, "ToDo: Unit test this code!");
      // If we have W_bar_op, the W_bar_adj_op was already created in
      // this->create_W_op()
      W_bar_adj_op = W_bar_op->getNonconstOp();
    }
    else if (!is_null(f_bar)) {
      TEST_FOR_EXCEPT_MSG(true, "ToDo: Unit test this code!");
      // If the user did not pass in W_bar or W_bar_op, then we need to create
      // our own local LOB form W_bar_adj_op of W_bar_adj in order to evaluate
      // the residual f_bar
      if (is_null(my_W_bar_adj_op_)) {
        my_W_bar_adj_op_ = fwdStateModel_->create_W_op();
      }
      W_bar_adj_op = my_W_bar_adj_op_;
    }
    
    // Set W_bar_adj or W_bar_adj_op on the OutArgs object
    if (!is_null(W_bar_adj)) {
      fwdOutArgs.set_W(W_bar_adj);
    }
    else if (!is_null(W_bar_adj_op)) {
      fwdOutArgs.set_W_op(W_bar_adj);
    }
    
    // Set alpha and beta on OutArgs object
    if (!is_null(W_bar_adj) || !is_null(W_bar_adj_op)) {
      fwdInArgs.set_alpha(alpha_bar);
      fwdInArgs.set_beta(beta_bar);
    }
    
    // Evaluate the model
    if (!is_null(W_bar_adj) || !is_null(W_bar_adj_op)) {
      fwdStateModel_->evalModel( fwdInArgs, fwdOutArgs );
    }
    
    // Print the objects if requested
    if (!is_null(W_bar_adj) && dumpAll)
      *out << "\nW_bar_adj = " << describe(*W_bar_adj, Teuchos::VERB_EXTREME);
    if (!is_null(W_bar_adj_op) && dumpAll)
      *out << "\nW_bar_adj_op = " << describe(*W_bar_adj_op, Teuchos::VERB_EXTREME);

  }
  
  // C.3) Evaluate d(f)/d(x_dot) if needed

  RCP<Thyra::LinearOpBase<Scalar> > d_f_d_x_dot_op;
  if (!is_null(f_bar)) {
    if (is_null(my_d_f_d_x_dot_op_)) {
      my_d_f_d_x_dot_op_ = fwdStateModel_->create_W_op();
    }
    d_f_d_x_dot_op = my_d_f_d_x_dot_op_;
    MEB::OutArgs<Scalar> fwdOutArgs = fwdStateModel_->createOutArgs();
    fwdOutArgs.set_W_op(d_f_d_x_dot_op);
    fwdInArgs.set_alpha(ST::one());
    fwdInArgs.set_beta(ST::zero());
    fwdStateModel_->evalModel( fwdInArgs, fwdOutArgs );
  }

  //
  // D) Evaluate the adjoint equation residual:
  //
  //   f_bar = d(f)/d(x_dot)^T * lambda_hat + 1/beta_bar * W_bar_adj^T * lambda
  //           - d(g)/d(x)^T
  //

  if (!is_null(f_bar)) {

    // D.1) lambda_hat = lambda_rev_dot - alpha_bar/beta_bar * lambda
    const RCP<Thyra::VectorBase<Scalar> >
      lambda_hat = createMember(lambda_rev_dot->space());
    Thyra::V_VpStV<Scalar>( outArg(*lambda_hat),
      *lambda_rev_dot, -alpha_bar/beta_bar, *lambda );
    if (dumpAll)
      *out << "\nlambda_hat = " << describe(*lambda_hat, Teuchos::VERB_EXTREME);

    // D.2) f_bar = d(f)/d(x_dot)^T * lambda_hat
    Thyra::apply<Scalar>( *d_f_d_x_dot_op, Thyra::CONJTRANS, *lambda_hat,
      outArg(*f_bar) );

    // D.3) f_bar += 1/beta_bar * W_bar_adj^T * lambda
    Thyra::apply<Scalar>( *W_bar_adj_op, Thyra::CONJTRANS, *lambda,
      outArg(*f_bar), 1.0/beta_bar, ST::one() );

    // D.4) f_bar += - d(g)/d(x)^T
    // 2008/05/15: rabart: ToDo: Implement once we add support for
    // distributed response functions

    if (dumpAll)
      *out << "\nf_bar = " << describe(*f_bar, Teuchos::VERB_EXTREME);

  }


  //
  // E) Do any remaining post processing
  //

  THYRA_MODEL_EVALUATOR_DECORATOR_EVAL_MODEL_END();

}


// private


template<class Scalar>
void AdjointModelEvaluator<Scalar>::initialize() const
{

  typedef Thyra::ModelEvaluatorBase MEB;

  //
  // A) Validate the that forward Model is of the correct form!
  //

  MEB::InArgs<Scalar> fwdStateModelInArgs = fwdStateModel_->createInArgs();
  MEB::OutArgs<Scalar> fwdStateModelOutArgs = fwdStateModel_->createOutArgs();

#ifdef TEUCHOS_DEBUG
  TEUCHOS_ASSERT( fwdStateModelInArgs.supports(MEB::IN_ARG_x_dot) );
  TEUCHOS_ASSERT( fwdStateModelInArgs.supports(MEB::IN_ARG_x) );
  TEUCHOS_ASSERT( fwdStateModelInArgs.supports(MEB::IN_ARG_t) );
  TEUCHOS_ASSERT( fwdStateModelInArgs.supports(MEB::IN_ARG_alpha) );
  TEUCHOS_ASSERT( fwdStateModelInArgs.supports(MEB::IN_ARG_beta) );
  TEUCHOS_ASSERT( fwdStateModelOutArgs.supports(MEB::OUT_ARG_f) );
  TEUCHOS_ASSERT( fwdStateModelOutArgs.supports(MEB::OUT_ARG_W) );
#endif

  //
  // B) Set up the prototypical InArgs and OutArgs
  //

  {
    MEB::InArgsSetup<Scalar> inArgs_bar;
    inArgs_bar.setModelEvalDescription(this->description());
    inArgs_bar.setSupports( MEB::IN_ARG_x_dot );
    inArgs_bar.setSupports( MEB::IN_ARG_x );
    inArgs_bar.setSupports( MEB::IN_ARG_t );
    inArgs_bar.setSupports( MEB::IN_ARG_alpha );
    inArgs_bar.setSupports( MEB::IN_ARG_beta );
    prototypeInArgs_bar_ = inArgs_bar;
  }

  {
    MEB::OutArgsSetup<Scalar> outArgs_bar;
    outArgs_bar.setModelEvalDescription(this->description());
    outArgs_bar.setSupports(MEB::OUT_ARG_f);
    if (fwdStateModelOutArgs.supports(MEB::OUT_ARG_W) ) {
      outArgs_bar.setSupports(MEB::OUT_ARG_W);
      outArgs_bar.set_W_properties(fwdStateModelOutArgs.get_W_properties());
    }
    if (fwdStateModelOutArgs.supports(MEB::OUT_ARG_W_op) ) {
      outArgs_bar.setSupports(MEB::OUT_ARG_W_op);
      outArgs_bar.set_W_properties(fwdStateModelOutArgs.get_W_properties());
    }
    prototypeOutArgs_bar_ = outArgs_bar;
  }

  //
  // D) Set up the nominal values for the adjoint
  //

  // Copy structure
  adjointNominalValues_ = prototypeInArgs_bar_;
  // Just set a zero initial condition for the adjoint
  const RCP<Thyra::VectorBase<Scalar> > zero_lambda_vec =
    createMember(fwdStateModel_->get_f_space());
  V_S( zero_lambda_vec.ptr(), ScalarTraits<Scalar>::zero() );
  adjointNominalValues_.set_x_dot(zero_lambda_vec);
  adjointNominalValues_.set_x(zero_lambda_vec);

  //
  // E) Wipe out other cached objects
  //

  my_W_bar_adj_op_ = Teuchos::null;
  my_d_f_d_x_dot_op_ = Teuchos::null;

}


} // namespace Rythmos


#endif // RYTHMOS_ADJOINT_MODEL_EVALUATOR_HPP