diff --git a/include/nanobind/eigen/dense.h b/include/nanobind/eigen/dense.h index 412ad4f7..7a4e6d42 100644 --- a/include/nanobind/eigen/dense.h +++ b/include/nanobind/eigen/dense.h @@ -26,46 +26,73 @@ template using DMap = Eigen::Map; NAMESPACE_BEGIN(detail) template -constexpr int NumDimensions = bool(T::IsVectorAtCompileTime) ? 1 : 2; +constexpr int num_dimensions = bool(T::IsVectorAtCompileTime) ? 1 : 2; -template -using array_for_eigen_t = ndarray< +template struct StrideExtr { + using Type = Eigen::Stride<0, 0>; +}; + +template struct StrideExtr> { + using Type = StrideType; +}; + +template using Stride = typename StrideExtr::Type; + +/// Is true for Eigen types that are known at compile-time to hold contiguous memory only, which includes all specializations of Matrix and Array, +/// and specializations of Map and Ref with according stride types and shapes. A (compile-time) stride of 0 means "contiguous" to Eigen. +template constexpr bool requires_contig_memory = + (Stride::InnerStrideAtCompileTime == 0 || Stride::InnerStrideAtCompileTime == 1) && + (num_dimensions == 1 || + Stride::OuterStrideAtCompileTime == 0 || + Stride::OuterStrideAtCompileTime != Eigen::Dynamic && Stride::OuterStrideAtCompileTime == T::InnerSizeAtCompileTime); + +/// Is true for StrideTypes that can describe the contiguous memory layout of the plain Eigen type T. +template constexpr bool can_map_contig_memory = + (StrideType::InnerStrideAtCompileTime == 0 || StrideType::InnerStrideAtCompileTime == 1 || StrideType::InnerStrideAtCompileTime == Eigen::Dynamic) && + (num_dimensions == 1 || + StrideType::OuterStrideAtCompileTime == 0 || + StrideType::OuterStrideAtCompileTime == Eigen::Dynamic || + StrideType::OuterStrideAtCompileTime == T::InnerSizeAtCompileTime); + +/// Alias ndarray for a given Eigen type, to be used by type_caster::from_python, which calls type_caster>::from_python. +/// If the Eigen type is known at compile-time to handle contiguous memory only, then this alias makes type_caster>::from_python +/// either fail or provide an ndarray with contiguous memory, triggering a conversion if necessary and supported by flags. +/// Otherwise, this alias makes type_caster>::from_python either fail or provide an ndarray with arbitrary strides, +/// which need to be checked for compatibility then. There is no way to ask type_caster for specific strides other than c_contig and f_contig. +/// Hence, if an Eigen type requires non-contiguous strides (at compile-time) and type_caster> provides an ndarray with unsuitable strides (at run-time), +/// then type_caster::from_python just fails. Note, however, that this is rather unusual, since the default stride type of Map requires contiguous memory, +/// and the one of Ref requires a contiguous inner stride, while it can handle any outer stride. +template using array_for_eigen_t = ndarray< typename T::Scalar, numpy, std::conditional_t< - NumDimensions == 1, + num_dimensions == 1, shape<(size_t) T::SizeAtCompileTime>, shape<(size_t) T::RowsAtCompileTime, (size_t) T::ColsAtCompileTime>>, std::conditional_t< - T::InnerStrideAtCompileTime == Eigen::Dynamic, - any_contig, + requires_contig_memory, std::conditional_t< - T::IsRowMajor || NumDimensions == 1, + num_dimensions == 1 || T::IsRowMajor, c_contig, - f_contig - > - > ->; + f_contig>, + any_contig>>; /// Any kind of Eigen class -template constexpr bool is_eigen_v = -is_base_of_template_v; +template constexpr bool is_eigen_v = is_base_of_template_v; /// Detects Eigen::Array, Eigen::Matrix, etc. -template constexpr bool is_eigen_plain_v = -is_base_of_template_v; +template constexpr bool is_eigen_plain_v = is_base_of_template_v; /// Detect Eigen::SparseMatrix -template constexpr bool is_eigen_sparse_v = -is_base_of_template_v; +template constexpr bool is_eigen_sparse_v = is_base_of_template_v; /// Detects expression templates template constexpr bool is_eigen_xpr_v = is_eigen_v && !is_eigen_plain_v && !is_eigen_sparse_v && !std::is_base_of_v, T>; -template struct type_caster>> { +template struct type_caster && is_ndarray_scalar_v>> { using Scalar = typename T::Scalar; using NDArray = array_for_eigen_t; using NDArrayCaster = make_caster; @@ -77,17 +104,12 @@ template struct type_caster>> { if (!caster.from_python(src, flags, cleanup)) return false; const NDArray &array = caster.value; - - if constexpr (NumDimensions == 1) { + if constexpr (num_dimensions == 1) value.resize(array.shape(0)); - memcpy(value.data(), array.data(), - array.shape(0) * sizeof(Scalar)); - } else { + else value.resize(array.shape(0), array.shape(1)); - memcpy(value.data(), array.data(), - array.shape(0) * array.shape(1) * sizeof(Scalar)); - } - + // array_for_eigen_t ensures that array holds contiguous memory. + memcpy(value.data(), array.data(), array.size() * sizeof(Scalar)); return true; } @@ -100,10 +122,10 @@ template struct type_caster>> { } static handle from_cpp(const T &v, rv_policy policy, cleanup_list *cleanup) noexcept { - size_t shape[NumDimensions]; - int64_t strides[NumDimensions]; + size_t shape[num_dimensions]; + int64_t strides[num_dimensions]; - if constexpr (NumDimensions == 1) { + if constexpr (num_dimensions == 1) { shape[0] = v.size(); strides[0] = v.innerStride(); } else { @@ -148,7 +170,7 @@ template struct type_caster>> { policy == rv_policy::move ? rv_policy::reference : policy; object o = steal(NDArrayCaster::from_cpp( - NDArray(ptr, NumDimensions, shape, owner, strides), + NDArray(ptr, num_dimensions, shape, owner, strides), array_rv_policy, cleanup)); return o.release(); @@ -156,10 +178,11 @@ template struct type_caster>> { }; /// Caster for Eigen expression templates -template struct type_caster>> { +template struct type_caster && is_ndarray_scalar_v>> { using Array = Eigen::Array; using Caster = make_caster; + static constexpr bool IsClass = false; static constexpr auto Name = Caster::Name; template using Cast = T; @@ -174,25 +197,60 @@ template struct type_caster>> { /// Caster for Eigen::Map template -struct type_caster, enable_if_t>> { +struct type_caster, enable_if_t && is_ndarray_scalar_v>> { using Map = Eigen::Map; using NDArray = array_for_eigen_t; using NDArrayCaster = type_caster; + static constexpr bool IsClass = false; static constexpr auto Name = NDArrayCaster::Name; template using Cast = Map; NDArrayCaster caster; - bool from_python(handle src, uint8_t flags, - cleanup_list *cleanup) noexcept { - return caster.from_python(src, flags, cleanup); + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + // Conversions result in an Eigen::Map pointing into a temporary ndarray. + // If src is not a bound function argument, but e.g. an argument of cast, then this temporary would be destroyed upon returning from cast. + // Hence, conversions cannot be supported in this case. + // If src is a bound function argument, then cleanup would keep alive this temporary until returning from the bound function. + // Hence, conversions could be supported in this case, resulting in a bound function altering the Map without an effect on the Python side. + // This behaviour would be surprising, however, as bound functions expecting a Map most probably expect that Map to point into the caller's data. + // Hence, do not support conversions in any case. + return from_python_(src, flags & ~(uint8_t)cast_flags::convert, cleanup); + } + + bool from_python_(handle src, uint8_t flags, cleanup_list* cleanup) noexcept { + if (!caster.from_python(src, flags, cleanup)) + return false; + + // Check if StrideType can cope with the strides of caster.value. Avoid this check if their types guarantee that, anyway. + + // If requires_contig_memory is true, then StrideType is known at compile-time to only cope with contiguous memory. + // Then since caster.from_python has succeeded, caster.value now surely provides contiguous memory, and so its strides surely fit. + if constexpr (!requires_contig_memory) { + // A stride that is dynamic at compile-time copes with any stride at run-time. + if constexpr (StrideType::InnerStrideAtCompileTime != Eigen::Dynamic) { + // A stride of 0 at compile-time means "contiguous" to Eigen, which is always 1 for the inner stride. + int64_t expected_inner_stride = StrideType::InnerStrideAtCompileTime == 0 ? 1 : StrideType::InnerStrideAtCompileTime; + if (expected_inner_stride != (num_dimensions == 1 || !T::IsRowMajor ? caster.value.stride(0) : caster.value.stride(1))) + return false; + } + if constexpr (num_dimensions == 2 && StrideType::OuterStrideAtCompileTime != Eigen::Dynamic) { + int64_t expected_outer_stride = + StrideType::OuterStrideAtCompileTime == 0 + ? T::IsRowMajor ? caster.value.shape(1) : caster.value.shape(0) + : StrideType::OuterStrideAtCompileTime; + if (expected_outer_stride != (T::IsRowMajor ? caster.value.stride(0) : caster.value.stride(1))) + return false; + } + } + return true; } static handle from_cpp(const Map &v, rv_policy, cleanup_list *cleanup) noexcept { - size_t shape[NumDimensions]; - int64_t strides[NumDimensions]; + size_t shape[num_dimensions]; + int64_t strides[num_dimensions]; - if constexpr (NumDimensions == 1) { + if constexpr (num_dimensions == 1) { shape[0] = v.size(); strides[0] = v.innerStride(); } else { @@ -203,24 +261,40 @@ struct type_caster, enable_if_t, shape, handle(), strides), + NDArray((void *) v.data(), num_dimensions, shape, handle(), strides), rv_policy::reference, cleanup); } StrideType strides() const { - constexpr int IS = StrideType::InnerStrideAtCompileTime, - OS = StrideType::OuterStrideAtCompileTime; + constexpr int is = StrideType::InnerStrideAtCompileTime, + os = StrideType::OuterStrideAtCompileTime; int64_t inner = caster.value.stride(0), - outer = caster.value.stride(1); - (void) outer; + outer; + if constexpr (num_dimensions == 1) + outer = caster.value.shape(0); + else + outer = caster.value.stride(1); - if constexpr (T::IsRowMajor) + if constexpr (num_dimensions == 2 && T::IsRowMajor) std::swap(inner, outer); - if constexpr (std::is_same_v>) + // Compile-time strides of 0 must be passed as such to constructors of StrideType, to avoid assertions in Eigen. + if constexpr (is == 0) { + // Ensured by stride checks in from_python_: + // assert(inner == 1); + inner = 0; + } + + if constexpr (os == 0) { + // Ensured by stride checks in from_python_: + // assert(num_dimensions == 1 || outer == (T::IsRowMajor ? int64_t(caster.value.shape(1)) : int64_t(caster.value.shape(0)))); + outer = 0; + } + + if constexpr (std::is_same_v>) return StrideType(inner); - else if constexpr (std::is_same_v>) + else if constexpr (std::is_same_v>) return StrideType(outer); else return StrideType(outer, inner); @@ -228,28 +302,85 @@ struct type_caster, enable_if_t == 1) + return Map(t.data(), t.shape(0), strides()); + else + return Map(t.data(), t.shape(0), t.shape(1), strides()); } }; + /// Caster for Eigen::Ref template -struct type_caster, enable_if_t>> { +struct type_caster, enable_if_t && is_ndarray_scalar_v>> { using Ref = Eigen::Ref; using Map = Eigen::Map; + using DMap = Eigen::Map; using MapCaster = make_caster; - static constexpr auto Name = MapCaster::Name; + using DMapCaster = make_caster; + using DmapMatches = typename Eigen::internal::traits::template match::type; + static constexpr bool can_map_contig_mem = can_map_contig_memory; + static constexpr bool IsClass = false; + static constexpr auto Name = const_name>(DMapCaster::Name, MapCaster::Name); template using Cast = Ref; MapCaster caster; + DMapCaster dcaster; + + + /// In short: + /// - type_caster> supports no conversions, independent of flags. + /// - type_caster> + /// + supports stride conversions, independent of flags, except for uncommon strides. + /// + It additionally supports conversions to T::Scalar if flags say so, + /// and if either a cleanup_list is passed, or if Ref is guaranteed to map its own data. + /// + /// type_caster> supports stride conversions independent of flags, because if the intention was to not allow them, + /// then the bound function would most probably expect a Map instead of a Ref. + /// + /// Both Ref and Ref map data. + /// Like for Map, type_caster>::from_python does not support conversions, and for the same reasons. + /// But unlike Ref, instead of mapping external data, Ref may alternatively map data that it owns itself. + /// Ref then maps its member variable m_object, having copy-constructed it from the passed Eigen type. + /// The primary use case of Ref is as function argument that either maps the caller's data, or a suitably converted copy thereof. + /// Hence, unlike with Map and Ref, a Ref that maps a (converted) copy is intended, + /// and thus, type_caster>::from_python may support conversions. + /// It first calls the type_caster for matching strides, not supporting conversions. + /// If that fails, it calls the one for arbitrary strides. Since conversions to T::Scalar create a temporary ndarray, + /// conversions are supported only if flags say so, and if either a cleanup_list is passed (that keeps the temporary alive), + /// or if Ref is guaranteed to map its own data (having copied the temporary), which is ensured only if DmapMatches::value is false. + /// + /// Unfortunately, if src's scalar type needs to be converted, then the latter means that e.g. + /// cast>(src) succeeds, while + /// cast< DRef>(src) fails - + /// even though DRef would be expected to support a superset of the types supported by Ref. + /// + /// Ref::m_object holds contiguous memory, which Ref silently fails to map if this is impossible given StrideType + /// and the passed object's shape. If mapping fails, then Ref is left with mapping nullptr. + /// While this could be considered below, it is not done for efficiency reasons: + /// due to Ref's missing move constructor, its unusual copy constructor, and since C++ does not guarantee named return value optimizations, + /// the Ref would need to be created only for checking it, and created a second time for returning it, + /// which seems too costly for a Ref that owns its data. + /// Instead of checking thoroughly after construction, conversion fails if it is known at compile-time that mapping may fail, + /// even though it may actually succeed in some of these cases at run-time (e.g. StrideType::OuterStrideAtCompileTime==4, + /// and a row-major Matrix with a dynamic number of columns and 4 columns at run-time). + /// Once Ref defines a move constructor https://gitlab.com/libeigen/eigen/-/issues/2668, this restriction may be lifted. + bool from_python(handle src, uint8_t flags, cleanup_list *cleanup) noexcept { + if constexpr (std::is_const_v) + return caster.from_python(src, flags, cleanup) || + can_map_contig_mem && + dcaster.from_python_(src, (!DmapMatches::value || cleanup) ? flags : flags & ~(uint8_t)cast_flags::convert, cleanup); + else + return caster.from_python(src, flags, cleanup); + } - bool from_python(handle src, uint8_t flags, - cleanup_list *cleanup) noexcept { - return caster.from_python(src, flags, cleanup); + operator Ref() { + if constexpr (std::is_const_v) + if (dcaster.caster.value.is_valid()) + return Ref(dcaster.operator DMap()); + return Ref(caster.operator Map()); } - operator Ref() { return Ref(caster.operator Map()); } }; NAMESPACE_END(detail) diff --git a/include/nanobind/ndarray.h b/include/nanobind/ndarray.h index ca4ed334..e760caa4 100644 --- a/include/nanobind/ndarray.h +++ b/include/nanobind/ndarray.h @@ -78,9 +78,16 @@ struct tensorflow { }; struct pytorch { }; struct jax { }; +NAMESPACE_BEGIN(detail) + +template constexpr bool is_ndarray_scalar_v = +std::is_floating_point_v || std::is_integral_v; + +NAMESPACE_END(detail) + template constexpr dlpack::dtype dtype() { static_assert( - std::is_floating_point_v || std::is_integral_v, + detail::is_ndarray_scalar_v, "nanobind::dtype: T must be a floating point or integer variable!" ); diff --git a/tests/test_eigen.cpp b/tests/test_eigen.cpp index bf8c5b57..9a6f5214 100644 --- a/tests/test_eigen.cpp +++ b/tests/test_eigen.cpp @@ -6,37 +6,35 @@ namespace nb = nanobind; using namespace nb::literals; NB_MODULE(test_eigen_ext, m) { - m.def( - "addV3i_1", - [](const Eigen::Vector3i &a, - const Eigen::Vector3i &b) -> Eigen::Vector3i { return a + b; }, - "a"_a, "b"_a.noconvert()); - - m.def( - "addV3i_2", - [](const Eigen::RowVector3i &a, - const Eigen::RowVector3i &b) -> Eigen::RowVector3i { return a + b; }, - "a"_a, "b"_a.noconvert()); - - m.def( - "addV3i_3", - [](const Eigen::Ref &a, - const Eigen::Ref &b) -> Eigen::Vector3i { - return a + b; - }, - "a"_a, "b"_a.noconvert()); - - m.def( - "addV3i_4", - [](const Eigen::Array3i &a, - const Eigen::Array3i &b) -> Eigen::Array3i { return a + b; }, - "a"_a, "b"_a.noconvert()); - - m.def( - "addV3i_5", - [](const Eigen::Array3i &a, - const Eigen::Array3i &b) { return a + b; }, - "a"_a, "b"_a.noconvert()); + m.def("addV3i", + [](const Eigen::Vector3i &a, + const Eigen::Vector3i &b) -> Eigen::Vector3i { return a + b; }, + "a"_a, "b"_a.noconvert()); + + m.def("addR3i", + [](const Eigen::RowVector3i &a, + const Eigen::RowVector3i &b) -> Eigen::RowVector3i { return a + b; }, + "a"_a, "b"_a.noconvert()); + + m.def("addRefCnstV3i", + [](const Eigen::Ref &a, + const Eigen::Ref &b) -> Eigen::Vector3i { return a + b; }, + "a"_a, "b"_a.noconvert()); + + m.def("addRefCnstR3i", + [](const Eigen::Ref& a, + const Eigen::Ref& b) -> Eigen::RowVector3i { return a + b; }, + "a"_a, "b"_a.noconvert()); + + m.def("addA3i", + [](const Eigen::Array3i &a, + const Eigen::Array3i &b) -> Eigen::Array3i { return a + b; }, + "a"_a, "b"_a.noconvert()); + + m.def("addA3i_retExpr", + [](const Eigen::Array3i &a, + const Eigen::Array3i &b) { return a + b; }, + "a"_a, "b"_a.noconvert()); m.def("addVXi", [](const Eigen::VectorXi &a, @@ -47,62 +45,111 @@ NB_MODULE(test_eigen_ext, m) { using MatrixXuC = Eigen::Matrix; using MatrixXuR = Eigen::Matrix; - m.def("addM4u_1", + m.def("addM4uCC", [](const Matrix4uC &a, const Matrix4uC &b) -> Matrix4uC { return a + b; }); - m.def("addMXu_1", + m.def("addMXuCC", [](const MatrixXuC &a, const MatrixXuC &b) -> MatrixXuC { return a + b; }); - m.def("addMXu_1_nc", + m.def("addMXuCC_nc", [](const MatrixXuC &a, const MatrixXuC &b) -> MatrixXuC { return a + b; }, "a"_a.noconvert(), "b"_a.noconvert()); - m.def("addM4u_2", + m.def("addM4uRR", [](const Matrix4uR &a, const Matrix4uR &b) -> Matrix4uR { return a + b; }); - m.def("addMXu_2", + m.def("addMXuRR", [](const MatrixXuR &a, const MatrixXuR &b) -> MatrixXuR { return a + b; }); - m.def("addMXu_2_nc", + m.def("addMXuRR_nc", [](const MatrixXuR &a, const MatrixXuR &b) -> MatrixXuR { return a + b; }, "a"_a.noconvert(), "b"_a.noconvert()); - m.def("addM4u_3", + m.def("addM4uCR", [](const Matrix4uC &a, const Matrix4uR &b) -> Matrix4uC { return a + b; }); - m.def("addMXu_3", + m.def("addMXuCR", [](const MatrixXuC &a, const MatrixXuR &b) -> MatrixXuC { return a + b; }); - m.def("addM4u_4", + m.def("addM4uRC", [](const Matrix4uR &a, const Matrix4uC &b) -> Matrix4uR { return a + b; }); - m.def("addMXu_4", + m.def("addMXuRC", [](const MatrixXuR &a, const MatrixXuC &b) -> MatrixXuR { return a + b; }); - m.def("addMXu_5", - [](const nb::DRef &a, - const nb::DRef &b) -> MatrixXuC { return a + b; }, + m.def("addMapMXuCC", + [](const Eigen::Map& a, + const Eigen::Map& b) -> MatrixXuC { return a + b; }); + + m.def("addMapCnstMXuCC", + [](const Eigen::Map& a, + const Eigen::Map& b) -> MatrixXuC { return a + b; }); + + m.def("addMapMXuRR", + [](const Eigen::Map& a, + const Eigen::Map& b) -> MatrixXuC { return a + b; }); + + m.def("addMapCnstMXuRR", + [](const Eigen::Map& a, + const Eigen::Map& b) -> MatrixXuC { return a + b; }); + + m.def("addRefMXuCC", + [](const Eigen::Ref& a, + const Eigen::Ref& b) -> MatrixXuC { return a + b; }); + + m.def("addRefCnstMXuCC", + [](const Eigen::Ref& a, + const Eigen::Ref& b) -> MatrixXuC { return a + b; }); + + m.def("addRefCnstMXuCC_nc", + [](const Eigen::Ref& a, + const Eigen::Ref& b) -> MatrixXuC { return a + b; }, "a"_a.noconvert(), "b"_a.noconvert()); - m.def("mutate_MXu", [](nb::DRef a) { a *= 2; }, nb::arg().noconvert()); + m.def("addRefMXuRR", + [](const Eigen::Ref& a, + const Eigen::Ref& b) -> MatrixXuC { return a + b; }); - m.def("updateV3i", [](Eigen::Ref a) { a[2] = 123; }); - m.def("updateVXi", [](Eigen::Ref a) { a[2] = 123; }); + m.def("addRefCnstMXuRR", + [](const Eigen::Ref& a, + const Eigen::Ref& b) -> MatrixXuC { return a + b; }); + + m.def("addRefCnstMXuRR_nc", + [](const Eigen::Ref& a, + const Eigen::Ref& b) -> MatrixXuC { return a + b; }, + "a"_a.noconvert(), "b"_a.noconvert()); + + m.def("addDRefMXuCC_nc", + [](const nb::DRef &a, + const nb::DRef &b) -> MatrixXuC { return a + b; }, + "a"_a.noconvert(), "b"_a.noconvert()); + + m.def("addDRefMXuRR_nc", + [](const nb::DRef& a, + const nb::DRef& b) -> MatrixXuC { return a + b; }, + "a"_a.noconvert(), "b"_a.noconvert()); + + m.def("mutate_DRefMXuC", [](nb::DRef a) { a *= 2; }, nb::arg().noconvert()); + + m.def("updateRefV3i", [](Eigen::Ref a) { a[2] = 123; }); + m.def("updateRefV3i_nc", [](Eigen::Ref a) { a[2] = 123; }, nb::arg().noconvert()); + m.def("updateRefVXi", [](Eigen::Ref a) { a[2] = 123; }); + m.def("updateRefVXi_nc", [](Eigen::Ref a) { a[2] = 123; }, nb::arg().noconvert()); using SparseMatrixR = Eigen::SparseMatrix; using SparseMatrixC = Eigen::SparseMatrix; Eigen::MatrixXf mat(5, 6); mat << - 0, 3, 0, 0, 0, 11, - 22, 0, 0, 0, 17, 11, - 7, 5, 0, 1, 0, 11, - 0, 0, 0, 0, 0, 11, - 0, 0, 14, 0, 8, 11; + 0, 3, 0, 0, 0, 11, + 22, 0, 0, 0, 17, 11, + 7, 5, 0, 1, 0, 11, + 0, 0, 0, 0, 0, 11, + 0, 0, 14, 0, 8, 11; m.def("sparse_r", [mat]() -> SparseMatrixR { return Eigen::SparseView(mat); }); @@ -114,7 +161,8 @@ NB_MODULE(test_eigen_ext, m) { m.def("sparse_r_uncompressed", []() -> SparseMatrixR { SparseMatrixR m(2,2); m.coeffRef(0,0) = 1.0f; - return m; + assert(!m.isCompressed()); + return m.markAsRValue(); }); /// issue #166 @@ -150,4 +198,21 @@ NB_MODULE(test_eigen_ext, m) { nb::class_(m, "ClassWithEigenMember") .def(nb::init<>()) .def_rw("member", &ClassWithEigenMember::member); + + m.def("castToMapVXi", [](nb::object obj) -> Eigen::Map { + return nb::cast>(obj); + }); + m.def("castToRefVXi", [](nb::object obj) -> Eigen::VectorXi { + return nb::cast>(obj); + }); + m.def("castToRefCnstVXi", [](nb::object obj) -> Eigen::VectorXi { + return nb::cast>(obj); + }); + m.def("castToDRefCnstVXi", [](nb::object obj) -> Eigen::VectorXi { + return nb::cast>(obj); + }); + m.def("castToRef03CnstVXi", [](nb::object obj) -> Eigen::VectorXi { + return nb::cast>>(obj); + }); + } diff --git a/tests/test_eigen.py b/tests/test_eigen.py index 4d3c8787..d0e78272 100644 --- a/tests/test_eigen.py +++ b/tests/test_eigen.py @@ -1,16 +1,19 @@ import pytest import gc +import itertools import re import sys try: import numpy as np + from numpy.testing import assert_array_equal import test_eigen_ext as t def needs_numpy_and_eigen(x): return x except: needs_numpy_and_eigen = pytest.mark.skip(reason="NumPy and Eigen are required") + @needs_numpy_and_eigen def test01_vector_fixed(): a = np.array([1, 2, 3], dtype=np.int32) @@ -20,63 +23,59 @@ def test01_vector_fixed(): af = np.float32(a) bf = np.float32(b) - assert np.all(t.addV3i_1(a, b) == c) - assert np.all(t.addV3i_2(a, b) == c) - assert np.all(t.addV3i_3(a, b) == c) - assert np.all(t.addV3i_4(a, b) == c) - assert np.all(t.addV3i_5(a, b) == c) + assert_array_equal(t.addV3i(a, b), c) + assert_array_equal(t.addR3i(a, b), c) + assert_array_equal(t.addRefCnstV3i(a, b), c) + assert_array_equal(t.addRefCnstR3i(a, b), c) + assert_array_equal(t.addA3i(a, b), c) + assert_array_equal(t.addA3i_retExpr(a, b), c) # Implicit conversion supported for first argument - assert np.all(t.addV3i_1(af, b) == c) - assert np.all(t.addV3i_2(af, b) == c) - assert np.all(t.addV3i_3(af, b) == c) - assert np.all(t.addV3i_4(af, b) == c) + assert_array_equal(t.addV3i(af, b), c) + assert_array_equal(t.addR3i(af, b), c) + assert_array_equal(t.addRefCnstV3i(af, b), c) + assert_array_equal(t.addRefCnstR3i(af, b), c) + assert_array_equal(t.addA3i(af, b), c) # But not the second one - with pytest.raises(TypeError) as e: - t.addV3i_1(a, bf) - assert 'incompatible function arguments' in str(e) - with pytest.raises(TypeError) as e: - t.addV3i_2(a, bf) - assert 'incompatible function arguments' in str(e) - with pytest.raises(TypeError) as e: - t.addV3i_3(a, bf) - assert 'incompatible function arguments' in str(e) - with pytest.raises(TypeError) as e: - t.addV3i_4(a, bf) - assert 'incompatible function arguments' in str(e) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addV3i(a, bf) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addR3i(a, bf) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addRefCnstV3i(a, bf) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addRefCnstR3i(a, bf) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addA3i(a, bf) # Catch size errors - with pytest.raises(TypeError) as e: - t.addV3i_1(x, b) - assert 'incompatible function arguments' in str(e) - with pytest.raises(TypeError) as e: - t.addV3i_2(x, b) - assert 'incompatible function arguments' in str(e) - with pytest.raises(TypeError) as e: - t.addV3i_3(x, b) - assert 'incompatible function arguments' in str(e) - with pytest.raises(TypeError) as e: - t.addV3i_4(x, b) - assert 'incompatible function arguments' in str(e) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addV3i(x, b) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addR3i(x, b) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addRefCnstV3i(x, b) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.addA3i(x, b) @needs_numpy_and_eigen def test02_vector_dynamic(): - a = np.array([1, 2, 3], dtype=np.int32) - b = np.array([0, 1, 2], dtype=np.int32) - c = np.array([1, 3, 5], dtype=np.int32) + a = np.array([1, 2, 3], dtype=np.int32) + b = np.array([0, 1, 2], dtype=np.int32) + c = np.array([1, 3, 5], dtype=np.int32) x = np.arange(10000, dtype=np.int32) af = np.float32(a) # Check call with dynamically sized arrays - assert np.all(t.addVXi(a, b) == c) + assert_array_equal(t.addVXi(a, b), c) # Implicit conversion - assert np.all(t.addVXi(af, b) == c) + assert_array_equal(t.addVXi(af, b), c) # Try with a big array. This will move the result to avoid a copy - r = np.all(t.addVXi(x, x) == 2*x) + assert_array_equal(t.addVXi(x, x), 2*x) @needs_numpy_and_eigen @@ -84,68 +83,124 @@ def test03_update_map(): a = np.array([1, 2, 3], dtype=np.int32) b = np.array([1, 2, 123], dtype=np.int32) c = a.copy() - t.updateV3i(c) - assert np.all(c == b) + t.updateRefV3i(c) + assert_array_equal(c, b) c = a.copy() - t.updateVXi(c) - assert np.all(c == b) + t.updateRefV3i_nc(c) + assert_array_equal(c, b) + + c = a.copy() + t.updateRefVXi(c) + assert_array_equal(c, b) + + c = a.copy() + t.updateRefVXi_nc(c) + assert_array_equal(c, b) + + c = np.float32(a) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.updateRefV3i(c) + + c = np.float32(a) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.updateRefV3i_nc(c) + + c = np.float32(a) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.updateRefVXi(c) + + c = np.float32(a) + with pytest.raises(TypeError, match='incompatible function arguments'): + t.updateRefVXi_nc(c) + @needs_numpy_and_eigen def test04_matrix(): A = np.vander((1, 2, 3, 4,)) At = A.T - A2 = 2*A - At2 = 2*At assert A.flags['C_CONTIGUOUS'] assert At.flags['F_CONTIGUOUS'] - assert np.all(t.addM4u_1(A, A) == A2) - assert np.all(t.addM4u_1(At, At) == At2) - assert np.all(t.addM4u_2(A, A) == A2) - assert np.all(t.addM4u_2(At, At) == At2) - assert np.all(t.addM4u_3(A, A) == A2) - assert np.all(t.addM4u_3(At, At) == At2) - assert np.all(t.addM4u_4(A, A) == A2) - assert np.all(t.addM4u_4(At, At) == At2) - assert np.all(t.addMXu_1(A, A) == A2) - assert np.all(t.addMXu_1(At, At) == At2) - assert np.all(t.addMXu_2(A, A) == A2) - assert np.all(t.addMXu_2(At, At) == At2) - assert np.all(t.addMXu_3(A, A) == A2) - assert np.all(t.addMXu_3(At, At) == At2) - assert np.all(t.addMXu_4(A, A) == A2) - assert np.all(t.addMXu_4(At, At) == At2) + base = np.zeros((A.shape[0] * 2, A.shape[1] * 2), A.dtype) + base[::2, ::2] = A + Av = base[-2::-2, -2::-2] + assert Av.base is base + Avt = Av.T + assert Avt.base is base + matrices = A, At, Av, Avt + for addM in (t.addM4uCC, t.addM4uRR, t.addM4uCR, t.addM4uRC, + t.addMXuCC, t.addMXuRR, t.addMXuCR, t.addMXuRC): + for left, right in itertools.product(matrices, matrices): + assert_array_equal(addM(left, right), left + right) @needs_numpy_and_eigen -@pytest.mark.parametrize("start", (0, 10)) -def test05_matrix_large_nonsymm(start): +@pytest.mark.parametrize("rowStart", (0, 1)) +@pytest.mark.parametrize("colStart", (0, 2)) +@pytest.mark.parametrize("rowStep", (1, 2, -2)) +@pytest.mark.parametrize("colStep", (1, 3, -3)) +@pytest.mark.parametrize("transpose", (False, True)) +def test05_matrix_large_nonsymm(rowStart, colStart, rowStep, colStep, transpose): A = np.uint32(np.vander(np.arange(80))) - A = A[:, start:] - A2 = A+A - out = t.addMXu_1(A, A) - assert np.all(t.addMXu_1(A, A) == A2) - assert np.all(t.addMXu_2(A, A) == A2) - assert np.all(t.addMXu_3(A, A) == A2) - assert np.all(t.addMXu_4(A, A) == A2) - assert np.all(t.addMXu_5(A, A) == A2) - + if rowStep < 0: + rowStart = -rowStart - 1 + if colStep < 0: + colStart = -colStart - 1 + A = A[rowStart::rowStep, colStart::colStep] + if transpose: + A = A.T + A2 = A + A + assert_array_equal(t.addMXuCC(A, A), A2) + assert_array_equal(t.addMXuRR(A, A), A2) + assert_array_equal(t.addMXuCR(A, A), A2) + assert_array_equal(t.addMXuRC(A, A), A2) + assert_array_equal(t.addDRefMXuCC_nc(A, A), A2) + assert_array_equal(t.addDRefMXuRR_nc(A, A), A2) + if A.flags['C_CONTIGUOUS']: + assert_array_equal(t.addMapMXuRR(A, A), A2) + assert_array_equal(t.addMapCnstMXuRR(A, A), A2) + else: + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addMapMXuRR(A, A) + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addMapCnstMXuRR(A, A) + + assert_array_equal(t.addRefCnstMXuRR(A, A), A2) + assert_array_equal(t.addRefCnstMXuRR(A.view(np.int32), A), A2) + assert_array_equal(t.addRefCnstMXuRR_nc(A, A), A2) + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addRefCnstMXuRR_nc(A.view(np.int32), A) + if A.strides[1] == A.itemsize: + assert_array_equal(t.addRefMXuRR(A, A), A2) + else: + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addRefMXuRR(A, A) + if A.flags['F_CONTIGUOUS']: + assert_array_equal(t.addMapMXuCC(A, A), A2) + assert_array_equal(t.addMapCnstMXuCC(A, A), A2) + else: + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addMapMXuCC(A, A) + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addMapCnstMXuCC(A, A) + + assert_array_equal(t.addRefCnstMXuCC(A, A), A2) + assert_array_equal(t.addRefCnstMXuCC(A.view(np.int32), A), A2) + assert_array_equal(t.addRefCnstMXuCC_nc(A, A), A2) + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addRefCnstMXuCC_nc(A.view(np.int32), A) + if A.strides[0] == A.itemsize: + assert_array_equal(t.addRefMXuCC(A, A), A2) + else: + with pytest.raises(TypeError, match="incompatible function arguments"): + t.addRefMXuCC(A, A) A = np.ascontiguousarray(A) assert A.flags['C_CONTIGUOUS'] - assert np.all(t.addMXu_2_nc(A, A) == A2) - + assert_array_equal(t.addMXuRR_nc(A, A), A2) A = np.asfortranarray(A) assert A.flags['F_CONTIGUOUS'] - assert np.all(t.addMXu_1_nc(A, A) == A2) - - A = A.T - A2 = A2.T - assert np.all(t.addMXu_1(A, A) == A2) - assert np.all(t.addMXu_2(A, A) == A2) - assert np.all(t.addMXu_3(A, A) == A2) - assert np.all(t.addMXu_4(A, A) == A2) - assert np.all(t.addMXu_5(A, A) == A2) + assert_array_equal(t.addMXuCC_nc(A, A), A2) @needs_numpy_and_eigen @@ -172,8 +227,8 @@ def test06_map(): def test07_mutate_arg(): A = np.uint32(np.vander(np.arange(10))) A2 = A.copy() - t.mutate_MXu(A) - assert np.all(A == 2*A2) + t.mutate_DRefMXuC(A) + assert_array_equal(A, 2*A2) @needs_numpy_and_eigen @@ -199,7 +254,7 @@ def assert_sparse_equal_ref(sparse_mat): [0, 0, 14, 0, 8, 11], ] ) - np.testing.assert_array_equal(sparse_mat.toarray(), ref) + assert_array_equal(sparse_mat.toarray(), ref) assert_sparse_equal_ref(t.sparse_r()) assert_sparse_equal_ref(t.sparse_c()) @@ -262,7 +317,7 @@ def test11_prop(): if j == 2 and i == 0: member[0, 0] = 10 ref[0, 0] = 10 - assert np.all(member == ref) + assert_array_equal(member, ref) del member gc.collect() gc.collect() @@ -271,4 +326,26 @@ def test11_prop(): del c gc.collect() gc.collect() - assert np.all(member == ref) + assert_array_equal(member, ref) + +@needs_numpy_and_eigen +def test12_cast(): + vec = np.arange(1000, dtype=np.int32) + vec2 = vec[::2] + vecf = np.float32(vec) + assert_array_equal(t.castToMapVXi(vec), vec) + assert_array_equal(t.castToRefVXi(vec), vec) + assert_array_equal(t.castToRefCnstVXi(vec), vec) + assert_array_equal(t.castToDRefCnstVXi(vec), vec) + for v in vec2, vecf: + with pytest.raises(RuntimeError, match="bad[_ ]cast"): + t.castToMapVXi(v) + with pytest.raises(RuntimeError, match="bad[_ ]cast"): + t.castToRefVXi(v) + assert_array_equal(t.castToRefCnstVXi(v), v) + assert_array_equal(t.castToDRefCnstVXi(vec2), vec2) + with pytest.raises(RuntimeError, match="bad[_ ]cast"): + t.castToDRefCnstVXi(vecf) + for v in vec, vec2, vecf: + with pytest.raises(RuntimeError, match='bad[_ ]cast'): + t.castToRef03CnstVXi(v)