variable.hpp

// Copyright (c) Sleipnir contributors
 
#pragma once
 
#include <algorithm>
#include <concepts>
#include <initializer_list>
#include <source_location>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <Eigen/Core>
 
#include "sleipnir/autodiff/expression.hpp"
#include "sleipnir/autodiff/expression_graph.hpp"
#include "sleipnir/util/assert.hpp"
#include "sleipnir/util/concepts.hpp"
#include "sleipnir/util/small_vector.hpp"
#include "sleipnir/util/symbol_exports.hpp"
 
#ifndef SLEIPNIR_DISABLE_DIAGNOSTICS
#include "sleipnir/util/print.hpp"
#endif
 
namespace slp {
 
// Forward declarations for friend declarations in Variable
namespace detail {
class AdjointExpressionGraph;
}  // namespace detail
template <int UpLo = Eigen::Lower | Eigen::Upper>
  requires(UpLo == Eigen::Lower) || (UpLo == (Eigen::Lower | Eigen::Upper))
class SLEIPNIR_DLLEXPORT Hessian;
class SLEIPNIR_DLLEXPORT Jacobian;
 
class SLEIPNIR_DLLEXPORT Variable {
 public:
  Variable() = default;
 
  explicit constexpr Variable(std::nullptr_t) : expr{nullptr} {}
 
  Variable(std::floating_point auto value)  // NOLINT
      : expr{detail::make_expression_ptr<detail::ConstExpression>(value)} {}
 
  Variable(std::integral auto value)  // NOLINT
      : expr{detail::make_expression_ptr<detail::ConstExpression>(value)} {}
 
  explicit Variable(const detail::ExpressionPtr& expr) : expr{expr} {}
 
  explicit constexpr Variable(detail::ExpressionPtr&& expr)
      : expr{std::move(expr)} {}
 
  Variable& operator=(double value) {
    expr = detail::make_expression_ptr<detail::ConstExpression>(value);
 
    return *this;
  }
 
  void set_value(double value) {
    if (expr->is_constant(0.0)) {
      expr = detail::make_expression_ptr<detail::ConstExpression>(value);
    } else {
#ifndef SLEIPNIR_DISABLE_DIAGNOSTICS
      // We only need to check the first argument since unary and binary
      // operators both use it
      if (expr->args[0] != nullptr) {
        auto location = std::source_location::current();
        slp::println(
            stderr,
            "WARNING: {}:{}: {}: Modified the value of a dependent variable",
            location.file_name(), location.line(), location.function_name());
      }
#endif
      expr->val = value;
    }
  }
 
  friend SLEIPNIR_DLLEXPORT Variable operator*(const Variable& lhs,
                                               const Variable& rhs) {
    return Variable{lhs.expr * rhs.expr};
  }
 
  Variable& operator*=(const Variable& rhs) {
    *this = *this * rhs;
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT Variable operator/(const Variable& lhs,
                                               const Variable& rhs) {
    return Variable{lhs.expr / rhs.expr};
  }
 
  Variable& operator/=(const Variable& rhs) {
    *this = *this / rhs;
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT Variable operator+(const Variable& lhs,
                                               const Variable& rhs) {
    return Variable{lhs.expr + rhs.expr};
  }
 
  Variable& operator+=(const Variable& rhs) {
    *this = *this + rhs;
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT Variable operator-(const Variable& lhs,
                                               const Variable& rhs) {
    return Variable{lhs.expr - rhs.expr};
  }
 
  Variable& operator-=(const Variable& rhs) {
    *this = *this - rhs;
    return *this;
  }
 
  friend SLEIPNIR_DLLEXPORT Variable operator-(const Variable& lhs) {
    return Variable{-lhs.expr};
  }
 
  friend SLEIPNIR_DLLEXPORT Variable operator+(const Variable& lhs) {
    return Variable{+lhs.expr};
  }
 
  double value() {
    if (!m_graph_initialized) {
      m_graph = detail::topological_sort(expr);
      m_graph_initialized = true;
    }
    detail::update_values(m_graph);
 
    return expr->val;
  }
 
  ExpressionType type() const { return expr->type(); }
 
 private:
  detail::ExpressionPtr expr =
      detail::make_expression_ptr<detail::DecisionVariableExpression>();
 
  small_vector<detail::Expression*> m_graph;
 
  bool m_graph_initialized = false;
 
  friend SLEIPNIR_DLLEXPORT Variable abs(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable acos(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable asin(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable atan(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable atan2(const Variable& y,
                                           const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable cos(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable cosh(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable erf(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable exp(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable hypot(const Variable& x,
                                           const Variable& y);
  friend SLEIPNIR_DLLEXPORT Variable log(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable log10(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable pow(const Variable& base,
                                         const Variable& power);
  friend SLEIPNIR_DLLEXPORT Variable sign(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable sin(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable sinh(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable sqrt(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable tan(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable tanh(const Variable& x);
  friend SLEIPNIR_DLLEXPORT Variable hypot(const Variable& x, const Variable& y,
                                           const Variable& z);
 
  friend class detail::AdjointExpressionGraph;
  template <int UpLo>
    requires(UpLo == Eigen::Lower) || (UpLo == (Eigen::Lower | Eigen::Upper))
  friend class SLEIPNIR_DLLEXPORT Hessian;
  friend class SLEIPNIR_DLLEXPORT Jacobian;
};
 
SLEIPNIR_DLLEXPORT inline Variable abs(const Variable& x) {
  return Variable{detail::abs(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable acos(const Variable& x) {
  return Variable{detail::acos(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable asin(const Variable& x) {
  return Variable{detail::asin(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable atan(const Variable& x) {
  return Variable{detail::atan(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable atan2(const Variable& y, const Variable& x) {
  return Variable{detail::atan2(y.expr, x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable cos(const Variable& x) {
  return Variable{detail::cos(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable cosh(const Variable& x) {
  return Variable{detail::cosh(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable erf(const Variable& x) {
  return Variable{detail::erf(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable exp(const Variable& x) {
  return Variable{detail::exp(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable hypot(const Variable& x, const Variable& y) {
  return Variable{detail::hypot(x.expr, y.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable pow(const Variable& base,
                                       const Variable& power) {
  return Variable{detail::pow(base.expr, power.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable log(const Variable& x) {
  return Variable{detail::log(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable log10(const Variable& x) {
  return Variable{detail::log10(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable sign(const Variable& x) {
  return Variable{detail::sign(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable sin(const Variable& x) {
  return Variable{detail::sin(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable sinh(const Variable& x) {
  return Variable{detail::sinh(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable sqrt(const Variable& x) {
  return Variable{detail::sqrt(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable tan(const Variable& x) {
  return Variable{detail::tan(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable tanh(const Variable& x) {
  return Variable{detail::tanh(x.expr)};
}
 
SLEIPNIR_DLLEXPORT inline Variable hypot(const Variable& x, const Variable& y,
                                         const Variable& z) {
  return Variable{slp::sqrt(slp::pow(x, 2) + slp::pow(y, 2) + slp::pow(z, 2))};
}
 
template <typename LHS, typename RHS>
  requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
          (ScalarLike<RHS> || MatrixLike<RHS>) &&
          (SleipnirType<LHS> || SleipnirType<RHS>)
small_vector<Variable> make_constraints(LHS&& lhs, RHS&& rhs) {
  small_vector<Variable> constraints;
 
  if constexpr (ScalarLike<LHS> && ScalarLike<RHS>) {
    constraints.emplace_back(lhs - rhs);
  } else if constexpr (ScalarLike<LHS> && MatrixLike<RHS>) {
    constraints.reserve(rhs.rows() * rhs.cols());
 
    for (int row = 0; row < rhs.rows(); ++row) {
      for (int col = 0; col < rhs.cols(); ++col) {
        // Make right-hand side zero
        if constexpr (EigenMatrixLike<std::decay_t<RHS>>) {
          constraints.emplace_back(lhs - rhs(row, col));
        } else {
          constraints.emplace_back(lhs - rhs[row, col]);
        }
      }
    }
  } else if constexpr (MatrixLike<LHS> && ScalarLike<RHS>) {
    constraints.reserve(lhs.rows() * lhs.cols());
 
    for (int row = 0; row < lhs.rows(); ++row) {
      for (int col = 0; col < lhs.cols(); ++col) {
        // Make right-hand side zero
        if constexpr (EigenMatrixLike<std::decay_t<LHS>>) {
          constraints.emplace_back(lhs(row, col) - rhs);
        } else {
          constraints.emplace_back(lhs[row, col] - rhs);
        }
      }
    }
  } else if constexpr (MatrixLike<LHS> && MatrixLike<RHS>) {
    slp_assert(lhs.rows() == rhs.rows() && lhs.cols() == rhs.cols());
    constraints.reserve(lhs.rows() * lhs.cols());
 
    for (int row = 0; row < lhs.rows(); ++row) {
      for (int col = 0; col < lhs.cols(); ++col) {
        // Make right-hand side zero
        if constexpr (EigenMatrixLike<std::decay_t<LHS>> &&
                      EigenMatrixLike<std::decay_t<RHS>>) {
          constraints.emplace_back(lhs(row, col) - rhs(row, col));
        } else if constexpr (EigenMatrixLike<std::decay_t<LHS>> &&
                             SleipnirMatrixLike<std::decay_t<RHS>>) {
          constraints.emplace_back(lhs(row, col) - rhs[row, col]);
        } else if constexpr (SleipnirMatrixLike<std::decay_t<LHS>> &&
                             EigenMatrixLike<std::decay_t<RHS>>) {
          constraints.emplace_back(lhs[row, col] - rhs(row, col));
        } else if constexpr (SleipnirMatrixLike<std::decay_t<LHS>> &&
                             SleipnirMatrixLike<std::decay_t<RHS>>) {
          constraints.emplace_back(lhs[row, col] - rhs[row, col]);
        }
      }
    }
  }
 
  return constraints;
}
 
struct SLEIPNIR_DLLEXPORT EqualityConstraints {
  small_vector<Variable> constraints;
 
  EqualityConstraints(
      std::initializer_list<EqualityConstraints> equality_constraints) {
    for (const auto& elem : equality_constraints) {
      constraints.insert(constraints.end(), elem.constraints.begin(),
                         elem.constraints.end());
    }
  }
 
  explicit EqualityConstraints(
      const std::vector<EqualityConstraints>& equality_constraints) {
    for (const auto& elem : equality_constraints) {
      constraints.insert(constraints.end(), elem.constraints.begin(),
                         elem.constraints.end());
    }
  }
 
  template <typename LHS, typename RHS>
    requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
            (ScalarLike<RHS> || MatrixLike<RHS>) &&
            (SleipnirType<LHS> || SleipnirType<RHS>)
  EqualityConstraints(LHS&& lhs, RHS&& rhs)
      : constraints{make_constraints(lhs, rhs)} {}
 
  operator bool() {  // NOLINT
    return std::ranges::all_of(constraints, [](auto& constraint) {
      return constraint.value() == 0.0;
    });
  }
};
 
struct SLEIPNIR_DLLEXPORT InequalityConstraints {
  small_vector<Variable> constraints;
 
  InequalityConstraints(
      std::initializer_list<InequalityConstraints> inequality_constraints) {
    for (const auto& elem : inequality_constraints) {
      constraints.insert(constraints.end(), elem.constraints.begin(),
                         elem.constraints.end());
    }
  }
 
  explicit InequalityConstraints(
      const std::vector<InequalityConstraints>& inequality_constraints) {
    for (const auto& elem : inequality_constraints) {
      constraints.insert(constraints.end(), elem.constraints.begin(),
                         elem.constraints.end());
    }
  }
 
  template <typename LHS, typename RHS>
    requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
            (ScalarLike<RHS> || MatrixLike<RHS>) &&
            (SleipnirType<LHS> || SleipnirType<RHS>)
  InequalityConstraints(LHS&& lhs, RHS&& rhs)
      : constraints{make_constraints(lhs, rhs)} {}
 
  operator bool() {  // NOLINT
    return std::ranges::all_of(constraints, [](auto& constraint) {
      return constraint.value() >= 0.0;
    });
  }
};
 
template <typename LHS, typename RHS>
  requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
          (ScalarLike<RHS> || MatrixLike<RHS>) &&
          (SleipnirType<LHS> || SleipnirType<RHS>)
EqualityConstraints operator==(LHS&& lhs, RHS&& rhs) {
  return EqualityConstraints{lhs, rhs};
}
 
template <typename LHS, typename RHS>
  requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
          (ScalarLike<RHS> || MatrixLike<RHS>) &&
          (SleipnirType<LHS> || SleipnirType<RHS>)
InequalityConstraints operator<(LHS&& lhs, RHS&& rhs) {
  return rhs >= lhs;
}
 
template <typename LHS, typename RHS>
  requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
          (ScalarLike<RHS> || MatrixLike<RHS>) &&
          (SleipnirType<LHS> || SleipnirType<RHS>)
InequalityConstraints operator<=(LHS&& lhs, RHS&& rhs) {
  return rhs >= lhs;
}
 
template <typename LHS, typename RHS>
  requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
          (ScalarLike<RHS> || MatrixLike<RHS>) &&
          (SleipnirType<LHS> || SleipnirType<RHS>)
InequalityConstraints operator>(LHS&& lhs, RHS&& rhs) {
  return lhs >= rhs;
}
 
template <typename LHS, typename RHS>
  requires(ScalarLike<LHS> || MatrixLike<LHS>) &&
          (ScalarLike<RHS> || MatrixLike<RHS>) &&
          (SleipnirType<LHS> || SleipnirType<RHS>)
InequalityConstraints operator>=(LHS&& lhs, RHS&& rhs) {
  return InequalityConstraints{lhs, rhs};
}
 
}  // namespace slp
 
namespace Eigen {
 
template <>
struct NumTraits<slp::Variable> : NumTraits<double> {
  using Real = slp::Variable;
  using NonInteger = slp::Variable;
  using Nested = slp::Variable;
 
  static constexpr int IsComplex = 0;
  static constexpr int IsInteger = 0;
  static constexpr int IsSigned = 1;
  static constexpr int RequireInitialization = 1;
  static constexpr int ReadCost = 1;
  static constexpr int AddCost = 3;
  static constexpr int MulCost = 3;
};
 
}  // namespace Eigen