[WIP] Array4: __array_interface__

src/Base/Array4.cpp

@@ -0,0 +1,160 @@
+/* Copyright 2021 The AMReX Community
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+
+#include <AMReX_Config.H>
+#include <AMReX_Array4.H>
+#include <AMReX_IntVect.H>
+
+#include <sstream>
+#include <type_traits>
+
+namespace py = pybind11;
+using namespace amrex;
+
+
+template< typename T >
+void make_Array4(py::module &m, std::string typestr)
+{
+    // dispatch simpler via: py::format_descriptor<T>::format() naming
+    auto const array_name = std::string("Array4_").append(typestr);
+    py::class_< Array4<T> >(m, array_name.c_str(), py::buffer_protocol())
+        .def("__repr__",
+             [](Array4<T> const & a4) {
+                 std::stringstream s;
+                 s << a4.size();
+                 return "<amrex.Array4 of size '" + s.str() + "'>";
+             }
+        )
+#if defined(AMREX_DEBUG) || defined(AMREX_BOUND_CHECK)
+        .def("index_assert", &Array4<T>::index_assert)
+#endif
+
+        .def_property_readonly("size", &Array4<T>::size)
+        .def_property_readonly("nComp", &Array4<T>::nComp)
+
+        .def(py::init< >())
+        .def(py::init< Array4<T> const & >())
+        .def(py::init< Array4<T> const &, int >())
+        .def(py::init< Array4<T> const &, int, int >())
+        //.def(py::init< T*, Dim3 const &, Dim3 const &, int >())
+
+        .def(py::init([](py::array_t<T> & arr) {
+            py::buffer_info buf = arr.request();
+
+            auto a4 = std::make_unique< Array4<T> >();
+            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<T> 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<T>::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<T> & 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 "<amrex.PolymorphicArray4 of size '" + s.str() + "'>";
+             }
+        )
+    ;
+     */
+
+    // free standing C++ functions:
+    /*
+    contains
+    lbound
+    ubound
+    length
+    makePolymorphic
+    */
+}

src/Base/CMakeLists.txt

@@ -1,6 +1,7 @@
 target_sources(pyAMReX
   PRIVATE
     AMReX.cpp
+    Array4.cpp
     Box.cpp
     Dim3.cpp
     IntVect.cpp

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

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)