Skip to content

Commit

Permalink
Enable numpy math functions to work with symbolic variables. (#15039)
Browse files Browse the repository at this point in the history 
* Enable numpy math functions to work with symbolic variables.

#15038

Co-authored-by: Soonho Kong <[email protected]>
  • Loading branch information
@soonho-tri
soonho-tri authored May 15, 2021
1 parent e14fc42 commit 2d57227
Show file tree
Hide file tree
Showing 6 changed files with 88 additions and 73 deletions.
5 changes: 4 additions & 1 deletion bindings/pydrake/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,7 +221,10 @@ drake_cc_library(
hdrs = ["symbolic_types_pybind.h"],
declare_installed_headers = 0,
visibility = ["//visibility:public"],
deps = ["//:drake_shared_library"],
deps = [
"//:drake_shared_library",
"//bindings/pydrake:documentation_pybind",
],
)

drake_pybind_library(
Expand Down
3 changes: 2 additions & 1 deletion bindings/pydrake/math_py.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ namespace pydrake {

using std::pow;
using symbolic::Expression;
using symbolic::Variable;

namespace {
template <typename T>
Expand Down Expand Up @@ -474,7 +475,7 @@ void DoScalarIndependentDefinitions(py::module m) {

// See TODO in corresponding header file - these should be removed soon!
pydrake::internal::BindAutoDiffMathOverloads(&m);
pydrake::internal::BindSymbolicMathOverloads(&m);
pydrake::internal::BindSymbolicMathOverloads<Expression>(&m);
}
} // namespace

Expand Down
33 changes: 5 additions & 28 deletions bindings/pydrake/symbolic_py.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,7 @@ PYBIND11_MODULE(symbolic, m) {
.def(py::self != Expression())
.def(py::self != py::self)
.def(py::self != double());
internal::BindSymbolicMathOverloads<Variable>(&var_cls);
DefCopyAndDeepCopy(&var_cls);

// Bind the free functions for MakeVectorXXXVariable
Expand Down Expand Up @@ -406,36 +407,12 @@ PYBIND11_MODULE(symbolic, m) {
.def(py::self != double())
.def("Differentiate", &Expression::Differentiate,
doc.Expression.Differentiate.doc)
.def("Jacobian", &Expression::Jacobian, doc.Expression.Jacobian.doc)
// TODO(eric.cousineau): Figure out how to consolidate with the `math`
// methods.
.def("log", &symbolic::log, doc.log.doc)
.def("__abs__", &symbolic::abs)
.def("exp", &symbolic::exp, doc.exp.doc)
.def("sqrt", &symbolic::sqrt, doc.sqrt.doc)
// TODO(eric.cousineau): Move `__pow__` here.
.def("sin", &symbolic::sin, doc.sin.doc)
.def("cos", &symbolic::cos, doc.cos.doc)
.def("tan", &symbolic::tan, doc.tan.doc)
.def("arcsin", &symbolic::asin, doc.asin.doc)
.def("arccos", &symbolic::acos, doc.acos.doc)
.def("arctan2", &symbolic::atan2, doc.atan2.doc)
.def("sinh", &symbolic::sinh, doc.sinh.doc)
.def("cosh", &symbolic::cosh, doc.cosh.doc)
.def("tanh", &symbolic::tanh, doc.tanh.doc)
.def("min", &symbolic::min, doc.min.doc)
.def("max", &symbolic::max, doc.max.doc)
.def("ceil", &symbolic::ceil, doc.ceil.doc)
.def("__ceil__", &symbolic::ceil, doc.ceil.doc)
.def("floor", &symbolic::floor, doc.floor.doc)
.def("__floor__", &symbolic::floor, doc.floor.doc);
.def("Jacobian", &Expression::Jacobian, doc.Expression.Jacobian.doc);
// TODO(eric.cousineau): Clean this overload stuff up (#15041).
pydrake::internal::BindSymbolicMathOverloads<Expression>(&expr_cls);
pydrake::internal::BindSymbolicMathOverloads<Expression>(&m);
DefCopyAndDeepCopy(&expr_cls);

// TODO(eric.cousineau): These should actually exist on the class, and should
// be should be consolidated with the above repeated definitions. This would
// yield the same parity with AutoDiff.
pydrake::internal::BindSymbolicMathOverloads(&m);

m.def("if_then_else", &symbolic::if_then_else);

m.def(
Expand Down
62 changes: 36 additions & 26 deletions bindings/pydrake/symbolic_types_pybind.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
#include "pybind11/eigen.h"
#include "pybind11/pybind11.h"

#include "drake/bindings/pydrake/documentation_pybind.h"
#include "drake/bindings/pydrake/pydrake_pybind.h"
#include "drake/common/symbolic.h"

Expand All @@ -25,44 +26,53 @@ namespace internal {
// TODO(eric.cousineau): Deprecate these methods once we support proper NumPy
// UFuncs.

// TODO(eric.cousineau): Per TODO in `symbolic_py.cc`, ensure that this binds
// directly to the class, so it is the same as AutoDiff.
// Adds math function overloads for Expression (ADL free functions from
// `symbolic_expression.h`) for both NumPy methods and `pydrake.math`.
// @param obj
// This is used to register functions or overloads in either (a)
// This is used to register functions or overloads in either
// `pydrake.symbolic` or `pydrake.math`.
template <typename PyObject>
template <typename Class, typename PyObject>
void BindSymbolicMathOverloads(PyObject* obj) {
using symbolic::Expression;
constexpr auto& doc = pydrake_doc.drake.symbolic;
(*obj) // BR
.def("log", &symbolic::log)
.def("abs", &symbolic::abs)
.def("exp", &symbolic::exp)
.def("sqrt", &symbolic::sqrt)
.def("log", &symbolic::log, doc.log.doc)
.def("abs", &symbolic::abs, doc.abs.doc)
.def("__abs__", &symbolic::abs, doc.abs.doc)
.def("exp", &symbolic::exp, doc.exp.doc)
.def("sqrt", &symbolic::sqrt, doc.sqrt.doc)
.def("pow",
[](const Expression& base, const Expression& exponent) {
[](const Class& base, const Expression& exponent) {
return pow(base, exponent);
})
.def("sin", &symbolic::sin)
.def("cos", &symbolic::cos)
.def("tan", &symbolic::tan)
.def("asin", &symbolic::asin)
.def("acos", &symbolic::acos)
.def("atan", &symbolic::atan)
.def("atan2", &symbolic::atan2)
.def("sinh", &symbolic::sinh)
.def("cosh", &symbolic::cosh)
.def("tanh", &symbolic::tanh)
.def("min", &symbolic::min)
.def("max", &symbolic::max)
.def("ceil", &symbolic::ceil)
.def("floor", &symbolic::floor)
.def("sin", &symbolic::sin, doc.sin.doc)
.def("cos", &symbolic::cos, doc.cos.doc)
.def("tan", &symbolic::tan, doc.tan.doc)
.def("arcsin", &symbolic::asin, doc.asin.doc)
.def("asin", &symbolic::asin, doc.asin.doc)
.def("arccos", &symbolic::acos, doc.acos.doc)
.def("acos", &symbolic::acos, doc.acos.doc)
.def("arctan", &symbolic::atan, doc.atan.doc)
.def("atan", &symbolic::atan, doc.atan.doc)
.def("arctan2", &symbolic::atan2, doc.atan2.doc)
.def("atan2", &symbolic::atan2, doc.atan2.doc)
.def("sinh", &symbolic::sinh, doc.sinh.doc)
.def("cosh", &symbolic::cosh, doc.cosh.doc)
.def("tanh", &symbolic::tanh, doc.tanh.doc)
.def("min", &symbolic::min, doc.min.doc)
.def("max", &symbolic::max, doc.max.doc)
.def("ceil", &symbolic::ceil, doc.ceil.doc)
.def("__ceil__", &symbolic::ceil, doc.ceil.doc)
.def("floor", &symbolic::floor, doc.floor.doc)
.def("__floor__", &symbolic::floor, doc.floor.doc)
// TODO(eric.cousineau): This is not a NumPy-overridable method using
// dtype=object. Deprecate and move solely into `pydrake.math`.
.def("inv", [](const MatrixX<Expression>& X) -> MatrixX<Expression> {
return X.inverse();
});
.def(
"inv",
[](const MatrixX<Expression>& X) -> MatrixX<Expression> {
return X.inverse();
},
"Symbolic matrix inverse");
}

} // namespace internal
Expand Down
47 changes: 30 additions & 17 deletions bindings/pydrake/test/math_overloads_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ def supports(self, func):
supported = backwards_compat
if func.__name__ in backwards_compat:
# Check backwards compatibility.
assert hasattr(self.m, func.__name__)
assert hasattr(self.m, func.__name__), self.m.__name__
return func.__name__ in supported

def to_float(self, y_T):
Expand Down Expand Up @@ -158,22 +158,35 @@ def check_eval(functions, nargs):
args_T = list(map(overload.to_type, args_float))
# Check each supported function.
for f_drake, f_builtin in functions:
if not overload.supports(f_drake):
continue
debug_print(
"- Functions: ", qualname(f_drake), qualname(f_builtin))
y_builtin = f_builtin(*args_float)
y_float = f_drake(*args_float)
debug_print(" - - Float Eval:", repr(y_builtin), repr(y_float))
self.assertEqual(y_float, y_builtin)
self.assertIsInstance(y_float, float)
# Test method current overload, and ensure value is accurate.
y_T = f_drake(*args_T)
y_T_float = overload.to_float(y_T)
debug_print(" - - Overload Eval:", repr(y_T), repr(y_T_float))
self.assertIsInstance(y_T, overload.T)
# - Ensure the translated value is accurate.
self.assertEqual(y_T_float, y_float)
with self.subTest(function=f_drake.__name__, nargs=nargs):
if not overload.supports(f_drake):
continue
debug_print(
"- Functions: ",
qualname(f_drake),
qualname(f_builtin),
)
y_builtin = f_builtin(*args_float)
y_float = f_drake(*args_float)
debug_print(
" - - Float Eval:",
repr(y_builtin),
repr(y_float),
)
self.assertEqual(y_float, y_builtin)
self.assertIsInstance(y_float, float)
# Test method current overload, and ensure value is
# accurate.
y_T = f_drake(*args_T)
y_T_float = overload.to_float(y_T)
debug_print(
" - - Overload Eval:",
repr(y_T),
repr(y_T_float),
)
self.assertIsInstance(y_T, overload.T)
# - Ensure the translated value is accurate.
self.assertEqual(y_T_float, y_float)

debug_print("\n\nOverload: ", qualname(type(overload)))
float_overload = FloatOverloads()
Expand Down
11 changes: 11 additions & 0 deletions bindings/pydrake/test/symbolic_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,18 +141,29 @@ def test_logical(self):
def test_functions_with_variable(self):
numpy_compare.assert_equal(sym.abs(x), "abs(x)")
numpy_compare.assert_equal(sym.exp(x), "exp(x)")
numpy_compare.assert_equal(np.exp(x), "exp(x)")
numpy_compare.assert_equal(sym.sqrt(x), "sqrt(x)")
numpy_compare.assert_equal(np.sqrt(x), "sqrt(x)")
numpy_compare.assert_equal(sym.pow(x, y), "pow(x, y)")
numpy_compare.assert_equal(sym.sin(x), "sin(x)")
numpy_compare.assert_equal(np.sin(x), "sin(x)")
numpy_compare.assert_equal(sym.cos(x), "cos(x)")
numpy_compare.assert_equal(np.cos(x), "cos(x)")
numpy_compare.assert_equal(sym.tan(x), "tan(x)")
numpy_compare.assert_equal(np.tan(x), "tan(x)")
numpy_compare.assert_equal(sym.asin(x), "asin(x)")
numpy_compare.assert_equal(np.arcsin(x), "asin(x)")
numpy_compare.assert_equal(sym.acos(x), "acos(x)")
numpy_compare.assert_equal(np.arccos(x), "acos(x)")
numpy_compare.assert_equal(sym.atan(x), "atan(x)")
numpy_compare.assert_equal(np.arctan(x), "atan(x)")
numpy_compare.assert_equal(sym.atan2(x, y), "atan2(x, y)")
numpy_compare.assert_equal(sym.sinh(x), "sinh(x)")
numpy_compare.assert_equal(np.sinh(x), "sinh(x)")
numpy_compare.assert_equal(sym.cosh(x), "cosh(x)")
numpy_compare.assert_equal(np.cosh(x), "cosh(x)")
numpy_compare.assert_equal(sym.tanh(x), "tanh(x)")
numpy_compare.assert_equal(np.tanh(x), "tanh(x)")
numpy_compare.assert_equal(sym.min(x, y), "min(x, y)")
numpy_compare.assert_equal(sym.max(x, y), "max(x, y)")
numpy_compare.assert_equal(sym.ceil(x), "ceil(x)")
Expand Down

0 comments on commit 2d57227

Please sign in to comment.