tapeHelper.hpp

/*
 * CoDiPack, a Code Differentiation Package
 *
 * Copyright (C) 2015-2024 Chair for Scientific Computing (SciComp), University of Kaiserslautern-Landau
 * Homepage: http://www.scicomp.uni-kl.de
 * Contact:  Prof. Nicolas R. Gauger (codi@scicomp.uni-kl.de)
 *
 * Lead developers: Max Sagebaum, Johannes Blühdorn (SciComp, University of Kaiserslautern-Landau)
 *
 * This file is part of CoDiPack (http://www.scicomp.uni-kl.de/software/codi).
 *
 * CoDiPack is free software: you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * CoDiPack 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 General Public License for more details.
 * You should have received a copy of the GNU
 * General Public License along with CoDiPack.
 * If not, see <http://www.gnu.org/licenses/>.
 *
 * For other licensing options please contact us.
 *
 * Authors:
 *  - SciComp, University of Kaiserslautern-Landau:
 *    - Max Sagebaum
 *    - Johannes Blühdorn
 *    - Former members:
 *      - Tim Albring
 */
#pragma once
 
#include <vector>
 
#include "../../config.h"
#include "../../expressions/lhsExpressionInterface.hpp"
#include "../../traits/realTraits.hpp"
#include "../../traits/tapeTraits.hpp"
#include "../algorithms.hpp"
#include "../data/hessian.hpp"
#include "../data/jacobian.hpp"
 
namespace codi {
 
  template<typename T_Type, typename T_Impl>
  struct TapeHelperBase {
    public:
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);
      using Impl = CODI_DD(T_Impl, TapeHelperBase);  
 
      using Real = typename Type::Real;              
      using Identifier = typename Type::Identifier;  
      using Gradient = typename Type::Gradient;      
 
      using PassiveReal = typename RealTraits::PassiveReal<Real>;  
 
      using JacobianType = Jacobian<PassiveReal>;  
      using HessianType = Hessian<PassiveReal>;    
 
    protected:
 
      using Tape = typename Type::Tape;  
 
      Tape& tape;  
 
      std::vector<Identifier> inputValues;   
      std::vector<Identifier> outputValues;  
 
      bool wasForwardEvaluated;  
 
    public:
 
      TapeHelperBase() : tape(Type::getTape()), inputValues(), outputValues(), wasForwardEvaluated(false) {}
 
      virtual ~TapeHelperBase() {}
 
      Gradient* createGradientVectorInput() {
        return new Gradient[getInputSize()];
      }
 
      Gradient* createGradientVectorOutput() {
        return new Gradient[getOutputSize()];
      }
 
      JacobianType& createJacobian() {
        JacobianType* jacPointer = new JacobianType(getOutputSize(), getInputSize());
 
        return *jacPointer;
      }
 
      HessianType& createHessian() {
        HessianType* hesPointer = new HessianType(getOutputSize(), getInputSize());
 
        return *hesPointer;
      }
 
      Real* createPrimalVectorInput() {
        return new Real[getInputSize()];
      }
 
      Real* createPrimalVectorOutput() {
        return new Real[getOutputSize()];
      }
 
      void deleteGradientVector(Gradient* vec) {
        delete[] vec;
 
        vec = nullptr;
      }
 
      void deleteJacobian(JacobianType& jac) {
        JacobianType* jacPointer = &jac;
 
        delete jacPointer;
      }
 
      void deleteHessian(HessianType& hes) {
        HessianType* hesPointer = &hes;
 
        delete hesPointer;
      }
 
      void deletePrimalVector(Real* vec) {
        delete[] vec;
 
        vec = nullptr;
      }
 
      size_t getInputSize() {
        return inputValues.size();
      }
 
      size_t getOutputSize() {
        return outputValues.size();
      }
 
      void registerInput(Type& value) {
        tape.registerInput(value);
        inputValues.push_back(value.getIdentifier());
      }
 
      void registerOutput(Type& value) {
        tape.registerOutput(value);
        outputValues.push_back(value.getIdentifier());
      }
 
      void startRecording() {
        tape.reset();
        inputValues.clear();
        outputValues.clear();
 
        tape.setActive();
      }
 
      void stopRecording() {
        tape.setPassive();
      }
 
      virtual void evalPrimal(Real const* x, Real* y = nullptr) = 0;
 
      CODI_INLINE void evalForward(Gradient const* x_d, Gradient* y_d) {
        changeStateToForwardEvaluation();
 
        for (size_t j = 0; j < inputValues.size(); j += 1) {
          tape.setGradient(inputValues[j], x_d[j]);
        }
 
        tape.evaluateForward();
 
        for (size_t i = 0; i < outputValues.size(); i += 1) {
          y_d[i] = tape.getGradient(outputValues[i]);
          tape.setGradient(outputValues[i], Gradient());
        }
      }
 
      CODI_INLINE void evalForwardAt(Real const* x, Gradient const* x_d, Gradient* y_d, Real* y = nullptr) {
        evalPrimal(x, y);
 
        evalForward(x_d, y_d);
      }
 
      CODI_INLINE void evalReverse(Gradient const* y_b, Gradient* x_b) {
        changeStateToReverseEvaluation();
 
        for (size_t i = 0; i < outputValues.size(); i += 1) {
          tape.setGradient(outputValues[i], y_b[i]);
        }
 
        tape.evaluate();
 
        for (size_t j = 0; j < inputValues.size(); j += 1) {
          x_b[j] = tape.getGradient(inputValues[j]);
          tape.setGradient(inputValues[j], Gradient());
        }
 
        if (!Config::ReversalZeroesAdjoints) {
          tape.clearAdjoints();
        }
      }
 
      CODI_INLINE void evalReverseAt(Real const* x, Gradient const* y_b, Gradient* x_b, Real* y = nullptr) {
        evalPrimal(x, y);
 
        evalReverse(y_b, x_b);
      }
 
      CODI_INLINE void evalJacobian(JacobianType& jac) {
        JacobianConvertWrapper<JacobianType> wrapper(jac);
 
        evalJacobianGen(wrapper);
      }
 
      CODI_INLINE void evalJacobianAt(Real const* x, JacobianType& jac, Real* y = nullptr) {
        evalPrimal(x, y);
 
        evalJacobian(jac);
      }
 
      template<typename Jac>
      CODI_INLINE void evalJacobianGen(Jac& jac) {
        using Algo = Algorithms<Type>;
        typename Algo::EvaluationType evalType = Algo::getEvaluationChoice(inputValues.size(), outputValues.size());
 
        if (Algo::EvaluationType::Forward == evalType) {
          changeStateToForwardEvaluation();
        } else if (Algo::EvaluationType::Reverse == evalType) {
          changeStateToReverseEvaluation();
        } else {
          CODI_EXCEPTION("Evaluation type not implemented.");
        }
 
        Algorithms<Type>::template computeJacobian<Jac, false>(tape, tape.getZeroPosition(), tape.getPosition(),
                                                               inputValues.data(), inputValues.size(),
                                                               outputValues.data(), outputValues.size(), jac);
      }
 
      template<typename Jac = DummyJacobian>
      void evalHessian(HessianType& hes, Jac& jac = StaticDummy<DummyJacobian>::dummy);
 
      template<typename Jac = DummyJacobian>
      CODI_INLINE void evalHessianAt(Real const* x, HessianType& hes, Real* y = nullptr,
                                     Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        evalPrimal(x, y);
 
        cast().evalHessian(hes, jac);
      }
 
    protected:
 
      CODI_INLINE Impl& cast() {
        return static_cast<Impl&>(*this);
      }
 
      void changeStateToForwardEvaluation() {
        wasForwardEvaluated = true;
 
        // No cleanup to do.
      }
 
      void changeStateToReverseEvaluation() {
        if (wasForwardEvaluated) {
          // Forward evaluation leaves the adjoint vector dirty.
 
          tape.clearAdjoints();
        }
 
        wasForwardEvaluated = false;
      }
  };
 
  // clang-format off
  template<typename T_Type>
  struct TapeHelperNoImpl : public TapeHelperBase<T_Type, TapeHelperNoImpl<T_Type>> {
    public:
 
      CODI_STATIC_ASSERT(false && std::is_void<T_Type>::value, "Tape helper not implemented for this tape.");
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);  
      using Real = typename Type::Real;                           
 
      using Base = TapeHelperBase<Type, TapeHelperNoImpl<Type>>;  
 
      virtual void evalPrimal(Real const* x, Real* y = nullptr) CODI_DD(= 0;, {})
 
      template<typename Jac = DummyJacobian>
      void evalHessian(typename Base::HessianType& hes, Jac& jac = StaticDummy<DummyJacobian>::dummy);
  };
  // clang-format on
 
  template<typename T_Type>
  struct TapeHelperJacobi : public TapeHelperBase<T_Type, TapeHelperJacobi<T_Type>> {
    public:
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);  
      using Real = typename Type::Real;                           
 
      using Base = TapeHelperBase<Type, TapeHelperJacobi<Type>>;  
 
      virtual void evalPrimal(Real const* x, Real* y = nullptr) {
        CODI_UNUSED(x, y);
 
        CODI_EXCEPTION(
            "No primal evaluation for Jacobian tapes. "
            "Please use codi::RealReversePrimal or codi::RealReversePrimalIndex types for this kind of functionality.");
      }
 
      template<typename Jac = DummyJacobian>
      void evalHessian(typename Base::HessianType& hes, Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        CODI_UNUSED(hes, jac);
 
        CODI_EXCEPTION(
            "No direct Hessian evaluation for Jacobian tapes. "
            "Please use codi::RealReversePrimal or codi::RealReversePrimalIndex types for this kind of functionality "
            "or the EvaluationHelper class.");
      }
  };
 
  template<typename T_Type>
  struct TapeHelperPrimal : public TapeHelperBase<T_Type, TapeHelperPrimal<T_Type>> {
    public:
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);  
      using Real = typename Type::Real;                           
 
      using Base = TapeHelperBase<Type, TapeHelperPrimal<Type>>;  
 
      virtual void evalPrimal(Real const* x, Real* y = nullptr) {
        for (size_t j = 0; j < this->inputValues.size(); j += 1) {
          this->tape.primal(this->inputValues[j]) = x[j];
        }
 
        this->tape.evaluatePrimal();
 
        if (nullptr != y) {
          for (size_t i = 0; i < this->outputValues.size(); i += 1) {
            y[i] = this->tape.primal(this->outputValues[i]);
          }
        }
      }
 
      template<typename Jac = DummyJacobian>
      void evalHessian(typename Base::HessianType& hes, Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        using Algo = Algorithms<Type>;
        typename Algo::EvaluationType evalType =
            Algo::getEvaluationChoice(this->inputValues.size(), this->outputValues.size());
 
        if (Algo::EvaluationType::Forward == evalType) {
          this->changeStateToForwardEvaluation();
        } else if (Algo::EvaluationType::Reverse == evalType) {
          this->changeStateToReverseEvaluation();
        } else {
          CODI_EXCEPTION("Evaluation type not implemented.");
        }
 
        Algorithms<Type>::computeHessianPrimalValueTape(
            this->tape, this->tape.getZeroPosition(), this->tape.getPosition(), this->inputValues.data(),
            this->inputValues.size(), this->outputValues.data(), this->outputValues.size(), hes, jac);
      }
  };
 
  template<typename Type, typename = void>
  struct TapeHelper : public TapeHelperNoImpl<Type> {};
 
#ifndef DOXYGEN_DISABLE
  template<typename Type>
  struct TapeHelper<Type, TapeTraits::EnableIfJacobianTape<typename Type::Tape>> : public TapeHelperJacobi<Type> {};
 
  template<typename Type>
  struct TapeHelper<Type, TapeTraits::EnableIfPrimalValueTape<typename Type::Tape>> : public TapeHelperPrimal<Type> {};
#endif
 
}