sleipnirgroup-jormungandr 0.6.7.dev13__tar.gz → 0.6.7.dev15__tar.gz

{sleipnirgroup_jormungandr-0.6.7.dev13 → sleipnirgroup_jormungandr-0.6.7.dev15}/CMakeLists.txt

@@ -111,6 +111,9 @@ if(PROJECT_IS_TOP_LEVEL AND BUILD_TESTING AND NOT CMAKE_CROSSCOMPILING)
         SYSTEM
     )
     FetchContent_MakeAvailable(Catch2)
+
+    # Enable StringMaker<std::string_view>
+    target_compile_features(Catch2 PRIVATE cxx_std_23)
 endif()
 
 install(

{sleipnirgroup_jormungandr-0.6.7.dev13 → sleipnirgroup_jormungandr-0.6.7.dev15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleipnirgroup-jormungandr
-Version: 0.6.7.dev13
+Version: 0.6.7.dev15
 Summary: Reverse mode autodiff library and NLP solver DSL
 License-Expression: BSD-3-Clause
 License-File: LICENSE.txt

{sleipnirgroup_jormungandr-0.6.7.dev13 → sleipnirgroup_jormungandr-0.6.7.dev15}/include/sleipnir/autodiff/expression.hpp

@@ -9,6 +9,7 @@
 #include <memory>
 #include <numbers>
 #include <string_view>
+#include <type_traits>
 #include <utility>
 
 #include <gch/small_vector.hpp>
@@ -1396,11 +1397,11 @@ template <typename Scalar>
 ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
                             const ExpressionPtr<Scalar>& y);
 
-/// Derived expression type for hypot().
+/// Derived expression type for hypot() with two arguments.
 ///
 /// @tparam Scalar Scalar type.
 template <typename Scalar>
-struct HypotExpression final : Expression<Scalar> {
+struct Hypot2Expression final : Expression<Scalar> {
   /// Binary operator's left operand.
   ExpressionPtr<Scalar> x;
 
@@ -1411,7 +1412,7 @@ struct HypotExpression final : Expression<Scalar> {
   ///
   /// @param x Binary operator's left operand.
   /// @param y Binary operator's right operand.
-  constexpr HypotExpression(ExpressionPtr<Scalar> x, ExpressionPtr<Scalar> y)
+  constexpr Hypot2Expression(ExpressionPtr<Scalar> x, ExpressionPtr<Scalar> y)
       : x{std::move(x)}, y{std::move(y)} {}
 
   void visit_args(
@@ -1430,32 +1431,16 @@ struct HypotExpression final : Expression<Scalar> {
   std::string_view name() const override { return "hypot"; }
 
   void accumulate_adjoints() const override {
-    x->adj += grad_l();
-    y->adj += grad_r();
-  }
-
-  void accumulate_adjoints_expr() const override {
-    x->adj_expr += grad_expr_l();
-    y->adj_expr += grad_expr_r();
-  }
-
- private:
-  Scalar grad_l() const {
-    using std::hypot;
-    return this->adj * x->val / hypot(x->val, y->val);
-  }
-
-  Scalar grad_r() const {
     using std::hypot;
-    return this->adj * y->val / hypot(x->val, y->val);
-  }
-
-  ExpressionPtr<Scalar> grad_expr_l() const {
-    return this->adj_expr * x / hypot(x, y);
+    Scalar norm = hypot(x->val, y->val);
+    x->adj += this->adj * x->val / norm;
+    y->adj += this->adj * y->val / norm;
   }
 
-  ExpressionPtr<Scalar> grad_expr_r() const {
-    return this->adj_expr * y / hypot(x, y);
+  void accumulate_adjoints_expr() const override {
+    auto norm = hypot(x, y);
+    x->adj_expr += this->adj_expr * x / norm;
+    y->adj_expr += this->adj_expr * y / norm;
   }
 };
 
@@ -1482,89 +1467,210 @@ ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
     return constant_ptr(hypot(x->val, y->val));
   }
 
-  return make_expression_ptr<HypotExpression<Scalar>>(x, y);
+  return make_expression_ptr<Hypot2Expression<Scalar>>(x, y);
 }
 
-/// Derived expression type for is_nonnegative().
+template <typename Scalar>
+ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
+                            const ExpressionPtr<Scalar>& y,
+                            const ExpressionPtr<Scalar>& z);
+
+/// Derived expression type for hypot() with three arguments.
 ///
 /// @tparam Scalar Scalar type.
 template <typename Scalar>
-struct IsNonnegativeExpression final : Expression<Scalar> {
-  /// Unary operator's operand.
+struct Hypot3Expression final : Expression<Scalar> {
+  /// Ternary operator's first operand.
   ExpressionPtr<Scalar> x;
 
-  /// Constructs an unary expression (an operator with one argument).
+  /// Ternary operator's second operand.
+  ExpressionPtr<Scalar> y;
+
+  /// Ternary operator's third operand.
+  ExpressionPtr<Scalar> z;
+
+  /// Constructs a ternary expression (an operator with three arguments).
   ///
-  /// @param x Unary operator's operand.
-  explicit constexpr IsNonnegativeExpression(ExpressionPtr<Scalar> x)
-      : x{std::move(x)} {}
+  /// @param x Ternary operator's first operand.
+  /// @param y Ternary operator's second operand.
+  /// @param z Ternary operator's third operand.
+  constexpr Hypot3Expression(ExpressionPtr<Scalar> x, ExpressionPtr<Scalar> y,
+                             ExpressionPtr<Scalar> z)
+      : x{std::move(x)}, y{std::move(y)}, z{std::move(z)} {}
 
   void visit_args(
       function_ref<void(Expression<Scalar>* arg)> func) const override {
     func(x.get());
+    func(y.get());
+    func(z.get());
   }
 
   Scalar value() const override {
-    return x->val >= Scalar(0) ? Scalar(1) : Scalar(0);
+    using std::hypot;
+    return hypot(x->val, y->val, z->val);
   }
 
   ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
-  std::string_view name() const override { return "is nonnegative"; }
+  std::string_view name() const override { return "hypot"; }
+
+  void accumulate_adjoints() const override {
+    using std::hypot;
+    Scalar norm = hypot(x->val, y->val, z->val);
+    x->adj += this->adj * x->val / norm;
+    y->adj += this->adj * y->val / norm;
+    z->adj += this->adj * z->val / norm;
+  }
+
+  void accumulate_adjoints_expr() const override {
+    auto norm = hypot(x, y, z);
+    x->adj_expr += this->adj_expr * x / norm;
+    y->adj_expr += this->adj_expr * y / norm;
+    z->adj_expr += this->adj_expr * z / norm;
+  }
 };
 
-/// Returns one if x is nonnegative and zero otherwise.
+/// hypot() for Expressions.
 ///
 /// @tparam Scalar Scalar type.
-/// @param x The argument.
+/// @param x The x argument.
+/// @param y The y argument.
+/// @param z The z argument.
 template <typename Scalar>
-ExpressionPtr<Scalar> is_nonnegative(const ExpressionPtr<Scalar>& x) {
-  if (x->type() == ExpressionType::CONSTANT) {
-    return constant_ptr(x->val >= Scalar(0) ? Scalar(1) : Scalar(0));
+ExpressionPtr<Scalar> hypot(const ExpressionPtr<Scalar>& x,
+                            const ExpressionPtr<Scalar>& y,
+                            const ExpressionPtr<Scalar>& z) {
+  using enum ExpressionType;
+  using std::hypot;
+
+  // 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 make_expression_ptr<IsNonnegativeExpression<Scalar>>(x);
+  // Evaluate constant
+  if (x->type() == CONSTANT && y->type() == CONSTANT && z->type() == CONSTANT) {
+    return constant_ptr(hypot(x->val, y->val, z->val));
+  }
+
+  return make_expression_ptr<Hypot3Expression<Scalar>>(x, y, z);
 }
 
-/// Derived expression type for is_positive().
+/// Derived expression type for if_else().
+///
+/// Returns t if cond(a, b) is true, otherwise f.
 ///
 /// @tparam Scalar Scalar type.
 template <typename Scalar>
-struct IsPositiveExpression final : Expression<Scalar> {
-  /// Unary operator's operand.
-  ExpressionPtr<Scalar> x;
+struct IfElseExpression final : Expression<Scalar> {
+  /// Condition evaluated on a and b.
+  bool (*cond)(Scalar a, Scalar b);
 
-  /// Constructs an unary expression (an operator with one argument).
+  /// The condition's first argument.
+  ExpressionPtr<Scalar> a;
+
+  /// The condition's second argument.
+  ExpressionPtr<Scalar> b;
+
+  /// Value selected when cond(a, b) is true.
+  ExpressionPtr<Scalar> t;
+
+  /// Value selected when cond(a, b) is false.
+  ExpressionPtr<Scalar> f;
+
+  /// Constructs an if-else expression.
   ///
-  /// @param x Unary operator's operand.
-  explicit constexpr IsPositiveExpression(ExpressionPtr<Scalar> x)
-      : x{std::move(x)} {}
+  /// @param cond Condition evaluated on a and b.
+  /// @param a The condition's first argument.
+  /// @param b The condition's second argument.
+  /// @param t Value selected when cond(a, b) is true.
+  /// @param f Value selected when cond(a, b) is false.
+  constexpr IfElseExpression(bool (*cond)(Scalar a, Scalar b),
+                             ExpressionPtr<Scalar> a, ExpressionPtr<Scalar> b,
+                             ExpressionPtr<Scalar> t, ExpressionPtr<Scalar> f)
+      : cond{cond},
+        a{std::move(a)},
+        b{std::move(b)},
+        t{std::move(t)},
+        f{std::move(f)} {}
 
   void visit_args(
       function_ref<void(Expression<Scalar>* arg)> func) const override {
-    func(x.get());
+    func(a.get());
+    func(b.get());
+    func(t.get());
+    func(f.get());
   }
 
   Scalar value() const override {
-    return x->val > Scalar(0) ? Scalar(1) : Scalar(0);
+    return cond(a->val, b->val) ? t->val : f->val;
   }
 
   ExpressionType type() const override { return ExpressionType::NONLINEAR; }
 
-  std::string_view name() const override { return "is positive"; }
+  std::string_view name() const override { return "if-else"; }
+
+  // a and b have zero gradient, so only t and f accumulate adjoints
+  void accumulate_adjoints() const override {
+    if (cond(a->val, b->val)) {
+      t->adj += this->adj;
+    } else {
+      f->adj += this->adj;
+    }
+  }
+
+  void accumulate_adjoints_expr() const override {
+    t->adj_expr += grad_expr_t();
+    f->adj_expr += grad_expr_f();
+  }
+
+ private:
+  ExpressionPtr<Scalar> grad_expr_t() const {
+    // adjoint if cond(a, b), otherwise 0
+    return if_else(cond, a, b, this->adj_expr, constant_ptr(Scalar(0)));
+  }
+
+  ExpressionPtr<Scalar> grad_expr_f() const {
+    // 0 if cond(a, b), otherwise adjoint
+    return if_else(cond, a, b, constant_ptr(Scalar(0)), this->adj_expr);
+  }
 };
 
-/// Returns one if x is positive and zero otherwise.
+/// Returns t if cond(a, b) is true, otherwise f.
 ///
 /// @tparam Scalar Scalar type.
-/// @param x The argument.
+/// @param cond Condition evaluated on a and b.
+/// @param a The condition's first argument.
+/// @param b The condition's second argument.
+/// @param t Value selected when cond(a, b) is true.
+/// @param f Value selected when cond(a, b) is false.
 template <typename Scalar>
-ExpressionPtr<Scalar> is_positive(const ExpressionPtr<Scalar>& x) {
-  if (x->type() == ExpressionType::CONSTANT) {
-    return constant_ptr(x->val > Scalar(0) ? Scalar(1) : Scalar(0));
+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<IsPositiveExpression<Scalar>>(x);
+  return make_expression_ptr<IfElseExpression<Scalar>>(cond, a, b, t, f);
 }
 
 /// Derived expression type for log().
@@ -1747,17 +1853,15 @@ struct MaxExpression final : Expression<Scalar> {
   Scalar grad_r() const { return a->val >= b->val ? Scalar(0) : this->adj; }
 
   ExpressionPtr<Scalar> grad_expr_l() const {
-    // adjoint * (a >= b)
-    // adjoint * (a - b >= 0)
-    return this->adj_expr * is_nonnegative(a - b);
+    // adjoint if a >= b, otherwise 0
+    return if_else([](Scalar a, Scalar b) { return a >= b; }, a, b,
+                   this->adj_expr, constant_ptr(Scalar(0)));
   }
 
   ExpressionPtr<Scalar> grad_expr_r() const {
-    // adjoint * !(a >= b)
-    // adjoint * (a < b)
-    // adjoint * (b > a)
-    // adjoint * (b - a > 0)
-    return this->adj_expr * is_positive(b - a);
+    // 0 if a >= b, otherwise adjoint
+    return if_else([](Scalar a, Scalar b) { return a >= b; }, a, b,
+                   constant_ptr(Scalar(0)), this->adj_expr);
   }
 };
 
@@ -1831,17 +1935,15 @@ struct MinExpression final : Expression<Scalar> {
   Scalar grad_r() const { return a->val <= b->val ? Scalar(0) : this->adj; }
 
   ExpressionPtr<Scalar> grad_expr_l() const {
-    // adjoint * (a <= b)
-    // adjoint * (b >= a)
-    // adjoint * (b - a >= 0)
-    return this->adj_expr * is_nonnegative(b - a);
+    // adjoint if a <= b, otherwise 0
+    return if_else([](Scalar a, Scalar b) { return a <= b; }, a, b,
+                   this->adj_expr, constant_ptr(Scalar(0)));
   }
 
   ExpressionPtr<Scalar> grad_expr_r() const {
-    // adjoint * !(a <= b)
-    // adjoint * (a > b)
-    // adjoint * (a - b > 0)
-    return this->adj_expr * is_positive(a - b);
+    // 0 if a <= b, otherwise adjoint
+    return if_else([](Scalar a, Scalar b) { return a <= b; }, a, b,
+                   constant_ptr(Scalar(0)), this->adj_expr);
   }
 };
 

{sleipnirgroup_jormungandr-0.6.7.dev13 → sleipnirgroup_jormungandr-0.6.7.dev15}/include/sleipnir/autodiff/variable.hpp

@@ -313,6 +313,11 @@ class Variable : public SleipnirBase {
   friend Variable<Scalar> hypot(const Variable<Scalar>& x,
                                 const Variable<Scalar>& y);
   template <typename Scalar>
+  friend Variable<Scalar> if_else(
+      std::type_identity_t<bool (*)(Scalar a, Scalar b)> cond,
+      const Variable<Scalar>& a, const Variable<Scalar>& b,
+      const Variable<Scalar>& t, const Variable<Scalar>& f);
+  template <typename Scalar>
   friend Variable<Scalar> log(const Variable<Scalar>& x);
   template <typename Scalar>
   friend Variable<Scalar> log10(const Variable<Scalar>& x);
@@ -522,6 +527,26 @@ Variable<Scalar> hypot(const Variable<Scalar>& x, const Variable<Scalar>& y) {
   return Variable{detail::hypot(x.expr, y.expr)};
 }
 
+/// if_else() for Variables.
+///
+/// Returns t if cond(a, b) is true, otherwise f. The condition is evaluated on
+/// the values of a and b each time the result is evaluated, and the gradient
+/// only flows through the selected branch.
+///
+/// @tparam Scalar Scalar type.
+/// @param cond Condition evaluated on a and b.
+/// @param a The condition's first argument.
+/// @param b The condition's second argument.
+/// @param t Value selected when cond(a, b) is true.
+/// @param f Value selected when cond(a, b) is false.
+template <typename Scalar>
+Variable<Scalar> if_else(
+    std::type_identity_t<bool (*)(Scalar a, Scalar b)> cond,
+    const Variable<Scalar>& a, const Variable<Scalar>& b,
+    const Variable<Scalar>& t, const Variable<Scalar>& f) {
+  return Variable{detail::if_else(cond, a.expr, b.expr, t.expr, f.expr)};
+}
+
 /// log() for Variables.
 ///
 /// @tparam Scalar Scalar type.
@@ -708,7 +733,7 @@ Variable<Scalar> tanh(const Variable<Scalar>& x) {
 template <typename Scalar>
 Variable<Scalar> hypot(const Variable<Scalar>& x, const Variable<Scalar>& y,
                        const Variable<Scalar>& z) {
-  return Variable{sqrt(pow(x, 2) + pow(y, 2) + pow(z, 2))};
+  return Variable{detail::hypot(x.expr, y.expr, z.expr)};
 }
 
 // The standard form for equality constraints is c(x) = 0, and the standard form

{sleipnirgroup_jormungandr-0.6.7.dev13 → sleipnirgroup_jormungandr-0.6.7.dev15}/python/cpp/docstrings.hpp

@@ -2811,6 +2811,25 @@ Template Args:
 
 )doc";
 
+static const char *mkd_doc_slp_if_else =
+R"doc(if_else() for Variables.
+
+Returns t if cond(a, b) is true, otherwise f. The condition is
+evaluated on the values of a and b each time the result is evaluated,
+and the gradient only flows through the selected branch.
+
+Args:
+    cond: Condition evaluated on a and b.
+    a: The condition's first argument.
+    b: The condition's second argument.
+    t: Value selected when cond(a, b) is true.
+    f: Value selected when cond(a, b) is false.
+
+Template Args:
+    Scalar: Scalar type.
+
+)doc";
+
 static const char *mkd_doc_slp_log =
 R"doc(log() for Variables.
 

{sleipnirgroup_jormungandr-0.6.7.dev13 → sleipnirgroup_jormungandr-0.6.7.dev15}/python/src/sleipnir/__init__.py

@@ -2,4 +2,4 @@
 
 from ._sleipnir import *
 
-__version__ = "0.6.7.dev13"
+__version__ = "0.6.7.dev15"