Skip to content

Commit

Permalink
Add ability to create object matrices
Browse files Browse the repository at this point in the history
  • Loading branch information
@EricCousineau-TRI
Mmanu Chaturvedi authored and EricCousineau-TRI committed May 23, 2018
1 parent e763f04 commit 0da7ac7
Show file tree
Hide file tree
Showing 4 changed files with 251 additions and 22 deletions.
163 changes: 144 additions & 19 deletions include/pybind11/eigen.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,10 @@ struct eigen_extract_stride<Eigen::Map<PlainObjectType, MapOptions, StrideType>>
template <typename PlainObjectType, int Options, typename StrideType>
struct eigen_extract_stride<Eigen::Ref<PlainObjectType, Options, StrideType>> { using type = StrideType; };

template <typename Scalar> bool is_pyobject_() {
return static_cast<pybind11::detail::npy_api::constants>(npy_format_descriptor<Scalar>::value) == npy_api::NPY_OBJECT_;
}

// Helper struct for extracting information from an Eigen type
template <typename Type_> struct EigenProps {
using Type = Type_;
Expand Down Expand Up @@ -145,14 +149,19 @@ template <typename Type_> struct EigenProps {
const auto dims = a.ndim();
if (dims < 1 || dims > 2)
return false;

bool is_pyobject = false;
if (is_pyobject_<Scalar>())
is_pyobject = true;
ssize_t scalar_size = (is_pyobject ? static_cast<ssize_t>(sizeof(PyObject*)) :
static_cast<ssize_t>(sizeof(Scalar)));
if (dims == 2) { // Matrix type: require exact match (or dynamic)

EigenIndex
np_rows = a.shape(0),
np_cols = a.shape(1),
np_rstride = a.strides(0) / static_cast<ssize_t>(sizeof(Scalar)),
np_cstride = a.strides(1) / static_cast<ssize_t>(sizeof(Scalar));
np_rstride = a.strides(0) / scalar_size,
np_cstride = a.strides(1) / scalar_size;

if ((fixed_rows && np_rows != rows) || (fixed_cols && np_cols != cols))
return false;

Expand All @@ -162,7 +171,7 @@ template <typename Type_> struct EigenProps {
// Otherwise we're storing an n-vector. Only one of the strides will be used, but whichever
// is used, we want the (single) numpy stride value.
const EigenIndex n = a.shape(0),
stride = a.strides(0) / static_cast<ssize_t>(sizeof(Scalar));
stride = a.strides(0) / scalar_size;

if (vector) { // Eigen type is a compile-time vector
if (fixed && size != n)
Expand Down Expand Up @@ -213,11 +222,51 @@ template <typename Type_> struct EigenProps {
template <typename props> handle eigen_array_cast(typename props::Type const &src, handle base = handle(), bool writeable = true) {
constexpr ssize_t elem_size = sizeof(typename props::Scalar);
array a;
if (props::vector)
a = array({ src.size() }, { elem_size * src.innerStride() }, src.data(), base);
else
a = array({ src.rows(), src.cols() }, { elem_size * src.rowStride(), elem_size * src.colStride() },
src.data(), base);
using Scalar = typename props::Type::Scalar;
bool is_pyoject = static_cast<pybind11::detail::npy_api::constants>(npy_format_descriptor<Scalar>::value) == npy_api::NPY_OBJECT_;

if (!is_pyoject) {
if (props::vector)
a = array({ src.size() }, { elem_size * src.innerStride() }, src.data(), base);
else
a = array({ src.rows(), src.cols() }, { elem_size * src.rowStride(), elem_size * src.colStride() },
src.data(), base);
}
else {
if (props::vector) {
a = array(
npy_format_descriptor<Scalar>::dtype(),
{ (size_t) src.size() },
nullptr,
base
);
auto policy = base ? return_value_policy::automatic_reference : return_value_policy::copy;
for (ssize_t i = 0; i < src.size(); ++i) {
const Scalar src_val = props::fixed_rows ? src(0, i) : src(i, 0);
auto value_ = reinterpret_steal<object>(make_caster<Scalar>::cast(src_val, policy, base));
if (!value_)
return handle();
a.attr("itemset")(i, value_);
}
}
else {
a = array(
npy_format_descriptor<Scalar>::dtype(),
{(size_t) src.rows(), (size_t) src.cols()},
nullptr,
base
);
auto policy = base ? return_value_policy::automatic_reference : return_value_policy::copy;
for (ssize_t i = 0; i < src.rows(); ++i) {
for (ssize_t j = 0; j < src.cols(); ++j) {
auto value_ = reinterpret_steal<object>(make_caster<Scalar>::cast(src(i, j), policy, base));
if (!value_)
return handle();
a.attr("itemset")(i, j, value_);
}
}
}
}

if (!writeable)
array_proxy(a.ptr())->flags &= ~detail::npy_api::NPY_ARRAY_WRITEABLE_;
Expand Down Expand Up @@ -271,14 +320,46 @@ struct type_caster<Type, enable_if_t<is_eigen_dense_plain<Type>::value>> {
auto fits = props::conformable(buf);
if (!fits)
return false;

int result = 0;
// Allocate the new type, then build a numpy reference into it
value = Type(fits.rows, fits.cols);
auto ref = reinterpret_steal<array>(eigen_ref_array<props>(value));
if (dims == 1) ref = ref.squeeze();
else if (ref.ndim() == 1) buf = buf.squeeze();

int result = detail::npy_api::get().PyArray_CopyInto_(ref.ptr(), buf.ptr());
bool is_pyobject = is_pyobject_<Scalar>();

if (!is_pyobject) {
auto ref = reinterpret_steal<array>(eigen_ref_array<props>(value));
if (dims == 1) ref = ref.squeeze();
else if (ref.ndim() == 1) buf = buf.squeeze();
result =
detail::npy_api::get().PyArray_CopyInto_(ref.ptr(), buf.ptr());
}
else {
if (dims == 1) {
if (Type::RowsAtCompileTime == Eigen::Dynamic)
value.resize(buf.shape(0), 1);
if (Type::ColsAtCompileTime == Eigen::Dynamic)
value.resize(1, buf.shape(0));

for (ssize_t i = 0; i < buf.shape(0); ++i) {
make_caster <Scalar> conv_val;
if (!conv_val.load(buf.attr("item")(i).cast<pybind11::object>(), convert))
return false;
value(i) = cast_op<Scalar>(conv_val);
}
} else {
if (Type::RowsAtCompileTime == Eigen::Dynamic || Type::ColsAtCompileTime == Eigen::Dynamic) {
value.resize(buf.shape(0), buf.shape(1));
}
for (ssize_t i = 0; i < buf.shape(0); ++i) {
for (ssize_t j = 0; j < buf.shape(1); ++j) {
// p is the const void pointer to the item
make_caster<Scalar> conv_val;
if (!conv_val.load(buf.attr("item")(i,j).cast<pybind11::object>(), convert))
return false;
value(i,j) = cast_op<Scalar>(conv_val);
}
}
}
}

if (result < 0) { // Copy failed!
PyErr_Clear();
Expand Down Expand Up @@ -430,13 +511,19 @@ struct type_caster<
// storage order conversion. (Note that we refuse to use this temporary copy when loading an
// argument for a Ref<M> with M non-const, i.e. a read-write reference).
Array copy_or_ref;
typename std::remove_cv<PlainObjectType>::type val;
public:
bool load(handle src, bool convert) {
// First check whether what we have is already an array of the right type. If not, we can't
// avoid a copy (because the copy is also going to do type conversion).
bool need_copy = !isinstance<Array>(src);

EigenConformable<props::row_major> fits;
bool is_pyobject = false;
if (is_pyobject_<Scalar>()) {
is_pyobject = true;
need_copy = true;
}
if (!need_copy) {
// We don't need a converting copy, but we also need to check whether the strides are
// compatible with the Ref's stride requirements
Expand All @@ -459,15 +546,53 @@ struct type_caster<
// We need to copy: If we need a mutable reference, or we're not supposed to convert
// (either because we're in the no-convert overload pass, or because we're explicitly
// instructed not to copy (via `py::arg().noconvert()`) we have to fail loading.
if (!convert || need_writeable) return false;
if (!is_pyobject && (!convert || need_writeable)) {
return false;
}

Array copy = Array::ensure(src);
if (!copy) return false;
fits = props::conformable(copy);
if (!fits || !fits.template stride_compatible<props>())
if (!fits || !fits.template stride_compatible<props>()) {
return false;
copy_or_ref = std::move(copy);
loader_life_support::add_patient(copy_or_ref);
}

if (!is_pyobject) {
copy_or_ref = std::move(copy);
loader_life_support::add_patient(copy_or_ref);
}
else {
auto dims = copy.ndim();
if (dims == 1) {
if (Type::RowsAtCompileTime == Eigen::Dynamic || Type::ColsAtCompileTime == Eigen::Dynamic) {
val.resize(copy.shape(0), 1);
}
for (ssize_t i = 0; i < copy.shape(0); ++i) {
make_caster <Scalar> conv_val;
if (!conv_val.load(copy.attr("item")(i).template cast<pybind11::object>(),
convert))
return false;
val(i) = cast_op<Scalar>(conv_val);

}
} else {
if (Type::RowsAtCompileTime == Eigen::Dynamic || Type::ColsAtCompileTime == Eigen::Dynamic) {
val.resize(copy.shape(0), copy.shape(1));
}
for (ssize_t i = 0; i < copy.shape(0); ++i) {
for (ssize_t j = 0; j < copy.shape(1); ++j) {
// p is the const void pointer to the item
make_caster <Scalar> conv_val;
if (!conv_val.load(copy.attr("item")(i, j).template cast<pybind11::object>(),
convert))
return false;
val(i, j) = cast_op<Scalar>(conv_val);
}
}
}
ref.reset(new Type(val));
return true;
}
}

ref.reset();
Expand Down
15 changes: 15 additions & 0 deletions include/pybind11/numpy.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1230,6 +1230,21 @@ template <typename T, typename SFINAE> struct npy_format_descriptor {
(::std::vector<::pybind11::detail::field_descriptor> \
{PYBIND11_MAP2_LIST (PYBIND11_FIELD_DESCRIPTOR_EX, Type, __VA_ARGS__)})

#define PYBIND11_NUMPY_OBJECT_DTYPE(Type) \
namespace pybind11 { namespace detail { \
template <> struct npy_format_descriptor<Type> { \
public: \
enum { value = npy_api::NPY_OBJECT_ }; \
static pybind11::dtype dtype() { \
if (auto ptr = npy_api::get().PyArray_DescrFromType_(value)) { \
return reinterpret_borrow<pybind11::dtype>(ptr); \
} \
pybind11_fail("Unsupported buffer format!"); \
} \
static constexpr auto name = _("object"); \
}; \
}}

#endif // __CLION_IDE__

template <class T>
Expand Down
42 changes: 39 additions & 3 deletions tests/test_eigen.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,14 @@
#endif

#include <Eigen/Cholesky>
#include <unsupported/Eigen/AutoDiff>
#include "Eigen/src/Core/util/DisableStupidWarnings.h"

using MatrixXdR = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;


typedef Eigen::AutoDiffScalar<Eigen::VectorXd> ADScalar;
typedef Eigen::Matrix<ADScalar, Eigen::Dynamic, 1> VectorXADScalar;
typedef Eigen::Matrix<ADScalar, 1, Eigen::Dynamic> VectorXADScalarR;
PYBIND11_NUMPY_OBJECT_DTYPE(ADScalar);

// Sets/resets a testing reference matrix to have values of 10*r + c, where r and c are the
// (1-based) row/column number.
Expand Down Expand Up @@ -79,7 +83,9 @@ TEST_SUBMODULE(eigen, m) {
using FixedMatrixR = Eigen::Matrix<float, 5, 6, Eigen::RowMajor>;
using FixedMatrixC = Eigen::Matrix<float, 5, 6>;
using DenseMatrixR = Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
using DenseADScalarMatrixR = Eigen::Matrix<ADScalar, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;
using DenseMatrixC = Eigen::Matrix<float, Eigen::Dynamic, Eigen::Dynamic>;
using DenseADScalarMatrixC = Eigen::Matrix<ADScalar, Eigen::Dynamic, Eigen::Dynamic>;
using FourRowMatrixC = Eigen::Matrix<float, 4, Eigen::Dynamic>;
using FourColMatrixC = Eigen::Matrix<float, Eigen::Dynamic, 4>;
using FourRowMatrixR = Eigen::Matrix<float, 4, Eigen::Dynamic>;
Expand All @@ -91,10 +97,14 @@ TEST_SUBMODULE(eigen, m) {

// various tests
m.def("double_col", [](const Eigen::VectorXf &x) -> Eigen::VectorXf { return 2.0f * x; });
m.def("double_adscalar_col", [](const VectorXADScalar &x) -> VectorXADScalar { return 2.0f * x; });
m.def("double_row", [](const Eigen::RowVectorXf &x) -> Eigen::RowVectorXf { return 2.0f * x; });
m.def("double_adscalar_row", [](const VectorXADScalarR &x) -> VectorXADScalarR { return 2.0f * x; });
m.def("double_complex", [](const Eigen::VectorXcf &x) -> Eigen::VectorXcf { return 2.0f * x; });
m.def("double_threec", [](py::EigenDRef<Eigen::Vector3f> x) { x *= 2; });
m.def("double_adscalarc", [](py::EigenDRef<VectorXADScalar> x) { x *= 2; });
m.def("double_threer", [](py::EigenDRef<Eigen::RowVector3f> x) { x *= 2; });
m.def("double_adscalarr", [](py::EigenDRef<VectorXADScalarR> x) { x *= 2; });
m.def("double_mat_cm", [](Eigen::MatrixXf x) -> Eigen::MatrixXf { return 2.0f * x; });
m.def("double_mat_rm", [](DenseMatrixR x) -> DenseMatrixR { return 2.0f * x; });

Expand Down Expand Up @@ -139,6 +149,12 @@ TEST_SUBMODULE(eigen, m) {
return m;
}, py::return_value_policy::reference);

// Increments ADScalar Matrix
m.def("incr_adscalar_matrix", [](Eigen::Ref<DenseADScalarMatrixC> m, double v) {
m += DenseADScalarMatrixC::Constant(m.rows(), m.cols(), v);
return m;
}, py::return_value_policy::reference);

// Same, but accepts a matrix of any strides
m.def("incr_matrix_any", [](py::EigenDRef<Eigen::MatrixXd> m, double v) {
m += Eigen::MatrixXd::Constant(m.rows(), m.cols(), v);
Expand Down Expand Up @@ -173,12 +189,16 @@ TEST_SUBMODULE(eigen, m) {
// return value referencing/copying tests:
class ReturnTester {
Eigen::MatrixXd mat = create();
DenseADScalarMatrixR ad_mat = create_ADScalar_mat();
public:
ReturnTester() { print_created(this); }
~ReturnTester() { print_destroyed(this); }
static Eigen::MatrixXd create() { return Eigen::MatrixXd::Ones(10, 10); }
static Eigen::MatrixXd create() { return Eigen::MatrixXd::Ones(10, 10); }
static DenseADScalarMatrixR create_ADScalar_mat() { DenseADScalarMatrixR ad_mat(2, 2);
ad_mat << 1, 2, 3, 7; return ad_mat; }
static const Eigen::MatrixXd createConst() { return Eigen::MatrixXd::Ones(10, 10); }
Eigen::MatrixXd &get() { return mat; }
DenseADScalarMatrixR& get_ADScalarMat() {return ad_mat;}
Eigen::MatrixXd *getPtr() { return &mat; }
const Eigen::MatrixXd &view() { return mat; }
const Eigen::MatrixXd *viewPtr() { return &mat; }
Expand All @@ -197,6 +217,7 @@ TEST_SUBMODULE(eigen, m) {
.def_static("create", &ReturnTester::create)
.def_static("create_const", &ReturnTester::createConst)
.def("get", &ReturnTester::get, rvp::reference_internal)
.def("get_ADScalarMat", &ReturnTester::get_ADScalarMat, rvp::reference_internal)
.def("get_ptr", &ReturnTester::getPtr, rvp::reference_internal)
.def("view", &ReturnTester::view, rvp::reference_internal)
.def("view_ptr", &ReturnTester::view, rvp::reference_internal)
Expand All @@ -216,6 +237,18 @@ TEST_SUBMODULE(eigen, m) {
.def("corners_const", &ReturnTester::cornersConst, rvp::reference_internal)
;

py::class_<ADScalar>(m, "AutoDiffXd")
.def("__init__",
[](ADScalar & self,
double value,
const Eigen::VectorXd& derivatives) {
new (&self) ADScalar(value, derivatives);
})
.def("value", [](const ADScalar & self) {
return self.value();
})
;

// test_special_matrix_objects
// Returns a DiagonalMatrix with diagonal (1,2,3,...)
m.def("incr_diag", [](int k) {
Expand Down Expand Up @@ -300,6 +333,9 @@ TEST_SUBMODULE(eigen, m) {
m.def("iss1105_col", [](Eigen::VectorXd) { return true; });
m.def("iss1105_row", [](Eigen::RowVectorXd) { return true; });

m.def("iss1105_col_obj", [](VectorXADScalar) { return true; });
m.def("iss1105_row_obj", [](VectorXADScalarR) { return true; });

// test_named_arguments
// Make sure named arguments are working properly:
m.def("matrix_multiply", [](const py::EigenDRef<const Eigen::MatrixXd> A, const py::EigenDRef<const Eigen::MatrixXd> B)
Expand Down
Loading

0 comments on commit 0da7ac7

Please sign in to comment.