sleipnirgroup-jormungandr 0.6.7.dev128__tar.gz → 0.6.7.dev129__tar.gz

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev129}/PKG-INFO

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

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev129}/include/sleipnir/autodiff/variable_matrix.hpp

@@ -9,6 +9,7 @@
 #include <initializer_list>
 #include <iterator>
 #include <span>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -566,24 +567,6 @@ class VariableMatrix : public SleipnirBase {
     return result;
   }
 
-  /// Matrix-scalar multiplication operator.
-  ///
-  /// @param lhs Operator left-hand side.
-  /// @param rhs Operator right-hand side.
-  template <EigenMatrixLike LHS>
-  friend VariableMatrix<Scalar> operator*(const LHS& lhs,
-                                          const Variable<Scalar>& rhs) {
-    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
-
-    for (int row = 0; row < result.rows(); ++row) {
-      for (int col = 0; col < result.cols(); ++col) {
-        result[row, col] = lhs[row, col] * rhs;
-      }
-    }
-
-    return result;
-  }
-
   /// Matrix-scalar multiplication operator.
   ///
   /// @param lhs Operator left-hand side.
@@ -601,24 +584,6 @@ class VariableMatrix : public SleipnirBase {
     return result;
   }
 
-  /// Scalar-matrix multiplication operator.
-  ///
-  /// @param lhs Operator left-hand side.
-  /// @param rhs Operator right-hand side.
-  template <EigenMatrixLike RHS>
-  friend VariableMatrix<Scalar> operator*(const Variable<Scalar>& lhs,
-                                          const RHS& rhs) {
-    VariableMatrix<Scalar> result(detail::empty, rhs.rows(), rhs.cols());
-
-    for (int row = 0; row < result.rows(); ++row) {
-      for (int col = 0; col < result.cols(); ++col) {
-        result[row, col] = rhs[row, col] * lhs;
-      }
-    }
-
-    return result;
-  }
-
   /// Scalar-matrix multiplication operator.
   ///
   /// @param lhs Operator left-hand side.
@@ -671,25 +636,6 @@ class VariableMatrix : public SleipnirBase {
     return *this;
   }
 
-  /// Binary division operator.
-  ///
-  /// @param lhs Operator left-hand side.
-  /// @param rhs Operator right-hand side.
-  /// @return Result of division.
-  template <EigenMatrixLike LHS>
-  friend VariableMatrix<Scalar> operator/(const LHS& lhs,
-                                          const Variable<Scalar>& rhs) {
-    VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
-
-    for (int row = 0; row < result.rows(); ++row) {
-      for (int col = 0; col < result.cols(); ++col) {
-        result[row, col] = lhs[row, col] / rhs;
-      }
-    }
-
-    return result;
-  }
-
   /// Binary division operator.
   ///
   /// @param lhs Operator left-hand side.
@@ -1373,6 +1319,64 @@ template <typename Derived>
 VariableMatrix(const Eigen::DiagonalBase<Derived>&)
     -> VariableMatrix<typename Derived::Scalar>;
 
+/// Matrix-scalar multiplication operator.
+///
+/// @tparam Scalar Scalar type.
+/// @param lhs Operator left-hand side.
+/// @param rhs Operator right-hand side.
+template <typename Scalar, EigenMatrixLike LHS>
+  requires std::same_as<typename std::decay_t<LHS>::Scalar, Scalar>
+VariableMatrix<Scalar> operator*(const LHS& lhs, const Variable<Scalar>& rhs) {
+  VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      result[row, col] = lhs[row, col] * rhs;
+    }
+  }
+
+  return result;
+}
+
+/// Scalar-matrix multiplication operator.
+///
+/// @tparam Scalar Scalar type.
+/// @param lhs Operator left-hand side.
+/// @param rhs Operator right-hand side.
+template <typename Scalar, EigenMatrixLike RHS>
+  requires std::same_as<typename std::decay_t<RHS>::Scalar, Scalar>
+VariableMatrix<Scalar> operator*(const Variable<Scalar>& lhs, const RHS& rhs) {
+  VariableMatrix<Scalar> result(detail::empty, rhs.rows(), rhs.cols());
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      result[row, col] = rhs[row, col] * lhs;
+    }
+  }
+
+  return result;
+}
+
+/// Binary division operator.
+///
+/// @tparam Scalar Scalar type.
+/// @param lhs Operator left-hand side.
+/// @param rhs Operator right-hand side.
+/// @return Result of division.
+template <typename Scalar, EigenMatrixLike LHS>
+  requires std::same_as<typename std::decay_t<LHS>::Scalar, Scalar>
+VariableMatrix<Scalar> operator/(const LHS& lhs, const Variable<Scalar>& rhs) {
+  VariableMatrix<Scalar> result(detail::empty, lhs.rows(), lhs.cols());
+
+  for (int row = 0; row < result.rows(); ++row) {
+    for (int col = 0; col < result.cols(); ++col) {
+      result[row, col] = lhs[row, col] / rhs;
+    }
+  }
+
+  return result;
+}
+
 /// Applies a coefficient-wise reduce operation to two matrices.
 ///
 /// @tparam Scalar Scalar type.

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev129}/python/cpp/autodiff/bind_variable.cpp

@@ -3,8 +3,12 @@
 #include <nanobind/eigen/dense.h>
 #include <nanobind/nanobind.h>
 #include <nanobind/operators.h>
+#include <nanobind/stl/string.h>
 
+#include <algorithm>
+#include <array>
 #include <concepts>
+#include <string>
 #include <utility>
 
 #include <sleipnir/autodiff/variable.hpp>
@@ -13,6 +17,7 @@
 
 #include "docstrings.hpp"
 #include "for_each_type.hpp"
+#include "try_cast.hpp"
 
 namespace nb = nanobind;
 
@@ -116,6 +121,142 @@ void bind_variable(nb::module_& autodiff, nb::class_<Variable<double>>& cls) {
     }
   });
 
+  // Scalar-matrix multiplication and division with NumPy arrays. These are also
+  // handled by __array_ufunc__ below, but type checkers only see these
+  // overloads.
+  cls.def(
+      "__mul__",
+      [](const Variable<double>& lhs, nb::DRef<Eigen::MatrixXd> rhs) {
+        return lhs * rhs;
+      },
+      nb::is_operator(), "rhs"_a);
+  cls.def(
+      "__imul__",
+      [](const Variable<double>& lhs, nb::DRef<Eigen::MatrixXd> rhs) {
+        return lhs * rhs;
+      },
+      nb::is_operator(), "rhs"_a);
+  cls.def(
+      "__rmul__",
+      [](const Variable<double>& rhs, nb::DRef<Eigen::MatrixXd> lhs) {
+        return lhs * rhs;
+      },
+      nb::is_operator(), "lhs"_a);
+  cls.def(
+      "__rtruediv__",
+      [](const Variable<double>& rhs, nb::DRef<Eigen::MatrixXd> lhs) {
+        return lhs / rhs;
+      },
+      nb::is_operator(), "lhs"_a);
+
+  // Map NumPy ufunc reprs to the equivalent function names in Python's
+  // operator module
+  static constexpr std::array scalar_operators{
+      std::pair{"<ufunc 'add'>", "add"},
+      std::pair{"<ufunc 'subtract'>", "sub"},
+      std::pair{"<ufunc 'multiply'>", "mul"},
+      std::pair{"<ufunc 'divide'>", "truediv"},
+      std::pair{"<ufunc 'equal'>", "eq"},
+      std::pair{"<ufunc 'less'>", "lt"},
+      std::pair{"<ufunc 'less_equal'>", "le"},
+      std::pair{"<ufunc 'greater'>", "gt"},
+      std::pair{"<ufunc 'greater_equal'>", "ge"}};
+
+  // https://numpy.org/doc/stable/user/basics.dispatch.html
+  cls.def(
+      "__array_ufunc__",
+      [](Variable<double>& self, nb::object ufunc, nb::str method,
+         nb::args inputs, const nb::kwargs&) -> nb::object {
+        std::string method_name = nb::cast<std::string>(method);
+        std::string ufunc_name =
+            nb::cast<std::string>(ufunc.attr("__repr__")());
+
+        // VariableMatrix only represents 2D arrays
+        auto is_2d_array = [](const nb::handle& obj) {
+          return nb::isinstance<nb::ndarray<>>(obj) &&
+                 nb::cast<nb::ndarray<>>(obj).ndim() == 2;
+        };
+
+        if (method_name == "__call__" && inputs.size() == 2) {
+          if (auto op = std::ranges::find(
+                  scalar_operators, ufunc_name,
+                  [](const auto& pair) { return pair.first; });
+              op != scalar_operators.end()) {
+            auto np = nb::module_::import_("numpy");
+
+            auto is_0d = [&](const nb::handle& obj) {
+              // np.ndim() returns 0 for any non-array object
+              return !nb::isinstance<Variable<double>>(obj) &&
+                     nb::cast<int>(np.attr("ndim")(obj)) == 0;
+            };
+
+            // Convert NumPy scalars and 0D arrays to the equivalent Python
+            // scalar (e.g., np.float64 to float) and return other objects
+            // unchanged
+            auto numpy_to_python_scalar =
+                [&](const nb::handle& obj) -> nb::object {
+              nb::object numpy_types =
+                  nb::make_tuple(np.attr("generic"), np.attr("ndarray"));
+              if (PyObject_IsInstance(obj.ptr(), numpy_types.ptr()) == 1) {
+                return obj.attr("item")();
+              } else {
+                return nb::borrow(obj);
+              }
+            };
+
+            // If an operand is 0D, apply the Python operator to its Python
+            // scalar and the other operand, which dispatches to Variable's
+            // scalar overloads
+            if (is_0d(inputs[0])) {
+              return nb::module_::import_("operator")
+                  .attr(op->second)(numpy_to_python_scalar(inputs[0]),
+                                    inputs[1]);
+            } else if (is_0d(inputs[1])) {
+              return nb::module_::import_("operator")
+                  .attr(op->second)(inputs[0],
+                                    numpy_to_python_scalar(inputs[1]));
+            }
+          }
+
+          if (ufunc_name == "<ufunc 'multiply'>") {
+            if (is_2d_array(inputs[0])) {
+              if (auto result = apply_eigen_op(
+                      [&](auto& lhs) {
+                        return lhs.template cast<double>() * self;
+                      },
+                      inputs[0])) {
+                return result.value();
+              }
+            } else if (is_2d_array(inputs[1])) {
+              if (auto result = apply_eigen_op(
+                      [&](auto& rhs) {
+                        return self * rhs.template cast<double>();
+                      },
+                      inputs[1])) {
+                return result.value();
+              }
+            }
+          } else if (ufunc_name == "<ufunc 'divide'>") {
+            if (is_2d_array(inputs[0])) {
+              if (auto result = apply_eigen_op(
+                      [&](auto& lhs) {
+                        return lhs.template cast<double>() / self;
+                      },
+                      inputs[0])) {
+                return result.value();
+              }
+            }
+          }
+        }
+
+        return nb::not_implemented();
+      },
+      // nanobind can't annotate nb::args or nb::kwargs, so the stub needs an
+      // explicit signature
+      nb::sig("def __array_ufunc__(self, ufunc: object, method: str, "
+              "*inputs: object, **kwargs: object) -> object"),
+      "ufunc"_a, "method"_a, "inputs"_a, "kwargs"_a);
+
   cls.def(
       "__pow__",
       [](const Variable<double>& self, int power) { return pow(self, power); },

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev129}/python/cpp/docstrings.hpp

@@ -2959,6 +2959,21 @@ Template Args:
 
 )doc";
 
+constexpr char mkd_doc_slp_operator_div[] =
+R"doc(Binary division operator.
+
+Args:
+    lhs: Operator left-hand side.
+    rhs: Operator right-hand side.
+
+Template Args:
+    Scalar: Scalar type.
+
+Returns:
+    Result of division.
+
+)doc";
+
 constexpr char mkd_doc_slp_operator_eq[] =
 R"doc(Equality operator that returns an equality constraint for two
 Variables.
@@ -3049,6 +3064,30 @@ Args:
 
 )doc";
 
+constexpr char mkd_doc_slp_operator_mul[] =
+R"doc(Matrix-scalar multiplication operator.
+
+Args:
+    lhs: Operator left-hand side.
+    rhs: Operator right-hand side.
+
+Template Args:
+    Scalar: Scalar type.
+
+)doc";
+
+constexpr char mkd_doc_slp_operator_mul_2[] =
+R"doc(Scalar-matrix multiplication operator.
+
+Args:
+    lhs: Operator left-hand side.
+    rhs: Operator right-hand side.
+
+Template Args:
+    Scalar: Scalar type.
+
+)doc";
+
 constexpr char mkd_doc_slp_pow[] =
 R"doc(pow() for Variables.
 

{sleipnirgroup_jormungandr-0.6.7.dev128 → sleipnirgroup_jormungandr-0.6.7.dev129}/python/src/sleipnir/__init__.py

@@ -2,4 +2,4 @@
 
 from ._sleipnir import *
 
-__version__ = "0.6.7.dev128"
+__version__ = "0.6.7.dev129"

sleipnirgroup_jormungandr-0.6.7.dev129/python/test/autodiff/variable_test.py

@@ -0,0 +1,85 @@
+import numpy as np
+import pytest
+from sleipnir.autodiff import ExpressionType, Variable, VariableMatrix
+
+
+def test_default_constructor() -> None:
+    a = Variable()
+    assert a.value() == 0.0
+    assert a.type() == ExpressionType.LINEAR
+
+
+def test_constant_constructor() -> None:
+    # float
+    a = Variable(1.0)
+    assert a.value() == 1.0
+    assert a.type() == ExpressionType.CONSTANT
+
+    # int
+    b = Variable(2)
+    assert b.value() == 2
+    assert b.type() == ExpressionType.CONSTANT
+
+
+def test_set_value() -> None:
+    a = Variable()
+
+    # float
+    a.set_value(1.0)
+    assert a.value() == 1.0
+
+    # int
+    a.set_value(2)
+    assert a.value() == 2
+
+
+def test_numpy_array_operators() -> None:
+    x = Variable(2.0)
+    A = np.array([[1.0, 2.0], [3.0, 4.0]])
+
+    # float arrays
+    for result, expected in [
+        (x * A, [[2.0, 4.0], [6.0, 8.0]]),
+        (A * x, [[2.0, 4.0], [6.0, 8.0]]),
+        (A / x, [[0.5, 1.0], [1.5, 2.0]]),
+    ]:
+        assert isinstance(result, VariableMatrix)
+        assert np.array_equal(result.value(), expected)
+
+    # int arrays
+    B = np.array([[1, 2], [3, 4]])
+    for result, expected in [
+        (x * B, [[2.0, 4.0], [6.0, 8.0]]),
+        (B * x, [[2.0, 4.0], [6.0, 8.0]]),
+        (B / x, [[0.5, 1.0], [1.5, 2.0]]),
+    ]:
+        assert isinstance(result, VariableMatrix)
+        assert np.array_equal(result.value(), expected)
+
+    # NumPy scalars
+    a = np.float64(3.0)
+    for result, expected in [
+        (a + x, 5.0),
+        (x + a, 5.0),
+        (a - x, 1.0),
+        (x - a, -1.0),
+        (a * x, 6.0),
+        (x * a, 6.0),
+        (a / x, 1.5),
+        (x / a, 2.0 / 3.0),
+    ]:
+        assert isinstance(result, Variable)
+        assert result.value() == pytest.approx(expected)
+    # mypy rejects this because NumPy's stubs don't let np.integer defer to
+    # Variable.__rmul__, but at runtime NumPy calls Variable.__array_ufunc__
+    assert isinstance(np.int64(3) * x, Variable)  # type: ignore[operator]
+    assert isinstance(a <= x, type(Variable() <= 0.0))
+    assert isinstance(x <= a, type(Variable() <= 0.0))
+
+    # Other operations with NumPy arrays raise TypeError
+    with pytest.raises(TypeError):
+        _ = A + x
+    with pytest.raises(TypeError):
+        _ = A <= x
+    with pytest.raises(TypeError):
+        _ = np.array([1.0, 2.0]) * x

sleipnirgroup_jormungandr-0.6.7.dev128/python/test/autodiff/variable_test.py

@@ -1,31 +0,0 @@
-from sleipnir.autodiff import ExpressionType, Variable
-
-
-def test_default_constructor() -> None:
-    a = Variable()
-    assert a.value() == 0.0
-    assert a.type() == ExpressionType.LINEAR
-
-
-def test_constant_constructor() -> None:
-    # float
-    a = Variable(1.0)
-    assert a.value() == 1.0
-    assert a.type() == ExpressionType.CONSTANT
-
-    # int
-    b = Variable(2)
-    assert b.value() == 2
-    assert b.type() == ExpressionType.CONSTANT
-
-
-def test_set_value() -> None:
-    a = Variable()
-
-    # float
-    a.set_value(1.0)
-    assert a.value() == 1.0
-
-    # int
-    a.set_value(2)
-    assert a.value() == 2