Experimental: NumPy Interoperability

.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')

pyproject.toml

@@ -1,6 +1,6 @@
 [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'

src/meson.build

@@ -3,7 +3,23 @@
 
 
 # Add Python installation details
-python = import('python').find_installation('python3', pure: false)
+pymod = import('python')
+with_numpy = get_option('PyORBIT_EXPERIMENTAL_WITH_NUMPY')
+
+if with_numpy
+  message('Compiling with Numpy support')
+  python = pymod.find_installation('python3', pure: false, modules: ['numpy'])
+  numpy_inc = run_command(
+    python, '-c', 'import numpy; print(numpy.get_include())',
+    check: true
+  ).stdout().strip()
+  add_project_arguments('-DPyORBIT_EXPERIMENTAL_WITH_NUMPY=1', language: 'cpp')
+  add_project_arguments('-I' + numpy_inc, language: 'cpp')
+
+else
+    python = pymod.find_installation('python3', pure: false)
+endif
+
 
 # Add C++ compiler details
 cpp = meson.get_compiler('cpp')
@@ -17,39 +33,37 @@ dependencies += dependency('fftw3', version: '>= 3.0.0', required: true)
 
 # Detecting if MPICH or OPENMPI are installed and enabling support if present
 
-
 mpi_use = get_option('USE_MPI')
-# message('MPI_USE is set to', mpi_use)
+
+cpp_args = ['-fPIC', '-std=c++11', '-O3', '-march=native']
 
 if mpi_use == 'mpich'
     message('Requested to use MPICH as the MPI implementation.')
     dependencies += dependency('mpich', version: '>= 4.0.0', required: true)
-    cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1']
+    cpp_args += ['-DUSE_MPI=1']
 
     # Configure dependencies or settings specific to MPICH
 elif mpi_use == 'ompi'
     message('Requested to use OpenMPI as the MPI implementation.')
     dependencies += dependency('ompi', version: '>= 4.0.0', required: true)
-    cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1']
+    cpp_args += ['-DUSE_MPI=1']
 
 elif mpi_use == 'none'
     message('Requested to not use MPI.')
-    cpp_args = ['-fPIC', '-std=c++11']
 
 else
    mpich_dependency = dependency('mpich', version: '>= 4.0.0', required: false)
    openmpi_dependency = dependency('ompi', version: '>= 4.0.0', required: false)
 
    if mpich_dependency.found()
-        cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1']
+        cpp_args = ['-DUSE_MPI=1']
         dependencies += mpich_dependency
         message('Using MPICH as the MPI implementation.')
    elif openmpi_dependency.found()
-        cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1']
+        cpp_args = ['-DUSE_MPI=1']
         dependencies += openmpi_dependency
         message('Using OpenMPI as the MPI implementation.')
    else
-        cpp_args = ['-fPIC', '-std=c++11']
         message('MPI will not be used.')
    endif
 
@@ -280,6 +294,7 @@ inc = include_directories([
 ])
 
 
+
 core_lib = library('core',
     sources: sources,
     include_directories: inc,
@@ -304,7 +319,7 @@ python.extension_module('orbit_mpi',
 python.extension_module('bunch',
     sources: [base + '/bunch_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args, #['-fPIC', '-std=c++11'],
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -313,7 +328,7 @@ python.extension_module('bunch',
 python.extension_module('spacecharge',
     sources: [base + '/spacecharge_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args, #['-fPIC', '-std=c++11'],
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -322,7 +337,7 @@ python.extension_module('spacecharge',
 python.extension_module('trackerrk4',
     sources: [base + '/trackerrk4_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args, #['-fPIC', '-std=c++11'],
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -331,7 +346,7 @@ python.extension_module('trackerrk4',
 python.extension_module('teapot_base',
     sources: [base + '/teapot_base_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args, #['-fPIC', '-std=c++11'],
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -340,7 +355,7 @@ python.extension_module('teapot_base',
 python.extension_module('linac',
     sources: [base + '/linac_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -349,7 +364,7 @@ python.extension_module('linac',
 python.extension_module('orbit_utils',
     sources: [base + '/utils_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -358,7 +373,7 @@ python.extension_module('orbit_utils',
 python.extension_module('aperture',
     sources: [base + '/aperture_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -367,7 +382,7 @@ python.extension_module('aperture',
 python.extension_module('foil',
     sources: [base + '/foil_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -376,7 +391,7 @@ python.extension_module('foil',
 python.extension_module('field_sources',
     sources: [base + '/field_sources_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -385,7 +400,7 @@ python.extension_module('field_sources',
 python.extension_module('rfcavities',
     sources: [base + '/rfcavities_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -394,7 +409,7 @@ python.extension_module('rfcavities',
 python.extension_module('impedances',
     sources: [base + '/impedances_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -403,7 +418,7 @@ python.extension_module('impedances',
 python.extension_module('fieldtracker',
     sources: [base + '/fieldtracker_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -412,7 +427,7 @@ python.extension_module('fieldtracker',
 python.extension_module('collimator',
     sources: [base + '/collimator_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
     subdir: 'orbit/core',
@@ -421,7 +436,7 @@ python.extension_module('collimator',
 python.extension_module('error_base',
    sources: [base + '/error_base_init.cc'],
     include_directories: inc,
-    cpp_args: ['-fPIC', '-std=c++11'],
+    cpp_args: cpp_args,
     dependencies: [core_dep],
     install: true,
    subdir: 'orbit/core',

src/orbit/wrap_bunch.cc

@@ -16,6 +16,19 @@
 
 #include "pyORBIT_Object.hh"
 
+#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY
+#include <numpy/arrayobject.h>
+
+static int ensure_numpy() {
+  static int numpy_initialized = 0;
+  if (!numpy_initialized) {
+    import_array1(-1);
+    numpy_initialized = 1;
+  }
+  return 0;
+}
+#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY
+
 #include "Bunch.hh"
 #include "ParticleAttributesFactory.hh"
 
@@ -623,7 +636,7 @@ namespace wrap_orbit_bunch{
         }
         std::string attr_name_str(attr_name);
         val = cpp_bunch->getBunchAttributeDouble(attr_name_str);
-        return Py_BuildValue("d",val);
+      return Py_BuildValue("d",val);
       }
       else{
         //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF()
@@ -1171,6 +1184,114 @@ namespace wrap_orbit_bunch{
         self->ob_base.ob_type->tp_free((PyObject*)self);
   }
 
+#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY
+static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args) {
+  pyORBIT_Object *pyBunch = (pyORBIT_Object*)self;
+  Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj;
+
+  if (!PyArg_ParseTuple(args, ":to_numpy")) {
+    ORBIT_MPI_Finalize("PyBunch - to_numpy() - no parameters are needed.");
+  }
+
+  const npy_intp nparts = (npy_intp)cpp_Bunch->getSize();
+  const npy_intp ncoords = 6;
+
+  npy_intp dims[2] = { nparts, ncoords };
+
+  PyObject *py_array = PyArray_SimpleNew(2, dims, NPY_FLOAT64);
+  PyArrayObject *arr_obj = (PyArrayObject*)py_array;
+  double *data_buffer = (double*)PyArray_DATA(arr_obj);
+  double **src = cpp_Bunch->coordArr();
+
+  for (npy_intp i = 0; i < nparts; ++i) {
+    for (npy_intp j = 0; j < ncoords; ++j) {
+      data_buffer[j + i*ncoords] = src[i][j];
+    }
+  }
+
+  return py_array;
+}
+
+  static int bunch_fill_from_numpy_args(Bunch *cpp_Bunch, PyObject *args) {
+    PyObject *arr_in = NULL;
+
+    if (!PyArg_ParseTuple(args, "O:from_numpy", &arr_in)) {
+        return -1;
+    }
+
+    PyArrayObject *arr =
+        (PyArrayObject*)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, NPY_ARRAY_IN_ARRAY);
+    if (!arr) return -1;
+
+    if (PyArray_NDIM(arr) != 2) {
+        Py_DECREF(arr);
+        PyErr_SetString(PyExc_ValueError,
+                        "from_numpy: array must be 2-dimensional with shape (nparts, 6)");
+        return -1;
+    }
+
+    const npy_intp nparts  = PyArray_DIM(arr, 0);
+    const npy_intp ncoords = PyArray_DIM(arr, 1);
+
+    if (ncoords != 6) {
+        Py_DECREF(arr);
+        PyErr_SetString(PyExc_ValueError,
+                        "from_numpy: expected coordinate dimension with shape 6 (x, px, y, py, z, dE)");
+        return -1;
+    }
+
+    const double *data = (const double*)PyArray_DATA(arr);
+
+    for (npy_intp i = 0; i < nparts; ++i) {
+        const npy_intp stride = i * ncoords;
+        cpp_Bunch->addParticle(
+            data[stride+0], data[stride+1], data[stride+2],
+            data[stride+3], data[stride+4], data[stride+5]
+        );
+    }
+
+    Py_DECREF(arr);
+    return 0;
+}
+
+static PyObject *Bunch_update_from_numpy(PyObject *self, PyObject *args) {
+  pyORBIT_Object *pyBunch = (pyORBIT_Object*)self;
+  Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj;
+
+  if (cpp_Bunch->getSizeGlobal() > 0) {
+    for (int i = 0; i < cpp_Bunch->getSizeGlobal(); ++i) {
+      cpp_Bunch->deleteParticleFast(i);
+    }
+    cpp_Bunch->compress();
+  }
+
+  if (bunch_fill_from_numpy_args(cpp_Bunch, args) < 0) return NULL;
+
+  Py_RETURN_NONE;
+}
+
+static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *args) {
+  PyObject *py_bunch_obj = PyObject_CallNoArgs(cls);
+  if (!py_bunch_obj) return NULL;
+
+  pyORBIT_Object *pyBunch = (pyORBIT_Object*)py_bunch_obj;
+  Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj;
+
+  if (!cpp_Bunch) {
+    Py_DECREF(py_bunch_obj);
+    PyErr_SetString(PyExc_RuntimeError, "from_numpy: Constructed bunch has NULL cpp_obj");
+    return NULL;
+  }
+
+  if (bunch_fill_from_numpy_args(cpp_Bunch, args) < 0) {
+    Py_DECREF(py_bunch_obj);
+    return NULL;
+  }
+
+  return py_bunch_obj;
+}
+#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY
+
   static PyMethodDef BunchClassMethods[] = {
     //--------------------------------------------------------
     // class Bunch wrapper                        START
@@ -1230,6 +1351,11 @@ namespace wrap_orbit_bunch{
     { "copyEmptyBunchTo",               Bunch_copyEmptyBunchTo              ,METH_VARARGS,"Copy bunch attrubutes and structure to another bunch"},
     { "copyBunchTo",                    Bunch_copyBunchTo                   ,METH_VARARGS,"Copy bunch all info including particles coordinates and attributes to another bunch"},
     { "addParticlesTo",                 Bunch_addParticlesTo                ,METH_VARARGS,"Copy particles coordinates from one bunch to another"},
+#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY
+    { "to_numpy",                       Bunch_to_numpy                      ,METH_VARARGS, "Convert bunch coordinates to a numpy array" },
+    { "update_from_numpy",              Bunch_update_from_numpy             ,METH_VARARGS, "Update bunch coordinates from a numpy array" },
+    { "from_numpy",                     Bunch_from_numpy                    ,METH_VARARGS | METH_CLASS, "Construct a new Bunch from a numpy array" },
+#endif
     {NULL,NULL}
     //--------------------------------------------------------
     // class Bunch wrapper                        STOP
@@ -1301,6 +1427,11 @@ extern "C" {
 
   /* The name of the function was changed to avoid collision with PyImport magic naming */
   PyMODINIT_FUNC initbunch(void) {
+  #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY
+    if (ensure_numpy() != 0) {
+      throw std::runtime_error("NumPy C-API init failed");
+    }
+  #endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY
       //check that the Bunch wrapper is ready
       if(PyType_Ready(&pyORBIT_Bunch_Type) < 0) return NULL;
       Py_INCREF(&pyORBIT_Bunch_Type);