Experimental: NumPy Interoperability

.github/workflows/compilation.yml

@@ -147,5 +147,81 @@ jobs:
            run: |
              .github/workflows/pip-versions.sh
 
+  python-tests:
+    name: Python tests (${{ matrix.mpi }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - mpi: none
+            mpi_packages: ""
+            mpiexec_preflags: ""
+          - mpi: mpich
+            mpi_packages: >-
+              mpich
+              libfabric-devel
+            mpiexec_preflags: ""
+          - mpi: ompi
+            mpi_packages: openmpi
+            mpiexec_preflags: --oversubscribe
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - uses: mamba-org/setup-micromamba@v3
+        with:
+          environment-name: pyorbit-test
+          create-args: >-
+            python=3.13
+            numpy
+            pytest
+            fftw
+            meson
+            ninja
+            pkg-config
+            setuptools
+            setuptools-scm
+            ${{ matrix.mpi_packages }}
+          cache-downloads: true
+
+      - name: Configure
+        shell: bash -el {0}
+        run: |
+          meson setup build-${{ matrix.mpi }} \
+            --prefix="$CONDA_PREFIX" \
+            --libdir=lib \
+            --buildtype=release \
+            -DUSE_MPI=${{ matrix.mpi }} \
+            -DPyORBIT_EXPERIMENTAL_WITH_NUMPY=true
+
+      - name: Compile
+        shell: bash -el {0}
+        run: |
+          meson compile -C build-${{ matrix.mpi }}
+
+      - name: Install
+        shell: bash -el {0}
+        run: |
+          meson install --no-rebuild -C build-${{ matrix.mpi }}
+
+      - name: Show MPI configuration
+        if: matrix.mpi != 'none'
+        shell: bash -el {0}
+        run: |
+          mpirun --version
+
+      - name: Run Python tests
+        shell: bash -el {0}
+        env:
+          PYORBIT_MPIEXEC_PREFLAGS: ${{ matrix.mpiexec_preflags }}
+        run: |
+          export LD_LIBRARY_PATH="$CONDA_PREFIX/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}"
+          python -m pytest tests/py/ -v
+
   build-docs:
     uses: ./.github/workflows/docs-build.yml

.github/workflows/pip-build.sh

@@ -4,4 +4,3 @@ pip install -U pip
 pip install -r requirements.txt
 pip install -U setuptools
 pip install --no-build-isolation --editable .
-pip install ".[numpy]"

.gitignore

@@ -17,3 +17,4 @@ _api/
 _cpp_api/
 doc/source/reference/*.rst
 doc/source/reference/*.rst.include
+.cache/

meson_options.txt

@@ -1 +1,2 @@
 option('USE_MPI', type: 'string', value: 'auto', description: 'Choose MPI implementation (mpich, openmpi, none, auto)')
+option('PyORBIT_EXPERIMENTAL_WITH_NUMPY', type: 'boolean', value: false, description: 'Use numpy')

py/orbit/bunch_utils/__init__.py

@@ -7,23 +7,20 @@
 
 from .particleidnumber import ParticleIdNumber
 
-# This guards against missing numpy.
-# Should be imporved with some meaningful (and MPI friendly?) warning printed out.
-try:
-    from .serialize import collect_bunch, save_bunch, load_bunch
-    from .serialize import BunchDict, SyncPartDict
-    from .serialize import FileHandler, NumPyHandler
-except ImportError:
-    pass
+__all__ = ["ParticleIdNumber"]
 
-__all__ = []
-# __all__.append("addParticleIdNumbers")  # doesn't exist
-__all__.append("ParticleIdNumber")
-__all__.append("collect_bunch")
-__all__.append("save_bunch")
-__all__.append("load_bunch")
-__all__.append("BunchDict")
-__all__.append("SyncPartDict")
-__all__.append("FileHandler")
-__all__.append("NumPyHandler")
+from .numpy_utils import bunch_from_shared_numpy
+from .serialize import collect_bunch, save_bunch, load_bunch
+from .serialize import BunchDict, SyncPartDict
+from .serialize import FileHandler, NumPyHandler
 
+__all__ += [
+    "bunch_from_shared_numpy",
+    "collect_bunch",
+    "save_bunch",
+    "load_bunch",
+    "BunchDict",
+    "SyncPartDict",
+    "FileHandler",
+    "NumPyHandler",
+]

py/orbit/bunch_utils/meson.build

@@ -3,6 +3,7 @@
 
 py_sources = files([
     '__init__.py',
+    'numpy_utils.py',
     'particleidnumber.py',
     'serialize.py',
 ])

py/orbit/bunch_utils/numpy_utils.py

@@ -0,0 +1,50 @@
+from orbit.core.bunch import Bunch
+from orbit.core import orbit_mpi
+
+import numpy as np
+
+
+def bunch_from_shared_numpy(global_coords: np.ndarray) -> Bunch:
+    """Construct a bunch from the local partition of a shared NumPy array.
+
+    The global particle coordinate array is divided into contiguous, approximately
+    equal partitions using the rank and size of ``MPI_COMM_WORLD``. Each rank reads
+    its partition directly from ``global_coords``; no particle data is
+    transferred through MPI. Consequently, every rank must be able to access
+    the complete input array, for example through a shared memory-mapped file.
+
+    Parameters
+    ----------
+    global_coords : numpy.ndarray
+        Global particle coordinates with shape ``(n_particles, 6)`` ordered as
+        ``(x, xp, y, yp, z, dE)``.
+
+    Returns
+    -------
+    Bunch
+        A bunch containing the calling rank's contiguous partition of the
+        global particle array.
+
+    Raises
+    ------
+    ValueError
+        If ``global_coords`` does not have shape ``(n_particles, 6)``.
+
+    Notes
+    -----
+    If the number of particles is not evenly divisible by the number of MPI
+    ranks, the lowest-numbered ranks receive one additional particle.
+    """
+    comm_world = orbit_mpi.mpi_comm.MPI_COMM_WORLD
+    mpi_size = orbit_mpi.MPI_Comm_size(comm_world)
+    rank = orbit_mpi.MPI_Comm_rank(comm_world)
+
+    global_size = global_coords.shape[0]
+    base = global_size // mpi_size
+    remainder = global_size % mpi_size
+
+    local_size = base + (1 if rank < remainder else 0)
+    start_row = rank * base + min(rank, remainder)
+    stop_row = start_row + local_size
+
+    return Bunch.from_numpy(global_coords[start_row:stop_row, :])

pyproject.toml

@@ -1,26 +1,19 @@
 [build-system]
 build-backend = 'mesonpy'
-requires = ['meson-python', "setuptools>=45", "wheel", "setuptools_scm"]
+requires = ['meson-python', "setuptools>=45", "wheel", "setuptools_scm", "numpy>=2.0"]
 
 [project]
 name = 'PyORBIT'
 dynamic = ["version"]
 description = 'Use meson-python to build c++ anf python modules.'
 requires-python = '>=3.9'
-#dependencies = [
-#'setuptools',
-#'setuptools-scm'
-#]
+dependencies = [
+    'numpy>=2.0',
+]
 authors = [
   {name = 'Alexander Zhukov', email = 'zhukovap@ornl.gov'},
 ]
 
 
 [tool.setuptools_scm]
 
-[project.optional-dependencies]
-numpy = [
-    "numpy",
-]
-
-