Merge pull request #1424 from OceanParcels/small_fix_mpi_partition_fu…

…nction Simplification of the MPI partition_function explanation
OceanParcels · Sep 4, 2023 · b7738e9 · b7738e9
2 parents bed2400 + fd1269e
commit b7738e9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/docs/examples/documentation_MPI.ipynb b/docs/examples/documentation_MPI.ipynb
@@ -67,10 +67,10 @@
    "outputs": [],
    "source": [
     "def simple_partition_function(coords, mpi_size=1):\n",
-    "    \"\"\"A very simple partition function that assigns particles to processors\"\"\"\n",
-    "    return np.array(\n",
-    "        [int(i) for i in np.linspace(0, mpi_size, coords.shape[0], endpoint=False)]\n",
-    "    )"
+    "    \"\"\"A very simple partition function\n",
+    "    that assigns particles to processors\n",
+    "    \"\"\"\n",
+    "    return np.linspace(0, mpi_size, coords.shape[0], endpoint=False, dtype=np.int32)"
    ]
   },
   {

diff --git a/docs/examples/example_stommel.py b/docs/examples/example_stommel.py
@@ -88,7 +88,7 @@ def AgeP(particle, fieldset, time):
 
 def simple_partition_function(coords, mpi_size=1):
     """A very simple partition function that assigns particles to processors (for MPI testing purposes))"""
-    return np.array([int(i) for i in np.linspace(0, mpi_size, coords.shape[0], endpoint=False)])
+    return np.linspace(0, mpi_size, coords.shape[0], endpoint=False, dtype=np.int32)
 
 
 def stommel_example(npart=1, mode='jit', verbose=False, method=AdvectionRK4, grid_type='A',