From 864a6ff0cd8dce120ef7f5e072576b27d8f73783 Mon Sep 17 00:00:00 2001 From: Axel Huebl Date: Mon, 15 Feb 2021 01:27:33 -0800 Subject: [PATCH] [Draft] Array4: __array_interface__ Implement the `__array_interface__` in Array for. --- src/Base/Array4.cpp | 160 ++++++++++++++++++++++++++++++++++++++++ src/Base/CMakeLists.txt | 1 + src/pyAMReX.cpp | 2 + tests/test_array4.py | 86 +++++++++++++++++++++ 4 files changed, 249 insertions(+) create mode 100644 src/Base/Array4.cpp create mode 100644 tests/test_array4.py diff --git a/src/Base/Array4.cpp b/src/Base/Array4.cpp new file mode 100644 index 00000000..caf53369 --- /dev/null +++ b/src/Base/Array4.cpp @@ -0,0 +1,160 @@ +/* Copyright 2021 The AMReX Community + * + * Authors: Axel Huebl + * License: BSD-3-Clause-LBNL + */ +#include +#include +#include + +#include +#include +#include + +#include +#include + +namespace py = pybind11; +using namespace amrex; + + +template< typename T > +void make_Array4(py::module &m, std::string typestr) +{ + // dispatch simpler via: py::format_descriptor::format() naming + auto const array_name = std::string("Array4_").append(typestr); + py::class_< Array4 >(m, array_name.c_str(), py::buffer_protocol()) + .def("__repr__", + [](Array4 const & a4) { + std::stringstream s; + s << a4.size(); + return ""; + } + ) +#if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK) + .def("index_assert", &Array4::index_assert) +#endif + + .def_property_readonly("size", &Array4::size) + .def_property_readonly("nComp", &Array4::nComp) + + .def(py::init< >()) + .def(py::init< Array4 const & >()) + .def(py::init< Array4 const &, int >()) + .def(py::init< Array4 const &, int, int >()) + //.def(py::init< T*, Dim3 const &, Dim3 const &, int >()) + + .def(py::init([](py::array_t & arr) { + py::buffer_info buf = arr.request(); + + auto a4 = std::make_unique< Array4 >(); + a4.get()->p = (T*)buf.ptr; + a4.get()->begin = Dim3{0, 0, 0}; + // TODO: likely C->F index conversion here + // p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride]; + a4.get()->end.x = (int)buf.shape.at(0); + a4.get()->end.y = (int)buf.shape.at(1); + a4.get()->end.z = (int)buf.shape.at(2); + a4.get()->ncomp = 1; + // buffer protocol strides are in bytes, AMReX strides are elements + a4.get()->jstride = (int)buf.strides.at(0) / sizeof(T); + a4.get()->kstride = (int)buf.strides.at(1) / sizeof(T); + a4.get()->nstride = (int)buf.strides.at(2) * (int)buf.shape.at(2) / sizeof(T); + return a4; + })) + + + .def_property_readonly("__array_interface__", [](Array4 const & a4) { + auto d = py::dict(); + auto const len = length(a4); + // TODO: likely F->C index conversion here + // p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride]; + auto shape = py::make_tuple( // Buffer dimensions + len.x < 0 ? 0 : len.x, + len.y < 0 ? 0 : len.y, + len.z < 0 ? 0 : len.z//, // zero-size shall not have negative dimension + //a4.ncomp + ); + // buffer protocol strides are in bytes, AMReX strides are elements + auto const strides = py::make_tuple( // Strides (in bytes) for each index + sizeof(T) * a4.jstride, + sizeof(T) * a4.kstride, + sizeof(T)//, + //sizeof(T) * a4.nstride + ); + d["data"] = py::make_tuple(long(a4.dataPtr()), false); + d["typestr"] = py::format_descriptor::format(); + d["shape"] = shape; + d["strides"] = strides; + // d["strides"] = py::none(); + d["version"] = 3; + return d; + }) + + // not sure if useful to have this implemented on top +/* + .def_buffer([](Array4 & a4) -> py::buffer_info { + auto const len = length(a4); + // TODO: likely F->C index conversion here + // p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride]; + auto shape = { // Buffer dimensions + len.x < 0 ? 0 : len.x, + len.y < 0 ? 0 : len.y, + len.z < 0 ? 0 : len.z//, // zero-size shall not have negative dimension + //a4.ncomp + }; + // buffer protocol strides are in bytes, AMReX strides are elements + auto const strides = { // Strides (in bytes) for each index + sizeof(T) * a4.jstride, + sizeof(T) * a4.kstride, + sizeof(T)//, + //sizeof(T) * a4.nstride + }; + return py::buffer_info( + a4.dataPtr(), + shape, + strides + ); + }) +*/ + ; +} + +void init_Array4(py::module &m) { + make_Array4< float >(m, "float"); + make_Array4< double >(m, "double"); + make_Array4< long double >(m, "longdouble"); + + make_Array4< short >(m, "short"); + make_Array4< int >(m, "int"); + make_Array4< long >(m, "long"); + make_Array4< long long >(m, "longlong"); + + make_Array4< unsigned short >(m, "ushort"); + make_Array4< unsigned int >(m, "uint"); + make_Array4< unsigned long >(m, "ulong"); + make_Array4< unsigned long long >(m, "ulonglong"); + + // std::complex< float|double|long double> ? + + /* + py::class_< PolymorphicArray4, Array4 >(m, "PolymorphicArray4") + .def("__repr__", + [](PolymorphicArray4 const & pa4) { + std::stringstream s; + s << pa4.size(); + return ""; + } + ) + ; + */ + + // free standing C++ functions: + /* + contains + lbound + ubound + length + makePolymorphic + */ +} diff --git a/src/Base/CMakeLists.txt b/src/Base/CMakeLists.txt index 024f8f0d..796ded82 100644 --- a/src/Base/CMakeLists.txt +++ b/src/Base/CMakeLists.txt @@ -1,6 +1,7 @@ target_sources(pyAMReX PRIVATE AMReX.cpp + Array4.cpp Box.cpp Dim3.cpp IntVect.cpp diff --git a/src/pyAMReX.cpp b/src/pyAMReX.cpp index f222863e..811a7f10 100644 --- a/src/pyAMReX.cpp +++ b/src/pyAMReX.cpp @@ -17,6 +17,7 @@ namespace py = pybind11; // forward declarations of exposed classes void init_AMReX(py::module&); +void init_Array4(py::module&); void init_Box(py::module &); void init_Dim3(py::module&); void init_IntVect(py::module &); @@ -37,6 +38,7 @@ PYBIND11_MODULE(amrex_pybind, m) { init_AMReX(m); init_Dim3(m); init_IntVect(m); + init_Array4(m); init_Box(m); // API runtime version diff --git a/tests/test_array4.py b/tests/test_array4.py new file mode 100644 index 00000000..0c8f0a42 --- /dev/null +++ b/tests/test_array4.py @@ -0,0 +1,86 @@ +# -*- coding: utf-8 -*- + +import pytest +import numpy as np +import amrex + +def test_array4_empty(): + empty = amrex.Array4_double() + + # Check properties + assert(empty.size == 0) + assert(empty.nComp == 0) + + # assign empty + emptyc = amrex.Array4_double(empty) + # Check properties + assert(emptyc.size == 0) + assert(emptyc.nComp == 0) + +def test_array4(): + # from numpy (also a non-owning view) + x = np.ones((2, 3, 4,)) + arr = amrex.Array4_double(x) + + x[1, 1, 1] = 42 + # TypeError: 'amrex.amrex_pybind.Array4_double' object is not subscriptable + # assert(arr[1, 1, 1] == 42) + + # copy to numpy + c_arr2np = np.array(arr, copy=True) + assert(c_arr2np.ndim == 3) + #assert(c_arr2np.ndim == 4) + assert(c_arr2np.dtype == np.dtype("double")) + np.testing.assert_array_equal(x, c_arr2np) + assert(c_arr2np[1, 1, 1] == 42) + + # view to numpy + v_arr2np = np.array(arr, copy=False) + assert(v_arr2np.ndim == 3) + #assert(c_arr2np.ndim == 4) + assert(v_arr2np.dtype == np.dtype("double")) + np.testing.assert_array_equal(x, v_arr2np) + assert(v_arr2np[1, 1, 1] == 42) + + # change original buffer once more + x[1, 1, 1] = 43 + assert(v_arr2np[1, 1, 1] == 43) + + # copy array4 (view) + c_arr = amrex.Array4_double(arr) + v_carr2np = np.array(c_arr, copy=False) + x[1, 1, 1] = 44 + assert(v_carr2np[1, 1, 1] == 44) + + # from cupy + + # to numpy + + # to cupy + + return + + # Check indexing + assert(obj[0] == 1) + assert(obj[1] == 2) + assert(obj[2] == 3) + assert(obj[-1] == 3) + assert(obj[-2] == 2) + assert(obj[-3] == 1) + with pytest.raises(IndexError): + obj[-4] + with pytest.raises(IndexError): + obj[3] + + # Check assignment + obj[0] = 2 + obj[1] = 3 + obj[2] = 4 + assert(obj[0] == 2) + assert(obj[1] == 3) + assert(obj[2] == 4) + +#def test_iv_conversions(): +# obj = amrex.IntVect.max_vector().numpy() +# assert(isinstance(obj, np.ndarray)) +# assert(obj.dtype == np.int32)