diff --git a/libsrc/core/python_ngcore.cpp b/libsrc/core/python_ngcore.cpp index 5e64d1b2..f7addf3c 100644 --- a/libsrc/core/python_ngcore.cpp +++ b/libsrc/core/python_ngcore.cpp @@ -7,7 +7,7 @@ using std::string; namespace ngcore { - + bool ngcore_have_numpy = false; void SetFlag(Flags &flags, string s, py::object value) { if (py::isinstance(value)) diff --git a/libsrc/core/python_ngcore.hpp b/libsrc/core/python_ngcore.hpp index f52c104e..21f610e3 100644 --- a/libsrc/core/python_ngcore.hpp +++ b/libsrc/core/python_ngcore.hpp @@ -4,6 +4,7 @@ #include "ngcore_api.hpp" // for operator new #include #include +#include #include "array.hpp" #include "archive.hpp" @@ -13,6 +14,57 @@ namespace py = pybind11; namespace ngcore { + NGCORE_API extern bool ngcore_have_numpy; + + // Python class name type traits + template + struct PyNameTraits { + static const std::string & GetName() + { + static const std::string name = + py::cast(py::cast(T()).attr("__class__").attr("__name__")); + return name; + } + }; + + template + std::string GetPyName(const char *prefix = 0) { + std::string s; + if(prefix) s = std::string(prefix); + s+= PyNameTraits::GetName(); + return s; + } + + template<> + struct PyNameTraits { + static std::string GetName() { return "I"; } + }; + + template<> + struct PyNameTraits { + static std::string GetName() { return "U"; } + }; + + template<> + struct PyNameTraits { + static std::string GetName() { return "F"; } + }; + + template<> + struct PyNameTraits { + static std::string GetName() { return "D"; } + }; + + template<> + struct PyNameTraits { + static std::string GetName() { return "S"; } + }; + + template + struct PyNameTraits> { + static std::string GetName() + { return std::string("sp_")+GetPyName(); } + }; template Array makeCArray(const py::object& obj) @@ -29,14 +81,30 @@ namespace ngcore return arr; } + namespace detail + { + template + struct HasPyFormat + { + private: + template + static auto check(T2*) -> std::enable_if_t>().format()), std::string>, std::true_type>; + static auto check(...) -> std::false_type; + public: + static constexpr bool value = decltype(check((T*) nullptr))::value; + }; + } // namespace detail + template ::index_type> void ExportArray (py::module &m) { using TFlat = FlatArray; using TArray = Array; - std::string suffix = std::string(typeid(T).name()) + "_" + typeid(TIND).name(); + std::string suffix = GetPyName() + "_" + + GetPyName(); std::string fname = std::string("FlatArray_") + suffix; - py::class_(m, fname.c_str()) + auto flatarray_class = py::class_(m, fname.c_str(), + py::buffer_protocol()) .def ("__len__", [] ( TFlat &self ) { return self.Size(); } ) .def ("__getitem__", [](TFlat & self, TIND i) -> T& @@ -54,7 +122,7 @@ namespace ngcore if (i < base || i >= self.Size()+base) throw py::index_error(); self[i] = val; - return self[i]; + return self[i]; }, py::return_value_policy::reference) @@ -66,17 +134,41 @@ namespace ngcore throw py::error_already_set(); static constexpr int base = IndexBASE(); if (start < base || start+(slicelength-1)*step >= self.Size()+base) - throw py::index_error(); - for (size_t i = 0; i < slicelength; i++, start+=step) - self[start] = val; + throw py::index_error(); + for (size_t i = 0; i < slicelength; i++, start+=step) + self[start] = val; }) - + .def("__iter__", [] ( TFlat & self) { return py::make_iterator (self.begin(),self.end()); }, py::keep_alive<0,1>()) // keep array alive while iterator is used ; + if constexpr (detail::HasPyFormat::value) + { + if(ngcore_have_numpy && !py::detail::npy_format_descriptor::dtype().is_none()) + { + flatarray_class + .def_buffer([](TFlat& self) + { + return py::buffer_info( + self.Addr(0), + sizeof(T), + py::format_descriptor::format(), + 1, + { self.Size() }, + { sizeof(T) * (self.Addr(1) - self.Addr(0)) }); + }) + .def("NumPy", [](py::object self) + { + return py::module::import("numpy") + .attr("frombuffer")(self, py::detail::npy_format_descriptor::dtype()); + }) + ; + } + } + std::string aname = std::string("Array_") + suffix; py::class_(m, aname.c_str()) .def(py::init([] (size_t n) { return new TArray(n); }),py::arg("n"), "Makes array of given length") @@ -100,7 +192,7 @@ namespace ngcore py::dict NGCORE_API CreateDictFromFlags(const Flags& flags); // *************** Archiving functionality ************** - + template Archive& Archive :: Shallow(T& val) { diff --git a/libsrc/core/python_ngcore_export.cpp b/libsrc/core/python_ngcore_export.cpp index 162f8163..d6cc9d6e 100644 --- a/libsrc/core/python_ngcore_export.cpp +++ b/libsrc/core/python_ngcore_export.cpp @@ -2,13 +2,18 @@ #include "bitarray.hpp" #include "taskmanager.hpp" - using namespace ngcore; using namespace std; using namespace pybind11::literals; PYBIND11_MODULE(pyngcore, m) // NOLINT { + try + { + auto numpy = py::module::import("numpy"); + ngcore_have_numpy = !numpy.is_none(); + } + catch(...) {} ExportArray(m); ExportArray(m); ExportArray(m); diff --git a/tests/pytest/test_array.py b/tests/pytest/test_array.py new file mode 100644 index 00000000..db73452e --- /dev/null +++ b/tests/pytest/test_array.py @@ -0,0 +1,17 @@ +from pyngcore import * +from numpy import sort, array + +def test_array_numpy(): + a = Array_I_S(5) + a[:] = 0 + a[3:] = 2 + assert(sum(a) == 4) + a[1] = 5 + b = sort(a) + assert(all(b == array([0,0,2,2,5]))) + assert(all(a == array([0,5,0,2,2]))) + a.NumPy().sort() + assert(all(a == array([0,0,2,2,5]))) + +if __name__ == "__main__": + test_array_numpy()