Python coupling patch

python/waLBerla_docs/modules/blockforest.rst

@@ -25,6 +25,21 @@ Reference
                                         on other processes). This information includes the process rank, the state, and
                                         the axis-aligned bounding box of any block (local or remote). [false by default]
 
+.. py:function:: createUniformBlockGrid(cells, dx=1, oneBlockPerProcess=True, periodic=(0,0,0), keepGlobalBlockInformation=False)
+
+   Creates a new uniform StructuredBlockForest. Similar to cpp function blockforest::createUniformBlockGrid.
+   Specify only number of cells. The distribution on the processes is then calculated automatically.
+
+   :param cells :                       3-tuple with total number of cells in x,y,z direction.
+
+   :param dx:                           Side length of a single cell.
+   :param oneBlockPerProcess:           If True, each process gets one block. If False, all blocks are put to one process.
+                                        The second option makes only sense for debugging or testing.
+   :param periodic:                     Periodicity of the domain in x,y,z direction
+   :param keepGlobalBlockInformation:   If true, each process keeps information about remote blocks (blocks that reside
+                                        on other processes). This information includes the process rank, the state, and
+                                        the axis-aligned bounding box of any block (local or remote). [false by default]
+
 
 
    
\ No newline at end of file

python/waLBerla_tests/test_field.py

@@ -47,6 +47,36 @@ class FieldModuleTest(unittest.TestCase):
         view4 = field.toArray(f, with_ghost_layers=[2, False, True])
         self.assertEqual(view4[:, :, :, 0].shape, tuple([size[0] + 2 * 2, size[1] + 2 * 0, size[2] + 2 * gl]))
 
+    def test_gather(self):
+        blocks = createUniformBlockGrid(blocks=(10, 4, 3), cellsPerBlock=(2, 2, 2),
+                                        periodic=(True, True, True), oneBlockPerProcess=False)
+        wlb.field.addToStorage(blocks, "test", dtype=np.int, fSize=3)
+
+        for block in blocks:
+            offset_in_global_domain = blocks.transformLocalToGlobal(block, wlb.Cell(0, 0, 0))[:]
+            f = wlb.field.toArray(block["test"])
+            s = f.shape
+            for i in range(0, s[0]):
+                for j in range(0, s[1]):
+                    for k in range(0, s[2]):
+                        f[i, j, k, 0] = i + offset_in_global_domain[0]
+                        f[i, j, k, 1] = j + offset_in_global_domain[1]
+                        f[i, j, k, 2] = k + offset_in_global_domain[2]
+
+        tp = tuple([slice(5, 15), slice(None, None), 0.5])
+        f = wlb.field.gather(blocks, "test", tp)
+
+        nparray = wlb.field.toArray(f)
+        self.assertEqual(nparray.shape, (11, 8, 1, 3))
+
+        s = nparray.shape
+        for i in range(0, s[0]):
+            for j in range(0, s[1]):
+                for k in range(0, s[2]):
+                    assert(nparray[i, j, k, 0] == i + 5)
+                    assert(nparray[i, j, k, 1] == j)
+                    assert(nparray[i, j, k, 2] == 2)
+
 
 if __name__ == '__main__':
     unittest.main()

src/python_coupling/export/BlockForestExport.cpp

@@ -332,6 +332,25 @@ void exportBlockForest(py::module_& m)
       },
       "blocks"_a, "cellsPerBlock"_a, "dx"_a = real_t(1), "oneBlockPerProcess"_a = true,
       "periodic"_a = std::array< bool, 3 >{ false, false, false }, "keepGlobalBlockInformation"_a = false);
+
+    m.def(
+            "createUniformBlockGrid",
+            [](std::array< uint_t, 3 > cells, real_t dx,
+               bool oneBlockPerProcess, std::array< bool, 3 > periodic, bool keepGlobalBlockInformation) {
+                Vector3<uint_t> cellsVec(cells[0], cells[1], cells[2]);
+                Vector3<uint_t> cellsPerBlock;
+                Vector3<uint_t> blocks;
+                uint_t nrOfProcesses = uint_c( MPIManager::instance()->numProcesses() );
+
+                calculateCellDistribution( cellsVec, nrOfProcesses, blocks, cellsPerBlock );
+
+
+                return blockforest::createUniformBlockGrid(blocks[0], blocks[1], blocks[2], cellsPerBlock[0], cellsPerBlock[1],
+                                                           cellsPerBlock[2], dx, oneBlockPerProcess, periodic[0], periodic[1], periodic[2],
+                                                           keepGlobalBlockInformation);
+            },
+            "cells"_a,"dx"_a = real_t(1), "oneBlockPerProcess"_a = true,
+            "periodic"_a = std::array< bool, 3 >{ false, false, false }, "keepGlobalBlockInformation"_a = false);
 }
 
 } // namespace blockforest

src/python_coupling/helper/SliceToCellInterval.h

@@ -38,11 +38,17 @@ inline cell_idx_t normalizeIdx(py::object pyIndex, uint_t coordinateSize)
 {
    cell_idx_t index;
 
-   try {
-      index = py::cast< cell_idx_t >(pyIndex);
+   try{
+       index = pyIndex.cast<cell_idx_t>();
    }
-   catch (py::error_already_set & ){
-      throw py::cast_error("Incompatible index data type");
+   catch (std::exception &){
+       try {
+           auto test = pyIndex.cast<real_t>();
+           index = cell_idx_c( test * real_t( coordinateSize) );
+       }
+       catch (std::exception &) {
+           throw py::cast_error("Incompatible index data type");
+       }
    }
 
    if (index < 0)