computeExpression.hpp

/*
 * CoDiPack, a Code Differentiation Package
 *
 * Copyright (C) 2015-2025 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 <tuple>
 
#include "../config.h"
#include "../misc/macros.hpp"
#include "../misc/tupleMemory.hpp"
#include "../traits/computationTraits.hpp"
#include "../traits/expressionTraits.hpp"
#include "expressionInterface.hpp"
#include "logic/compileTimeTraversalLogic.hpp"
#include "logic/nodeInterface.hpp"
#include "logic/traversalLogic.hpp"

namespace codi {

  template<typename T_Real>
  struct ComputeOperation {
    public:
 
      using Real = CODI_DD(T_Real, double);  

      template<typename... Args>
      static CODI_INLINE Real primal(Args const&... args);

      template<size_t argNumber, typename Tangent, typename... Args>
      static CODI_INLINE auto applyTangent(Tangent const& tangent, Real const& result, Args const&... args);

      template<size_t argNumber, typename Adjoint, typename... Args>
      static CODI_INLINE auto applyAdjoint(Adjoint const& adjoint, Real const& result, Args const&... args);

      static CODI_INLINE std::string getMathRep();
  };

  template<typename T_Real, typename T_Impl>
  struct UnaryOperation : public ComputeOperation<T_Real> {
    public:
 
      using Real = CODI_DD(T_Real, double);  

      template<typename Arg>
      static CODI_INLINE Real primal(Arg const& arg);

      template<typename Tangent, typename Arg>
      static CODI_INLINE auto applyTangentArg(Tangent const& tangent, Real const& result, Arg const& arg);

      template<typename Adjoint, typename Arg>
      static CODI_INLINE auto applyAdjointArg(Adjoint const& adjoint, Real const& result, Arg const& arg);
 
      // Interface implementation
 
      using Impl = CODI_DD(T_Impl, UnaryOperation);  

      template<size_t argNumber, typename Tangent, typename Arg>
      static CODI_INLINE auto applyTangent(Tangent const& tangent, Real const& result, Arg const& arg) {
        return Impl::applyTangentArg(tangent, result, arg);
      }

      template<size_t argNumber, typename Adjoint, typename Arg>
      static CODI_INLINE auto applyAdjoint(Adjoint const& adjoint, Real const& result, Arg const& arg) {
        return Impl::applyAdjointArg(adjoint, result, arg);
      }
  };

  template<typename T_Real, typename T_Impl>
  struct UnaryJacobianOperation : UnaryOperation<T_Real, T_Impl> {
    public:
 
      using Real = CODI_DD(T_Real, double);  

      template<typename Arg>
      static CODI_INLINE Real primal(Arg const& arg);

      template<typename Arg>
      static CODI_INLINE auto gradient(Arg const& arg, Real const& result);
 
      // Interface implementation
 
      using Impl = CODI_DD(T_Impl, UnaryJacobianOperation);  

      template<typename Tangent, typename Arg>
      static CODI_INLINE auto applyTangentArg(Tangent const& tangent, Real const& result, Arg const& arg) {
        return Impl::gradient(arg, result) * tangent;
      }

      template<typename Adjoint, typename Arg>
      static CODI_INLINE Arg applyAdjointArg(Adjoint const& adjoint, Real const& result, Arg const& arg) {
        return ComputationTraits::transpose(Impl::gradient(arg, result)) * adjoint;
      }
  };

  template<typename T_Real, typename T_Impl>
  struct BinaryOperation : public ComputeOperation<T_Real> {
    public:
 
      using Real = CODI_DD(T_Real, double);  

      template<typename ArgA, typename ArgB>
      static CODI_INLINE Real primal(ArgA const& argA, ArgB const& argB);

      template<typename Tangent, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyTangentArgA(Tangent const& tangent, Real const& result, ArgA const& argA,
                                               ArgB const& argB);

      template<typename Adjoint, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyAdjointArgA(Adjoint const& adjoint, Real const& result, ArgA const& argA,
                                               ArgB const& argB);

      template<typename Tangent, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyTangentArgB(Tangent const& tangent, Real const& result, ArgA const& argA,
                                               ArgB const& argB);

      template<typename Adjoint, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyAdjointArgB(Adjoint const& adjoint, Real const& result, ArgA const& argA,
                                               ArgB const& argB);
 
      // Interface implementation
 
      using Impl = CODI_DD(T_Impl, BinaryOperation);  
 
#ifndef DOXYGEN_DISABLE
      template<typename Func, bool derivB = false>
      struct CallSwitchTangent {
        public:
          template<typename... Args>
          CODI_INLINE static auto call(Args&&... args) {
            return Impl::applyTangentArgA(std::forward<Args>(args)...);
          }
      };
 
      template<typename Temp>
      struct CallSwitchTangent<Temp, true> {
        public:
          template<typename... Args>
          CODI_INLINE static auto call(Args&&... args) {
            return Impl::applyTangentArgB(std::forward<Args>(args)...);
          }
      };
 
      template<typename Func, bool derivB = false>
      struct CallSwitchAdjoint {
        public:
          template<typename... Args>
          CODI_INLINE static auto call(Args&&... args) {
            return Impl::applyAdjointArgA(std::forward<Args>(args)...);
          }
      };
 
      template<typename Temp>
      struct CallSwitchAdjoint<Temp, true> {
        public:
          template<typename... Args>
          CODI_INLINE static auto call(Args&&... args) {
            return Impl::applyAdjointArgB(std::forward<Args>(args)...);
          }
      };
#endif

      template<size_t argNumber, typename Tangent, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyTangent(Tangent const& tangent, Real const& result, ArgA const& argA,
                                           ArgB const& argB) {
        return CallSwitchTangent < void, argNumber == 1 > ::call(tangent, result, argA, argB);
      }

      template<size_t argNumber, typename Adjoint, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyAdjoint(Adjoint const& adjoint, Real const& result, ArgA const& argA,
                                           ArgB const& argB) {
        return CallSwitchAdjoint < void, argNumber == 1 > ::call(adjoint, result, argA, argB);
      }
  };

  template<typename T_Real, typename T_Impl>
  struct BinaryJacobianOperation : public BinaryOperation<T_Real, T_Impl> {
    public:
 
      using Real = CODI_DD(T_Real, double);  

      template<typename ArgA, typename ArgB>
      static CODI_INLINE Real primal(ArgA const& argA, ArgB const& argB);

      template<typename ArgA, typename ArgB>
      static CODI_INLINE auto gradientA(ArgA const& argA, ArgB const& argB, Real const& result);

      template<typename ArgA, typename ArgB>
      static CODI_INLINE auto gradientB(ArgA const& argA, ArgB const& argB, Real const& result);
 
      // Interface implementation
 
      using Impl = CODI_DD(T_Impl, BinaryJacobianOperation);  

      template<typename Tangent, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyTangentArgA(Tangent const& tangent, Real const& result, ArgA const& argA,
                                               ArgB const& argB) {
        return Impl::gradientA(argA, argB, result) * tangent;
      }

      template<typename Adjoint, typename ArgA, typename ArgB>
      static CODI_INLINE ArgA applyAdjointArgA(Adjoint const& adjoint, Real const& result, ArgA const& argA,
                                               ArgB const& argB) {
        return ComputationTraits::transpose(Impl::gradientA(argA, argB, result)) * adjoint;
      }

      template<typename Tangent, typename ArgA, typename ArgB>
      static CODI_INLINE auto applyTangentArgB(Tangent const& tangent, Real const& result, ArgA const& argA,
                                               ArgB const& argB) {
        return Impl::gradientB(argA, argB, result) * tangent;
      }

      template<typename Adjoint, typename ArgA, typename ArgB>
      static CODI_INLINE ArgB applyAdjointArgB(Adjoint const& adjoint, Real const& result, ArgA const& argA,
                                               ArgB const& argB) {
        return ComputationTraits::transpose(Impl::gradientB(argA, argB, result)) * adjoint;
      }
  };

  template<typename T_Real, template<typename> class T_Operation, typename... T_ArgExprs>
  struct ComputeExpression
      : public ExpressionInterface<T_Real, ComputeExpression<T_Real, T_Operation, T_ArgExprs...> > {
    public:
      using Real = CODI_DD(T_Real, double);  
      template<typename T>
      using Operation = CODI_DD(CODI_T(T_Operation<T>), CODI_T(ComputeOperation<T>));  
      using ArgExprs = std::tuple<T_ArgExprs...>;                                      
 
      using ArgReals = std::tuple<typename T_ArgExprs::Real...>;  
      using ArgStoresTypes =
          std::tuple<typename T_ArgExprs::StoreAs...>;  
      using ArgStores = TupleMemory<typename T_ArgExprs::StoreAs...>;  
 
      ArgStores args;  
 
      Real result;  

      CODI_INLINE explicit ComputeExpression(ExpressionInterface<typename T_ArgExprs::Real, T_ArgExprs> const&... args)
          : args(args.cast()...), result(Operation<Real>::primal(args.cast().getValue()...)) {}

      CODI_INLINE ComputeExpression(ComputeExpression const& v) = default;
 
      /*******************************************************************************/

      CODI_INLINE std::string getMathRep() const {
        return Operation<Real>::getMathRep();
      }

      /*******************************************************************************/
 
      using StoreAs = ComputeExpression;  
      using ADLogic = typename ExpressionTraits::ValidateADLogic<
          typename T_ArgExprs::ADLogic...>::ADLogic;  

      CODI_INLINE Real const& getValue() const {
        return result;
      }

      template<size_t argNumber, typename Tangent>
      CODI_INLINE Real applyTangent(Tangent const& tangent) const {
        return callTangent<argNumber>(tangent, result, args);
      }

      template<size_t argNumber, typename Adjoint>
      CODI_INLINE auto applyAdjoint(Adjoint const& adjoint) const {
        return callAdjoint<argNumber>(adjoint, result, args);
      }

      /*******************************************************************************/
 
      static size_t constexpr LinkCount = std::tuple_size<ArgReals>::value;  

      template<size_t argNumber>
      CODI_INLINE std::tuple_element_t<argNumber, ArgStoresTypes> const& getLink() const {
        return args.template get<argNumber>();
      }

 
    private:
 
      // Unpack helpers
 
      template<size_t argNumber, typename EvalArg, typename Tuple, size_t... I>
      static CODI_INLINE auto callAdjoint(EvalArg const& evalArg, Real const& result, Tuple t,
                                          std::index_sequence<I...>) {
        return Operation<Real>::template applyAdjoint<argNumber>(evalArg, result, t.template get<I>().getValue()...);
      }
 
      template<size_t argNumber, typename EvalArg, typename Tuple>
      static CODI_INLINE auto callAdjoint(EvalArg const& evalArg, Real const& result, Tuple const& t) {
        static constexpr auto size = std::tuple_size<ArgStoresTypes>::value;
        return callAdjoint<argNumber>(evalArg, result, t, std::make_index_sequence<size>{});
      }
 
      template<size_t argNumber, typename EvalArg, typename Tuple, size_t... I>
      static CODI_INLINE auto callTangent(EvalArg const& evalArg, Real const& result, Tuple t,
                                          std::index_sequence<I...>) {
        return Operation<Real>::template applyTangent<argNumber>(evalArg, result, t.template get<I>().getValue()...);
      }
 
      template<size_t argNumber, typename EvalArg, typename Tuple>
      static CODI_INLINE auto callTangent(EvalArg const& evalArg, Real const& result, Tuple const& t) {
        static constexpr auto size = std::tuple_size<ArgStoresTypes>::value;
        return callTangent<argNumber>(evalArg, result, t, std::make_index_sequence<size>{});
      }
  };
}