MultiFab: Iterate via Array4

AMReX-Codes · Mar 9, 2022 · 24584b0 · 24584b0
1 parent e524ba9
commit 24584b0
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 9 deletions.
diff --git a/src/Base/MultiFab.cpp b/src/Base/MultiFab.cpp
@@ -28,13 +28,13 @@ void init_MultiFab(py::module &m) {
         .def("is_nodal",
              py::overload_cast< int >(&FabArrayBase::is_nodal, py::const_))
 
+        .def_property_readonly("nComp", &FabArrayBase::nComp)
         .def_property_readonly("num_comp", &FabArrayBase::nComp)
+        .def_property_readonly("size", &FabArrayBase::size)
     ;
 
     py::class_< FArrayBox >(m, "FArrayBox");
 
-    py::class_< FabArray<FArrayBox>, FabArrayBase >(m, "FabArray_FArrayBox");
-
     //py::class_< MFIterWrapper >(m, "MFIterWrapper");
     py::class_< MFIter >(m, "MFIter")
         .def("__repr__",
@@ -94,6 +94,13 @@ void init_MultiFab(py::module &m) {
         */
     ;
 
+    py::class_< FabArray<FArrayBox>, FabArrayBase >(m, "FabArray_FArrayBox")
+        //.def("array", py::overload_cast< const MFIter& >(&FabArray<FArrayBox>::array))
+        //.def("const_array", &FabArray<FArrayBox>::const_array)
+        .def("array", [](FabArray<FArrayBox> & fa, MFIter const & mfi) {return fa.array(mfi);})
+        .def("const_array", [](FabArray<FArrayBox> & fa, MFIter const & mfi) {return fa.const_array(mfi);})
+    ;
+
     py::class_< MultiFab, FabArray<FArrayBox> >(m, "MultiFab")
         .def("__repr__",
              [](MultiFab const & mf) {

diff --git a/tests/test_multifab.py b/tests/test_multifab.py
@@ -10,13 +10,28 @@ def test_mfab_loop(mfab, nghost):
     for mfi in mfab:
         print(mfi)
         bx = mfi.tilebox()
-        #marr = mfab.array(mfi) # Array4: add __array_interface__ here
-
-        #for i, j, k in bx:
-        #   mar[i, j, k, 0] = 10.0 * i
-        #    mar[i, j, k, 1] = 10.0 * j
-        #    mar[i, j, k, 2] = 10.0 * k
-        #    print(i,j,k)
+        marr = mfab.array(mfi) # Array4: add __array_interface__ here
+
+        print(mfab)
+        print(mfab.num_comp)
+        print(mfab.size)
+        print(marr.size)
+        print(marr.nComp)
+
+        # slow, index by index assignment
+        for i, j, k in bx:
+            #print(i,j,k)
+            marr[i, j, k] = 10.0 * i
+            #marr[i, j, k, 0] = 10.0 * i
+            #marr[i, j, k, 1] = 10.0 * j
+            #marr[i, j, k, 2] = 10.0 * k
+
+        # fast, range based assignment
+        # TODO
+
+        # challenge: offset from our index space for...
+        # extra test: numpy assignment
+        # extra test: cupy assignment
 
 
 def test_mfab_simple(mfab):