algorithms.hpp

/*
 * CoDiPack, a Code Differentiation Package
 *
 * Copyright (C) 2015-2024 Chair for Scientific Computing (SciComp), University of Kaiserslautern-Landau
 * Homepage: http://scicomp.rptu.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://scicomp.rptu.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 "../config.h"
#include "../expressions/lhsExpressionInterface.hpp"
#include "../misc/exceptions.hpp"
#include "../tapes/misc/tapeParameters.hpp"
#include "../traits/adjointVectorTraits.hpp"
#include "../traits/gradientTraits.hpp"
#include "data/dummy.hpp"
#include "data/jacobian.hpp"
#include "data/staticDummy.hpp"
 
namespace codi {
 
  template<typename T_Type, bool T_ActiveChecks = true>
  struct Algorithms {
    public:
 
      using Type = CODI_DD(T_Type, CODI_DEFAULT_LHS_EXPRESSION);
 
      static bool constexpr ActiveChecks = T_ActiveChecks;  
 
      using Tape = typename Type::Tape;              
      using Position = typename Tape::Position;      
      using Real = typename Type::Real;              
      using Identifier = typename Type::Identifier;  
      using Gradient = typename Type::Gradient;      
 
      using GT = GradientTraits::TraitsImplementation<Gradient>;  
 
      enum class EvaluationType {
        Forward,
        Reverse
      };
 
      static CODI_INLINE EvaluationType getEvaluationChoice(size_t const inputs, size_t const outputs) {
        if (inputs <= outputs) {
          return EvaluationType::Forward;
        } else {
          return EvaluationType::Reverse;
        }
      }
 
      template<typename Jac, bool keepState = true>
      static CODI_INLINE void computeJacobian(Tape& tape, Position const& start, Position const& end,
                                              Identifier const* input, size_t const inputSize, Identifier const* output,
                                              size_t const outputSize, Jac& jac,
                                              AdjointsManagement adjointsManagement = AdjointsManagement::Automatic) {
        // internally, automatic management is implemented in an optimized way that uses manual management
        if (AdjointsManagement::Automatic == adjointsManagement) {
          tape.resizeAdjointVector();
          tape.beginUseAdjointVector();
        }
 
        computeJacobianCustomAdjoints(tape, start, end, input, inputSize, output, outputSize, jac,
                                      tape.getInternalAdjoints());
 
        if (AdjointsManagement::Automatic == adjointsManagement) {
          tape.endUseAdjointVector();
        }
      }
 
      // clang-format off
      // clang-format on
      template<typename Jac, typename AdjointVector, bool keepState = true>
      static CODI_INLINE void computeJacobianCustomAdjoints(Tape& tape, Position const& start, Position const& end,
                                                            Identifier const* input, size_t const inputSize,
                                                            Identifier const* output, size_t const outputSize, Jac& jac,
                                                            AdjointVector&& adjoints) {
        using Adjoint = AdjointVectorTraits::Gradient<AdjointVector>;
 
        using CustomGT = GradientTraits::TraitsImplementation<Adjoint>;
 
        size_t constexpr gradDim = CustomGT::dim;
 
        EvaluationType evalType = getEvaluationChoice(inputSize, outputSize);
        if (EvaluationType::Forward == evalType) {
          for (size_t j = 0; j < inputSize; j += gradDim) {
            setGradientOnIdentifierCustomAdjoints(tape, j, input, inputSize, typename CustomGT::Real(1.0), adjoints);
 
            if (keepState) {
              tape.evaluateForwardKeepState(start, end, std::forward<AdjointVector>(adjoints));
            } else {
              tape.evaluateForward(start, end, std::forward<AdjointVector>(adjoints));
            }
 
            for (size_t i = 0; i < outputSize; i += 1) {
              for (size_t curDim = 0; curDim < gradDim && j + curDim < inputSize; curDim += 1) {
                jac(outputSize - i - 1, j + curDim) = CustomGT::at(adjoints[output[outputSize - i - 1]], curDim);
                if (tape.isIdentifierActive(output[i])) {
                  CustomGT::at(adjoints[output[outputSize - i - 1]], curDim) = typename GT::Real();
                }
              }
            }
 
            setGradientOnIdentifierCustomAdjoints(tape, j, input, inputSize, typename CustomGT::Real(), adjoints);
          }
 
          tape.clearCustomAdjoints(end, start, std::forward<AdjointVector>(adjoints));
 
        } else if (EvaluationType::Reverse == evalType) {
          for (size_t i = 0; i < outputSize; i += gradDim) {
            setGradientOnIdentifierCustomAdjoints(tape, i, output, outputSize, typename CustomGT::Real(1.0), adjoints);
 
            if (keepState) {
              tape.evaluateKeepState(end, start, std::forward<AdjointVector>(adjoints));
            } else {
              tape.evaluate(end, start, std::forward<AdjointVector>(adjoints));
            }
 
            for (size_t j = 0; j < inputSize; j += 1) {
              for (size_t curDim = 0; curDim < gradDim && i + curDim < outputSize; curDim += 1) {
                jac(i + curDim, j) = CustomGT::at(adjoints[input[j]], curDim);
                CustomGT::at(adjoints[input[j]], curDim) = typename CustomGT::Real();
              }
            }
 
            setGradientOnIdentifierCustomAdjoints(tape, i, output, outputSize, typename CustomGT::Real(), adjoints);
 
            if (!Config::ReversalZeroesAdjoints) {
              tape.clearCustomAdjoints(end, start, std::forward<AdjointVector>(adjoints));
            }
          }
        } else {
          CODI_EXCEPTION("Evaluation mode not implemented. Mode is: %d.", (int)evalType);
        }
      }
 
      // clang-format off
      // clang-format on
      template<typename Jac>
      static CODI_INLINE void computeJacobian(Position const& start, Position const& end, Identifier const* input,
                                              size_t const inputSize, Identifier const* output, size_t const outputSize,
                                              Jac& jac,
                                              AdjointsManagement adjointsManagement = AdjointsManagement::Automatic) {
        computeJacobian(Type::getTape(), start, end, input, inputSize, output, outputSize, jac, adjointsManagement);
      }
 
      // clang-format off
      // clang-format on
      template<typename Jac, typename AdjointVector>
      static CODI_INLINE void computeJacobianCustomAdjoints(Position const& start, Position const& end,
                                                            Identifier const* input, size_t const inputSize,
                                                            Identifier const* output, size_t const outputSize, Jac& jac,
                                                            AdjointVector&& adjoints) {
        computeJacobianCustomAdjoints<Jac, AdjointVector>(Type::getTape(), start, end, input, inputSize, output,
                                                          outputSize, jac, std::forward<AdjointVector>(adjoints));
      }
 
      template<typename Hes, typename Jac = DummyJacobian>
      static CODI_INLINE void computeHessianPrimalValueTape(Tape& tape, Position const& start, Position const& end,
                                                            Identifier const* input, size_t const inputSize,
                                                            Identifier const* output, size_t const outputSize, Hes& hes,
                                                            Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        EvaluationType evalType = getEvaluationChoice(inputSize, outputSize);
        if (EvaluationType::Forward == evalType) {
          computeHessianPrimalValueTapeForward(tape, start, end, input, inputSize, output, outputSize, hes, jac);
        } else if (EvaluationType::Reverse == evalType) {
          computeHessianPrimalValueTapeReverse(tape, start, end, input, inputSize, output, outputSize, hes, jac);
        } else {
          CODI_EXCEPTION("Evaluation mode not implemented. Mode is: %d.", (int)evalType);
        }
      }
 
      template<typename Hes, typename Jac = DummyJacobian>
      static CODI_INLINE void computeHessianPrimalValueTapeForward(Tape& tape, Position const& start,
                                                                   Position const& end, Identifier const* input,
                                                                   size_t const inputSize, Identifier const* output,
                                                                   size_t const outputSize, Hes& hes,
                                                                   Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        using GT1st = GT;
        size_t constexpr gradDim1st = GT1st::dim;
        using GT2nd = GradientTraits::TraitsImplementation<CODI_DD(typename Real::Gradient, double)>;
        size_t constexpr gradDim2nd = GT2nd::dim;
 
        // Assume that the tape was just recorded.
        tape.revertPrimals(start);
 
        for (size_t j = 0; j < inputSize; j += gradDim2nd) {
          setGradient2ndOnIdentifier(tape, j, input, inputSize, typename GT2nd::Real(1.0));
 
          // The k = j init is no problem, it will evaluate slightly more elements around the diagonal.
          for (size_t k = j; k < inputSize; k += gradDim1st) {
            setGradientOnIdentifier(tape, k, input, inputSize, typename GT1st::Real(1.0));
 
            tape.evaluateForward(start, end);
 
            for (size_t i = 0; i < outputSize; i += 1) {
              for (size_t vecPos1st = 0; vecPos1st < gradDim1st && k + vecPos1st < inputSize; vecPos1st += 1) {
                for (size_t vecPos2nd = 0; vecPos2nd < gradDim2nd && j + vecPos2nd < inputSize; vecPos2nd += 1) {
                  hes(i, j + vecPos2nd, k + vecPos1st) =
                      GT2nd::at(GT1st::at(tape.getGradient(output[i]), vecPos1st).gradient(), vecPos2nd);
                  hes(i, k + vecPos1st, j + vecPos2nd) = hes(i, j + vecPos2nd, k + vecPos1st);  // Symmetry
                }
              }
 
              if (j == 0) {
                for (size_t vecPos = 0; vecPos < gradDim1st && k + vecPos < inputSize; vecPos += 1) {
                  jac(i, k + vecPos) = GT1st::at(tape.getGradient(output[i]), vecPos).value();
                }
              }
            }
 
            setGradientOnIdentifier(tape, k, input, inputSize, typename GT1st::Real());
          }
 
          setGradient2ndOnIdentifier(tape, j, input, inputSize, typename GT2nd::Real());
        }
      }
 
      template<typename Hes, typename Jac = DummyJacobian>
      static CODI_INLINE void computeHessianPrimalValueTapeReverse(Tape& tape, Position const& start,
                                                                   Position const& end, Identifier const* input,
                                                                   size_t const inputSize, Identifier const* output,
                                                                   size_t const outputSize, Hes& hes,
                                                                   Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        using GT1st = GT;
        size_t constexpr gradDim1st = GT1st::dim;
        using GT2nd = GradientTraits::TraitsImplementation<CODI_DD(typename Real::Gradient, double)>;
        size_t constexpr gradDim2nd = GT2nd::dim;
 
        // Assume that the tape was just recorded.
        tape.revertPrimals(start);
 
        for (size_t j = 0; j < inputSize; j += gradDim2nd) {
          setGradient2ndOnIdentifier(tape, j, input, inputSize, typename GT2nd::Real(1.0));
 
          // Propagate the new derivative information.
          tape.evaluatePrimal(start, end);
 
          for (size_t i = 0; i < outputSize; i += gradDim1st) {
            setGradientOnIdentifier(tape, i, output, outputSize, typename GT1st::Real(1.0));
 
            // Propagate the derivatives backward for second order derivatives.
            tape.evaluateKeepState(end, start);
 
            for (size_t k = 0; k < inputSize; k += 1) {
              for (size_t vecPos1st = 0; vecPos1st < gradDim1st && i + vecPos1st < outputSize; vecPos1st += 1) {
                for (size_t vecPos2nd = 0; vecPos2nd < gradDim2nd && j + vecPos2nd < inputSize; vecPos2nd += 1) {
                  hes(i + vecPos1st, j + vecPos2nd, k) =
                      GT2nd::at(GT1st::at(tape.gradient(input[k]), vecPos1st).gradient(), vecPos2nd);
                }
              }
 
              if (j == 0) {
                for (size_t vecPos1st = 0; vecPos1st < gradDim1st && i + vecPos1st < outputSize; vecPos1st += 1) {
                  jac(i + vecPos1st, k) = GT1st::at(tape.getGradient(input[k]), vecPos1st).value();
                }
              }
 
              tape.gradient(input[k]) = Gradient();
            }
 
            setGradientOnIdentifier(tape, i, output, outputSize, typename GT1st::Real());
 
            if (!Config::ReversalZeroesAdjoints) {
              tape.clearAdjoints(end, start);
            }
          }
 
          setGradient2ndOnIdentifier(tape, j, input, inputSize, typename GT2nd::Real());
 
          if (j + gradDim2nd < inputSize) {
            tape.revertPrimals(start);
          }
        }
      }
 
      template<typename Func, typename VecIn, typename VecOut, typename Hes, typename Jac = DummyJacobian>
      static CODI_INLINE void computeHessian(Func func, VecIn& input, VecOut& output, Hes& hes,
                                             Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        EvaluationType evalType = getEvaluationChoice(input.size(), output.size());
        if (EvaluationType::Forward == evalType) {
          computeHessianForward(func, input, output, hes, jac);
        } else if (EvaluationType::Reverse == evalType) {
          computeHessianReverse(func, input, output, hes, jac);
        } else {
          CODI_EXCEPTION("Evaluation mode not implemented. Mode is: %d.", (int)evalType);
        }
      }
 
      template<typename Func, typename VecIn, typename VecOut, typename Hes, typename Jac = DummyJacobian>
      static CODI_INLINE void computeHessianForward(Func func, VecIn& input, VecOut& output, Hes& hes,
                                                    Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        using GT1st = GT;
        size_t constexpr gradDim1st = GT1st::dim;
        using GT2nd = GradientTraits::TraitsImplementation<CODI_DD(typename Real::Gradient, double)>;
        size_t constexpr gradDim2nd = GT2nd::dim;
 
        Tape& tape = Type::getTape();
 
        for (size_t j = 0; j < input.size(); j += gradDim2nd) {
          setGradient2ndOnCoDiValue(j, input.data(), input.size(), typename GT2nd::Real(1.0));
 
          // Propagate the new derivative information.
          recordTape(func, input, output);
 
          // The k = j init is no problem, it will evaluate slightly more elements around the diagonal.
          for (size_t k = j; k < input.size(); k += gradDim1st) {
            setGradientOnCoDiValue(tape, k, input.data(), input.size(), typename GT1st::Real(1.0));
 
            // Propagate the derivatives forward for second order derivatives.
            tape.evaluateForwardKeepState(tape.getZeroPosition(), tape.getPosition());
 
            for (size_t i = 0; i < output.size(); i += 1) {
              for (size_t vecPos1st = 0; vecPos1st < gradDim1st && k + vecPos1st < input.size(); vecPos1st += 1) {
                for (size_t vecPos2nd = 0; vecPos2nd < gradDim2nd && j + vecPos2nd < input.size(); vecPos2nd += 1) {
                  hes(i, j + vecPos2nd, k + vecPos1st) = GT2nd::at(
                      GT1st::at(tape.getGradient(output[i].getIdentifier()), vecPos1st).gradient(), vecPos2nd);
                  hes(i, k + vecPos1st, j + vecPos2nd) = hes(i, j + vecPos2nd, k + vecPos1st);  // Symmetry
                }
              }
 
              if (j == 0) {
                for (size_t vecPos = 0; vecPos < gradDim1st && k + vecPos < input.size(); vecPos += 1) {
                  jac(i, k + vecPos) = GT1st::at(tape.getGradient(output[i].getIdentifier()), vecPos).value();
                }
              }
            }
 
            setGradientOnCoDiValue(tape, k, input.data(), input.size(), typename GT1st::Real());
          }
 
          setGradient2ndOnCoDiValue(j, input.data(), input.size(), typename GT2nd::Real());
 
          tape.reset();
        }
      }
 
      template<typename Func, typename VecIn, typename VecOut, typename Hes, typename Jac = DummyJacobian>
      static CODI_INLINE void computeHessianReverse(Func func, VecIn& input, VecOut& output, Hes& hes,
                                                    Jac& jac = StaticDummy<DummyJacobian>::dummy) {
        using GT1st = GT;
        size_t constexpr gradDim1st = GT1st::dim;
        using GT2nd = GradientTraits::TraitsImplementation<CODI_DD(typename Real::Gradient, double)>;
        size_t constexpr gradDim2nd = GT2nd::dim;
 
        Tape& tape = Type::getTape();
 
        for (size_t j = 0; j < input.size(); j += gradDim2nd) {
          setGradient2ndOnCoDiValue(j, input.data(), input.size(), typename GT2nd::Real(1.0));
 
          // Propagate the new derivative information.
          recordTape(func, input, output);
 
          for (size_t i = 0; i < output.size(); i += gradDim1st) {
            setGradientOnCoDiValue(tape, i, output.data(), output.size(), typename GT1st::Real(1.0));
 
            // Propagate the derivatives backward for second order derivatives.
            tape.evaluateKeepState(tape.getPosition(), tape.getZeroPosition());
 
            for (size_t k = 0; k < input.size(); k += 1) {
              for (size_t vecPos1st = 0; vecPos1st < gradDim1st && i + vecPos1st < output.size(); vecPos1st += 1) {
                for (size_t vecPos2nd = 0; vecPos2nd < gradDim2nd && j + vecPos2nd < input.size(); vecPos2nd += 1) {
                  hes(i + vecPos1st, j + vecPos2nd, k) =
                      GT2nd::at(GT1st::at(tape.gradient(input[k].getIdentifier()), vecPos1st).gradient(), vecPos2nd);
                }
              }
 
              if (j == 0) {
                for (size_t vecPos1st = 0; vecPos1st < gradDim1st && i + vecPos1st < output.size(); vecPos1st += 1) {
                  jac(i + vecPos1st, k) = GT1st::at(tape.getGradient(input[k].getIdentifier()), vecPos1st).value();
                }
              }
 
              tape.gradient(input[k].getIdentifier()) = Gradient();
            }
 
            setGradientOnCoDiValue(tape, i, output.data(), output.size(), typename GT1st::Real());
 
            if (!Config::ReversalZeroesAdjoints) {
              tape.clearAdjoints(tape.getPosition(), tape.getZeroPosition());
            }
          }
 
          setGradient2ndOnCoDiValue(j, input.data(), input.size(), typename GT2nd::Real());
 
          tape.reset();
        }
      }
 
    private:
 
      template<typename T>
      static CODI_INLINE void setGradientOnIdentifier(
          Tape& tape, size_t const pos, Identifier const* identifiers, size_t const size, T value,
          AdjointsManagement adjointsManagement = AdjointsManagement::Automatic) {
        size_t constexpr gradDim = GT::dim;
 
        for (size_t curDim = 0; curDim < gradDim && pos + curDim < size; curDim += 1) {
          if (CODI_ENABLE_CHECK(ActiveChecks, tape.isIdentifierActive(identifiers[pos + curDim]))) {
            GT::at(tape.gradient(identifiers[pos + curDim], adjointsManagement), curDim) = value;
          }
        }
      }
 
      template<typename T, typename Adjoints>
      static CODI_INLINE void setGradientOnIdentifierCustomAdjoints(Tape& tape, size_t const pos,
                                                                    Identifier const* identifiers, size_t const size,
                                                                    T value, Adjoints& adjoints) {
        size_t constexpr gradDim = GT::dim;
 
        for (size_t curDim = 0; curDim < gradDim && pos + curDim < size; curDim += 1) {
          if (CODI_ENABLE_CHECK(ActiveChecks, tape.isIdentifierActive(identifiers[pos + curDim]))) {
            GT::at(adjoints[identifiers[pos + curDim]], curDim) = value;
          }
        }
      }
 
      template<typename T>
      static CODI_INLINE void setGradient2ndOnIdentifier(Tape& tape, size_t const pos, Identifier const* identifiers,
                                                         size_t const size, T value) {
        using GT2nd = GradientTraits::TraitsImplementation<CODI_DD(typename Real::Gradient, double)>;
        size_t constexpr gradDim2nd = GT2nd::dim;
 
        for (size_t curDim = 0; curDim < gradDim2nd && pos + curDim < size; curDim += 1) {
          // No activity check on the identifier required since forward types are used.
          GT2nd::at(tape.primal(identifiers[pos + curDim]).gradient(), curDim) = value;
        }
      }
 
      template<typename T>
      static CODI_INLINE void setGradientOnCoDiValue(Tape& tape, size_t const pos, Type* identifiers, size_t const size,
                                                     T value) {
        size_t constexpr gradDim = GT::dim;
 
        for (size_t curDim = 0; curDim < gradDim && pos + curDim < size; curDim += 1) {
          if (CODI_ENABLE_CHECK(ActiveChecks, tape.isIdentifierActive(identifiers[pos + curDim].getIdentifier()))) {
            GT::at(tape.gradient(identifiers[pos + curDim].getIdentifier()), curDim) = value;
          }
        }
      }
 
      template<typename T>
      static CODI_INLINE void setGradient2ndOnCoDiValue(size_t const pos, Type* identifiers, size_t const size,
                                                        T value) {
        using GT2nd = GradientTraits::TraitsImplementation<CODI_DD(typename Real::Gradient, double)>;
        size_t constexpr gradDim2nd = GT2nd::dim;
 
        for (size_t curDim = 0; curDim < gradDim2nd && pos + curDim < size; curDim += 1) {
          // No activity check on the identifier required since forward types are used.
          GT2nd::at(identifiers[pos + curDim].value().gradient(), curDim) = value;
        }
      }
 
      template<typename Func, typename VecIn, typename VecOut>
      static CODI_INLINE void recordTape(Func func, VecIn& input, VecOut& output) {
        Tape& tape = Type::getTape();
        tape.setActive();
        for (size_t curIn = 0; curIn < input.size(); curIn += 1) {
          tape.registerInput(input[curIn]);
        }
 
        func(input, output);
 
        for (size_t curOut = 0; curOut < output.size(); curOut += 1) {
          tape.registerOutput(output[curOut]);
        }
        tape.setPassive();
      }
  };
 
}