MultiFab & Array4

.github/workflows/windows.yml

@@ -24,7 +24,7 @@ jobs:
   clang:
     name: Clang w/o MPI
     runs-on: windows-latest
-    env: {PYAMREX_LIBDIR: build/lib/site-packages/amrex/Release/}
+    env: {PYAMREX_LIBDIR: build/lib/site-packages/amrex/Debug/}
     steps:
     - uses: actions/checkout@v2
     - uses: seanmiddleditch/gha-setup-ninja@master
@@ -38,11 +38,11 @@ jobs:
               -DAMReX_MPI=OFF             `
               -DPython_EXECUTABLE=C:\\hostedtoolcache\\windows\\python\\3.9.10\\x64\\python3.exe
 
-        cmake --build build --config Release -j 2
+        cmake --build build --config Debug -j 2
         python3 -m pip install -v .
     - name: Unit tests
       run: |
-        ctest --test-dir build -C Release --output-on-failure
+        ctest --test-dir build -C Debug --output-on-failure
 
 # C:\\hostedtoolcache\\windows\\python\\3.9.10\\x64\\python3.exe
 # C:\\hostedtoolcache\\windows\\Python\\3.10.2\\x64\\python3.exe

README.md

@@ -147,6 +147,9 @@ python3 -m pytest tests/test_intvect.py
 
 # Run a single test (useful during debugging)
 python3 -m pytest tests/test_intvect.py::test_iv_conversions
+
+# Run all tests, do not capture "print" output and be verbose
+python3 -m pytest -s -vvvv tests/
 ```
 
 ### Build Options

src/Base/Array4.cpp

@@ -0,0 +1,196 @@
+/* Copyright 2021-2022 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_Array4.H>
+#include <AMReX_BLassert.H>
+#include <AMReX_IntVect.H>
+
+#include <cstdint>
+#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())
+        .def("__repr__",
+             [typestr](Array4<T> const & a4) {
+                 std::stringstream s;
+                 s << a4.size();
+                 return "<amrex.Array4 of type '" + typestr +
+                        "' and 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_property_readonly("num_comp", &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 >())
+
+        /* init from a numpy or other buffer protocol array: non-owning view
+         */
+        .def(py::init([](py::array_t<T> & arr) {
+            py::buffer_info buf = arr.request();
+
+            AMREX_ALWAYS_ASSERT_WITH_MESSAGE(buf.ndim == 3,
+                "We can only create amrex::Array4 views into 3D Python arrays at the moment.");
+            // TODO:
+            //   In 2D, Array4 still needs to be accessed with (i,j,k) or (i,j,k,n), with k = 0.
+            //   Likewise in 1D.
+            //   We could also add support for 4D numpy arrays, treating the slowest
+            //   varying index as component "n".
+
+            if (buf.format != py::format_descriptor<T>::format())
+                throw std::runtime_error("Incompatible format: expected '" +
+                    py::format_descriptor<T>::format() +
+                    "' and received '" + buf.format + "'!");
+
+            auto a4 = std::make_unique< Array4<T> >();
+            a4.get()->p = static_cast<T*>(buf.ptr);
+            a4.get()->begin = Dim3{0, 0, 0};
+            // 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(2); // fastest varying index
+            a4.get()->end.y = (int)buf.shape.at(1);
+            a4.get()->end.z = (int)buf.shape.at(0);
+            a4.get()->ncomp = 1;
+            // buffer protocol strides are in bytes, AMReX strides are elements
+            a4.get()->jstride = (int)buf.strides.at(1) / sizeof(T); // fastest varying index
+            a4.get()->kstride = (int)buf.strides.at(0) / sizeof(T);
+            // 3D == no component: stride here should not matter
+            a4.get()->nstride = a4.get()->kstride * (int)buf.shape.at(0);
+
+            // todo: we could check and store here if the array buffer we got is read-only
+
+            return a4;
+        }))
+
+        .def_property_readonly("__array_interface__", [](Array4<T> const & a4) {
+            auto d = py::dict();
+            auto const len = length(a4);
+            // F->C index conversion here
+            // p[(i-begin.x)+(j-begin.y)*jstride+(k-begin.z)*kstride+n*nstride];
+            // Buffer dimensions: zero-size shall not skip dimension
+            auto shape = py::make_tuple(
+                    a4.ncomp,
+                    len.z <= 0 ? 1 : len.z,
+                    len.y <= 0 ? 1 : len.y,
+                    len.x <= 0 ? 1 : len.x  // fastest varying index
+            );
+            // buffer protocol strides are in bytes, AMReX strides are elements
+            auto const strides = py::make_tuple(
+                    sizeof(T) * a4.nstride,
+                    sizeof(T) * a4.kstride,
+                    sizeof(T) * a4.jstride,
+                    sizeof(T)  // fastest varying index
+            );
+            bool const read_only = false;
+            d["data"] = py::make_tuple(std::intptr_t(a4.dataPtr()), read_only);
+            // note: if we want to keep the same global indexing with non-zero
+            //       box small_end as in AMReX, then we can explore playing with
+            //       this offset as well
+            //d["offset"] = 0;         // default
+            //d["mask"] = py::none();  // default
+
+            d["shape"] = shape;
+            // we could also set this after checking the strides are C-style contiguous:
+            //if (is_contiguous<T>(shape, strides))
+            //    d["strides"] = py::none();  // C-style contiguous
+            //else
+            d["strides"] = strides;
+
+            d["typestr"] = py::format_descriptor<T>::format();
+            d["version"] = 3;
+            return d;
+        })
+
+
+        // TODO: __cuda_array_interface__
+        // https://numba.readthedocs.io/en/latest/cuda/cuda_array_interface.html
+
+
+        // TODO: __dlpack__
+        // DLPack protocol (CPU, NVIDIA GPU, AMD GPU, Intel GPU, etc.)
+        // https://dmlc.github.io/dlpack/latest/
+        // https://data-apis.org/array-api/latest/design_topics/data_interchange.html
+        // https://github.com/data-apis/consortium-feedback/issues/1
+        // https://github.com/dmlc/dlpack/blob/master/include/dlpack/dlpack.h
+
+
+        .def("contains", &Array4<T>::contains)
+        //.def("__contains__", &Array4<T>::contains)
+
+        // setter & getter
+        .def("__setitem__", [](Array4<T> & a4, IntVect const & v, T const value){ a4(v) = value; })
+        .def("__setitem__", [](Array4<T> & a4, std::array<int, 4> const key, T const value){
+            a4(key[0], key[1], key[2], key[3]) = value;
+        })
+        .def("__setitem__", [](Array4<T> & a4, std::array<int, 3> const key, T const value){
+            a4(key[0], key[1], key[2]) = value;
+        })
+
+        .def("__getitem__", [](Array4<T> & a4, IntVect const & v){ return a4(v); })
+        .def("__getitem__", [](Array4<T> & a4, std::array<int, 4> const key){
+            return a4(key[0], key[1], key[2], key[3]);
+        })
+        .def("__getitem__", [](Array4<T> & a4, std::array<int, 3> const key){
+            return a4(key[0], key[1], key[2]);
+        })
+    ;
+
+    // free standing C++ functions:
+    m.def("lbound", &lbound< Array4<T> >);
+    m.def("ubound", &ubound< Array4<T> >);
+    m.def("length", &length< Array4<T> >);
+    m.def("makePolymorphic", &makePolymorphic< Array4<T> >);
+}
+
+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() + "'>";
+             }
+        )
+    ;
+     */
+}

src/Base/Box.cpp

@@ -1,22 +1,75 @@
-/* Copyright 2021 The AMReX Community
+/* Copyright 2021-2022 The AMReX Community
  *
  * Authors: Axel Huebl
  * License: BSD-3-Clause-LBNL
  */
 #include <pybind11/pybind11.h>
+#include <pybind11/operators.h>
 #include <pybind11/stl.h>
 
 #include <AMReX_Config.H>
 #include <AMReX_Box.H>
 #include <AMReX_IntVect.H>
 
 #include <sstream>
+#include <optional>
 
 namespace py = pybind11;
 using namespace amrex;
 
+namespace
+{
+    /** A little Wrapper class to iterate an amrex::Box via
+     *  amrex::Box::next().
+     */
+    struct Box3DConstIter {
+        Box m_box;
+        std::optional<IntVect> m_it;
+
+        Box3DConstIter(Box const & bx) : m_box(bx) {
+            m_it = m_box.smallEnd();
+        }
+
+        Box3DConstIter& operator++() {
+            // from FABio_ascii::write
+            if (m_it < m_box.bigEnd()) {
+                m_box.next(m_it.value());
+                return *this;
+            }
+            else
+            {
+                m_it = std::nullopt;
+                return *this;
+            }
+        }
+
+        bool operator==(Box3DConstIter const & other) const {
+            return other.m_it == m_it;
+        }
+
+        Box3DConstIter begin() const
+        {
+            return Box3DConstIter(m_box);
+        }
+        Box3DConstIter end() const
+        {
+            auto it = Box3DConstIter(m_box);
+            it.m_it = std::nullopt;
+            return it;
+        }
+
+        IntVect operator*() const
+        {
+            return m_it.value();
+        }
+    };
+}
 
 void init_Box(py::module &m) {
+    py::class_< Direction >(m, "Direction");
+
+
+
     py::class_< Box >(m, "Box")
         .def("__repr__",
             [](Box const & b) {
@@ -30,14 +83,16 @@ void init_Box(py::module &m) {
         .def(py::init< IntVect const &, IntVect const &, IntVect const & >())
         //.def(py::init< IntVect const &, IntVect const &, IndexType >())
 
+        .def_property_readonly("small_end", [](Box const & bx){ return bx.smallEnd(); })
+        .def_property_readonly("big_end", [](Box const & bx){ return bx.bigEnd(); })
         /*
         .def_property("small_end",
             &Box::smallEnd,
             py::overload_cast< IntVect const & >(&Box::setSmall))
         .def_property("big_end",
             &Box::bigEnd,
             &Box::setBig)
-        */
+
         .def_property("type",
             py::overload_cast<>(&Box::type, py::const_),
             &Box::setType)
@@ -57,8 +112,8 @@ void init_Box(py::module &m) {
 
         // loVect3d
         // hiVect3d
-        // loVect
-        // hiVect
+        .def_property_readonly("lo_vect", &Box::loVect)
+        .def_property_readonly("hi_vect", &Box::hiVect)
 
         .def("contains",
             py::overload_cast< IntVect const & >(&Box::contains, py::const_))
@@ -74,10 +129,31 @@ void init_Box(py::module &m) {
         // atOffset
         // atOffset3d
         // setRange
-        // shift
         // shiftHalf
-        // convert
-        // surroundingNodes
+        */
+        .def("shift", [](Box & bx, IntVect const& iv) { return bx.shift(iv); })
+
+        .def(py::self + IntVect())
+        .def(py::self - IntVect())
+        .def(py::self += IntVect())
+        .def(py::self -= IntVect())
+
+        .def("convert",
+             py::overload_cast< IndexType >(&Box::convert))
+        .def("convert",
+             py::overload_cast< IntVect const & >(&Box::convert))
+
+        .def("grow", [](Box & bx, IntVect const& iv) { return bx.grow(iv); })
+
+        //.def("surrounding_nodes",
+        //     py::overload_cast< >(&Box::surroundingNodes))
+        //.def("surrounding_nodes",
+        //     py::overload_cast< int >(&Box::surroundingNodes),
+        //     py::arg("dir"))
+        //.def("surrounding_nodes",
+        //     py::overload_cast< Direction >(&Box::surroundingNodes),
+        //     py::arg("d"))
+
         // enclosedCells
         // minBox
         // chop
@@ -91,5 +167,28 @@ void init_Box(py::module &m) {
 
         // __getitem__
 
+        /* iterate Box index space */
+        .def("__iter__",
+             [](Box const & bx) {
+                 auto box_iter = Box3DConstIter(bx);
+                 return py::make_iterator(box_iter.begin(), box_iter.end());
+             },
+            // Essential: keep object alive while iterator exists
+             py::keep_alive<0, 1>()
+        )
+
+        .def("lbound", [](Box const &, Box const & other){ return lbound(other); })
+        .def("ubound", [](Box const &, Box const & other){ return ubound(other); })
+        .def("begin", [](Box const &, Box const & other){ return begin(other); })
+        .def("end", [](Box const &, Box const & other){ return end(other); })
+        // already an attribute
+        //.def("length", [](Box const &, Box const & other){ return length(other); })
     ;
+
+    // free standing C++ functions:
+    m.def("lbound", [](Box const & other){ return lbound(other); });
+    m.def("ubound", [](Box const & other){ return ubound(other); });
+    m.def("begin", [](Box const & other){ return begin(other); });
+    m.def("end", [](Box const & other){ return end(other); });
+    m.def("length", [](Box const & other){ return length(other); });
 }