Expose parrsb_part_mesh() in Python wrapper

bin/test.sh

@@ -100,10 +100,10 @@ echo "  uv    : `which ${UV}`"
 echo "  cmake : `which ${CMAKE}`"
 echo "  mpirun: `which ${MPIRUN}`, opts: \"${MPIOPTS} -np ${NP}\""
 echo -e "${RESET}"
-build_genbox
-build_meshes
 init_venv
 build_parrsb
 build_parrsb_py
+build_genbox
+build_meshes
 run_tests
 echo -e "${GREEN}Tests passed.${RESET}"

src/parrsb.pyx

@@ -5,9 +5,26 @@ from parrsb cimport parrsb_conn_mesh, parrsb_read_mesh
 import numpy as np
 
 
+cdef class Options:
+    cdef parrsb_options_t opts
+
+    def __cinit__(self):
+        parrsb_options_get_default(&self.opts)
+
+    def partitioner(self, int partitioner):
+        parrsb_options_set_partitioner(self.opts, partitioner)
+
+    def rsb_algo(self, int rsb_algo):
+        parrsb_options_set_rsb_algo(self.opts, rsb_algo)
+
+    def __dealloc__(self):
+        parrsb_options_free(&self.opts)
+
+
 cdef class Mesh:
     cdef unsigned nel, nv, nbcs, ndim
     cdef long long *bcs
+    cdef long long *vl
     cdef double *coord
     cdef public MPI.Comm c
 
@@ -20,20 +37,50 @@ cdef class Mesh:
             raise RuntimeError("Only 3D meshes are supported!")
         self.ndim = 3
 
-    def find_connectivity(self, double tol):
         cdef unsigned long ndof = <cython.ulong>self.nel * <cython.ulong>self.nv
-        cdef long long *vl = <long long *>malloc(ndof * cython.sizeof(cython.longlong))
+        self.vl = <long long *>malloc(ndof * cython.sizeof(cython.longlong))
+
+    cdef _parrsb_conn_mesh(self, double tol=0.2):
+        return parrsb_conn_mesh(self.vl, self.coord, self.nel, self.ndim, self.bcs,
+                                self.nbcs, tol, self.c.ob_mpi)
+
+    def partition(self, long long[:, :] vtx=None, Options opts=None):
+        cdef long long [:, :] vtx_
+        cdef long long *vl_
 
-        cdef int err = parrsb_conn_mesh(vl, self.coord, self.nel, self.ndim, self.bcs,
-                                        self.nbcs, tol, self.c.ob_mpi)
+        if vtx is None:
+            self._parrsb_conn_mesh()
+            vl_ = self.vl
+        else:
+            vtx_= np.ascontiguousarray(vtx)
+            vl_ = &vtx_[0, 0]
 
-        arr = np.zeros((self.nel, self.nv), dtype=int)
+        opts_ = Options() if opts is None else opts
+
+        cdef int *part = <int *>malloc(self.nel * cython.sizeof(cython.int))
+        cdef int err = parrsb_part_mesh(part, vl_, self.coord, NULL, self.nel, self.nv,
+                                        opts_.opts, self.c.ob_mpi)
+
+        arr = np.zeros(self.nel, dtype=np.int32)
         if err == 0:
             for e in range(self.nel):
-                for v in range(self.nv):
-                    arr[e, v] = vl[e * self.nv + v]
-        free(vl)
+                arr[e] = part[e]
+        free(part)
+
+        if err != 0:
+            raise RuntimeError(f"parrsb_part_mesh() failed with error = {err}")
+        return arr
 
+    def connectivity(self, double tol):
+        cdef int err = self._parrsb_conn_mesh(tol)
+
+        arr = np.zeros((self.nel, self.nv), dtype=np.int64)
+        if err == 0:
+            for e in range(self.nel):
+                for v in range(self.nv):
+                    arr[e, v] = self.vl[e * self.nv + v]
+        else:
+            raise RuntimeError(f"parrsb_conn_mesh() failed with error = {err}")
         return arr
 
     @property
@@ -55,3 +102,4 @@ cdef class Mesh:
     def __dealloc__(self):
         free(self.bcs)
         free(self.coord)
+        free(self.vl)

tests/10_find_connectivity.py → tests/10_connectivity.py

@@ -10,7 +10,7 @@
     comm = MPI.COMM_WORLD
     m = Mesh(sys.argv[1], comm)
 
-    vtx = m.find_connectivity(0.2)
+    vtx = m.connectivity(0.2)
 
     minv = comm.allreduce(np.min(vtx), op=MPI.MIN)
     maxv = comm.allreduce(np.max(vtx), op=MPI.MAX)

tests/20_partition.py

@@ -0,0 +1,18 @@
+import numpy as np
+from mpi4py import MPI
+from parrsb import Mesh
+import sys
+
+if __name__ == "__main__":
+    if len(sys.argv) == 1:
+        sys.exit(1)
+
+    comm = MPI.COMM_WORLD
+    m = Mesh(sys.argv[1], comm)
+
+    # Let parRSB.py calculate the vertices.
+    part = m.partition()
+
+    # Use vertices already calculated.
+    vtx = m.connectivity(0.2)
+    part = m.partition(vtx)

tests/box/10_find_connectivity.box → tests/box/10_connectivity.box

@@ -1,19 +1,6 @@
 base.rea
 -3                     spatial dimension  ( < 0 --> generate .rea/.re2 pair)
 1                      number of fields
-#=======================================================================
-#
-#    Example of .box file for Taylor-Green
-#
-#    If nelx (y or z) < 0, then genbox automatically generates the
-#                          grid spacing in the x (y or z) direction
-#                          with a geometric ratio given by "ratio".
-#                          ( ratio=1 implies uniform spacing )
-#
-#    Note that the character bcs _must_ have 3 spaces.
-#
-#=======================================================================
-#
 Box
 -8  -1  -1                                           nelx,nely,nelz for Box ! multiple of 3
 0 1 1.                                               x0,x1,gain  (rescaled in usrdat)

tests/box/20_partition.box

@@ -0,0 +1,9 @@
+base.rea
+-3                     spatial dimension  ( < 0 --> generate .rea/.re2 pair)
+1                      number of fields
+Box
+-4  -2  -1                                           nelx,nely,nelz for Box ! multiple of 3
+0 1 1.                                               x0,x1,gain  (rescaled in usrdat)
+0 1 1.                                               y0,y1,gain  (rescaled in usrdat)
+0 1 1.                                               z0,z1,gain
+W  ,W  ,W  ,W  ,W  ,W                                bc's  (3 chars each!)