diff --git a/nanobind_test/CMakeLists.txt b/nanobind_test/CMakeLists.txt index fe2064f..8c16009 100644 --- a/nanobind_test/CMakeLists.txt +++ b/nanobind_test/CMakeLists.txt @@ -6,6 +6,8 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS include(cmake/CPM.cmake) +project(nanobind_test LANGUAGES CXX) + # Set C++ standard set(CMAKE_CXX_STANDARD 23) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/nanobind_test/mdspan_wrapper.cpp b/nanobind_test/mdspan_wrapper.cpp index 55d7e4a..1f7006d 100644 --- a/nanobind_test/mdspan_wrapper.cpp +++ b/nanobind_test/mdspan_wrapper.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -106,6 +107,12 @@ template struct PauliOp { } fmt::print("]\n"); } + + void return_coeffs(std::mdspan> &out) { + for (size_t i = 0; i < coeffs.extent(0); ++i) { + out[i] = coeffs[i]; + } + } }; NB_MODULE(mdspan_wrapper, m) { @@ -115,7 +122,62 @@ NB_MODULE(mdspan_wrapper, m) { .def("scale", &PauliOp::scale, "scale"_a) .def("print", &PauliOp::print) .def("__eq__", &PauliOp::operator==) - .def("multiply_coeff", [](PauliOp &op, nb::ndarray &c) { - op.multiply_coeff(ndarray_cast_from_py(c)); + .def("multiply_coeff", + [](PauliOp &op, nb::ndarray &c) { + op.multiply_coeff(ndarray_cast_from_py(c)); + }) + .def("return_coeffs", + [](PauliOp &op, nb::ndarray &out) { + auto out_mdspan = ndarray_cast_from_py(out); + op.return_coeffs(out_mdspan); + }) + .def("return_coeffs_owning", [](PauliOp &op) { + struct Temp { + std::vector data; + }; + + Temp *tmp = new Temp{op._coeffs}; + + fmt::println("copied data: [{}]", fmt::join(tmp->data, ", ")); + std::cout << std::flush; + + nb::capsule deleter( + tmp, [](void *data) noexcept { delete static_cast(data); }); + + return nb::ndarray( + /*data*/ tmp->data.data(), + /*shape */ {tmp->data.size()}, + /*deleter*/ deleter); }); -} \ No newline at end of file + + m.def("return_coeffs", + [](size_t n) { + std::vector data(n); + for (size_t i = 0; i < n; ++i) { + data[i] = i; + } + + struct Temp { + std::vector data; + }; + + Temp *tmp = new Temp{data}; + + nb::capsule deleter(tmp, [](void *data) noexcept { + delete static_cast(data); + }); + + return nb::ndarray( + /*data*/ tmp->data.data(), + /*shape */ {tmp->data.size()}, + /*deleter*/ deleter); + + // nb::ndarray out{data.data(), {data.size()}, deleter}; + + // nb::object res = nb::cast(out, nb::rv_policy::copy); + // return res; + } + + /**/ + ); +} diff --git a/nanobind_test/test.py b/nanobind_test/test.py index de7e143..dec05d0 100644 --- a/nanobind_test/test.py +++ b/nanobind_test/test.py @@ -65,3 +65,22 @@ def test_multiply_coeff_diff_type(my_op) -> None: coeffs = np.array([1.0, 2.0, 3.0], dtype=np.complex128) with pytest.raises(TypeError): my_op.multiply_coeff(coeffs) + + +def test_return_coeffs_non_owning(my_op) -> None: + c = np.zeros(3) + print(c) + my_op.return_coeffs(c) + + print(c) + np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0])) + + +def test_return_coeffs_owning(my_op) -> None: + c = my_op.return_coeffs_owning() + print("returned array", c) + np.testing.assert_allclose(c, np.array([1.0, 2.0, 3.0])) + + +def test_standalone() -> None: + np.testing.assert_allclose(mdspan_wrapper.return_coeffs(3), np.array([0, 1.0, 2.0]))