expression_functions.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <algorithm>
#include <cmath>
#include <concepts>
#include <numbers>
#include <string_view>
#include <type_traits>
#include <utility>
 
#include "sleipnir/autodiff/expression.hpp"
#include "sleipnir/autodiff/expression_type.hpp"
 
namespace slp::detail {
template <typename Scalar>
struct AbsExpression final
    : NonLeafExpression<AbsExpression<Scalar>, Scalar, 1> {
  using AbsExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "abs";
 
  static Scalar eval(Scalar x) {
    using std::abs;
    return abs(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    if constexpr (std::same_as<Value, Scalar>) {
      if (x < Scalar(0)) {
        g[0] = Scalar(-1);
      } else if (x > Scalar(0)) {
        g[0] = Scalar(1);
      } else {
        g[0] = Scalar(0);
      }
    } else {
      g[0] = sign(x);
    }
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> abs(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<AbsExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct AcosExpression final
    : NonLeafExpression<AcosExpression<Scalar>, Scalar, 1> {
  using AcosExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "acos";
 
  static Scalar eval(Scalar x) {
    using std::acos;
    return acos(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::sqrt;
    g[0] = lit<Value>(Scalar(-1)) / sqrt(lit<Value>(Scalar(1)) - x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> acos(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<AcosExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct AsinExpression final
    : NonLeafExpression<AsinExpression<Scalar>, Scalar, 1> {
  using AsinExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "asin";
 
  static Scalar eval(Scalar x) {
    using std::asin;
    return asin(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::sqrt;
    g[0] = lit<Value>(Scalar(1)) / sqrt(lit<Value>(Scalar(1)) - x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> asin(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<AsinExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct AtanExpression final
    : NonLeafExpression<AtanExpression<Scalar>, Scalar, 1> {
  using AtanExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "atan";
 
  static Scalar eval(Scalar x) {
    using std::atan;
    return atan(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    g[0] = lit<Value>(Scalar(1)) / (lit<Value>(Scalar(1)) + x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> atan(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<AtanExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct Atan2Expression final
    : NonLeafExpression<Atan2Expression<Scalar>, Scalar, 2> {
  using Atan2Expression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "atan2";
 
  static Scalar eval(Scalar y, Scalar x) {
    using std::atan2;
    return atan2(y, x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& y, const Value& x) {
    Value d = y * y + x * x;
    g[0] = x / d;
    g[1] = -y / d;
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> atan2(const ExpressionPtr<Scalar>& y,
                            const ExpressionPtr<Scalar>& x) {
  return fold_or_make<Atan2Expression<Scalar>>(y, x);
}
 
template <typename Scalar>
struct CbrtExpression final
    : NonLeafExpression<CbrtExpression<Scalar>, Scalar, 1> {
  using CbrtExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "cbrt";
 
  static Scalar eval(Scalar x) {
    using std::cbrt;
    return cbrt(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, const Value& val,
                        [[maybe_unused]] const Value& x) {
    g[0] = lit<Value>(Scalar(1)) / (lit<Value>(Scalar(3)) * val * val);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> cbrt(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<CbrtExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct CosExpression final
    : NonLeafExpression<CosExpression<Scalar>, Scalar, 1> {
  using CosExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "cos";
 
  static Scalar eval(Scalar x) {
    using std::cos;
    return cos(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::sin;
    g[0] = -sin(x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> cos(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<CosExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct CoshExpression final
    : NonLeafExpression<CoshExpression<Scalar>, Scalar, 1> {
  using CoshExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "cosh";
 
  static Scalar eval(Scalar x) {
    using std::cosh;
    return cosh(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::sinh;
    g[0] = sinh(x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> cosh(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<CoshExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct ErfExpression final
    : NonLeafExpression<ErfExpression<Scalar>, Scalar, 1> {
  using ErfExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "erf";
 
  static Scalar eval(Scalar x) {
    using std::erf;
    return erf(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::exp;
    g[0] = lit<Value>(Scalar(2.0 * std::numbers::inv_sqrtpi)) * exp(-x * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> erf(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<ErfExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct ExpExpression final
    : NonLeafExpression<ExpExpression<Scalar>, Scalar, 1> {
  using ExpExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "exp";
 
  static Scalar eval(Scalar x) {
    using std::exp;
    return exp(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, const Value& val,
                        [[maybe_unused]] const Value& x) {
    g[0] = val;
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> exp(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<ExpExpression<Scalar>>(x);
}
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y);
 
template <typename Scalar>
struct Hypot2Expression final
    : NonLeafExpression<Hypot2Expression<Scalar>, Scalar, 2> {
  using Hypot2Expression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "hypot";
 
  static Scalar eval(Scalar x, Scalar y) {
    using std::hypot;
    return hypot(x, y);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, const Value& val, const Value& x,
                        const Value& y) {
    g[0] = x / val;
    g[1] = y / val;
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y) {
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return abs(y);
  } else if (y->is_constant(Scalar(0))) {
    return abs(x);
  }
 
  return fold_or_make<Hypot2Expression<Scalar>>(x, y);
}
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y,
                            const ExpressionPtr<Scalar>& z);
 
template <typename Scalar>
struct Hypot3Expression final
    : NonLeafExpression<Hypot3Expression<Scalar>, Scalar, 3> {
  using Hypot3Expression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "hypot";
 
  static Scalar eval(Scalar x, Scalar y, Scalar z) {
    using std::hypot;
    return hypot(x, y, z);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, const Value& val, const Value& x,
                        const Value& y, const Value& z) {
    g[0] = x / val;
    g[1] = y / val;
    g[2] = z / val;
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                            const ExpressionPtr<Scalar>& y,
                            const ExpressionPtr<Scalar>& z) {
  // Prune expression
  if (x->is_constant(Scalar(0))) {
    return hypot(y, z);
  } else if (y->is_constant(Scalar(0))) {
    return hypot(x, z);
  } else if (z->is_constant(Scalar(0))) {
    return hypot(x, y);
  }
 
  return fold_or_make<Hypot3Expression<Scalar>>(x, y, z);
}
 
template <typename Scalar>
struct IfElseExpression final
    : NonLeafExpression<IfElseExpression<Scalar>, Scalar, 4> {
  static constexpr std::string_view NAME = "if-else";
 
  bool (*cond)(Scalar a, Scalar b);
 
  constexpr IfElseExpression(bool (*cond)(Scalar a, Scalar b),
                             ExpressionPtr<Scalar> a, ExpressionPtr<Scalar> b,
                             ExpressionPtr<Scalar> t, ExpressionPtr<Scalar> f)
      : IfElseExpression::NonLeafExpression{std::move(a), std::move(b),
                                            std::move(t), std::move(f)},
        cond{cond} {}
 
  Scalar eval(Scalar a, Scalar b, Scalar t, Scalar f) const {
    return cond(a, b) ? t : f;
  }
 
  template <typename Value>
  void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                 const Value& a, const Value& b,
                 [[maybe_unused]] const Value& t,
                 [[maybe_unused]] const Value& f) const {
    // a and b have zero gradient
    if constexpr (std::same_as<Value, Scalar>) {
      bool c = cond(a, b);
      g[2] = c ? Scalar(1) : Scalar(0);
      g[3] = c ? Scalar(0) : Scalar(1);
    } else {
      auto zero = constant_ptr(Scalar(0));
      auto one = constant_ptr(Scalar(1));
      g[2] = if_else(cond, a, b, one, zero);
      g[3] = if_else(cond, a, b, zero, one);
    }
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> if_else(
    std::type_identity_t<bool (*)(Scalar a, Scalar b)> cond,
    const ExpressionPtr<Scalar>& a, const ExpressionPtr<Scalar>& b,
    const ExpressionPtr<Scalar>& t, const ExpressionPtr<Scalar>& f) {
  using enum ExpressionType;
 
  // Prune expression
  if (t == f) {
    // Return t, which both branches currently are
    return t;
  } else if (t->type() == CONSTANT && f->type() == CONSTANT &&
             t->val == f->val) {
    // Return t, which both branches currently equal
    return t;
  }
 
  // Evaluate constant condition
  if (a->type() == CONSTANT && b->type() == CONSTANT) {
    return cond(a->val, b->val) ? t : f;
  }
 
  return make_expression_ptr<IfElseExpression<Scalar>>(cond, a, b, t, f);
}
 
template <typename Scalar>
struct LogExpression final
    : NonLeafExpression<LogExpression<Scalar>, Scalar, 1> {
  using LogExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "log";
 
  static Scalar eval(Scalar x) {
    using std::log;
    return log(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    g[0] = lit<Value>(Scalar(1)) / x;
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> log(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<LogExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct Log10Expression final
    : NonLeafExpression<Log10Expression<Scalar>, Scalar, 1> {
  using Log10Expression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "log10";
 
  static Scalar eval(Scalar x) {
    using std::log10;
    return log10(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    g[0] = lit<Value>(Scalar(1)) / (lit<Value>(Scalar(std::numbers::ln10)) * x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> log10(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<Log10Expression<Scalar>>(x);
}
 
template <typename Scalar>
struct MaxExpression final
    : NonLeafExpression<MaxExpression<Scalar>, Scalar, 2> {
  using MaxExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "max";
 
  static Scalar eval(Scalar a, Scalar b) {
    using std::max;
    return max(a, b);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& a, const Value& b) {
    if constexpr (std::same_as<Value, Scalar>) {
      bool c = a >= b;
      g[0] = c ? Scalar(1) : Scalar(0);
      g[1] = c ? Scalar(0) : Scalar(1);
    } else {
      auto zero = constant_ptr(Scalar(0));
      auto one = constant_ptr(Scalar(1));
      auto cond = [](Scalar a, Scalar b) { return a >= b; };
      g[0] = if_else(cond, a, b, one, zero);
      g[1] = if_else(cond, a, b, zero, one);
    }
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> max(const ExpressionPtr<Scalar>& a,
                          const ExpressionPtr<Scalar>& b) {
  return fold_or_make<MaxExpression<Scalar>>(a, b);
}
 
template <typename Scalar>
struct MinExpression final
    : NonLeafExpression<MinExpression<Scalar>, Scalar, 2> {
  using MinExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "min";
 
  static Scalar eval(Scalar a, Scalar b) {
    using std::min;
    return min(a, b);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& a, const Value& b) {
    if constexpr (std::same_as<Value, Scalar>) {
      bool c = a <= b;
      g[0] = c ? Scalar(1) : Scalar(0);
      g[1] = c ? Scalar(0) : Scalar(1);
    } else {
      auto zero = constant_ptr(Scalar(0));
      auto one = constant_ptr(Scalar(1));
      auto cond = [](Scalar a, Scalar b) { return a <= b; };
      g[0] = if_else(cond, a, b, one, zero);
      g[1] = if_else(cond, a, b, zero, one);
    }
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> min(const ExpressionPtr<Scalar>& a,
                          const ExpressionPtr<Scalar>& b) {
  return fold_or_make<MinExpression<Scalar>>(a, b);
}
 
template <typename Scalar>
ExpressionPtr<Scalar> pow(const ExpressionPtr<Scalar>& base,
                          const ExpressionPtr<Scalar>& power);
 
template <typename Scalar>
ExpressionPtr<Scalar> sqrt(const ExpressionPtr<Scalar>& x);
 
template <typename Scalar, ExpressionType T>
struct PowExpression final
    : NonLeafExpression<PowExpression<Scalar, T>, Scalar, 2, T> {
  using PowExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "pow";
 
  static Scalar eval(Scalar base, Scalar power) {
    using std::pow;
    return pow(base, power);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, const Value& val, const Value& base,
                        const Value& power) {
    using std::log;
    using std::pow;
 
    g[0] = pow(base, power - lit<Value>(Scalar(1))) * power;
    g[1] = val * log(base);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> pow(const ExpressionPtr<Scalar>& base,
                          const ExpressionPtr<Scalar>& power) {
  using enum ExpressionType;
  using std::pow;
 
  // Prune expression
  if (base->is_constant(Scalar(0))) {
    // Return zero, which base currently is
    return base;
  } else if (base->is_constant(Scalar(1))) {
    // Return one, which base currently is
    return base;
  }
  if (power->is_constant(Scalar(0))) {
    return constant_ptr(Scalar(1));
  } else if (power->is_constant(Scalar(1))) {
    // Return base unmodified
    return base;
  }
 
  // Evaluate constant
  if (base->type() == CONSTANT && power->type() == CONSTANT) {
    return constant_ptr(pow(base->val, power->val));
  }
 
  if (power->is_constant(Scalar(2))) {
    auto type = base->type() == LINEAR ? QUADRATIC : NONLINEAR;
    return make_typed_expression_ptr<MultExpression, Scalar>(type, base, base);
  } else if (power->is_constant(Scalar(0.5))) {
    return sqrt(base);
  }
 
  return make_expression_ptr<PowExpression<Scalar, NONLINEAR>>(base, power);
}
 
template <typename Scalar>
struct SignExpression final
    : NonLeafExpression<SignExpression<Scalar>, Scalar, 1> {
  using SignExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "sign";
 
  static Scalar eval(Scalar x) {
    if (x < Scalar(0)) {
      return Scalar(-1);
    } else if (x == Scalar(0)) {
      return Scalar(0);
    } else {
      return Scalar(1);
    }
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sign(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<SignExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct SinExpression final
    : NonLeafExpression<SinExpression<Scalar>, Scalar, 1> {
  using SinExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "sin";
 
  static Scalar eval(Scalar x) {
    using std::sin;
    return sin(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::cos;
    g[0] = cos(x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sin(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<SinExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct SinhExpression final
    : NonLeafExpression<SinhExpression<Scalar>, Scalar, 1> {
  using SinhExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "sinh";
 
  static Scalar eval(Scalar x) {
    using std::sinh;
    return sinh(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::cosh;
    g[0] = cosh(x);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sinh(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<SinhExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct SqrtExpression final
    : NonLeafExpression<SqrtExpression<Scalar>, Scalar, 1> {
  using SqrtExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "sqrt";
 
  static Scalar eval(Scalar x) {
    using std::sqrt;
    return sqrt(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, const Value& val,
                        [[maybe_unused]] const Value& x) {
    g[0] = lit<Value>(Scalar(1)) / (lit<Value>(Scalar(2)) * val);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> sqrt(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<SqrtExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct TanExpression final
    : NonLeafExpression<TanExpression<Scalar>, Scalar, 1> {
  using TanExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "tan";
 
  static Scalar eval(Scalar x) {
    using std::tan;
    return tan(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::cos;
 
    Value c = cos(x);
    g[0] = lit<Value>(Scalar(1)) / (c * c);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> tan(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<TanExpression<Scalar>>(x);
}
 
template <typename Scalar>
struct TanhExpression final
    : NonLeafExpression<TanhExpression<Scalar>, Scalar, 1> {
  using TanhExpression::NonLeafExpression::NonLeafExpression;
 
  static constexpr std::string_view NAME = "tanh";
 
  static Scalar eval(Scalar x) {
    using std::tanh;
    return tanh(x);
  }
 
  template <typename Value>
  static void eval_grad(ArgArray<Value>& g, [[maybe_unused]] const Value& val,
                        const Value& x) {
    using std::cosh;
 
    Value c = cosh(x);
    g[0] = lit<Value>(Scalar(1)) / (c * c);
  }
};
 
template <typename Scalar>
ExpressionPtr<Scalar> tanh(const ExpressionPtr<Scalar>& x) {
  return fold_or_make<TanhExpression<Scalar>>(x);
}
 
}  // namespace slp::detail