Skip to content

Commit

Permalink
Use user-defined NumPy dtypes for AutoDiffXd, Expression, etc.
Browse files Browse the repository at this point in the history
  • Loading branch information
@EricCousineau-TRI
EricCousineau-TRI committed Apr 10, 2018
1 parent e949912 commit 5b552db
Show file tree
Hide file tree
Showing 12 changed files with 602 additions and 457 deletions.
13 changes: 13 additions & 0 deletions bindings/pydrake/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,11 +207,24 @@ drake_py_library(
srcs = ["test/algebra_test_util.py"],
)

drake_pybind_library(
name = "math_test_util_py",
testonly = 1,
add_install = False,
cc_deps = [
":autodiff_types_pybind",
],
cc_so_name = "test/math_test_util",
cc_srcs = ["test/math_test_util_py.cc"],
package_info = PACKAGE_INFO,
)

drake_py_unittest(
name = "autodiffutils_test",
deps = [
":algebra_test_util_py",
":autodiffutils_py",
":math_test_util_py",
],
)

Expand Down
5 changes: 3 additions & 2 deletions bindings/pydrake/autodiff_types_pybind.h
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
#pragma once

#include <Eigen/Core>
#include "pybind11/numpy_dtypes_user.h"
#include "pybind11/pybind11.h"

#include "drake/bindings/pydrake/pydrake_pybind.h"
#include "drake/common/autodiff.h"

// The macro `PYBIND11_NUMPY_OBJECT_DTYPE` place symbols into the namespace
// The macro `PYBIND11_NUMPY_DTYPE_USER` place symbols into the namespace
// `pybind11::detail`, so we should not place these in `drake::pydrake`.

PYBIND11_NUMPY_OBJECT_DTYPE(drake::AutoDiffXd);
PYBIND11_NUMPY_DTYPE_USER(drake::AutoDiffXd);
136 changes: 72 additions & 64 deletions bindings/pydrake/autodiffutils_py.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ PYBIND11_MODULE(_autodiffutils_py, m) {
py::module::import("pydrake.util.deprecation")
.attr("install_numpy_warning_filters")();

py::class_<AutoDiffXd> autodiff(m, "AutoDiffXd");
py::dtype_user<AutoDiffXd> autodiff(m, "AutoDiffXd");
autodiff
.def(py::init<const double&, const Eigen::VectorXd&>())
.def("value", [](const AutoDiffXd& self) {
Expand All @@ -41,72 +41,80 @@ PYBIND11_MODULE(_autodiffutils_py, m) {
self.value(), self.derivatives().size());
})
// Arithmetic
.def(-py::self)
.def(py::self + py::self)
.def(py::self + double())
.def(double() + py::self)
.def(py::self - py::self)
.def(py::self - double())
.def(double() - py::self)
.def(py::self * py::self)
.def(py::self * double())
.def(double() * py::self)
.def(py::self / py::self)
.def(py::self / double())
.def(double() / py::self)
.def_loop(-py::self)
.def_loop(py::self + py::self)
.def_loop(py::self + double())
.def_loop(double() + py::self)
.def_loop(py::self - py::self)
.def_loop(py::self - double())
.def_loop(double() - py::self)
.def_loop(py::self * py::self)
.def_loop(py::self * double())
.def_loop(double() * py::self)
.def_loop(py::self / py::self)
.def_loop(py::self / double())
.def_loop(double() / py::self)
// Logical comparison
.def(py::self == py::self)
.def(py::self == double())
.def(py::self != py::self)
.def(py::self != double())
.def(py::self < py::self)
.def(py::self < double())
.def(py::self <= py::self)
.def(py::self <= double())
.def(py::self > py::self)
.def(py::self > double())
.def(py::self >= py::self)
.def(py::self >= double())
// Additional math
.def("__pow__",
.def_loop(py::self == py::self)
.def_loop(py::self == double())
.def_loop(py::self != py::self)
.def_loop(py::self != double())
// .def_loop(double() != py::self)
.def_loop(py::self < py::self)
.def_loop(py::self < double())
.def_loop(double() < py::self)
.def_loop(py::self <= py::self)
.def_loop(py::self <= double())
.def_loop(double() <= py::self)
.def_loop(py::self > py::self)
.def_loop(py::self > double())
.def_loop(double() > py::self)
.def_loop(py::self >= py::self)
.def_loop(py::self >= double())
.def_loop(double() >= py::self)
// Dot-product
.def_loop(py::dtype_method::dot())
// Casting
// - Downcasting must be explicit, to prevent inadvertent information loss.
.def_loop_cast([](const AutoDiffXd& self) -> double {
return self.value();
})
// - Upcasting can be implicit, especially for matrix multiplication.
.def_loop_cast([](double x) -> AutoDiffXd { return x; }, true);

// Add overloads for `math` functions.
auto math = py::module::import("pydrake.math");
UfuncMirrorDef<decltype(autodiff)>(&autodiff, math)
.def_loop("__pow__", "pow",
[](const AutoDiffXd& base, int exponent) {
return pow(base, exponent);
}, py::is_operator())
.def("__abs__", [](const AutoDiffXd& x) { return abs(x); });

// Add overloads for `math` functions.
auto math = py::module::import("pydrake.math");
MirrorDef<py::module, decltype(autodiff)>(&math, &autodiff)
.def("log", [](const AutoDiffXd& x) { return log(x); })
.def("abs", [](const AutoDiffXd& x) { return abs(x); })
.def("exp", [](const AutoDiffXd& x) { return exp(x); })
.def("sqrt", [](const AutoDiffXd& x) { return sqrt(x); })
.def("pow", [](const AutoDiffXd& x, int y) {
return pow(x, y);
})
.def("sin", [](const AutoDiffXd& x) { return sin(x); })
.def("cos", [](const AutoDiffXd& x) { return cos(x); })
.def("tan", [](const AutoDiffXd& x) { return tan(x); })
.def("asin", [](const AutoDiffXd& x) { return asin(x); })
.def("acos", [](const AutoDiffXd& x) { return acos(x); })
.def("atan2", [](const AutoDiffXd& y, const AutoDiffXd& x) {
return atan2(y, x);
})
.def("sinh", [](const AutoDiffXd& x) { return sinh(x); })
.def("cosh", [](const AutoDiffXd& x) { return cosh(x); })
.def("tanh", [](const AutoDiffXd& x) { return tanh(x); })
.def("min", [](const AutoDiffXd& x, const AutoDiffXd& y) {
return min(x, y);
})
.def("max", [](const AutoDiffXd& x, const AutoDiffXd& y) {
return max(x, y);
})
.def("ceil", [](const AutoDiffXd& x) { return ceil(x); })
.def("floor", [](const AutoDiffXd& x) { return floor(x); });
// Mirror for numpy.
autodiff.attr("arcsin") = autodiff.attr("asin");
autodiff.attr("arccos") = autodiff.attr("acos");
autodiff.attr("arctan2") = autodiff.attr("atan2");
})
.def_loop("__abs__", "abs", [](const AutoDiffXd& x) { return abs(x); })
.def_loop("log", [](const AutoDiffXd& x) { return log(x); })
.def_loop("exp", [](const AutoDiffXd& x) { return exp(x); })
.def_loop("sqrt", [](const AutoDiffXd& x) { return sqrt(x); })
.def_loop("sin", [](const AutoDiffXd& x) { return sin(x); })
.def_loop("cos", [](const AutoDiffXd& x) { return cos(x); })
.def_loop("tan", [](const AutoDiffXd& x) { return tan(x); })
.def_loop("arcsin", "asin", [](const AutoDiffXd& x) { return asin(x); })
.def_loop("arccos", "acos", [](const AutoDiffXd& x) { return acos(x); })
.def_loop("arctan2", "atan2",
[](const AutoDiffXd& y, const AutoDiffXd& x) {
return atan2(y, x);
})
.def_loop("sinh", [](const AutoDiffXd& x) { return sinh(x); })
.def_loop("cosh", [](const AutoDiffXd& x) { return cosh(x); })
.def_loop("tanh", [](const AutoDiffXd& x) { return tanh(x); })
.def_loop("fmin", "min",
[](const AutoDiffXd& x, const AutoDiffXd& y) {
return min(x, y);
})
.def_loop("fmax", "max",
[](const AutoDiffXd& x, const AutoDiffXd& y) {
return max(x, y);
})
.def_loop("ceil", [](const AutoDiffXd& x) { return ceil(x); })
.def_loop("floor", [](const AutoDiffXd& x) { return floor(x); });
}

} // namespace pydrake
Expand Down
Loading

0 comments on commit 5b552db

Please sign in to comment.