From 224c71477165f6ec557deaa965f51a46b96bb274 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 21 May 2026 16:52:07 -0400 Subject: [PATCH 01/17] NumPy+Grid3D interop Add optional compile-time flag that will compile wrap_grid3D against numpy, enabling to/from_numpy methods for Grid3D instances. When flag is disabled, to/from_numpy don't do anything. --- meson_options.txt | 1 + pyproject.toml | 2 +- src/meson.build | 20 +++++- src/spacecharge/wrap_grid3D.cc | 128 ++++++++++++++++++++++++++++++--- 4 files changed, 139 insertions(+), 12 deletions(-) diff --git a/meson_options.txt b/meson_options.txt index 3ae5f6ab..2e76a4d0 100644 --- a/meson_options.txt +++ b/meson_options.txt @@ -1 +1,2 @@ option('USE_MPI', type: 'string', value: 'auto', description: 'Choose MPI implementation (mpich, openmpi, none, auto)') +option('WITH_NUMPY', type: 'boolean', value: true, description: 'Use numpy') diff --git a/pyproject.toml b/pyproject.toml index a65921e6..a53ad656 100644 --- a/pyproject.toml +++ b/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' diff --git a/src/meson.build b/src/meson.build index c223553d..a70ecf74 100644 --- a/src/meson.build +++ b/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('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('-DWITH_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,7 +33,6 @@ 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) @@ -280,6 +295,7 @@ inc = include_directories([ ]) + core_lib = library('core', sources: sources, include_directories: inc, diff --git a/src/spacecharge/wrap_grid3D.cc b/src/spacecharge/wrap_grid3D.cc index 5dcc32b0..e51356f5 100644 --- a/src/spacecharge/wrap_grid3D.cc +++ b/src/spacecharge/wrap_grid3D.cc @@ -7,6 +7,20 @@ #include +#ifdef WITH_NUMPY + #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION + #include + +static int ensure_numpy() { + static int numpy_initialized = 0; + if (!numpy_initialized) { + import_array1(-1); + numpy_initialized = 1; + } + return 0; +} +#endif + #include "Grid3D.hh" using namespace OrbitUtils; @@ -318,6 +332,95 @@ extern "C" { self->ob_base.ob_type->tp_free((PyObject*)self); } + static PyObject* Grid3D_to_numpy(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object*)self; + Grid3D *cpp_Grid3D = (Grid3D*)pyGrid3D->cpp_obj; + +#ifndef WITH_NUMPY + PyErr_SetString(PyExc_RuntimeError, + "Grid3D.from_numpy is unavailable: built without NumPy support (compile with -DWITH_NUMPY)."); + return NULL; +#else + + if(!PyArg_ParseTuple(args, ":to_numpy")) { + ORBIT_MPI_Finalize("PyGrid3D - to_numpy() - no parameters are needed."); + } + + const npy_intp nz = (npy_intp)cpp_Grid3D->getSizeZ(); + const npy_intp nx = (npy_intp)cpp_Grid3D->getSizeX(); + const npy_intp ny = (npy_intp)cpp_Grid3D->getSizeY(); + + npy_intp dims[3] = {nz, nx, ny}; + PyObject *arr_obj = PyArray_SimpleNew(3, dims, NPY_FLOAT64); + if(NULL == arr_obj) { return NULL; } + + PyArrayObject *arr = (PyArrayObject*)arr_obj; + double *out_buffer = (double*)PyArray_DATA(arr); + double ***src = cpp_Grid3D->getArr3D(); + + for(npy_intp iz = 0; iz < nz; ++iz) { + for(npy_intp ix = 0; ix < ny; ++ix) { + for(npy_intp iy = 0; iy < nx; ++iy) { + out_buffer[iy + ix*ny + iz*nx*ny] = src[iz][ix][iy]; + } + } + } + + return arr_obj; +#endif + } + + static PyObject* Grid3D_from_numpy(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object*)self; + Grid3D *cpp_Grid3D = (Grid3D*)pyGrid3D->cpp_obj; +#ifndef WITH_NUMPY + PyErr_SetString(PyExc_RuntimeError, + "Grid3D.from_numpy is unavailable: built without NumPy support (compile with -DWITH_NUMPY)."); + return NULL; +#else + + PyObject* arr_in = NULL; + if(!PyArg_ParseTuple(args, "O:from_numpy", &arr_in)) { + ORBIT_MPI_Finalize("PyGrid3D - from_numpy() - ndarray is needed."); + } + + PyArrayObject *arr = (PyArrayObject*)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, NPY_ARRAY_IN_ARRAY); + if(NULL == arr) { return NULL; } + + if(PyArray_NDIM(arr) != 3) { + Py_DECREF(arr); + PyErr_SetString(PyExc_ValueError, "from_numpy: array must be 3-dimensional with shape (nz,nx,ny)"); + return NULL; + } + + const npy_intp nz_in = PyArray_DIM(arr, 0); + const npy_intp nx_in = PyArray_DIM(arr, 1); + const npy_intp ny_in = PyArray_DIM(arr, 2); + const npy_intp nz_grid = (npy_intp)cpp_Grid3D->getSizeZ(); + const npy_intp nx_grid = (npy_intp)cpp_Grid3D->getSizeX(); + const npy_intp ny_grid = (npy_intp)cpp_Grid3D->getSizeY(); + + if(nz_in != nz_grid || nx_in != nx_grid || ny_in != ny_grid){ + Py_DECREF(arr); + PyErr_SetString(PyExc_ValueError, "from_numpy: shape mismatch; expected (zSize, xSize, ySize)"); + return NULL; + } + const double *in_buffer = (double*)PyArray_DATA(arr); + double ***dst = cpp_Grid3D->getArr3D(); + + for(npy_intp iz = 0; iz < nz_grid; ++iz) { + for(npy_intp ix = 0; ix < nx_grid; ++ix) { + for(npy_intp iy = 0; iy < ny_grid; ++iy) { + dst[iz][ix][iy] = in_buffer[iy + ix*ny_grid + iz*nx_grid*ny_grid]; + } + } + } + + Py_DECREF(arr); + Py_RETURN_NONE; +#endif + } + // defenition of the methods of the python Grid3D wrapper class // they will be vailable from python level static PyMethodDef Grid3DClassMethods[] = { @@ -345,6 +448,8 @@ extern "C" { { "calcGradient", Grid3D_calcGradient, METH_VARARGS,"returns gradient as (gx,gy,gz) for point (x,y,z)"}, { "longWrapping", Grid3D_longWrapping, METH_VARARGS,"set/get isWrapping variable defining long. wrapping policy"}, { "synchronizeMPI", Grid3D_synchronizeMPI, METH_VARARGS,"synchronize through the MPI communicator"}, + { "to_numpy", Grid3D_to_numpy, METH_VARARGS,"converts the 3D grid to a numpy array"}, + { "from_numpy", Grid3D_from_numpy, METH_VARARGS,"converts the numpy array to a 3D grid"}, {NULL} }; @@ -396,15 +501,20 @@ extern "C" { Grid3D_new, /* tp_new */ }; - //-------------------------------------------------- - //Initialization function of the pyGrid3D class - //It will be called from SpaceCharge wrapper initialization - //-------------------------------------------------- - void initGrid3D(PyObject* module){ - if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) return; - Py_INCREF(&pyORBIT_Grid3D_Type); - PyModule_AddObject(module, "Grid3D", (PyObject *)&pyORBIT_Grid3D_Type); - } +//-------------------------------------------------- +//Initialization function of the pyGrid3D class +//It will be called from SpaceCharge wrapper initialization +//-------------------------------------------------- +void initGrid3D(PyObject* module) { +#ifdef WITH_NUMPY + if (ensure_numpy() != 0) { + throw std::runtime_error("NumPy C-API init failed"); + } +#endif + if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) return; + Py_INCREF(&pyORBIT_Grid3D_Type); + PyModule_AddObject(module, "Grid3D", (PyObject *)&pyORBIT_Grid3D_Type); +} #ifdef __cplusplus } From 6259ab290b77af85fc56026eb8ffde3f562126d2 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 21 May 2026 17:08:35 -0400 Subject: [PATCH 02/17] Apply clang-format --style=llvm --- src/spacecharge/wrap_grid3D.cc | 932 +++++++++++++++++---------------- 1 file changed, 479 insertions(+), 453 deletions(-) diff --git a/src/spacecharge/wrap_grid3D.cc b/src/spacecharge/wrap_grid3D.cc index e51356f5..40f40f9a 100644 --- a/src/spacecharge/wrap_grid3D.cc +++ b/src/spacecharge/wrap_grid3D.cc @@ -1,15 +1,15 @@ -#include "orbit_mpi.hh" -#include "pyORBIT_Object.hh" - #include "wrap_grid3D.hh" -#include "wrap_spacecharge.hh" -#include "wrap_bunch.hh" #include +#include "orbit_mpi.hh" +#include "pyORBIT_Object.hh" +#include "wrap_bunch.hh" +#include "wrap_spacecharge.hh" + #ifdef WITH_NUMPY - #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION - #include +#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION +#include static int ensure_numpy() { static int numpy_initialized = 0; @@ -25,493 +25,519 @@ static int ensure_numpy() { using namespace OrbitUtils; -namespace wrap_spacecharge{ +namespace wrap_spacecharge { #ifdef __cplusplus extern "C" { #endif - //--------------------------------------------------------- - //Python Grid3D class definition - //--------------------------------------------------------- - - //constructor for python class wrapping Grid3D instance - //It never will be called directly - static PyObject* Grid3D_new(PyTypeObject *type, PyObject *args, PyObject *kwds) - { - pyORBIT_Object* self; - self = (pyORBIT_Object *) type->tp_alloc(type, 0); - self->cpp_obj = NULL; - return (PyObject *) self; - } - - //initializator for python Grid3D class - //this is implementation of the __init__ method - static int Grid3D_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds){ - int binX, binY, binZ; - if(!PyArg_ParseTuple(args,"iii:__init__",&binX,&binY,&binZ)){ - ORBIT_MPI_Finalize("PyGrid3D - Grid3D(nX,nY,nZ) - constructor needs parameters."); - } - self->cpp_obj = new Grid3D(binX,binY,binZ); - ((Grid3D*) self->cpp_obj)->setPyWrapper((PyObject*) self); - return 0; +//--------------------------------------------------------- +// Python Grid3D class definition +//--------------------------------------------------------- + +// constructor for python class wrapping Grid3D instance +// It never will be called directly +static PyObject *Grid3D_new(PyTypeObject *type, PyObject *args, + PyObject *kwds) { + pyORBIT_Object *self; + self = (pyORBIT_Object *)type->tp_alloc(type, 0); + self->cpp_obj = NULL; + return (PyObject *)self; +} + +// initializator for python Grid3D class +// this is implementation of the __init__ method +static int Grid3D_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds) { + int binX, binY, binZ; + if (!PyArg_ParseTuple(args, "iii:__init__", &binX, &binY, &binZ)) { + ORBIT_MPI_Finalize( + "PyGrid3D - Grid3D(nX,nY,nZ) - constructor needs parameters."); + } + self->cpp_obj = new Grid3D(binX, binY, binZ); + ((Grid3D *)self->cpp_obj)->setPyWrapper((PyObject *)self); + return 0; +} + +// setZero() +static PyObject *Grid3D_setZero(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + cpp_Grid3D->setZero(); + Py_INCREF(Py_None); + return Py_None; +} + +// getValue(double x, double y, double z) +static PyObject *Grid3D_getValue(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + double x, y, z; + if (!PyArg_ParseTuple(args, "ddd:getValue", &x, &y, &z)) { + ORBIT_MPI_Finalize("PyGrid3D - getValue(x,y,z) - parameters are needed."); + } + return Py_BuildValue("d", cpp_Grid3D->getValue(x, y, z)); +} + +// setValue(double value, int ix, int iy) +static PyObject *Grid3D_setValue(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + double val; + int ix, iy, iz; + if (!PyArg_ParseTuple(args, "diii:setValue", &val, &ix, &iy, &iz)) { + ORBIT_MPI_Finalize( + "PyGrid3D - setValue(val,ix,iy,iz) - parameters are needed."); + } + cpp_Grid3D->setValue(val, ix, iy, iz); + Py_INCREF(Py_None); + return Py_None; +} + +// getValueOnGrid(int ix, int iy, int iz) +static PyObject *Grid3D_getValueOnGrid(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int ix, iy, iz; + if (!PyArg_ParseTuple(args, "iii:getValueOnGrid", &ix, &iy, &iz)) { + ORBIT_MPI_Finalize( + "PyGrid3D - getValueOnGrid(ix,iy,iz) - parameters are needed."); + } + return Py_BuildValue("d", cpp_Grid3D->getValueOnGrid(ix, iy, iz)); +} + +// setGridX(double min, double max) +static PyObject *Grid3D_setGridX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + double min, max; + if (!PyArg_ParseTuple(args, "dd:setGridX", &min, &max)) { + ORBIT_MPI_Finalize("PyGrid3D - setGridX(min,max) - parameters are needed."); + } + cpp_Grid3D->setGridX(min, max); + Py_INCREF(Py_None); + return Py_None; +} + +// setGridY(double min, double max, int n) +static PyObject *Grid3D_setGridY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + double min, max; + if (!PyArg_ParseTuple(args, "dd:setGridY", &min, &max)) { + ORBIT_MPI_Finalize("PyGrid3D - setGridY(min,max) - parameters are needed."); + } + cpp_Grid3D->setGridY(min, max); + Py_INCREF(Py_None); + return Py_None; +} + +// setGridZ(double min, double max) +static PyObject *Grid3D_setGridZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + double min, max; + if (!PyArg_ParseTuple(args, "dd:setGridZ", &min, &max)) { + ORBIT_MPI_Finalize("PyGrid3D - setGridZ(min,max) - parameters are needed."); + } + cpp_Grid3D->setGridZ(min, max); + Py_INCREF(Py_None); + return Py_None; +} + +// getGridX(ix) +static PyObject *Grid3D_getGridX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int ind = -1; + if (!PyArg_ParseTuple(args, "i:getGridX", &ind) || ind < 0 || + ind >= cpp_Grid3D->getSizeX()) { + ORBIT_MPI_Finalize( + "PyGrid3D - getGridX(ix) - parameter is needed. [0 - sizeX["); + } + return Py_BuildValue("d", cpp_Grid3D->getGridX(ind)); +} + +// getGridY(iy) +static PyObject *Grid3D_getGridY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int ind = -1; + if (!PyArg_ParseTuple(args, "i:getGridY", &ind) || ind < 0 || + ind >= cpp_Grid3D->getSizeY()) { + ORBIT_MPI_Finalize( + "PyGrid3D - getGridY(iy) - parameter is needed. [0 - sizeY["); + } + return Py_BuildValue("d", cpp_Grid3D->getGridY(ind)); +} + +// getGridZ(iy) +static PyObject *Grid3D_getGridZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int ind = -1; + if (!PyArg_ParseTuple(args, "i:getGridZ", &ind) || ind < 0 || + ind >= cpp_Grid3D->getSizeZ()) { + ORBIT_MPI_Finalize( + "PyGrid3D - getGridZ(iz) - parameter is needed. [0 - sizeZ["); + } + return Py_BuildValue("d", cpp_Grid3D->getGridZ(ind)); +} + +// getSizeX() +static PyObject *Grid3D_getSizeX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("i", cpp_Grid3D->getSizeX()); +} + +// getSizeY() +static PyObject *Grid3D_getSizeY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("i", cpp_Grid3D->getSizeY()); +} + +// getSizeZ() +static PyObject *Grid3D_getSizeZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("i", cpp_Grid3D->getSizeZ()); +} + +// It will synchronize through the MPI communicator +static PyObject *Grid3D_synchronizeMPI(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int nVars = PyTuple_Size(args); + if (nVars == 0) { + cpp_Grid3D->synchronizeMPI(NULL); + } else { + PyObject *py_mpi_comm_type = + wrap_orbit_mpi_comm::getMPI_CommType("MPI_Comm"); + PyObject *pyMPIComm = PyTuple_GetItem(args, 0); + if ((!PyObject_IsInstance(pyMPIComm, py_mpi_comm_type))) { + ORBIT_MPI_Finalize("Grid3D.synchronizeMPI(MPI_Comm) - input " + "parameter is not MPI_Comm"); + } + cpp_Grid3D->synchronizeMPI((pyORBIT_MPI_Comm *)pyMPIComm); + } + Py_INCREF(Py_None); + return Py_None; +} + +// getMinX() +static PyObject *Grid3D_getMinX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMinX()); +} + +// getMaxX() +static PyObject *Grid3D_getMaxX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMaxX()); +} + +// getMinY() +static PyObject *Grid3D_getMinY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMinY()); +} + +// getMaxY() +static PyObject *Grid3D_getMaxY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMaxY()); +} + +// getMinZ() +static PyObject *Grid3D_getMinZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMinZ()); +} + +// getMaxZ() +static PyObject *Grid3D_getMaxZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMaxZ()); +} + +// longWrapping([isWrapped]) set or return the longitudinal wrapping policy +static PyObject *Grid3D_longWrapping(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + PyObject *pyIsWrapped = NULL; + if (!PyArg_ParseTuple(args, "|O:longWrapping", &pyIsWrapped)) { + ORBIT_MPI_Finalize("PyGrid3D - longWrapping([True/False]) - " + "parameter may be needed."); + } + if (pyIsWrapped != NULL) { + int isWrapped = PyObject_IsTrue(pyIsWrapped); + if (isWrapped != 1) + isWrapped = 0; + cpp_Grid3D->setLongWrapping(isWrapped); + } + return Py_BuildValue("i", cpp_Grid3D->getLongWrapping()); +} + +// binBunch(Bunch* bunch) +static PyObject *Grid3D_binBunch(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + PyObject *pyBunch; + double lambda = -1.0; + if (!PyArg_ParseTuple(args, "O|d:binBunch", &pyBunch, &lambda)) { + ORBIT_MPI_Finalize("PyGrid3D - binBunch(Bunch* bunch [,lambda]) - " + "parameters are needed."); + } + PyObject *pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); + if (!PyObject_IsInstance(pyBunch, pyORBIT_Bunch_Type)) { + ORBIT_MPI_Finalize("PyGrid3D - binBunch(Bunch* bunch [,lambda]) - " + "method needs a Bunch."); + } + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)pyBunch)->cpp_obj; + if (lambda > 0.) { + cpp_Grid3D->binBunch(cpp_bunch, lambda); + } else { + cpp_Grid3D->binBunch(cpp_bunch); } + Py_INCREF(Py_None); + return Py_None; +} - //setZero() - static PyObject* Grid3D_setZero(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - cpp_Grid3D->setZero(); - Py_INCREF(Py_None); - return Py_None; - } - - //getValue(double x, double y, double z) - static PyObject* Grid3D_getValue(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - double x,y,z; - if(!PyArg_ParseTuple(args,"ddd:getValue",&x,&y,&z)){ - ORBIT_MPI_Finalize("PyGrid3D - getValue(x,y,z) - parameters are needed."); - } - return Py_BuildValue("d",cpp_Grid3D->getValue(x,y,z)); - } - - //setValue(double value, int ix, int iy) - static PyObject* Grid3D_setValue(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - double val; - int ix,iy,iz; - if(!PyArg_ParseTuple(args,"diii:setValue",&val,&ix,&iy,&iz)){ - ORBIT_MPI_Finalize("PyGrid3D - setValue(val,ix,iy,iz) - parameters are needed."); - } - cpp_Grid3D->setValue(val,ix,iy,iz); - Py_INCREF(Py_None); - return Py_None; - } - - //getValueOnGrid(int ix, int iy, int iz) - static PyObject* Grid3D_getValueOnGrid(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - int ix,iy,iz; - if(!PyArg_ParseTuple(args,"iii:getValueOnGrid",&ix,&iy,&iz)){ - ORBIT_MPI_Finalize("PyGrid3D - getValueOnGrid(ix,iy,iz) - parameters are needed."); - } - return Py_BuildValue("d",cpp_Grid3D->getValueOnGrid(ix,iy,iz)); - } - - //setGridX(double min, double max) - static PyObject* Grid3D_setGridX(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - double min,max; - if(!PyArg_ParseTuple(args,"dd:setGridX",&min,&max)){ - ORBIT_MPI_Finalize("PyGrid3D - setGridX(min,max) - parameters are needed."); - } - cpp_Grid3D->setGridX(min,max); - Py_INCREF(Py_None); - return Py_None; - } - - //setGridY(double min, double max, int n) - static PyObject* Grid3D_setGridY(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - double min,max; - if(!PyArg_ParseTuple(args,"dd:setGridY",&min,&max)){ - ORBIT_MPI_Finalize("PyGrid3D - setGridY(min,max) - parameters are needed."); - } - cpp_Grid3D->setGridY(min,max); - Py_INCREF(Py_None); - return Py_None; - } - - //setGridZ(double min, double max) - static PyObject* Grid3D_setGridZ(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - double min,max; - if(!PyArg_ParseTuple(args,"dd:setGridZ",&min,&max)){ - ORBIT_MPI_Finalize("PyGrid3D - setGridZ(min,max) - parameters are needed."); - } - cpp_Grid3D->setGridZ(min,max); - Py_INCREF(Py_None); - return Py_None; - } - - //getGridX(ix) - static PyObject* Grid3D_getGridX(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - int ind = -1; - if(!PyArg_ParseTuple(args,"i:getGridX",&ind) || ind < 0 || ind >= cpp_Grid3D->getSizeX()){ - ORBIT_MPI_Finalize("PyGrid3D - getGridX(ix) - parameter is needed. [0 - sizeX["); - } - return Py_BuildValue("d",cpp_Grid3D->getGridX(ind)); - } - - //getGridY(iy) - static PyObject* Grid3D_getGridY(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - int ind = -1; - if(!PyArg_ParseTuple(args,"i:getGridY",&ind) || ind < 0 || ind >= cpp_Grid3D->getSizeY()){ - ORBIT_MPI_Finalize("PyGrid3D - getGridY(iy) - parameter is needed. [0 - sizeY["); - } - return Py_BuildValue("d",cpp_Grid3D->getGridY(ind)); - } - - //getGridZ(iy) - static PyObject* Grid3D_getGridZ(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - int ind = -1; - if(!PyArg_ParseTuple(args,"i:getGridZ",&ind) || ind < 0 || ind >= cpp_Grid3D->getSizeZ()){ - ORBIT_MPI_Finalize("PyGrid3D - getGridZ(iz) - parameter is needed. [0 - sizeZ["); - } - return Py_BuildValue("d",cpp_Grid3D->getGridZ(ind)); - } - - //getSizeX() - static PyObject* Grid3D_getSizeX(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("i",cpp_Grid3D->getSizeX()); - } - - //getSizeY() - static PyObject* Grid3D_getSizeY(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("i",cpp_Grid3D->getSizeY()); - } - - //getSizeZ() - static PyObject* Grid3D_getSizeZ(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("i",cpp_Grid3D->getSizeZ()); - } - - //It will synchronize through the MPI communicator - static PyObject* Grid3D_synchronizeMPI(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - int nVars = PyTuple_Size(args); - if(nVars == 0){ - cpp_Grid3D->synchronizeMPI(NULL); - } - else { - PyObject* py_mpi_comm_type = wrap_orbit_mpi_comm::getMPI_CommType("MPI_Comm"); - PyObject* pyMPIComm = PyTuple_GetItem(args,0); - if((!PyObject_IsInstance(pyMPIComm,py_mpi_comm_type))){ - ORBIT_MPI_Finalize("Grid3D.synchronizeMPI(MPI_Comm) - input parameter is not MPI_Comm"); - } - cpp_Grid3D->synchronizeMPI((pyORBIT_MPI_Comm*) pyMPIComm); - } - Py_INCREF(Py_None); - return Py_None; +// binValue(double value, double x, double y, double z) +static PyObject *Grid3D_binValue(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + double val, x, y, z; + if (!PyArg_ParseTuple(args, "dddd:binValue", &val, &x, &y, &z)) { + ORBIT_MPI_Finalize( + "PyGrid3D - binValue(val,x,y,z) - parameters are needed."); } + cpp_Grid3D->binValue(val, x, y, z); + Py_INCREF(Py_None); + return Py_None; +} - //getMinX() - static PyObject* Grid3D_getMinX(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("d",cpp_Grid3D->getMinX()); - } - - //getMaxX() - static PyObject* Grid3D_getMaxX(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("d",cpp_Grid3D->getMaxX()); - } - - //getMinY() - static PyObject* Grid3D_getMinY(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("d",cpp_Grid3D->getMinY()); - } - - //getMaxY() - static PyObject* Grid3D_getMaxY(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("d",cpp_Grid3D->getMaxY()); - } - - //getMinZ() - static PyObject* Grid3D_getMinZ(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("d",cpp_Grid3D->getMinZ()); - } - - //getMaxZ() - static PyObject* Grid3D_getMaxZ(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - return Py_BuildValue("d",cpp_Grid3D->getMaxZ()); - } - - //longWrapping([isWrapped]) set or return the longitudinal wrapping policy - static PyObject* Grid3D_longWrapping(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - PyObject* pyIsWrapped = NULL; - if(!PyArg_ParseTuple(args,"|O:longWrapping",&pyIsWrapped)){ - ORBIT_MPI_Finalize("PyGrid3D - longWrapping([True/False]) - parameter may be needed."); - } - if(pyIsWrapped != NULL){ - int isWrapped = PyObject_IsTrue(pyIsWrapped); - if(isWrapped != 1) isWrapped = 0; - cpp_Grid3D->setLongWrapping(isWrapped); - } - return Py_BuildValue("i",cpp_Grid3D->getLongWrapping()); - } - - //binBunch(Bunch* bunch) - static PyObject* Grid3D_binBunch(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - PyObject* pyBunch; - double lambda = -1.0; - if(!PyArg_ParseTuple(args,"O|d:binBunch",&pyBunch,&lambda)){ - ORBIT_MPI_Finalize("PyGrid3D - binBunch(Bunch* bunch [,lambda]) - parameters are needed."); - } - PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); - if(!PyObject_IsInstance(pyBunch,pyORBIT_Bunch_Type)){ - ORBIT_MPI_Finalize("PyGrid3D - binBunch(Bunch* bunch [,lambda]) - method needs a Bunch."); - } - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object*)pyBunch)->cpp_obj; - if(lambda > 0.){ - cpp_Grid3D->binBunch(cpp_bunch,lambda); - } else { - cpp_Grid3D->binBunch(cpp_bunch); - } - Py_INCREF(Py_None); - return Py_None; - } - - //binValue(double value, double x, double y, double z) - static PyObject* Grid3D_binValue(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - double val,x,y,z; - if(!PyArg_ParseTuple(args,"dddd:binValue",&val,&x,&y,&z)){ - ORBIT_MPI_Finalize("PyGrid3D - binValue(val,x,y,z) - parameters are needed."); - } - cpp_Grid3D->binValue(val,x,y,z); - Py_INCREF(Py_None); - return Py_None; - } - - //calcGradient(double x, double y, double y) - static PyObject* Grid3D_calcGradient(PyObject *self, PyObject *args){ - pyORBIT_Object* pyGrid3D = (pyORBIT_Object*) self; - Grid3D* cpp_Grid3D = (Grid3D*) pyGrid3D->cpp_obj; - double x,y,z; - double ex, ey, ez; - if(!PyArg_ParseTuple(args,"ddd:calcGradient",&x,&y,&z)){ - ORBIT_MPI_Finalize("PyGrid3D - calcGradient(x,y,z) - parameters are needed."); - } - cpp_Grid3D->calcGradient(x,ex,y,ey,z,ez); - return Py_BuildValue("(ddd)",ex,ey,ez); - } - - //----------------------------------------------------- - //destructor for python Grid3D class (__del__ method). - //----------------------------------------------------- - static void Grid3D_del(pyORBIT_Object* self){ - //std::cerr<<"The Grid3D __del__ has been called!"<cpp_obj; - delete cpp_Grid3D; - self->ob_base.ob_type->tp_free((PyObject*)self); +// calcGradient(double x, double y, double y) +static PyObject *Grid3D_calcGradient(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + double x, y, z; + double ex, ey, ez; + if (!PyArg_ParseTuple(args, "ddd:calcGradient", &x, &y, &z)) { + ORBIT_MPI_Finalize( + "PyGrid3D - calcGradient(x,y,z) - parameters are needed."); } + cpp_Grid3D->calcGradient(x, ex, y, ey, z, ez); + return Py_BuildValue("(ddd)", ex, ey, ez); +} - static PyObject* Grid3D_to_numpy(PyObject *self, PyObject *args) { - pyORBIT_Object *pyGrid3D = (pyORBIT_Object*)self; - Grid3D *cpp_Grid3D = (Grid3D*)pyGrid3D->cpp_obj; +//----------------------------------------------------- +// destructor for python Grid3D class (__del__ method). +//----------------------------------------------------- +static void Grid3D_del(pyORBIT_Object *self) { + // std::cerr<<"The Grid3D __del__ has been called!"<cpp_obj; + delete cpp_Grid3D; + self->ob_base.ob_type->tp_free((PyObject *)self); +} + +static PyObject *Grid3D_to_numpy(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; #ifndef WITH_NUMPY - PyErr_SetString(PyExc_RuntimeError, - "Grid3D.from_numpy is unavailable: built without NumPy support (compile with -DWITH_NUMPY)."); - return NULL; + PyErr_SetString(PyExc_RuntimeError, + "Grid3D.from_numpy is unavailable: built without NumPy " + "support (compile with -DWITH_NUMPY)."); + return NULL; #else - if(!PyArg_ParseTuple(args, ":to_numpy")) { - ORBIT_MPI_Finalize("PyGrid3D - to_numpy() - no parameters are needed."); - } + if (!PyArg_ParseTuple(args, ":to_numpy")) { + ORBIT_MPI_Finalize("PyGrid3D - to_numpy() - no parameters are needed."); + } - const npy_intp nz = (npy_intp)cpp_Grid3D->getSizeZ(); - const npy_intp nx = (npy_intp)cpp_Grid3D->getSizeX(); - const npy_intp ny = (npy_intp)cpp_Grid3D->getSizeY(); + const npy_intp nz = (npy_intp)cpp_Grid3D->getSizeZ(); + const npy_intp nx = (npy_intp)cpp_Grid3D->getSizeX(); + const npy_intp ny = (npy_intp)cpp_Grid3D->getSizeY(); - npy_intp dims[3] = {nz, nx, ny}; - PyObject *arr_obj = PyArray_SimpleNew(3, dims, NPY_FLOAT64); - if(NULL == arr_obj) { return NULL; } + npy_intp dims[3] = {nz, nx, ny}; + PyObject *arr_obj = PyArray_SimpleNew(3, dims, NPY_FLOAT64); + if (NULL == arr_obj) { + return NULL; + } - PyArrayObject *arr = (PyArrayObject*)arr_obj; - double *out_buffer = (double*)PyArray_DATA(arr); - double ***src = cpp_Grid3D->getArr3D(); + PyArrayObject *arr = (PyArrayObject *)arr_obj; + double *out_buffer = (double *)PyArray_DATA(arr); + double ***src = cpp_Grid3D->getArr3D(); - for(npy_intp iz = 0; iz < nz; ++iz) { - for(npy_intp ix = 0; ix < ny; ++ix) { - for(npy_intp iy = 0; iy < nx; ++iy) { - out_buffer[iy + ix*ny + iz*nx*ny] = src[iz][ix][iy]; - } + for (npy_intp iz = 0; iz < nz; ++iz) { + for (npy_intp ix = 0; ix < ny; ++ix) { + for (npy_intp iy = 0; iy < nx; ++iy) { + out_buffer[iy + ix * ny + iz * nx * ny] = src[iz][ix][iy]; } } + } - return arr_obj; + return arr_obj; #endif - } +} + +static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; - static PyObject* Grid3D_from_numpy(PyObject *self, PyObject *args) { - pyORBIT_Object *pyGrid3D = (pyORBIT_Object*)self; - Grid3D *cpp_Grid3D = (Grid3D*)pyGrid3D->cpp_obj; #ifndef WITH_NUMPY PyErr_SetString(PyExc_RuntimeError, - "Grid3D.from_numpy is unavailable: built without NumPy support (compile with -DWITH_NUMPY)."); + "Grid3D.from_numpy is unavailable: built without NumPy " + "support (compile with -DWITH_NUMPY)."); return NULL; #else - PyObject* arr_in = NULL; - if(!PyArg_ParseTuple(args, "O:from_numpy", &arr_in)) { - ORBIT_MPI_Finalize("PyGrid3D - from_numpy() - ndarray is needed."); - } + PyObject *arr_in = NULL; + if (!PyArg_ParseTuple(args, "O:from_numpy", &arr_in)) { + ORBIT_MPI_Finalize("PyGrid3D - from_numpy() - ndarray is needed."); + } - PyArrayObject *arr = (PyArrayObject*)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, NPY_ARRAY_IN_ARRAY); - if(NULL == arr) { return NULL; } + PyArrayObject *arr = (PyArrayObject *)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, + NPY_ARRAY_IN_ARRAY); + if (NULL == arr) { + return NULL; + } - if(PyArray_NDIM(arr) != 3) { - Py_DECREF(arr); - PyErr_SetString(PyExc_ValueError, "from_numpy: array must be 3-dimensional with shape (nz,nx,ny)"); - return NULL; - } + if (PyArray_NDIM(arr) != 3) { + Py_DECREF(arr); + PyErr_SetString( + PyExc_ValueError, + "from_numpy: array must be 3-dimensional with shape (nz,nx,ny)"); + return NULL; + } - const npy_intp nz_in = PyArray_DIM(arr, 0); - const npy_intp nx_in = PyArray_DIM(arr, 1); - const npy_intp ny_in = PyArray_DIM(arr, 2); - const npy_intp nz_grid = (npy_intp)cpp_Grid3D->getSizeZ(); - const npy_intp nx_grid = (npy_intp)cpp_Grid3D->getSizeX(); - const npy_intp ny_grid = (npy_intp)cpp_Grid3D->getSizeY(); - - if(nz_in != nz_grid || nx_in != nx_grid || ny_in != ny_grid){ - Py_DECREF(arr); - PyErr_SetString(PyExc_ValueError, "from_numpy: shape mismatch; expected (zSize, xSize, ySize)"); - return NULL; - } - const double *in_buffer = (double*)PyArray_DATA(arr); - double ***dst = cpp_Grid3D->getArr3D(); - - for(npy_intp iz = 0; iz < nz_grid; ++iz) { - for(npy_intp ix = 0; ix < nx_grid; ++ix) { - for(npy_intp iy = 0; iy < ny_grid; ++iy) { - dst[iz][ix][iy] = in_buffer[iy + ix*ny_grid + iz*nx_grid*ny_grid]; - } + const npy_intp nz_in = PyArray_DIM(arr, 0); + const npy_intp nx_in = PyArray_DIM(arr, 1); + const npy_intp ny_in = PyArray_DIM(arr, 2); + const npy_intp nz_grid = (npy_intp)cpp_Grid3D->getSizeZ(); + const npy_intp nx_grid = (npy_intp)cpp_Grid3D->getSizeX(); + const npy_intp ny_grid = (npy_intp)cpp_Grid3D->getSizeY(); + + if (nz_in != nz_grid || nx_in != nx_grid || ny_in != ny_grid) { + Py_DECREF(arr); + PyErr_SetString( + PyExc_ValueError, + "from_numpy: shape mismatch; expected (zSize, xSize, ySize)"); + return NULL; + } + + const double *in_buffer = (double *)PyArray_DATA(arr); + double ***dst = cpp_Grid3D->getArr3D(); + + for (npy_intp iz = 0; iz < nz_grid; ++iz) { + for (npy_intp ix = 0; ix < nx_grid; ++ix) { + for (npy_intp iy = 0; iy < ny_grid; ++iy) { + dst[iz][ix][iy] = in_buffer[iy + ix * ny_grid + iz * nx_grid * ny_grid]; } } + } - Py_DECREF(arr); - Py_RETURN_NONE; + Py_DECREF(arr); + Py_RETURN_NONE; #endif - } +} - // defenition of the methods of the python Grid3D wrapper class - // they will be vailable from python level - static PyMethodDef Grid3DClassMethods[] = { - { "setZero", Grid3D_setZero, METH_VARARGS,"sets all points on the grid to zero"}, - { "getValue", Grid3D_getValue, METH_VARARGS,"returns value for (x,y,z) point"}, - { "getValueOnGrid", Grid3D_getValueOnGrid, METH_VARARGS,"returns value for indeces (ix,iy,iz) "}, - { "setValue", Grid3D_setValue, METH_VARARGS,"sets value for (ix,iy,iz) point - (val,ix,iy,iz)"}, - { "setGridX", Grid3D_setGridX, METH_VARARGS,"sets the X grid with min,max"}, - { "setGridY", Grid3D_setGridY, METH_VARARGS,"sets the Y grid with min,max"}, - { "setGridZ", Grid3D_setGridZ, METH_VARARGS,"sets the Z grid with min,max"}, - { "getGridX", Grid3D_getGridX, METH_VARARGS,"returns the x-grid point with index ind"}, - { "getGridY", Grid3D_getGridY, METH_VARARGS,"returns the y-grid point with index ind"}, - { "getGridZ", Grid3D_getGridZ, METH_VARARGS,"returns the z-grid point with index ind"}, - { "getSizeX", Grid3D_getSizeX, METH_VARARGS,"returns the size of grid in X dir."}, - { "getSizeY", Grid3D_getSizeY, METH_VARARGS,"returns the size of grid in Y dir."}, - { "getSizeZ", Grid3D_getSizeZ, METH_VARARGS,"returns the size of grid in Z dir."}, - { "getMinX", Grid3D_getMinX, METH_VARARGS,"returns the min grid point in X dir."}, - { "getMaxX", Grid3D_getMaxX, METH_VARARGS,"returns the max grid point in X dir."}, - { "getMinY", Grid3D_getMinY, METH_VARARGS,"returns the min grid point in Y dir."}, - { "getMaxY", Grid3D_getMaxY, METH_VARARGS,"returns the max grid point in Y dir."}, - { "getMinZ", Grid3D_getMinZ, METH_VARARGS,"returns the min grid point in Z dir."}, - { "getMaxZ", Grid3D_getMaxZ, METH_VARARGS,"returns the max grid point in Z dir."}, - { "binValue", Grid3D_binValue, METH_VARARGS,"bins the value into the 3D mesh"}, - { "binBunch", Grid3D_binBunch, METH_VARARGS,"bins the Bunch into the 3D mesh"}, - { "calcGradient", Grid3D_calcGradient, METH_VARARGS,"returns gradient as (gx,gy,gz) for point (x,y,z)"}, - { "longWrapping", Grid3D_longWrapping, METH_VARARGS,"set/get isWrapping variable defining long. wrapping policy"}, - { "synchronizeMPI", Grid3D_synchronizeMPI, METH_VARARGS,"synchronize through the MPI communicator"}, - { "to_numpy", Grid3D_to_numpy, METH_VARARGS,"converts the 3D grid to a numpy array"}, - { "from_numpy", Grid3D_from_numpy, METH_VARARGS,"converts the numpy array to a 3D grid"}, - {NULL} - }; - - // defenition of the memebers of the python Grid3D wrapper class - // they will be vailable from python level - static PyMemberDef Grid3DClassMembers [] = { - {NULL} - }; - - //new python Grid3D wrapper type definition - static PyTypeObject pyORBIT_Grid3D_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - "Grid3D", /*tp_name*/ - sizeof(pyORBIT_Object), /*tp_basicsize*/ - 0, /*tp_itemsize*/ - (destructor) Grid3D_del , /*tp_dealloc*/ - 0, /*tp_print*/ - 0, /*tp_getattr*/ - 0, /*tp_setattr*/ - 0, /*tp_compare*/ - 0, /*tp_repr*/ - 0, /*tp_as_number*/ - 0, /*tp_as_sequence*/ - 0, /*tp_as_mapping*/ - 0, /*tp_hash */ - 0, /*tp_call*/ - 0, /*tp_str*/ - 0, /*tp_getattro*/ - 0, /*tp_setattro*/ - 0, /*tp_as_buffer*/ - Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ - "The Grid3D python wrapper", /* tp_doc */ - 0, /* tp_traverse */ - 0, /* tp_clear */ - 0, /* tp_richcompare */ - 0, /* tp_weaklistoffset */ - 0, /* tp_iter */ - 0, /* tp_iternext */ - Grid3DClassMethods, /* tp_methods */ - Grid3DClassMembers, /* tp_members */ - 0, /* tp_getset */ - 0, /* tp_base */ - 0, /* tp_dict */ - 0, /* tp_descr_get */ - 0, /* tp_descr_set */ - 0, /* tp_dictoffset */ - (initproc) Grid3D_init, /* tp_init */ - 0, /* tp_alloc */ - Grid3D_new, /* tp_new */ - }; +// defenition of the methods of the python Grid3D wrapper class +// they will be vailable from python level +static PyMethodDef Grid3DClassMethods[] = { + {"setZero", Grid3D_setZero, METH_VARARGS, "sets all points on the grid to zero"}, + {"getValue", Grid3D_getValue, METH_VARARGS, "returns value for (x,y,z) point"}, + {"getValueOnGrid", Grid3D_getValueOnGrid, METH_VARARGS, "returns value for indeces (ix,iy,iz) "}, + {"setValue", Grid3D_setValue, METH_VARARGS, "sets value for (ix,iy,iz) point - (val,ix,iy,iz)"}, + {"setGridX", Grid3D_setGridX, METH_VARARGS, "sets the X grid with min,max"}, + {"setGridY", Grid3D_setGridY, METH_VARARGS, "sets the Y grid with min,max"}, + {"setGridZ", Grid3D_setGridZ, METH_VARARGS, "sets the Z grid with min,max"}, + {"getGridX", Grid3D_getGridX, METH_VARARGS, "returns the x-grid point with index ind"}, + {"getGridY", Grid3D_getGridY, METH_VARARGS, "returns the y-grid point with index ind"}, + {"getGridZ", Grid3D_getGridZ, METH_VARARGS, "returns the z-grid point with index ind"}, + {"getSizeX", Grid3D_getSizeX, METH_VARARGS, "returns the size of grid in X dir."}, + {"getSizeY", Grid3D_getSizeY, METH_VARARGS, "returns the size of grid in Y dir."}, + {"getSizeZ", Grid3D_getSizeZ, METH_VARARGS, "returns the size of grid in Z dir."}, + {"getMinX", Grid3D_getMinX, METH_VARARGS, "returns the min grid point in X dir."}, + {"getMaxX", Grid3D_getMaxX, METH_VARARGS, "returns the max grid point in X dir."}, + {"getMinY", Grid3D_getMinY, METH_VARARGS, "returns the min grid point in Y dir."}, + {"getMaxY", Grid3D_getMaxY, METH_VARARGS, "returns the max grid point in Y dir."}, + {"getMinZ", Grid3D_getMinZ, METH_VARARGS, "returns the min grid point in Z dir."}, + {"getMaxZ", Grid3D_getMaxZ, METH_VARARGS, "returns the max grid point in Z dir."}, + {"binValue", Grid3D_binValue, METH_VARARGS, "bins the value into the 3D mesh"}, + {"binBunch", Grid3D_binBunch, METH_VARARGS, "bins the Bunch into the 3D mesh"}, + {"calcGradient", Grid3D_calcGradient, METH_VARARGS, "returns gradient as (gx,gy,gz) for point (x,y,z)"}, + {"longWrapping", Grid3D_longWrapping, METH_VARARGS, "set/get isWrapping variable defining long. wrapping policy"}, + {"synchronizeMPI", Grid3D_synchronizeMPI, METH_VARARGS, "synchronize through the MPI communicator"}, + {"to_numpy", Grid3D_to_numpy, METH_VARARGS, "converts the 3D grid to a numpy array"}, + {"from_numpy", Grid3D_from_numpy, METH_VARARGS, "converts the numpy array to a 3D grid"}, + {NULL}}; + +// defenition of the memebers of the python Grid3D wrapper class +// they will be vailable from python level +static PyMemberDef Grid3DClassMembers[] = {{NULL}}; + +// new python Grid3D wrapper type definition +static PyTypeObject pyORBIT_Grid3D_Type = { + PyVarObject_HEAD_INIT(NULL, 0) "Grid3D", /*tp_name*/ + sizeof(pyORBIT_Object), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor)Grid3D_del, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ + "The Grid3D python wrapper", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + Grid3DClassMethods, /* tp_methods */ + Grid3DClassMembers, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc)Grid3D_init, /* tp_init */ + 0, /* tp_alloc */ + Grid3D_new, /* tp_new */ +}; //-------------------------------------------------- -//Initialization function of the pyGrid3D class -//It will be called from SpaceCharge wrapper initialization +// Initialization function of the pyGrid3D class +// It will be called from SpaceCharge wrapper initialization //-------------------------------------------------- -void initGrid3D(PyObject* module) { +void initGrid3D(PyObject *module) { #ifdef WITH_NUMPY - if (ensure_numpy() != 0) { + if (ensure_numpy() != 0) { throw std::runtime_error("NumPy C-API init failed"); - } + } #endif - if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) return; + if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) + return; Py_INCREF(&pyORBIT_Grid3D_Type); PyModule_AddObject(module, "Grid3D", (PyObject *)&pyORBIT_Grid3D_Type); } @@ -520,5 +546,5 @@ void initGrid3D(PyObject* module) { } #endif -//end of namespace wrap_spacecharge -} +// end of namespace wrap_spacecharge +} // namespace wrap_spacecharge From eacf2d1f4f132d14288f8c10a232b9acd2a6add6 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Fri, 29 May 2026 14:31:21 -0400 Subject: [PATCH 03/17] wip --- src/orbit/wrap_bunch.cc | 133 ++++++++++++++++++++++++++++++++- src/spacecharge/wrap_grid3D.cc | 29 +++---- 2 files changed, 142 insertions(+), 20 deletions(-) diff --git a/src/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index 68824b4e..bb12f911 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -16,6 +16,19 @@ #include "pyORBIT_Object.hh" +#ifdef WITH_NUMPY +#include + +static int ensure_numpy() { + static int numpy_initialized = 0; + if (!numpy_initialized) { + import_array1(-1); + numpy_initialized = 1; + } + return 0; +} +#endif // 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 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 // 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 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 WITH_NUMPY + if (ensure_numpy() != 0) { + throw std::runtime_error("NumPy C-API init failed"); + } + #endif // WITH_NUMPY //check that the Bunch wrapper is ready if(PyType_Ready(&pyORBIT_Bunch_Type) < 0) return NULL; Py_INCREF(&pyORBIT_Bunch_Type); diff --git a/src/spacecharge/wrap_grid3D.cc b/src/spacecharge/wrap_grid3D.cc index 40f40f9a..e55f4780 100644 --- a/src/spacecharge/wrap_grid3D.cc +++ b/src/spacecharge/wrap_grid3D.cc @@ -19,7 +19,7 @@ static int ensure_numpy() { } return 0; } -#endif +#endif // WITH_NUMPY #include "Grid3D.hh" @@ -348,17 +348,11 @@ static void Grid3D_del(pyORBIT_Object *self) { self->ob_base.ob_type->tp_free((PyObject *)self); } +#ifdef WITH_NUMPY static PyObject *Grid3D_to_numpy(PyObject *self, PyObject *args) { pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; -#ifndef WITH_NUMPY - PyErr_SetString(PyExc_RuntimeError, - "Grid3D.from_numpy is unavailable: built without NumPy " - "support (compile with -DWITH_NUMPY)."); - return NULL; -#else - if (!PyArg_ParseTuple(args, ":to_numpy")) { ORBIT_MPI_Finalize("PyGrid3D - to_numpy() - no parameters are needed."); } @@ -386,22 +380,16 @@ static PyObject *Grid3D_to_numpy(PyObject *self, PyObject *args) { } return arr_obj; -#endif } static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; -#ifndef WITH_NUMPY - PyErr_SetString(PyExc_RuntimeError, - "Grid3D.from_numpy is unavailable: built without NumPy " - "support (compile with -DWITH_NUMPY)."); - return NULL; -#else - PyObject *arr_in = NULL; - if (!PyArg_ParseTuple(args, "O:from_numpy", &arr_in)) { + const char* order = "zxy"; + + if (!PyArg_ParseTuple(args, "O|s:from_numpy", &arr_in, &order)) { ORBIT_MPI_Finalize("PyGrid3D - from_numpy() - ndarray is needed."); } @@ -422,6 +410,7 @@ static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { const npy_intp nz_in = PyArray_DIM(arr, 0); const npy_intp nx_in = PyArray_DIM(arr, 1); const npy_intp ny_in = PyArray_DIM(arr, 2); + const npy_intp nz_grid = (npy_intp)cpp_Grid3D->getSizeZ(); const npy_intp nx_grid = (npy_intp)cpp_Grid3D->getSizeX(); const npy_intp ny_grid = (npy_intp)cpp_Grid3D->getSizeY(); @@ -447,8 +436,8 @@ static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { Py_DECREF(arr); Py_RETURN_NONE; -#endif } +#endif // WITH_NUMPY // defenition of the methods of the python Grid3D wrapper class // they will be vailable from python level @@ -477,8 +466,10 @@ static PyMethodDef Grid3DClassMethods[] = { {"calcGradient", Grid3D_calcGradient, METH_VARARGS, "returns gradient as (gx,gy,gz) for point (x,y,z)"}, {"longWrapping", Grid3D_longWrapping, METH_VARARGS, "set/get isWrapping variable defining long. wrapping policy"}, {"synchronizeMPI", Grid3D_synchronizeMPI, METH_VARARGS, "synchronize through the MPI communicator"}, +#ifdef WITH_NUMPY {"to_numpy", Grid3D_to_numpy, METH_VARARGS, "converts the 3D grid to a numpy array"}, {"from_numpy", Grid3D_from_numpy, METH_VARARGS, "converts the numpy array to a 3D grid"}, +#endif // WITH_NUMPY {NULL}}; // defenition of the memebers of the python Grid3D wrapper class @@ -535,7 +526,7 @@ void initGrid3D(PyObject *module) { if (ensure_numpy() != 0) { throw std::runtime_error("NumPy C-API init failed"); } -#endif +#endif // WITH_NUMPY if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) return; Py_INCREF(&pyORBIT_Grid3D_Type); From 0e4e9d643f8d104197772ce25a47fa19bec89912 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Fri, 5 Jun 2026 12:43:47 -0400 Subject: [PATCH 04/17] fix gitignore conflict --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 4bd2411b..78c05689 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,4 @@ _api/ _cpp_api/ doc/source/reference/*.rst doc/source/reference/*.rst.include +.cache/ From 6d29eabd3892ef1ad799599112d748198098b092 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Fri, 5 Jun 2026 12:19:44 -0400 Subject: [PATCH 05/17] Rename numpy option, default to false, mark experimental --- meson_options.txt | 2 +- src/meson.build | 45 +++++++++++++++++----------------- src/orbit/wrap_bunch.cc | 14 +++++------ src/spacecharge/wrap_grid3D.cc | 16 ++++++------ 4 files changed, 38 insertions(+), 39 deletions(-) diff --git a/meson_options.txt b/meson_options.txt index 2e76a4d0..90e09823 100644 --- a/meson_options.txt +++ b/meson_options.txt @@ -1,2 +1,2 @@ option('USE_MPI', type: 'string', value: 'auto', description: 'Choose MPI implementation (mpich, openmpi, none, auto)') -option('WITH_NUMPY', type: 'boolean', value: true, description: 'Use numpy') +option('PyORBIT_EXPERIMENTAL_WITH_NUMPY', type: 'boolean', value: false, description: 'Use numpy') diff --git a/src/meson.build b/src/meson.build index a70ecf74..a3e86e34 100644 --- a/src/meson.build +++ b/src/meson.build @@ -4,7 +4,7 @@ # Add Python installation details pymod = import('python') -with_numpy = get_option('WITH_NUMPY') +with_numpy = get_option('PyORBIT_EXPERIMENTAL_WITH_NUMPY') if with_numpy message('Compiling with Numpy support') @@ -13,7 +13,7 @@ if with_numpy python, '-c', 'import numpy; print(numpy.get_include())', check: true ).stdout().strip() - add_project_arguments('-DWITH_NUMPY=1', language: 'cpp') + add_project_arguments('-DPyORBIT_EXPERIMENTAL_WITH_NUMPY=1', language: 'cpp') add_project_arguments('-I' + numpy_inc, language: 'cpp') else @@ -34,37 +34,36 @@ 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 @@ -320,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', @@ -329,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', @@ -338,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', @@ -347,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', @@ -356,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', @@ -365,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', @@ -374,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', @@ -383,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', @@ -392,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', @@ -401,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', @@ -410,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', @@ -419,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', @@ -428,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', @@ -437,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', diff --git a/src/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index bb12f911..4293dd78 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -16,7 +16,7 @@ #include "pyORBIT_Object.hh" -#ifdef WITH_NUMPY +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY #include static int ensure_numpy() { @@ -27,7 +27,7 @@ static int ensure_numpy() { } return 0; } -#endif // WITH_NUMPY +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY #include "Bunch.hh" #include "ParticleAttributesFactory.hh" @@ -1184,7 +1184,7 @@ namespace wrap_orbit_bunch{ self->ob_base.ob_type->tp_free((PyObject*)self); } -#ifdef WITH_NUMPY +#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; @@ -1290,7 +1290,7 @@ static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *args) { return py_bunch_obj; } -#endif // WITH_NUMPY +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY static PyMethodDef BunchClassMethods[] = { //-------------------------------------------------------- @@ -1351,7 +1351,7 @@ static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *args) { { "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 WITH_NUMPY +#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" }, @@ -1427,11 +1427,11 @@ extern "C" { /* The name of the function was changed to avoid collision with PyImport magic naming */ PyMODINIT_FUNC initbunch(void) { - #ifdef WITH_NUMPY + #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY if (ensure_numpy() != 0) { throw std::runtime_error("NumPy C-API init failed"); } - #endif // WITH_NUMPY + #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); diff --git a/src/spacecharge/wrap_grid3D.cc b/src/spacecharge/wrap_grid3D.cc index e55f4780..3100109a 100644 --- a/src/spacecharge/wrap_grid3D.cc +++ b/src/spacecharge/wrap_grid3D.cc @@ -7,7 +7,7 @@ #include "wrap_bunch.hh" #include "wrap_spacecharge.hh" -#ifdef WITH_NUMPY +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION #include @@ -19,7 +19,7 @@ static int ensure_numpy() { } return 0; } -#endif // WITH_NUMPY +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY #include "Grid3D.hh" @@ -348,7 +348,7 @@ static void Grid3D_del(pyORBIT_Object *self) { self->ob_base.ob_type->tp_free((PyObject *)self); } -#ifdef WITH_NUMPY +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY static PyObject *Grid3D_to_numpy(PyObject *self, PyObject *args) { pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; @@ -437,7 +437,7 @@ static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { Py_DECREF(arr); Py_RETURN_NONE; } -#endif // WITH_NUMPY +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY // defenition of the methods of the python Grid3D wrapper class // they will be vailable from python level @@ -466,10 +466,10 @@ static PyMethodDef Grid3DClassMethods[] = { {"calcGradient", Grid3D_calcGradient, METH_VARARGS, "returns gradient as (gx,gy,gz) for point (x,y,z)"}, {"longWrapping", Grid3D_longWrapping, METH_VARARGS, "set/get isWrapping variable defining long. wrapping policy"}, {"synchronizeMPI", Grid3D_synchronizeMPI, METH_VARARGS, "synchronize through the MPI communicator"}, -#ifdef WITH_NUMPY +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY {"to_numpy", Grid3D_to_numpy, METH_VARARGS, "converts the 3D grid to a numpy array"}, {"from_numpy", Grid3D_from_numpy, METH_VARARGS, "converts the numpy array to a 3D grid"}, -#endif // WITH_NUMPY +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY {NULL}}; // defenition of the memebers of the python Grid3D wrapper class @@ -522,11 +522,11 @@ static PyTypeObject pyORBIT_Grid3D_Type = { // It will be called from SpaceCharge wrapper initialization //-------------------------------------------------- void initGrid3D(PyObject *module) { -#ifdef WITH_NUMPY +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY if (ensure_numpy() != 0) { throw std::runtime_error("NumPy C-API init failed"); } -#endif // WITH_NUMPY +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) return; Py_INCREF(&pyORBIT_Grid3D_Type); From 62ebe2d61e453385c96df55f297abb64b625ef1e Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Wed, 26 Aug 2026 10:47:11 -0400 Subject: [PATCH 06/17] more tests --- tests/py/orbit/core/test_numpy_interop.py | 286 ++++++++++++++++++++++ 1 file changed, 286 insertions(+) create mode 100644 tests/py/orbit/core/test_numpy_interop.py diff --git a/tests/py/orbit/core/test_numpy_interop.py b/tests/py/orbit/core/test_numpy_interop.py new file mode 100644 index 00000000..3b5ab814 --- /dev/null +++ b/tests/py/orbit/core/test_numpy_interop.py @@ -0,0 +1,286 @@ +import os +import shutil +import subprocess +import sys +import textwrap + +import numpy as np +import pytest + +from orbit.core.bunch import Bunch +from orbit.core.spacecharge import Grid3D + + +NUMPY_INTEROP_ENABLED = hasattr(Bunch, "to_numpy") and hasattr( + Grid3D, "to_numpy" +) + +pytestmark = pytest.mark.skipif( + not NUMPY_INTEROP_ENABLED, + reason="PyORBIT was built without experimental NumPy interoperability", +) + + +def _process_failure(result): + return textwrap.dedent( + f""" + subprocess exited with status {result.returncode} + stdout: + {result.stdout} + stderr: + {result.stderr} + """ + ) + + +def test_grid3d_to_numpy_round_trips_rectangular_grid(): + """Exporting a non-square grid must not access memory out of bounds.""" + script = textwrap.dedent( + """ + import numpy as np + + from orbit.core.spacecharge import Grid3D + + expected = np.arange(4 * 5 * 6, dtype=np.float64).reshape(4, 5, 6) + grid = Grid3D(4, 5, 6) + grid.from_numpy(expected) + np.testing.assert_array_equal(grid.to_numpy(), expected) + """ + ) + + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=30, + check=False, + ) + + assert result.returncode == 0, _process_failure(result) + + +def test_bunch_update_from_numpy_is_atomic_on_invalid_input(): + """A rejected update must leave the original particles untouched.""" + original = np.arange(18, dtype=np.float64).reshape(3, 6) + bunch = Bunch.from_numpy(original) + + with pytest.raises(ValueError, match=r"shape \(nparts, 6\)"): + bunch.update_from_numpy(np.zeros((4, 5), dtype=np.float64)) + + np.testing.assert_array_equal(bunch.to_numpy(), original) + + +def test_bunch_to_numpy_gather_on_single_rank(): + """Gathering on a serial communicator returns the local array.""" + expected = np.arange(18, dtype=np.float64).reshape(3, 6) + bunch = Bunch.from_numpy(expected) + + np.testing.assert_array_equal(bunch.to_numpy(gather=True), expected) + + +def test_grid3d_from_numpy_uses_xyz_order(): + values_xyz = np.arange(4 * 5 * 6, dtype=np.float64).reshape(4, 5, 6) + grid = Grid3D(4, 5, 6) + + grid.from_numpy(values_xyz) + + for ix, iy, iz in ((0, 0, 0), (1, 2, 3), (3, 4, 5)): + assert grid.getValueOnGrid(ix, iy, iz) == values_xyz[ix, iy, iz] + + +@pytest.mark.parametrize( + "extension_name", + ["bunch", "spacecharge"], +) +def test_missing_numpy_raises_import_error_instead_of_aborting(extension_name): + """NumPy C-API initialization failure must not terminate Python.""" + script = textwrap.dedent( + f""" + import sys + + class BlockNumpyImports: + @staticmethod + def find_spec(fullname, path=None, target=None): + if fullname == "numpy" or fullname.startswith("numpy."): + raise ModuleNotFoundError("NumPy is unavailable for this test") + return None + + sys.meta_path.insert(0, BlockNumpyImports()) + try: + from orbit.core import {extension_name} + except ImportError: + raise SystemExit(0) + raise SystemExit("expected importing without NumPy to fail") + """ + ) + env = os.environ.copy() + env.pop("PYTHONPATH", None) + + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=30, + check=False, + env=env, + ) + + assert result.returncode == 0, _process_failure(result) + + +def test_bunch_update_from_numpy_uses_only_local_particle_count(tmp_path): + """Each MPI rank must clear only its locally allocated particles.""" + mpirun = shutil.which("mpirun") or shutil.which("mpiexec") + if mpirun is None: + pytest.skip("MPI launcher is not available") + + script_path = tmp_path / "numpy_bunch_mpi_update.py" + script_path.write_text( + textwrap.dedent( + """ + import numpy as np + + from orbit.core.bunch import Bunch + from orbit.core.orbit_mpi import ( + MPI_Comm_rank, + MPI_Comm_size, + mpi_comm, + ) + + comm = mpi_comm.MPI_COMM_WORLD + rank = MPI_Comm_rank(comm) + size = MPI_Comm_size(comm) + if size != 2: + print("PYORBIT_MPI_DISABLED", flush=True) + raise SystemExit(0) + + # Rank zero deliberately has only the minimum allocation while + # rank one makes the global count much larger than that allocation. + local_size = 0 if rank == 0 else 100_000 + bunch = Bunch.from_numpy( + np.ones((local_size, 6), dtype=np.float64) + ) + bunch.update_from_numpy(np.full((1, 6), rank, dtype=np.float64)) + + assert bunch.getSize() == 1 + np.testing.assert_array_equal( + bunch.to_numpy(), np.full((1, 6), rank, dtype=np.float64) + ) + """ + ) + ) + + env = os.environ.copy() + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") + result = subprocess.run( + [mpirun, "-np", "2", sys.executable, str(script_path)], + capture_output=True, + text=True, + timeout=60, + check=False, + env=env, + ) + + if "PYORBIT_MPI_DISABLED" in result.stdout: + pytest.skip("PyORBIT was built without MPI support") + + assert result.returncode == 0, _process_failure(result) + + +def test_bunch_to_numpy_gathers_particles_across_mpi_ranks(tmp_path): + """Gathered rows are rank ordered and returned only on the chosen root.""" + mpirun = shutil.which("mpirun") or shutil.which("mpiexec") + if mpirun is None: + pytest.skip("MPI launcher is not available") + + script_path = tmp_path / "numpy_bunch_mpi_gather.py" + script_path.write_text( + textwrap.dedent( + """ + import numpy as np + import pytest + + from orbit.core.bunch import Bunch + from orbit.core.orbit_mpi import ( + MPI_Comm_rank, + MPI_Comm_size, + mpi_comm, + ) + + comm = mpi_comm.MPI_COMM_WORLD + rank = MPI_Comm_rank(comm) + size = MPI_Comm_size(comm) + if size != 3: + print("PYORBIT_MPI_DISABLED", flush=True) + raise SystemExit(0) + + local_sizes = (2, 0, 3) + + def particles_for_rank(mpi_rank): + count = local_sizes[mpi_rank] + return ( + np.arange(count * 6, dtype=np.float64).reshape(count, 6) + + 100.0 * mpi_rank + ) + + local_particles = particles_for_rank(rank) + expected_global = np.concatenate( + [particles_for_rank(mpi_rank) for mpi_rank in range(size)] + ) + bunch = Bunch.from_numpy(local_particles) + + # Complete every collective before making assertions so an + # assertion failure on one rank cannot strand another rank. + default_root_result = bunch.to_numpy(gather=True) + empty_root_result = bunch.to_numpy(gather=True, root=1) + + empty_bunch = Bunch.from_numpy(np.empty((0, 6), dtype=np.float64)) + empty_global_result = empty_bunch.to_numpy(gather=True, root=2) + + # Gathering must not change the ordinary local conversion. + np.testing.assert_array_equal(bunch.to_numpy(), local_particles) + + if rank == 0: + np.testing.assert_array_equal( + default_root_result, expected_global + ) + else: + assert default_root_result is None + + # Rank one owns no local particles but can still be the root. + if rank == 1: + np.testing.assert_array_equal( + empty_root_result, expected_global + ) + else: + assert empty_root_result is None + + if rank == 2: + assert empty_global_result.shape == (0, 6) + assert empty_global_result.dtype == np.float64 + else: + assert empty_global_result is None + + with pytest.raises(ValueError, match="root must be between 0 and 2"): + bunch.to_numpy(gather=True, root=3) + """ + ) + ) + + env = os.environ.copy() + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") + result = subprocess.run( + [mpirun, "-np", "3", sys.executable, str(script_path)], + capture_output=True, + text=True, + timeout=60, + check=False, + env=env, + ) + + if "PYORBIT_MPI_DISABLED" in result.stdout: + pytest.skip("PyORBIT was built without MPI support") + + assert result.returncode == 0, _process_failure(result) From 441915a61239a5b075ab15913be0ee46fa3be56a Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Wed, 26 Aug 2026 10:50:20 -0400 Subject: [PATCH 07/17] Create workflow for running tests with pytest --- .github/workflows/compilation.yml | 67 +++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/.github/workflows/compilation.yml b/.github/workflows/compilation.yml index c40130af..bda30bc4 100644 --- a/.github/workflows/compilation.yml +++ b/.github/workflows/compilation.yml @@ -147,5 +147,72 @@ 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_package: "" + - mpi: mpich + mpi_package: mpich + - mpi: ompi + mpi_package: openmpi + + 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_package }} + cache-downloads: true + + - name: Configure + shell: bash -el {0} + run: | + meson setup build-${{ matrix.mpi }} \ + --prefix="$CONDA_PREFIX" \ + --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} + run: | + python -m pytest tests/py/ -v + build-docs: uses: ./.github/workflows/docs-build.yml From f26631d319eab000bd18fb527257febbb94b7398 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Wed, 26 Aug 2026 10:57:45 -0400 Subject: [PATCH 08/17] let's try that again --- .github/workflows/compilation.yml | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/.github/workflows/compilation.yml b/.github/workflows/compilation.yml index bda30bc4..4598907a 100644 --- a/.github/workflows/compilation.yml +++ b/.github/workflows/compilation.yml @@ -157,11 +157,13 @@ jobs: matrix: include: - mpi: none - mpi_package: "" + mpi_packages: "" - mpi: mpich - mpi_package: mpich + mpi_packages: >- + mpich + libfabric-devel - mpi: ompi - mpi_package: openmpi + mpi_packages: openmpi steps: - uses: actions/checkout@v4 @@ -181,7 +183,7 @@ jobs: pkg-config setuptools setuptools-scm - ${{ matrix.mpi_package }} + ${{ matrix.mpi_packages }} cache-downloads: true - name: Configure @@ -189,6 +191,7 @@ jobs: run: | meson setup build-${{ matrix.mpi }} \ --prefix="$CONDA_PREFIX" \ + --libdir=lib \ --buildtype=release \ -DUSE_MPI=${{ matrix.mpi }} \ -DPyORBIT_EXPERIMENTAL_WITH_NUMPY=true @@ -212,6 +215,7 @@ jobs: - name: Run Python tests shell: bash -el {0} run: | + export LD_LIBRARY_PATH="$CONDA_PREFIX/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" python -m pytest tests/py/ -v build-docs: From e4743f3777f01782e519cc83765441fbf8800fba Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Wed, 26 Aug 2026 11:22:15 -0400 Subject: [PATCH 09/17] Test to ensure exit status is preserved. Workflows are passing even though tests are failing, this makes sure that the exit status is preserved when calling Py_AtExit. --- src/mpi/orbit_mpi.cc | 19 ++++++++++++++++++- tests/py/orbit/core/test_orbit_mpi.py | 16 ++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/mpi/orbit_mpi.cc b/src/mpi/orbit_mpi.cc index 4965a9e1..5d7191c9 100644 --- a/src/mpi/orbit_mpi.cc +++ b/src/mpi/orbit_mpi.cc @@ -28,6 +28,23 @@ static std::size_t ORBIT_MPI_Type_size(MPI_Datatype data) { } #endif +#if USE_MPI > 0 +/** Finalize MPI during interpreter shutdown without changing its exit status. */ +static void ORBIT_MPI_Finalize_AtExit(){ + int initialized = 0; + if(MPI_Initialized(&initialized) != MPI_SUCCESS || !initialized){ + return; + } + + int finalized = 0; + if(MPI_Finalized(&finalized) != MPI_SUCCESS || finalized){ + return; + } + + MPI_Finalize(); +} +#endif + /** A C wrapper around MPI_Init. */ int ORBIT_MPI_Init(){ #if USE_MPI > 0 @@ -35,7 +52,7 @@ int ORBIT_MPI_Init(){ MPI_Init(NULL, NULL); // Registering MPI finalize method at cleanup stage - Py_AtExit(ORBIT_MPI_Finalize); + Py_AtExit(ORBIT_MPI_Finalize_AtExit); #endif return MPI_SUCCESS; } diff --git a/tests/py/orbit/core/test_orbit_mpi.py b/tests/py/orbit/core/test_orbit_mpi.py index 07c70936..05483aec 100644 --- a/tests/py/orbit/core/test_orbit_mpi.py +++ b/tests/py/orbit/core/test_orbit_mpi.py @@ -1,3 +1,6 @@ +import subprocess +import sys + import pytest from orbit.core.orbit_mpi import ( mpi_comm, @@ -17,6 +20,19 @@ ] +def test_mpi_cleanup_preserves_interpreter_exit_status(): + result = subprocess.run( + [ + sys.executable, + "-c", + "from orbit.core import orbit_mpi; raise SystemExit(7)", + ], + check=False, + ) + + assert result.returncode == 7 + + def _expected(val, op_name, size): if op_name == "sum": return val * size From 1d84fcb55e67457c2d4f215583621e2d96494271 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 10:02:52 -0400 Subject: [PATCH 10/17] Organize tests into classes --- tests/py/orbit/core/test_numpy_interop.py | 382 +++++++++++----------- 1 file changed, 195 insertions(+), 187 deletions(-) diff --git a/tests/py/orbit/core/test_numpy_interop.py b/tests/py/orbit/core/test_numpy_interop.py index 3b5ab814..4e1d9a58 100644 --- a/tests/py/orbit/core/test_numpy_interop.py +++ b/tests/py/orbit/core/test_numpy_interop.py @@ -33,66 +33,7 @@ def _process_failure(result): ) -def test_grid3d_to_numpy_round_trips_rectangular_grid(): - """Exporting a non-square grid must not access memory out of bounds.""" - script = textwrap.dedent( - """ - import numpy as np - - from orbit.core.spacecharge import Grid3D - - expected = np.arange(4 * 5 * 6, dtype=np.float64).reshape(4, 5, 6) - grid = Grid3D(4, 5, 6) - grid.from_numpy(expected) - np.testing.assert_array_equal(grid.to_numpy(), expected) - """ - ) - - result = subprocess.run( - [sys.executable, "-c", script], - capture_output=True, - text=True, - timeout=30, - check=False, - ) - - assert result.returncode == 0, _process_failure(result) - - -def test_bunch_update_from_numpy_is_atomic_on_invalid_input(): - """A rejected update must leave the original particles untouched.""" - original = np.arange(18, dtype=np.float64).reshape(3, 6) - bunch = Bunch.from_numpy(original) - - with pytest.raises(ValueError, match=r"shape \(nparts, 6\)"): - bunch.update_from_numpy(np.zeros((4, 5), dtype=np.float64)) - - np.testing.assert_array_equal(bunch.to_numpy(), original) - - -def test_bunch_to_numpy_gather_on_single_rank(): - """Gathering on a serial communicator returns the local array.""" - expected = np.arange(18, dtype=np.float64).reshape(3, 6) - bunch = Bunch.from_numpy(expected) - - np.testing.assert_array_equal(bunch.to_numpy(gather=True), expected) - - -def test_grid3d_from_numpy_uses_xyz_order(): - values_xyz = np.arange(4 * 5 * 6, dtype=np.float64).reshape(4, 5, 6) - grid = Grid3D(4, 5, 6) - - grid.from_numpy(values_xyz) - - for ix, iy, iz in ((0, 0, 0), (1, 2, 3), (3, 4, 5)): - assert grid.getValueOnGrid(ix, iy, iz) == values_xyz[ix, iy, iz] - - -@pytest.mark.parametrize( - "extension_name", - ["bunch", "spacecharge"], -) -def test_missing_numpy_raises_import_error_instead_of_aborting(extension_name): +def _assert_missing_numpy_raises_import_error(extension_name): """NumPy C-API initialization failure must not terminate Python.""" script = textwrap.dedent( f""" @@ -128,159 +69,226 @@ def find_spec(fullname, path=None, target=None): assert result.returncode == 0, _process_failure(result) -def test_bunch_update_from_numpy_uses_only_local_particle_count(tmp_path): - """Each MPI rank must clear only its locally allocated particles.""" - mpirun = shutil.which("mpirun") or shutil.which("mpiexec") - if mpirun is None: - pytest.skip("MPI launcher is not available") +class TestBunchNumpyInterop: + def test_update_from_numpy_is_atomic_on_invalid_input(self): + """A rejected update must leave the original particles untouched.""" + original = np.arange(18, dtype=np.float64).reshape(3, 6) + bunch = Bunch.from_numpy(original) - script_path = tmp_path / "numpy_bunch_mpi_update.py" - script_path.write_text( - textwrap.dedent( - """ - import numpy as np + with pytest.raises(ValueError, match=r"shape \(nparts, 6\)"): + bunch.update_from_numpy(np.zeros((4, 5), dtype=np.float64)) - from orbit.core.bunch import Bunch - from orbit.core.orbit_mpi import ( - MPI_Comm_rank, - MPI_Comm_size, - mpi_comm, - ) - - comm = mpi_comm.MPI_COMM_WORLD - rank = MPI_Comm_rank(comm) - size = MPI_Comm_size(comm) - if size != 2: - print("PYORBIT_MPI_DISABLED", flush=True) - raise SystemExit(0) - - # Rank zero deliberately has only the minimum allocation while - # rank one makes the global count much larger than that allocation. - local_size = 0 if rank == 0 else 100_000 - bunch = Bunch.from_numpy( - np.ones((local_size, 6), dtype=np.float64) - ) - bunch.update_from_numpy(np.full((1, 6), rank, dtype=np.float64)) + np.testing.assert_array_equal(bunch.to_numpy(), original) - assert bunch.getSize() == 1 - np.testing.assert_array_equal( - bunch.to_numpy(), np.full((1, 6), rank, dtype=np.float64) - ) - """ - ) - ) + def test_to_numpy_gather_on_single_rank(self): + """Gathering on a serial communicator returns the local array.""" + expected = np.arange(18, dtype=np.float64).reshape(3, 6) + bunch = Bunch.from_numpy(expected) - env = os.environ.copy() - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") - result = subprocess.run( - [mpirun, "-np", "2", sys.executable, str(script_path)], - capture_output=True, - text=True, - timeout=60, - check=False, - env=env, - ) + np.testing.assert_array_equal(bunch.to_numpy(gather=True), expected) - if "PYORBIT_MPI_DISABLED" in result.stdout: - pytest.skip("PyORBIT was built without MPI support") + def test_missing_numpy_raises_import_error_instead_of_aborting(self): + _assert_missing_numpy_raises_import_error("bunch") - assert result.returncode == 0, _process_failure(result) + def test_update_from_numpy_uses_only_local_particle_count(self, tmp_path): + """Each MPI rank must clear only its locally allocated particles.""" + mpirun = shutil.which("mpirun") or shutil.which("mpiexec") + if mpirun is None: + pytest.skip("MPI launcher is not available") + script_path = tmp_path / "numpy_bunch_mpi_update.py" + script_path.write_text( + textwrap.dedent( + """ + import numpy as np -def test_bunch_to_numpy_gathers_particles_across_mpi_ranks(tmp_path): - """Gathered rows are rank ordered and returned only on the chosen root.""" - mpirun = shutil.which("mpirun") or shutil.which("mpiexec") - if mpirun is None: - pytest.skip("MPI launcher is not available") + from orbit.core.bunch import Bunch + from orbit.core.orbit_mpi import ( + MPI_Comm_rank, + MPI_Comm_size, + mpi_comm, + ) - script_path = tmp_path / "numpy_bunch_mpi_gather.py" - script_path.write_text( - textwrap.dedent( - """ - import numpy as np - import pytest + comm = mpi_comm.MPI_COMM_WORLD + rank = MPI_Comm_rank(comm) + size = MPI_Comm_size(comm) + if size != 2: + print("PYORBIT_MPI_DISABLED", flush=True) + raise SystemExit(0) + + # Rank zero deliberately has only the minimum allocation while + # rank one makes the global count much larger than that allocation. + local_size = 0 if rank == 0 else 100_000 + bunch = Bunch.from_numpy( + np.ones((local_size, 6), dtype=np.float64) + ) + bunch.update_from_numpy( + np.full((1, 6), rank, dtype=np.float64) + ) - from orbit.core.bunch import Bunch - from orbit.core.orbit_mpi import ( - MPI_Comm_rank, - MPI_Comm_size, - mpi_comm, + assert bunch.getSize() == 1 + np.testing.assert_array_equal( + bunch.to_numpy(), np.full((1, 6), rank, dtype=np.float64) + ) + """ ) + ) - comm = mpi_comm.MPI_COMM_WORLD - rank = MPI_Comm_rank(comm) - size = MPI_Comm_size(comm) - if size != 3: - print("PYORBIT_MPI_DISABLED", flush=True) - raise SystemExit(0) + env = os.environ.copy() + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") + result = subprocess.run( + [mpirun, "-np", "2", sys.executable, str(script_path)], + capture_output=True, + text=True, + timeout=60, + check=False, + env=env, + ) - local_sizes = (2, 0, 3) + if "PYORBIT_MPI_DISABLED" in result.stdout: + pytest.skip("PyORBIT was built without MPI support") + + assert result.returncode == 0, _process_failure(result) + + def test_to_numpy_gathers_particles_across_mpi_ranks(self, tmp_path): + """Gathered rows are rank ordered and returned only on the chosen root.""" + mpirun = shutil.which("mpirun") or shutil.which("mpiexec") + if mpirun is None: + pytest.skip("MPI launcher is not available") + + script_path = tmp_path / "numpy_bunch_mpi_gather.py" + script_path.write_text( + textwrap.dedent( + """ + import numpy as np + import pytest + + from orbit.core.bunch import Bunch + from orbit.core.orbit_mpi import ( + MPI_Comm_rank, + MPI_Comm_size, + mpi_comm, + ) - def particles_for_rank(mpi_rank): - count = local_sizes[mpi_rank] - return ( - np.arange(count * 6, dtype=np.float64).reshape(count, 6) - + 100.0 * mpi_rank + comm = mpi_comm.MPI_COMM_WORLD + rank = MPI_Comm_rank(comm) + size = MPI_Comm_size(comm) + if size != 3: + print("PYORBIT_MPI_DISABLED", flush=True) + raise SystemExit(0) + + local_sizes = (2, 0, 3) + + def particles_for_rank(mpi_rank): + count = local_sizes[mpi_rank] + return ( + np.arange(count * 6, dtype=np.float64).reshape(count, 6) + + 100.0 * mpi_rank + ) + + local_particles = particles_for_rank(rank) + expected_global = np.concatenate( + [particles_for_rank(mpi_rank) for mpi_rank in range(size)] ) + bunch = Bunch.from_numpy(local_particles) + + # Complete every collective before making assertions so an + # assertion failure on one rank cannot strand another rank. + default_root_result = bunch.to_numpy(gather=True) + empty_root_result = bunch.to_numpy(gather=True, root=1) - local_particles = particles_for_rank(rank) - expected_global = np.concatenate( - [particles_for_rank(mpi_rank) for mpi_rank in range(size)] + empty_bunch = Bunch.from_numpy( + np.empty((0, 6), dtype=np.float64) + ) + empty_global_result = empty_bunch.to_numpy(gather=True, root=2) + + # Gathering must not change the ordinary local conversion. + np.testing.assert_array_equal(bunch.to_numpy(), local_particles) + + if rank == 0: + np.testing.assert_array_equal( + default_root_result, expected_global + ) + else: + assert default_root_result is None + + # Rank one owns no local particles but can still be the root. + if rank == 1: + np.testing.assert_array_equal( + empty_root_result, expected_global + ) + else: + assert empty_root_result is None + + if rank == 2: + assert empty_global_result.shape == (0, 6) + assert empty_global_result.dtype == np.float64 + else: + assert empty_global_result is None + + with pytest.raises( + ValueError, match="root must be between 0 and 2" + ): + bunch.to_numpy(gather=True, root=3) + """ ) - bunch = Bunch.from_numpy(local_particles) + ) - # Complete every collective before making assertions so an - # assertion failure on one rank cannot strand another rank. - default_root_result = bunch.to_numpy(gather=True) - empty_root_result = bunch.to_numpy(gather=True, root=1) + env = os.environ.copy() + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") + result = subprocess.run( + [mpirun, "-np", "3", sys.executable, str(script_path)], + capture_output=True, + text=True, + timeout=60, + check=False, + env=env, + ) - empty_bunch = Bunch.from_numpy(np.empty((0, 6), dtype=np.float64)) - empty_global_result = empty_bunch.to_numpy(gather=True, root=2) + if "PYORBIT_MPI_DISABLED" in result.stdout: + pytest.skip("PyORBIT was built without MPI support") - # Gathering must not change the ordinary local conversion. - np.testing.assert_array_equal(bunch.to_numpy(), local_particles) + assert result.returncode == 0, _process_failure(result) - if rank == 0: - np.testing.assert_array_equal( - default_root_result, expected_global - ) - else: - assert default_root_result is None - # Rank one owns no local particles but can still be the root. - if rank == 1: - np.testing.assert_array_equal( - empty_root_result, expected_global - ) - else: - assert empty_root_result is None +class TestGrid3DNumpyInterop: + def test_to_numpy_round_trips_rectangular_grid(self): + """Exporting a non-square grid must not access memory out of bounds.""" + script = textwrap.dedent( + """ + import numpy as np - if rank == 2: - assert empty_global_result.shape == (0, 6) - assert empty_global_result.dtype == np.float64 - else: - assert empty_global_result is None + from orbit.core.spacecharge import Grid3D - with pytest.raises(ValueError, match="root must be between 0 and 2"): - bunch.to_numpy(gather=True, root=3) + expected = np.arange( + 4 * 5 * 6, dtype=np.float64 + ).reshape(4, 5, 6) + grid = Grid3D(4, 5, 6) + grid.from_numpy(expected) + np.testing.assert_array_equal(grid.to_numpy(), expected) """ ) - ) - env = os.environ.copy() - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") - result = subprocess.run( - [mpirun, "-np", "3", sys.executable, str(script_path)], - capture_output=True, - text=True, - timeout=60, - check=False, - env=env, - ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=30, + check=False, + ) - if "PYORBIT_MPI_DISABLED" in result.stdout: - pytest.skip("PyORBIT was built without MPI support") + assert result.returncode == 0, _process_failure(result) - assert result.returncode == 0, _process_failure(result) + def test_from_numpy_uses_xyz_order(self): + values_xyz = np.arange(4 * 5 * 6, dtype=np.float64).reshape(4, 5, 6) + grid = Grid3D(4, 5, 6) + + grid.from_numpy(values_xyz) + + for ix, iy, iz in ((0, 0, 0), (1, 2, 3), (3, 4, 5)): + assert grid.getValueOnGrid(ix, iy, iz) == values_xyz[ix, iy, iz] + + def test_missing_numpy_raises_import_error_instead_of_aborting(self): + _assert_missing_numpy_raises_import_error("spacecharge") From b6604c36d45bce1a1e88e72e263b9ca23987aed1 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 13:56:39 -0400 Subject: [PATCH 11/17] Easier interface for exporting bunch to root MPI rank. Also, includes a helper for distributing the coordinates during bunch initialization with several MPI ranks. --- py/orbit/bunch_utils/__init__.py | 2 + py/orbit/bunch_utils/numpy_utils.py | 50 ++++ src/orbit/wrap_bunch.cc | 342 ++++++++++++++++++++++------ 3 files changed, 324 insertions(+), 70 deletions(-) create mode 100644 py/orbit/bunch_utils/numpy_utils.py diff --git a/py/orbit/bunch_utils/__init__.py b/py/orbit/bunch_utils/__init__.py index 42ea363e..58c92160 100644 --- a/py/orbit/bunch_utils/__init__.py +++ b/py/orbit/bunch_utils/__init__.py @@ -10,6 +10,7 @@ # This guards against missing numpy. # Should be imporved with some meaningful (and MPI friendly?) warning printed out. try: + 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 @@ -19,6 +20,7 @@ __all__ = [] # __all__.append("addParticleIdNumbers") # doesn't exist __all__.append("ParticleIdNumber") +__all__.append("bunch_from_shared_numpy") __all__.append("collect_bunch") __all__.append("save_bunch") __all__.append("load_bunch") diff --git a/py/orbit/bunch_utils/numpy_utils.py b/py/orbit/bunch_utils/numpy_utils.py new file mode 100644 index 00000000..191d170a --- /dev/null +++ b/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, :]) diff --git a/src/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index 82f91caf..c4409e38 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -9,6 +9,8 @@ // /////////////////////////////////////////////////////////////////////////// #include "wrap_bunch.hh" +// #include "modsupport.h" +// #include "pyerrors.h" #include "wrap_syncpart.hh" #include "wrap_bunch_twiss_analysis.hh" #include "wrap_bunch_tune_analysis.hh" @@ -1198,109 +1200,309 @@ namespace wrap_orbit_bunch{ } #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; +static PyArrayObject *parse_bunch_array(PyObject *input) { + PyArrayObject *array = (PyArrayObject *)PyArray_FROM_OTF( + input, + NPY_FLOAT64, + NPY_ARRAY_IN_ARRAY + ); + + if (array == NULL) { + return NULL; + } + + if (PyArray_NDIM(array) != 2) { + PyErr_SetString(PyExc_ValueError, + "array must be 2-dimensional with shape (nparts, 6)" + ); + Py_DECREF(array); + return NULL; + } + + if (PyArray_DIM(array, 1) != 6) { + PyErr_SetString(PyExc_ValueError, + "expected shape (nparts, 6) for (x, xp, y, yp, z, dE)" + ); + Py_DECREF(array); + return NULL; + } - if (!PyArg_ParseTuple(args, ":to_numpy")) { - ORBIT_MPI_Finalize("PyBunch - to_numpy() - no parameters are needed."); + return array; // new ref; caller MUST Py_DECREF(). +} + +static void append_bunch_with_PyArray(Bunch *bunch, PyArrayObject *array) { + const npy_intp nparts = PyArray_DIM(array, 0); + const double *data = (const double *)PyArray_DATA(array); + + for (npy_intp i = 0; i < nparts; ++i) { + const double *coords = data + i*6; + + bunch->addParticle(coords[0], coords[1], coords[2], coords[3], coords[4], coords[5]); + } +} + +PyDoc_STRVAR( + Bunch_to_numpy_doc, + "to_numpy($self, /, gather=False, root=0)\n" + "--\n" + "\n" + "Return the particle coordinates as a NumPy array.\n" + "\n" + "Parameters\n" + "----------\n" + "gather : bool, optional\n" + " Collect particles from all MPI ranks. The default is False.\n" + "root : int, optional\n" + " Rank that receives the collected array. The default is 0.\n" + "\n" + "Returns\n" + "-------\n" + "numpy.ndarray or None\n" + " Array with shape ``(n_particles, 6)`` and dtype ``float64``.\n" + " When gathering, non-root ranks return None.\n" + "\n" + "Raises\n" + "------\n" + "ValueError\n" + " If root is not a valid MPI rank.\n" + "OverflowError\n" + " If the bunch is too large for the MPI gather operation.\n" +); + +static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args, PyObject *kwargs) { + static const char* kwlist[] = {"gather", "root", NULL}; + + int gather = 0; + int root = 0; + + if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|pi:to_numpy", kwlist, &gather, &root)) { + return NULL; } - const npy_intp nparts = (npy_intp)cpp_Bunch->getSize(); + Bunch *bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + + const int rank = bunch->getMPI_Rank(); + const int size = bunch->getMPI_Size(); + const npy_intp nparts = (npy_intp)bunch->getSize(); const npy_intp ncoords = 6; + if (gather && (root < 0 || root >= size)) { + PyErr_Format(PyExc_ValueError, "root must be between 0 and %d, got %d", size - 1, rank); + return NULL; + } + 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(); + PyObject *local_array = PyArray_SimpleNew(2, dims, NPY_FLOAT64); + + // In the unlikely event that we try to gather a bunch with size larger than INT_MAX + // or if we try to allocate a buffer to hold the bunch data and it fails on one rank + // then propagate that result to the other ranks and raise. + if (gather && size > 1) { + int local_counts_ok = nparts <= INT_MAX / ncoords; + int all_counts_ok = 0; + + ORBIT_MPI_Allreduce(&local_counts_ok, &all_counts_ok, 1, MPI_INT, MPI_MIN, bunch->getMPI_Comm_Local()->comm); + + if(!all_counts_ok) { + PyErr_SetString(PyExc_OverflowError, "local bunch is too large for MPI_Gatherv"); + return NULL; + } + + int local_alloc_ok = local_array != NULL; + int all_alloc_ok = 0; + + ORBIT_MPI_Allreduce(&local_alloc_ok, &all_alloc_ok, 1, MPI_INT, MPI_MIN, bunch->getMPI_Comm_Local()->comm); + + if(!all_alloc_ok) { + Py_XDECREF(local_array); + + if(!PyErr_Occurred()) { + PyErr_NoMemory(); + } + + return NULL; + } + } else { + if (local_array == NULL) { + return NULL; + } + } + + PyArrayObject *arr_obj = (PyArrayObject*)local_array; + double *local_data = (double*)PyArray_DATA(arr_obj); + double **src = 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]; + local_data[j + i*ncoords] = src[i][j]; } } - return py_array; -} + if (!gather || size == 1) { + return local_array; + } - static int bunch_fill_from_numpy_args(Bunch *cpp_Bunch, PyObject *args) { - PyObject *arr_in = NULL; +#if USE_MPI > 0 + MPI_Comm comm = bunch->getMPI_Comm_Local()->comm; - if (!PyArg_ParseTuple(args, "O:from_numpy", &arr_in)) { - return -1; - } + std::vector global_nparts(size); + std::vector displacements(size); + std::vector recv_counts(size); - PyArrayObject *arr = - (PyArrayObject*)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, NPY_ARRAY_IN_ARRAY); - if (!arr) return -1; + const int local_nparts = (int)nparts; - 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; - } + if (MPI_SUCCESS != MPI_Gather(&local_nparts, 1, MPI_INT, global_nparts.data(), 1, MPI_INT, root, comm)) { + Py_DECREF(local_array); + PyErr_SetString(PyExc_RuntimeError, "MPI_Gather failed to collect the bunch sizes across ranks"); + return NULL; + } - const npy_intp nparts = PyArray_DIM(arr, 0); - const npy_intp ncoords = PyArray_DIM(arr, 1); + int total = 0; + int layout_ok = 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; - } + if (rank == root) { + for (int mpi_rank = 0; mpi_rank < size; ++mpi_rank) { + const int rank_nvalues = global_nparts[mpi_rank] * ncoords; - const double *data = (const double*)PyArray_DATA(arr); + if (total + rank_nvalues > INT_MAX) { + layout_ok = 0; + break; + } - 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] - ); - } + displacements[mpi_rank] = total; + recv_counts[mpi_rank] = rank_nvalues; + total += rank_nvalues; + } + } - Py_DECREF(arr); - return 0; -} + ORBIT_MPI_Bcast(&layout_ok, 1, MPI_INT, root, comm); -static PyObject *Bunch_update_from_numpy(PyObject *self, PyObject *args) { - pyORBIT_Object *pyBunch = (pyORBIT_Object*)self; - Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj; + if(!layout_ok) { + Py_DECREF(local_array); + PyErr_SetString(PyExc_OverflowError, "global bunch is too large to collect"); + return NULL; + } - if (cpp_Bunch->getSizeGlobal() > 0) { - for (int i = 0; i < cpp_Bunch->getSizeGlobal(); ++i) { - cpp_Bunch->deleteParticleFast(i); - } - cpp_Bunch->compress(); + PyObject *global_array = NULL; + + if (rank == root) { + npy_intp global_dims[2] = {(npy_intp)(total / ncoords), ncoords}; + global_array = PyArray_SimpleNew(2, global_dims, NPY_FLOAT64); + } + + int global_alloc_ok = rank != root || global_array != NULL; + + ORBIT_MPI_Bcast(&global_alloc_ok, 1, MPI_INT, root, comm); + + if(!global_alloc_ok) { + Py_DECREF(local_array); + + if(!PyErr_Occurred()) { + PyErr_NoMemory(); + } + + return NULL; } - if (bunch_fill_from_numpy_args(cpp_Bunch, args) < 0) return NULL; + double* global_data = rank == root ? (double *)PyArray_DATA((PyArrayObject*)global_array) : NULL; + if(MPI_SUCCESS != MPI_Gatherv(local_data, nparts*ncoords, MPI_DOUBLE, global_data, recv_counts.data(), displacements.data(), MPI_DOUBLE, root, comm)) { + Py_XDECREF(global_array); + PyErr_SetString(PyExc_RuntimeError, "MPI_Gatherv failed to collect bunch."); + return NULL; + } + + if (rank == root) { + return global_array; + } + +#endif // USE_MPI > 0 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; +PyDoc_STRVAR( + Bunch_update_from_numpy_doc, + "update_from_numpy($self, array, /)\n" + "--\n" + "\n" + "Replace the local particle coordinates from an array.\n" + "\n" + "Parameters\n" + "----------\n" + "array : array_like\n" + " Particle coordinates with shape ``(n_particles, 6)`` ordered as\n" + " ``(x, xp, y, yp, z, dE)``. Values are converted to ``float64``.\n" + "\n" + "Returns\n" + "-------\n" + "None\n" + "\n" + "Raises\n" + "------\n" + "ValueError\n" + " If the array does not have shape ``(n_particles, 6)``.\n" +); + +static PyObject *Bunch_update_from_numpy(PyObject *self, PyObject *arg) { + Bunch *bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + + PyArrayObject *array = parse_bunch_array(arg); + + if (array == NULL) { + return NULL; + } - pyORBIT_Object *pyBunch = (pyORBIT_Object*)py_bunch_obj; - Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj; + bunch->deleteAllParticles(); + append_bunch_with_PyArray(bunch, array); - if (!cpp_Bunch) { - Py_DECREF(py_bunch_obj); - PyErr_SetString(PyExc_RuntimeError, "from_numpy: Constructed bunch has NULL cpp_obj"); + Py_DECREF(array); + Py_RETURN_NONE; +} + +PyDoc_STRVAR( + Bunch_from_numpy_doc, + "from_numpy($type, array, /)\n" + "--\n" + "\n" + "Construct a bunch from particle coordinates.\n" + "\n" + "Parameters\n" + "----------\n" + "array : array_like\n" + " Particle coordinates with shape ``(n_particles, 6)`` ordered as\n" + " ``(x, xp, y, yp, z, dE)``. Values are converted to ``float64``.\n" + "\n" + "Returns\n" + "-------\n" + "Bunch\n" + " A new bunch containing the supplied particles.\n" + "\n" + "Raises\n" + "------\n" + "ValueError\n" + " If the array does not have shape ``(n_particles, 6)``.\n" +); + +static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *arg) { + PyArrayObject *array = parse_bunch_array(arg); + + if (array == NULL) { return NULL; } - if (bunch_fill_from_numpy_args(cpp_Bunch, args) < 0) { - Py_DECREF(py_bunch_obj); - return NULL; + PyObject *py_bunch_obj = PyObject_CallNoArgs(cls); + + if (py_bunch_obj == NULL) { + Py_DECREF(array); + return NULL; } + Bunch *bunch = (Bunch*)((pyORBIT_Object*)py_bunch_obj)->cpp_obj; + + append_bunch_with_PyArray(bunch, array); + + Py_DECREF(array); return py_bunch_obj; } #endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY @@ -1366,9 +1568,9 @@ static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *args) { { "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" }, + { "to_numpy", _PyCFunction_CAST(Bunch_to_numpy) ,METH_VARARGS | METH_KEYWORDS, Bunch_to_numpy_doc }, + { "update_from_numpy", Bunch_update_from_numpy ,METH_O, Bunch_update_from_numpy_doc }, + { "from_numpy", Bunch_from_numpy ,METH_O | METH_CLASS, Bunch_from_numpy_doc }, #endif {NULL,NULL} //-------------------------------------------------------- @@ -1443,7 +1645,7 @@ extern "C" { PyMODINIT_FUNC initbunch(void) { #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY if (ensure_numpy() != 0) { - throw std::runtime_error("NumPy C-API init failed"); + return NULL; } #endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY //check that the Bunch wrapper is ready From f9dc50f3f5314480b11fddeccb19ff47aaa3df7e Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 13:58:10 -0400 Subject: [PATCH 12/17] Grid3D->numpy bug fixes. Make initGrid3D return an int if numpy init fails so that we can catch it. --- src/spacecharge/wrap_grid3D.cc | 75 ++++++++++++++--------------- src/spacecharge/wrap_grid3D.hh | 2 +- src/spacecharge/wrap_spacecharge.cc | 18 ++++--- 3 files changed, 48 insertions(+), 47 deletions(-) diff --git a/src/spacecharge/wrap_grid3D.cc b/src/spacecharge/wrap_grid3D.cc index 3100109a..956cfecc 100644 --- a/src/spacecharge/wrap_grid3D.cc +++ b/src/spacecharge/wrap_grid3D.cc @@ -349,20 +349,18 @@ static void Grid3D_del(pyORBIT_Object *self) { } #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY -static PyObject *Grid3D_to_numpy(PyObject *self, PyObject *args) { - pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; - Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; +static PyObject *Grid3D_to_numpy(PyObject *self, + PyObject *Py_UNUSED(ignored)) { + Grid3D *cpp_Grid3D = (Grid3D *)((pyORBIT_Object *)self)->cpp_obj; - if (!PyArg_ParseTuple(args, ":to_numpy")) { - ORBIT_MPI_Finalize("PyGrid3D - to_numpy() - no parameters are needed."); - } - - const npy_intp nz = (npy_intp)cpp_Grid3D->getSizeZ(); const npy_intp nx = (npy_intp)cpp_Grid3D->getSizeX(); const npy_intp ny = (npy_intp)cpp_Grid3D->getSizeY(); + const npy_intp nz = (npy_intp)cpp_Grid3D->getSizeZ(); + + npy_intp dims[3] = {nx, ny, nz}; - npy_intp dims[3] = {nz, nx, ny}; PyObject *arr_obj = PyArray_SimpleNew(3, dims, NPY_FLOAT64); + if (NULL == arr_obj) { return NULL; } @@ -371,10 +369,10 @@ static PyObject *Grid3D_to_numpy(PyObject *self, PyObject *args) { double *out_buffer = (double *)PyArray_DATA(arr); double ***src = cpp_Grid3D->getArr3D(); - for (npy_intp iz = 0; iz < nz; ++iz) { - for (npy_intp ix = 0; ix < ny; ++ix) { - for (npy_intp iy = 0; iy < nx; ++iy) { - out_buffer[iy + ix * ny + iz * nx * ny] = src[iz][ix][iy]; + for (npy_intp ix = 0; ix < nx; ++ix) { + for (npy_intp iy = 0; iy < ny; ++iy) { + for (npy_intp iz = 0; iz < nz; ++iz) { + out_buffer[(ix * ny + iy) * nz + iz] = src[iz][ix][iy]; } } } @@ -382,18 +380,10 @@ static PyObject *Grid3D_to_numpy(PyObject *self, PyObject *args) { return arr_obj; } -static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { - pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; - Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; - - PyObject *arr_in = NULL; - const char* order = "zxy"; +static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *arg) { + Grid3D *cpp_Grid3D = (Grid3D *)((pyORBIT_Object *)self)->cpp_obj; - if (!PyArg_ParseTuple(args, "O|s:from_numpy", &arr_in, &order)) { - ORBIT_MPI_Finalize("PyGrid3D - from_numpy() - ndarray is needed."); - } - - PyArrayObject *arr = (PyArrayObject *)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, + PyArrayObject *arr = (PyArrayObject *)PyArray_FROM_OTF(arg, NPY_FLOAT64, NPY_ARRAY_IN_ARRAY); if (NULL == arr) { return NULL; @@ -403,23 +393,23 @@ static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { Py_DECREF(arr); PyErr_SetString( PyExc_ValueError, - "from_numpy: array must be 3-dimensional with shape (nz,nx,ny)"); + "from_numpy: array must be 3-dimensional with shape (nx,ny,nz)"); return NULL; } - const npy_intp nz_in = PyArray_DIM(arr, 0); - const npy_intp nx_in = PyArray_DIM(arr, 1); - const npy_intp ny_in = PyArray_DIM(arr, 2); + const npy_intp nx_in = PyArray_DIM(arr, 0); + const npy_intp ny_in = PyArray_DIM(arr, 1); + const npy_intp nz_in = PyArray_DIM(arr, 2); - const npy_intp nz_grid = (npy_intp)cpp_Grid3D->getSizeZ(); const npy_intp nx_grid = (npy_intp)cpp_Grid3D->getSizeX(); const npy_intp ny_grid = (npy_intp)cpp_Grid3D->getSizeY(); + const npy_intp nz_grid = (npy_intp)cpp_Grid3D->getSizeZ(); - if (nz_in != nz_grid || nx_in != nx_grid || ny_in != ny_grid) { + if (nx_in != nx_grid || ny_in != ny_grid || nz_in != nz_grid) { Py_DECREF(arr); PyErr_SetString( PyExc_ValueError, - "from_numpy: shape mismatch; expected (zSize, xSize, ySize)"); + "from_numpy: shape mismatch; expected (xSize, ySize, zSize)"); return NULL; } @@ -429,7 +419,8 @@ static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { for (npy_intp iz = 0; iz < nz_grid; ++iz) { for (npy_intp ix = 0; ix < nx_grid; ++ix) { for (npy_intp iy = 0; iy < ny_grid; ++iy) { - dst[iz][ix][iy] = in_buffer[iy + ix * ny_grid + iz * nx_grid * ny_grid]; + dst[iz][ix][iy] = + in_buffer[(ix * ny_grid + iy) * nz_grid + iz]; } } } @@ -467,8 +458,8 @@ static PyMethodDef Grid3DClassMethods[] = { {"longWrapping", Grid3D_longWrapping, METH_VARARGS, "set/get isWrapping variable defining long. wrapping policy"}, {"synchronizeMPI", Grid3D_synchronizeMPI, METH_VARARGS, "synchronize through the MPI communicator"}, #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY - {"to_numpy", Grid3D_to_numpy, METH_VARARGS, "converts the 3D grid to a numpy array"}, - {"from_numpy", Grid3D_from_numpy, METH_VARARGS, "converts the numpy array to a 3D grid"}, + {"to_numpy", Grid3D_to_numpy, METH_NOARGS, "converts the 3D grid to a numpy array in (x,y,z) order"}, + {"from_numpy", Grid3D_from_numpy, METH_O, "converts a numpy array in (x,y,z) order to a 3D grid"}, #endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY {NULL}}; @@ -521,16 +512,22 @@ static PyTypeObject pyORBIT_Grid3D_Type = { // Initialization function of the pyGrid3D class // It will be called from SpaceCharge wrapper initialization //-------------------------------------------------- -void initGrid3D(PyObject *module) { +int initGrid3D(PyObject *module) { #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY if (ensure_numpy() != 0) { - throw std::runtime_error("NumPy C-API init failed"); + return -1; } #endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY - if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) - return; + if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) { + return -1; + } Py_INCREF(&pyORBIT_Grid3D_Type); - PyModule_AddObject(module, "Grid3D", (PyObject *)&pyORBIT_Grid3D_Type); + if (PyModule_AddObject(module, "Grid3D", + (PyObject *)&pyORBIT_Grid3D_Type) < 0) { + Py_DECREF(&pyORBIT_Grid3D_Type); + return -1; + } + return 0; } #ifdef __cplusplus diff --git a/src/spacecharge/wrap_grid3D.hh b/src/spacecharge/wrap_grid3D.hh index 4a465399..24c4278d 100644 --- a/src/spacecharge/wrap_grid3D.hh +++ b/src/spacecharge/wrap_grid3D.hh @@ -8,7 +8,7 @@ extern "C" { #endif namespace wrap_spacecharge{ - void initGrid3D(PyObject* module); + int initGrid3D(PyObject* module); } #ifdef __cplusplus diff --git a/src/spacecharge/wrap_spacecharge.cc b/src/spacecharge/wrap_spacecharge.cc index ab8d9ed5..390cd177 100644 --- a/src/spacecharge/wrap_spacecharge.cc +++ b/src/spacecharge/wrap_spacecharge.cc @@ -31,13 +31,17 @@ extern "C" { spacechargeMethods }; - PyMODINIT_FUNC initspacecharge(){ - //create new module - PyObject* module = PyModule_Create(&cModPyDem); - //add the other classes init - wrap_spacecharge::initGrid1D(module); - wrap_spacecharge::initGrid2D(module); - wrap_spacecharge::initGrid3D(module); + PyMODINIT_FUNC initspacecharge(){ + //create new module + PyObject* module = PyModule_Create(&cModPyDem); + if(module == NULL) return NULL; + //add the other classes init + wrap_spacecharge::initGrid1D(module); + wrap_spacecharge::initGrid2D(module); + if(wrap_spacecharge::initGrid3D(module) < 0){ + Py_DECREF(module); + return NULL; + } wrap_spacecharge::initUniformEllipsoidFieldCalculator(module); wrap_spacecharge::initSpaceChargeCalcUniformEllipse(module); wrap_spacecharge::initPoissonSolverFFT2D(module); From 870bc4e3b446a3cc21a522d74054c6d92c70ca1d Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 14:20:43 -0400 Subject: [PATCH 13/17] I propose we just make numpy>=2.0 a required dependency. --- .github/workflows/pip-build.sh | 1 - py/orbit/bunch_utils/__init__.py | 35 +++++++++++++---------------- py/orbit/bunch_utils/meson.build | 1 + py/orbit/bunch_utils/numpy_utils.py | 2 +- pyproject.toml | 13 +++-------- 5 files changed, 20 insertions(+), 32 deletions(-) diff --git a/.github/workflows/pip-build.sh b/.github/workflows/pip-build.sh index 9c29f800..7bec7708 100755 --- a/.github/workflows/pip-build.sh +++ b/.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]" \ No newline at end of file diff --git a/py/orbit/bunch_utils/__init__.py b/py/orbit/bunch_utils/__init__.py index 58c92160..3096e0d9 100644 --- a/py/orbit/bunch_utils/__init__.py +++ b/py/orbit/bunch_utils/__init__.py @@ -7,25 +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 .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 -except ImportError: - pass +__all__ = ["ParticleIdNumber"] -__all__ = [] -# __all__.append("addParticleIdNumbers") # doesn't exist -__all__.append("ParticleIdNumber") -__all__.append("bunch_from_shared_numpy") -__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", +] diff --git a/py/orbit/bunch_utils/meson.build b/py/orbit/bunch_utils/meson.build index 256d6f5e..341bed29 100644 --- a/py/orbit/bunch_utils/meson.build +++ b/py/orbit/bunch_utils/meson.build @@ -3,6 +3,7 @@ py_sources = files([ '__init__.py', + 'numpy_utils.py', 'particleidnumber.py', 'serialize.py', ]) diff --git a/py/orbit/bunch_utils/numpy_utils.py b/py/orbit/bunch_utils/numpy_utils.py index 191d170a..76ff52c2 100644 --- a/py/orbit/bunch_utils/numpy_utils.py +++ b/py/orbit/bunch_utils/numpy_utils.py @@ -7,7 +7,7 @@ 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 + 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 diff --git a/pyproject.toml b/pyproject.toml index 7a9cb3de..2a8c6a06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,10 +7,9 @@ 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'}, ] @@ -18,9 +17,3 @@ authors = [ [tool.setuptools_scm] -[project.optional-dependencies] -numpy = [ - "numpy", -] - - From 90ace4bb731582fa9f7e87b0894ff5b4e5001d73 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 14:32:06 -0400 Subject: [PATCH 14/17] Inject --oversubscribe flag into the ompi MPI tests. --- .github/workflows/compilation.yml | 5 ++++ tests/py/orbit/core/test_numpy_interop.py | 34 ++++++++++++++++------- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/.github/workflows/compilation.yml b/.github/workflows/compilation.yml index 4598907a..76102b1a 100644 --- a/.github/workflows/compilation.yml +++ b/.github/workflows/compilation.yml @@ -158,12 +158,15 @@ jobs: 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 @@ -214,6 +217,8 @@ jobs: - 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 diff --git a/tests/py/orbit/core/test_numpy_interop.py b/tests/py/orbit/core/test_numpy_interop.py index 4e1d9a58..cb47a5d3 100644 --- a/tests/py/orbit/core/test_numpy_interop.py +++ b/tests/py/orbit/core/test_numpy_interop.py @@ -1,4 +1,5 @@ import os +import shlex import shutil import subprocess import sys @@ -33,6 +34,25 @@ def _process_failure(result): ) +def _mpi_test_env(): + env = os.environ.copy() + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") + env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") + return env + + +def _mpi_test_command(mpirun, nprocs, script_path): + preflags = shlex.split(os.environ.get("PYORBIT_MPIEXEC_PREFLAGS", "")) + return [ + mpirun, + *preflags, + "-np", + str(nprocs), + sys.executable, + str(script_path), + ] + + def _assert_missing_numpy_raises_import_error(extension_name): """NumPy C-API initialization failure must not terminate Python.""" script = textwrap.dedent( @@ -134,16 +154,13 @@ def test_update_from_numpy_uses_only_local_particle_count(self, tmp_path): ) ) - env = os.environ.copy() - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") result = subprocess.run( - [mpirun, "-np", "2", sys.executable, str(script_path)], + _mpi_test_command(mpirun, 2, script_path), capture_output=True, text=True, timeout=60, check=False, - env=env, + env=_mpi_test_env(), ) if "PYORBIT_MPI_DISABLED" in result.stdout: @@ -235,16 +252,13 @@ def particles_for_rank(mpi_rank): ) ) - env = os.environ.copy() - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT", "1") - env.setdefault("OMPI_ALLOW_RUN_AS_ROOT_CONFIRM", "1") result = subprocess.run( - [mpirun, "-np", "3", sys.executable, str(script_path)], + _mpi_test_command(mpirun, 3, script_path), capture_output=True, text=True, timeout=60, check=False, - env=env, + env=_mpi_test_env(), ) if "PYORBIT_MPI_DISABLED" in result.stdout: From b14e0d5e71a3a036e2b6a3dbb8412b37b8c808c0 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 14:57:47 -0400 Subject: [PATCH 15/17] Mark methods METH_O and METH_NOARGS where appropriate. --- src/orbit/wrap_bunch.cc | 271 ++++++++++++++++++---------------------- 1 file changed, 119 insertions(+), 152 deletions(-) diff --git a/src/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index c4409e38..427d0c28 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -77,7 +77,7 @@ namespace wrap_orbit_bunch{ //---------------------------------------------------------------- //returns the SyncPart python class wrapper instance - static PyObject* Bunch_getSyncParticle(PyObject *self, PyObject *args){ + static PyObject* Bunch_getSyncParticle(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; PyObject* pySyncPart = cpp_bunch->getSyncPart()->getPyWrapper(); Py_INCREF(pySyncPart); @@ -91,7 +91,7 @@ namespace wrap_orbit_bunch{ //---------------------------------------------------------------- //returns the local MPI Comm for this bunch - static PyObject* Bunch_getMPIComm(PyObject *self, PyObject *args){ + static PyObject* Bunch_getMPIComm(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; PyObject* pyMPIComm = (PyObject*) cpp_bunch->getMPI_Comm_Local(); Py_INCREF(pyMPIComm); @@ -99,19 +99,9 @@ namespace wrap_orbit_bunch{ } //sets a new local MPI Comm for this bunch - static PyObject* Bunch_setMPIComm(PyObject *self, PyObject *args){ + static PyObject* Bunch_setMPIComm(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - int nVars = PyTuple_Size(args); - PyObject* pyMPIComm; - if(nVars == 1){ - if(!PyArg_ParseTuple(args,"O:setMPIComm",&pyMPIComm)){ - error("The Bunch method setMPIComm(mpi_comm) - mpi_comm is needed."); - } - cpp_bunch->setMPI_Comm_Local( (pyORBIT_MPI_Comm*) pyMPIComm); - } - else{ - error("The Bunch method should be setMPIComm(mpi_comm)."); - } + cpp_bunch->setMPI_Comm_Local( (pyORBIT_MPI_Comm*) arg); Py_INCREF(Py_None); return Py_None; } @@ -142,13 +132,13 @@ namespace wrap_orbit_bunch{ //removes a particle to the Bunch object //returns the number of particles in the bunch //this is implementation of the deleteParticle(int index) method - static PyObject* Bunch_deleteParticle(PyObject *self, PyObject *args){ + static PyObject* Bunch_deleteParticle(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; int ind; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:deleteParticle",&ind)){ - error("PyBunch - deleteParticle - needs index of particle for deleting"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"i:deleteParticle",&ind)){ + return NULL; } cpp_bunch->deleteParticle(ind); @@ -160,25 +150,25 @@ namespace wrap_orbit_bunch{ //removes a particle to the Bunch object //returns the index of removed macro-particle //this is implementation of the deleteParticleFast(int index) method - static PyObject* Bunch_deleteParticleFast(PyObject *self, PyObject *args){ + static PyObject* Bunch_deleteParticleFast(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; int ind; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:deleteParticleFast",&ind)){ - error("PyBunch - deleteParticleFast - needs index of particle for deleting"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"i:deleteParticleFast",&ind)){ + return NULL; } cpp_bunch->deleteParticleFast(ind); return Py_BuildValue("i",ind); } - static PyObject* Bunch_recoverParticle(PyObject *self, PyObject *args){ + static PyObject* Bunch_recoverParticle(PyObject *self, PyObject *arg){ Bunch *cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; int ind; - if(!PyArg_ParseTuple(args,"i:recoverParticle",&ind)){ - error("PyBunch - recoverParticle - needs index of particle for recovering"); + if(!PyArg_Parse(arg,"i:recoverParticle",&ind)){ + return NULL; } cpp_bunch->recoverParticle(ind); @@ -188,7 +178,7 @@ namespace wrap_orbit_bunch{ //removes all particles from the Bunch object //this is implementation of the deleteAllParticles() method - static PyObject* Bunch_deleteAllParticles(PyObject *self, PyObject *args){ + static PyObject* Bunch_deleteAllParticles(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; cpp_bunch->deleteAllParticles(); Py_INCREF(Py_None); @@ -198,7 +188,7 @@ namespace wrap_orbit_bunch{ //compress the bunch. This method should be called after deleting one // or more macro-particles //this is implementation of the compress() method - static PyObject* Bunch_compress(PyObject *self, PyObject *args){ + static PyObject* Bunch_compress(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; cpp_bunch->compress(); Py_INCREF(Py_None); @@ -450,35 +440,25 @@ namespace wrap_orbit_bunch{ // the action is depended on the number of arguments // (index) - returns flag //this is implementation of the flag(int index) method - static PyObject* Bunch_flag(PyObject *self, PyObject *args){ + static PyObject* Bunch_flag(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 get flag - int nVars = PyTuple_Size(args); int index = 0; - int flag = 0; - if(nVars == 1){ - if(!PyArg_ParseTuple(args,"i:flag",&index)){ - error("PyBunch - flag(index) - index is needed"); - } - flag = cpp_bunch->flag(index); - return Py_BuildValue("i",flag); - } - else{ - error("PyBunch. You should call bunch.flag(index)"); + if(!PyArg_Parse(arg,"i:flag",&index)){ + return NULL; } - Py_INCREF(Py_None); - return Py_None; + int flag = cpp_bunch->flag(index); + return Py_BuildValue("i",flag); } //Wraps long. coords in the bunch //ringwrap(ring_length) - static PyObject* Bunch_ringwrap(PyObject *self, PyObject *args) { + static PyObject* Bunch_ringwrap(PyObject *self, PyObject *arg) { Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; double ring_length = 0.; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! //NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"d:py",&ring_length)){ - error("PyBunch - ringwrap(ring_length) - pyBunch object needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! //NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"d:ringwrap",&ring_length)){ + return NULL; } cpp_bunch->ringwrap(ring_length); @@ -527,14 +507,14 @@ namespace wrap_orbit_bunch{ } //Returns classicalRadius of the particle in meters - static PyObject* Bunch_classicalRadius(PyObject *self, PyObject *args){ + static PyObject* Bunch_classicalRadius(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; double val = cpp_bunch->getClassicalRadius(); return Py_BuildValue("d",val); } //Returns B_Rho of the particle in [Tesla*meter]. Parameter is used in TEAPOT - static PyObject* Bunch_B_Rho(PyObject *self, PyObject *args){ + static PyObject* Bunch_B_Rho(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; double val = cpp_bunch->getB_Rho(); return Py_BuildValue("d",val); @@ -615,12 +595,12 @@ namespace wrap_orbit_bunch{ //---------------------------------------------------------------- //initilizes bunch attributes from the bunch file - static PyObject* Bunch_initBunchAttr(PyObject *self, PyObject *args){ + static PyObject* Bunch_initBunchAttr(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* file_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:initBunchAttr",&file_name)){ - error("PyBunch - initBunchAttr(fileName) - the file name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:initBunchAttr",&file_name)){ + return NULL; } cpp_bunch->initBunchAttributes(file_name); Py_INCREF(Py_None); @@ -718,7 +698,7 @@ namespace wrap_orbit_bunch{ } //Returns a list (tuple) of ther double bunch attribute names - static PyObject* Bunch_bunchAttrDoubleNames(PyObject *self, PyObject *args){ + static PyObject* Bunch_bunchAttrDoubleNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; std::vector names; cpp_bunch->getDoubleBunchAttributeNames(names); @@ -734,7 +714,7 @@ namespace wrap_orbit_bunch{ } //Returns a list (tuple) of ther integer bunch attribute names - static PyObject* Bunch_bunchAttrIntNames(PyObject *self, PyObject *args){ + static PyObject* Bunch_bunchAttrIntNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; std::vector names; cpp_bunch->getIntBunchAttributeNames(names); @@ -750,12 +730,12 @@ namespace wrap_orbit_bunch{ } //Returns 0 or 1. The result is 1 if the bunch has an attribute with a particular name - static PyObject* Bunch_hasBunchAttrDouble(PyObject *self, PyObject *args){ + static PyObject* Bunch_hasBunchAttrDouble(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:hasBunchAttrDouble",&attr_name)){ - error("PyBunch - hasBunchAttrDouble(name) - a bunch attr. name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:hasBunchAttrDouble",&attr_name)){ + return NULL; } std::string attr_name_str(attr_name); int res = cpp_bunch->getBunchAttributes()->hasDoubleAttribute(attr_name_str); @@ -763,12 +743,12 @@ namespace wrap_orbit_bunch{ } //Returns 0 or 1. The result is 1 if the bunch has an attribute with a particular name - static PyObject* Bunch_hasBunchAttrInt(PyObject *self, PyObject *args){ + static PyObject* Bunch_hasBunchAttrInt(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:hasBunchAttrInt",&attr_name)){ - error("PyBunch - hasBunchAttrInt(name) - a bunch attr. name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:hasBunchAttrInt",&attr_name)){ + return NULL; } std::string attr_name_str(attr_name); int res = cpp_bunch->getBunchAttributes()->hasIntAttribute(attr_name_str); @@ -813,12 +793,12 @@ namespace wrap_orbit_bunch{ } //Removes a particles' attributes with a particular name from the bunch - static PyObject* Bunch_removePartAttr(PyObject *self, PyObject *args){ + static PyObject* Bunch_removePartAttr(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:removePartAttr",&attr_name)){ - error("PyBunch - removePartAttr(name) - pyBunch object and a particle attr. name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:removePartAttr",&attr_name)){ + return NULL; } std::string attr_name_str(attr_name); cpp_bunch->removeParticleAttributes(attr_name_str); @@ -827,7 +807,7 @@ namespace wrap_orbit_bunch{ } //Removes all particles' attributes from the bunch - static PyObject* Bunch_removeAllPartAttr(PyObject *self, PyObject *args){ + static PyObject* Bunch_removeAllPartAttr(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; cpp_bunch->removeAllParticleAttributes(); Py_INCREF(Py_None); @@ -835,7 +815,7 @@ namespace wrap_orbit_bunch{ } //Returns a list (tuple) of the particles' attributes names - static PyObject* Bunch_getPartAttrNames(PyObject *self, PyObject *args){ + static PyObject* Bunch_getPartAttrNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; std::vector names; cpp_bunch->getParticleAttributesNames(names); @@ -851,7 +831,7 @@ namespace wrap_orbit_bunch{ } //Returns a dict{"part. attribute name":dict{"key":val}} - static PyObject* Bunch_getPartAttrDicts(PyObject *self, PyObject *args){ + static PyObject* Bunch_getPartAttrDicts(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; std::vector names; cpp_bunch->getParticleAttributesNames(names); @@ -872,7 +852,7 @@ namespace wrap_orbit_bunch{ } //Returns a list (tuple) of the possible particles' attributes names - static PyObject* Bunch_getPossiblePartAttrNames(PyObject *self, PyObject *args){ + static PyObject* Bunch_getPossiblePartAttrNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ std::vector names; ParticleAttributesFactory::getParticleAttributesNames(names); //create tuple with names @@ -887,7 +867,7 @@ namespace wrap_orbit_bunch{ } //temporary removes and memorizes all particles' attributes names - static PyObject* Bunch_clearAllPartAttrAndMemorize(PyObject *self, PyObject *args){ + static PyObject* Bunch_clearAllPartAttrAndMemorize(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; cpp_bunch->clearAllParticleAttributesAndMemorize(); Py_INCREF(Py_None); @@ -895,7 +875,7 @@ namespace wrap_orbit_bunch{ } //restores all particles' attributes names from memory - static PyObject* Bunch_restoreAllPartAttrFromMemory(PyObject *self, PyObject *args){ + static PyObject* Bunch_restoreAllPartAttrFromMemory(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; cpp_bunch->restoreAllParticleAttributesFromMemory(); Py_INCREF(Py_None); @@ -903,12 +883,12 @@ namespace wrap_orbit_bunch{ } //Returns 0 or 1. The result is 1 if the bunch has a particles' attributes with a particular name - static PyObject* Bunch_hasPartAttr(PyObject *self, PyObject *args){ + static PyObject* Bunch_hasPartAttr(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:hasPartAttr",&attr_name)){ - error("PyBunch - hasPartAttr(name) - a particles' attr. name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:hasPartAttr",&attr_name)){ + return NULL; } std::string attr_name_str(attr_name); int res = cpp_bunch->hasParticleAttributes(attr_name_str); @@ -916,14 +896,14 @@ namespace wrap_orbit_bunch{ } //Returns a list (tuple) of their bunch particles attribute names specified in the bunch file - static PyObject* Bunch_readPartAttrNames(PyObject *self, PyObject *args){ + static PyObject* Bunch_readPartAttrNames(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* file_name = NULL; std::vector names; std::map > part_attr_dicts; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:readPartAttrNames",&file_name)){ - error("PyBunch - readPartAttrNames(fileName) - a file name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:readPartAttrNames",&file_name)){ + return NULL; } cpp_bunch->readParticleAttributesNames(file_name,names,part_attr_dicts); //create tuple with names @@ -939,13 +919,13 @@ namespace wrap_orbit_bunch{ //Returns a dictionary with the bunch particles attribute names as keys and //dictionaries with parameter:value for each attribute - static PyObject* Bunch_readPartAttrDicts(PyObject *self, PyObject *args){ + static PyObject* Bunch_readPartAttrDicts(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* file_name = NULL; std::vector names; std::map > part_attr_dicts; - if(!PyArg_ParseTuple(args,"s:readPartAttrDicts",&file_name)){ - error("PyBunch - readPartAttrDicts(fileName) - a file name are needed"); + if(!PyArg_Parse(arg,"s:readPartAttrDicts",&file_name)){ + return NULL; } cpp_bunch->readParticleAttributesNames(file_name,names,part_attr_dicts); PyObject* resDict = PyDict_New(); @@ -966,12 +946,12 @@ namespace wrap_orbit_bunch{ } //initilizes particles' attributes from the bunch file - static PyObject* Bunch_readPartAttr(PyObject *self, PyObject *args){ + static PyObject* Bunch_readPartAttr(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* file_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:readPartAttr",&file_name)){ - error("PyBunch - readPartAttr(fileName) - pyBunch object and file name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:readPartAttr",&file_name)){ + return NULL; } cpp_bunch->readParticleAttributes(file_name); Py_INCREF(Py_None); @@ -979,12 +959,12 @@ namespace wrap_orbit_bunch{ } //Returns the number of variables in the particles' attributes with a particular name - static PyObject* Bunch_getPartAttrSize(PyObject *self, PyObject *args){ + static PyObject* Bunch_getPartAttrSize(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:getPartAttrSize",&attr_name)){ - error("PyBunch - getPartAttrSize(name) - a particles' attr. name are needed"); + //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if(!PyArg_Parse(arg,"s:getPartAttrSize",&attr_name)){ + return NULL; } std::string attr_name_str(attr_name); int size = cpp_bunch->getParticleAttributes(attr_name_str)->getAttSize(); @@ -1044,14 +1024,14 @@ namespace wrap_orbit_bunch{ //returns the number of macro-particles in the bunch //this is implementation of the "getSize()" method - static PyObject* Bunch_getSize(PyObject *self, PyObject *args){ + static PyObject* Bunch_getSize(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; return Py_BuildValue("i",cpp_bunch->getSize()); } //returns the number of macro-particles in the bunch in all CPUs //this is implementation of the "getSizeGlobal()" method - static PyObject* Bunch_getSizeGlobal(PyObject *self, PyObject *args){ + static PyObject* Bunch_getSizeGlobal(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; return Py_BuildValue("i",cpp_bunch->getSizeGlobal()); } @@ -1059,21 +1039,21 @@ namespace wrap_orbit_bunch{ //returns the number of macro-particles in the bunch in all CPUs // that was calculated in the previous call of getSizeGlobal() //this is implementation of the "getSizeGlobalFromMemory()" method - static PyObject* Bunch_getSizeGlobalFromMemory(PyObject *self, PyObject *args){ + static PyObject* Bunch_getSizeGlobalFromMemory(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; return Py_BuildValue("i",cpp_bunch->getSizeGlobalFromMemory()); } //returns the number of all macro-particles - alive, dead, new //this is implementation of the "getTotalCount()" method - static PyObject* Bunch_getTotalCount(PyObject *self, PyObject *args){ + static PyObject* Bunch_getTotalCount(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; return Py_BuildValue("i",cpp_bunch->getTotalCount()); } //returns the capacity of the bunch-container. It could be changed. //this is implementation of the "getCapacity()" method - static PyObject* Bunch_getCapacity(PyObject *self, PyObject *args){ + static PyObject* Bunch_getCapacity(PyObject *self, PyObject *Py_UNUSED(ignored)){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; return Py_BuildValue("i",cpp_bunch->getCapacity()); } @@ -1147,13 +1127,9 @@ namespace wrap_orbit_bunch{ } //Copy bunch attrubutes and structure to another bunch - static PyObject* Bunch_copyEmptyBunchTo(PyObject *self, PyObject *args){ + static PyObject* Bunch_copyEmptyBunchTo(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyBunch_Target; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"O:copyEmptyBunchTo",&pyBunch_Target)){ - error("PyBunch - copyEmptyBunchTo(pyBunch) - target pyBunch object is needed"); - } + PyObject* pyBunch_Target = arg; Bunch* cpp_target_bunch = (Bunch*) ((pyORBIT_Object *) pyBunch_Target)->cpp_obj; cpp_bunch->copyEmptyBunchTo(cpp_target_bunch); Py_INCREF(Py_None); @@ -1161,13 +1137,9 @@ namespace wrap_orbit_bunch{ } //Copy bunch all info including particles coordinates and attributes to another bunch - static PyObject* Bunch_copyBunchTo(PyObject *self, PyObject *args){ + static PyObject* Bunch_copyBunchTo(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyBunch_Target; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"O:copyBunchTo",&pyBunch_Target)){ - error("PyBunch - copyBunchTo(pyBunch) - target pyBunch object is needed"); - } + PyObject* pyBunch_Target = arg; Bunch* cpp_target_bunch = (Bunch*) ((pyORBIT_Object *) pyBunch_Target)->cpp_obj; cpp_bunch->copyBunchTo(cpp_target_bunch); Py_INCREF(Py_None); @@ -1175,14 +1147,9 @@ namespace wrap_orbit_bunch{ } //Copy particles coordinates from one bunch to another - static PyObject* Bunch_addParticlesTo(PyObject *self, PyObject *args){ + static PyObject* Bunch_addParticlesTo(PyObject *self, PyObject *arg){ Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyBunch_Target; - - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"O:addParticlesTo",&pyBunch_Target)){ - error("PyBunch - addParticlesTo(pyBunch) - target pyBunch object is needed"); - } + PyObject* pyBunch_Target = arg; Bunch* cpp_target_bunch =(Bunch*) ((pyORBIT_Object *) pyBunch_Target)->cpp_obj ; cpp_bunch->addParticlesTo(cpp_target_bunch); Py_INCREF(Py_None); @@ -1511,15 +1478,15 @@ static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *arg) { //-------------------------------------------------------- // class Bunch wrapper START //-------------------------------------------------------- - { "getMPIComm", Bunch_getMPIComm ,METH_VARARGS,"Returns MPI Comm of this bunch"}, - { "setMPIComm", Bunch_setMPIComm ,METH_VARARGS,"Sets a new MPI Comm for this bunch"}, - { "getSyncParticle", Bunch_getSyncParticle ,METH_VARARGS,"Returns syncParticle class instance"}, + { "getMPIComm", Bunch_getMPIComm ,METH_NOARGS,"Returns MPI Comm of this bunch"}, + { "setMPIComm", Bunch_setMPIComm ,METH_O,"Sets a new MPI Comm for this bunch"}, + { "getSyncParticle", Bunch_getSyncParticle ,METH_NOARGS,"Returns syncParticle class instance"}, { "addParticle", Bunch_addParticle ,METH_VARARGS,"Adds a macro-particle to the bunch"}, - { "deleteParticle", Bunch_deleteParticle ,METH_VARARGS,"Removes macro-particle from the bunch and call compress inside"}, - { "deleteParticleFast", Bunch_deleteParticleFast ,METH_VARARGS,"Removes macro-particle from the bunch very fast"}, - { "recoverParticle", Bunch_recoverParticle ,METH_VARARGS,"Recovers a particle marked for removal"}, - { "deleteAllParticles", Bunch_deleteAllParticles ,METH_VARARGS,"Removes all macro-particles from the bunch"}, - { "compress", Bunch_compress ,METH_VARARGS,"Compress the bunch"}, + { "deleteParticle", Bunch_deleteParticle ,METH_O,"Removes macro-particle from the bunch and call compress inside"}, + { "deleteParticleFast", Bunch_deleteParticleFast ,METH_O,"Removes macro-particle from the bunch very fast"}, + { "recoverParticle", Bunch_recoverParticle ,METH_O,"Recovers a particle marked for removal"}, + { "deleteAllParticles", Bunch_deleteAllParticles ,METH_NOARGS,"Removes all macro-particles from the bunch"}, + { "compress", Bunch_compress ,METH_NOARGS,"Compress the bunch"}, { "x", Bunch_x ,METH_VARARGS,"Set x(index,value) or get x(index) coordinate"}, { "y", Bunch_y ,METH_VARARGS,"Set y(index,value) or get y(index) coordinate"}, { "z", Bunch_z ,METH_VARARGS,"Set z(index,value) or get z(index) coordinate"}, @@ -1529,44 +1496,44 @@ static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *arg) { { "dE", Bunch_pz ,METH_VARARGS,"Set dE(index,value) or get dE(index) coordinate"}, { "xp", Bunch_px ,METH_VARARGS,"Set xp(index,value) or get xp(index) coordinate"}, { "yp", Bunch_py ,METH_VARARGS,"Set yp(index,value) or get yp(index) coordinate"}, - { "flag", Bunch_flag ,METH_VARARGS,"Returns flag(index) for particle with index"}, - { "ringwrap", Bunch_ringwrap ,METH_VARARGS,"Perform the ring wrap. Usage: ringwrap(ring_length)"}, + { "flag", Bunch_flag ,METH_O,"Returns flag(index) for particle with index"}, + { "ringwrap", Bunch_ringwrap ,METH_O,"Perform the ring wrap. Usage: ringwrap(ring_length)"}, { "mass", Bunch_mass ,METH_VARARGS,"Set mass(value) or get mass() the mass of particle in MeV"}, - { "classicalRadius", Bunch_classicalRadius ,METH_VARARGS,"Returns a classical radius of particle in [m]"}, - { "B_Rho", Bunch_B_Rho ,METH_VARARGS,"Returns B*Rho parameter of particle in [T*m]"}, + { "classicalRadius", Bunch_classicalRadius ,METH_NOARGS,"Returns a classical radius of particle in [m]"}, + { "B_Rho", Bunch_B_Rho ,METH_NOARGS,"Returns B*Rho parameter of particle in [T*m]"}, { "charge", Bunch_charge ,METH_VARARGS,"Set charge(value) or get charge() the charge of particle in e-charge"}, { "macroSize", Bunch_macroSize ,METH_VARARGS,"Set macroSize(value) or get macroSize() the charge of particle in e-charge"}, - { "initBunchAttr", Bunch_initBunchAttr ,METH_VARARGS,"Reads and initilizes bunch attributes from a bunch file"}, + { "initBunchAttr", Bunch_initBunchAttr ,METH_O,"Reads and initilizes bunch attributes from a bunch file"}, { "bunchAttrDouble", Bunch_bunchAttrDouble ,METH_VARARGS,"Returns and sets a double bunch attribute"}, { "bunchAttrInt", Bunch_bunchAttrInt ,METH_VARARGS,"Returns and sets an integer bunch attribute"}, - { "bunchAttrDoubleNames", Bunch_bunchAttrDoubleNames ,METH_VARARGS,"Returns a list of double bunch attribute names"}, - { "bunchAttrIntNames", Bunch_bunchAttrIntNames ,METH_VARARGS,"Returns a list of integer bunch attribute names"}, - { "hasBunchAttrDouble", Bunch_hasBunchAttrDouble ,METH_VARARGS,"Returns 1 if there is a double bunch attr. with this name, 0 - otherwise"}, - { "hasBunchAttrInt", Bunch_hasBunchAttrInt ,METH_VARARGS,"Returns 1 if there is a int bunch attr. with this name, 0 - otherwise"}, + { "bunchAttrDoubleNames", Bunch_bunchAttrDoubleNames ,METH_NOARGS,"Returns a list of double bunch attribute names"}, + { "bunchAttrIntNames", Bunch_bunchAttrIntNames ,METH_NOARGS,"Returns a list of integer bunch attribute names"}, + { "hasBunchAttrDouble", Bunch_hasBunchAttrDouble ,METH_O,"Returns 1 if there is a double bunch attr. with this name, 0 - otherwise"}, + { "hasBunchAttrInt", Bunch_hasBunchAttrInt ,METH_O,"Returns 1 if there is a int bunch attr. with this name, 0 - otherwise"}, { "addPartAttr", Bunch_addPartAttr ,METH_VARARGS,"Adds a particles' attributes to the bunch"}, - { "removePartAttr", Bunch_removePartAttr ,METH_VARARGS,"Removes a particles' attributes from the bunch"}, - { "removeAllPartAttr", Bunch_removeAllPartAttr ,METH_VARARGS,"Removes all particles' attributes from the bunch"}, - { "getPartAttrNames", Bunch_getPartAttrNames ,METH_VARARGS,"Returns all particles' attributes names in the bunch at this moment"}, - { "getPartAttrDicts", Bunch_getPartAttrDicts ,METH_VARARGS,"Returns dict{part. attribute name:dict{key:val}}"}, - { "getPossiblePartAttrNames", Bunch_getPossiblePartAttrNames ,METH_VARARGS,"Returns all possible particles' attributes names"}, - { "clearAllPartAttrAndMemorize", Bunch_clearAllPartAttrAndMemorize ,METH_VARARGS,"Temporary removes and memorizes all particles' attributes names"}, - { "restoreAllPartAttrFromMemory", Bunch_restoreAllPartAttrFromMemory ,METH_VARARGS,"Restores all particles' attributes names from memory"}, - { "hasPartAttr", Bunch_hasPartAttr ,METH_VARARGS,"Returns 1 if there is a particles' attr. with this name, 0 - otherwis"}, - { "readPartAttrNames", Bunch_readPartAttrNames ,METH_VARARGS,"Returns a tuple with particles' attr. names in the bunch file"}, - { "readPartAttrDicts", Bunch_readPartAttrDicts ,METH_VARARGS,"Returns a dict{attr_name:dicts{param_name:val}} in the bunch file"}, - { "readPartAttr", Bunch_readPartAttr ,METH_VARARGS,"Initializes the particles' attr. from the bunch file"}, - { "getPartAttrSize", Bunch_getPartAttrSize ,METH_VARARGS,"Returns the number of variables in the particles' attributes with a particular name"}, + { "removePartAttr", Bunch_removePartAttr ,METH_O,"Removes a particles' attributes from the bunch"}, + { "removeAllPartAttr", Bunch_removeAllPartAttr ,METH_NOARGS,"Removes all particles' attributes from the bunch"}, + { "getPartAttrNames", Bunch_getPartAttrNames ,METH_NOARGS,"Returns all particles' attributes names in the bunch at this moment"}, + { "getPartAttrDicts", Bunch_getPartAttrDicts ,METH_NOARGS,"Returns dict{part. attribute name:dict{key:val}}"}, + { "getPossiblePartAttrNames", Bunch_getPossiblePartAttrNames ,METH_NOARGS,"Returns all possible particles' attributes names"}, + { "clearAllPartAttrAndMemorize", Bunch_clearAllPartAttrAndMemorize ,METH_NOARGS,"Temporary removes and memorizes all particles' attributes names"}, + { "restoreAllPartAttrFromMemory", Bunch_restoreAllPartAttrFromMemory ,METH_NOARGS,"Restores all particles' attributes names from memory"}, + { "hasPartAttr", Bunch_hasPartAttr ,METH_O,"Returns 1 if there is a particles' attr. with this name, 0 - otherwis"}, + { "readPartAttrNames", Bunch_readPartAttrNames ,METH_O,"Returns a tuple with particles' attr. names in the bunch file"}, + { "readPartAttrDicts", Bunch_readPartAttrDicts ,METH_O,"Returns a dict{attr_name:dicts{param_name:val}} in the bunch file"}, + { "readPartAttr", Bunch_readPartAttr ,METH_O,"Initializes the particles' attr. from the bunch file"}, + { "getPartAttrSize", Bunch_getPartAttrSize ,METH_O,"Returns the number of variables in the particles' attributes with a particular name"}, { "partAttrValue", Bunch_partAttrValue ,METH_VARARGS,"Sets or returns a particles' attribute value"}, - { "getSize", Bunch_getSize ,METH_VARARGS,"Returns number of macro-particles"}, - { "getSizeGlobal", Bunch_getSizeGlobal ,METH_VARARGS,"Returns number of macro-particles in all CPUs"}, - { "getSizeGlobalFromMemory", Bunch_getSizeGlobalFromMemory ,METH_VARARGS,"Returns number of macro-particles in all CPUs from memory"}, - { "getTotalCount", Bunch_getTotalCount ,METH_VARARGS,"Returns number of all particles - alive,dead,new"}, - { "getCapacity", Bunch_getCapacity ,METH_VARARGS,"Returns the capacity of the bunch-contaiter"}, + { "getSize", Bunch_getSize ,METH_NOARGS,"Returns number of macro-particles"}, + { "getSizeGlobal", Bunch_getSizeGlobal ,METH_NOARGS,"Returns number of macro-particles in all CPUs"}, + { "getSizeGlobalFromMemory", Bunch_getSizeGlobalFromMemory ,METH_NOARGS,"Returns number of macro-particles in all CPUs from memory"}, + { "getTotalCount", Bunch_getTotalCount ,METH_NOARGS,"Returns number of all particles - alive,dead,new"}, + { "getCapacity", Bunch_getCapacity ,METH_NOARGS,"Returns the capacity of the bunch-contaiter"}, { "dumpBunch", Bunch_dumpBunch ,METH_VARARGS,"Prints the bunch info into a standart output stream or file"}, { "readBunch", Bunch_readBunch ,METH_VARARGS,"Reads the bunch info from a file"}, - { "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"}, + { "copyEmptyBunchTo", Bunch_copyEmptyBunchTo ,METH_O,"Copy bunch attrubutes and structure to another bunch"}, + { "copyBunchTo", Bunch_copyBunchTo ,METH_O,"Copy bunch all info including particles coordinates and attributes to another bunch"}, + { "addParticlesTo", Bunch_addParticlesTo ,METH_O,"Copy particles coordinates from one bunch to another"}, #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY { "to_numpy", _PyCFunction_CAST(Bunch_to_numpy) ,METH_VARARGS | METH_KEYWORDS, Bunch_to_numpy_doc }, { "update_from_numpy", Bunch_update_from_numpy ,METH_O, Bunch_update_from_numpy_doc }, From 06e1491214dff4801ddd9846b0556abb1583080c Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 15:18:27 -0400 Subject: [PATCH 16/17] Bunch_mass simplification --- src/orbit/wrap_bunch.cc | 31 +++++++++---------------------- 1 file changed, 9 insertions(+), 22 deletions(-) diff --git a/src/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index 427d0c28..129515d1 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -478,32 +478,19 @@ namespace wrap_orbit_bunch{ // mass(value) - sets the new value //this is implementation of the getMass() and setMass methods of the Bunch class static PyObject* Bunch_mass(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 0 get mass - //if nVars == 1 set mass - int nVars = PyTuple_Size(args); + Bunch* bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - double val = 0.; + if(0 == PyTuple_GET_SIZE(args)) { + return PyFloat_FromDouble(bunch->getMass()); + } - if(nVars == 0 || nVars == 1){ - if(nVars == 0){ - val = cpp_bunch->getMass(); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"d:mass",&val)){ - error("PyBunch - mass(value) - value is needed"); + double value; + if (!PyArg_ParseTuple(args, "d:mass", &value)) { + return NULL; } - cpp_bunch->setMass(val); - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call mass() or mass(value)"); - } - Py_INCREF(Py_None); - return Py_None; + bunch->setMass(value); + return PyFloat_FromDouble(value); } //Returns classicalRadius of the particle in meters From ad19735c471d0b517fac9447bbc70cba98c9da7a Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 27 Aug 2026 15:52:07 -0400 Subject: [PATCH 17/17] clang format --- src/orbit/wrap_bunch.cc | 2785 ++++++++++++++++++++------------------- 1 file changed, 1419 insertions(+), 1366 deletions(-) diff --git a/src/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index 129515d1..cdeb9563 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -9,19 +9,18 @@ // /////////////////////////////////////////////////////////////////////////// #include "wrap_bunch.hh" -// #include "modsupport.h" -// #include "pyerrors.h" -#include "wrap_syncpart.hh" -#include "wrap_bunch_twiss_analysis.hh" -#include "wrap_bunch_tune_analysis.hh" -#include "wrap_synch_part_redefinition_z_de.hh" #include "pyORBIT_Object.hh" +#include "wrap_bunch_tune_analysis.hh" +#include "wrap_bunch_twiss_analysis.hh" +#include "wrap_synch_part_redefinition_z_de.hh" +#include "wrap_syncpart.hh" #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY #include -static int ensure_numpy() { +static int ensure_numpy() +{ static int numpy_initialized = 0; if (!numpy_initialized) { import_array1(-1); @@ -34,1218 +33,1148 @@ static int ensure_numpy() { #include "Bunch.hh" #include "ParticleAttributesFactory.hh" -namespace wrap_orbit_bunch{ +namespace wrap_orbit_bunch +{ - void error(const char* msg){ ORBIT_MPI_Finalize(msg); } - //--------------------------------------------------------- - //Python Bunch class definition - //--------------------------------------------------------- +void error(const char *msg) +{ + ORBIT_MPI_Finalize(msg); +} - //constructor for python class wrapping Bunch instance - //It never will be called directly - static PyObject* Bunch_new(PyTypeObject *type, PyObject *args, PyObject *kwds){ - pyORBIT_Object* self; - self = (pyORBIT_Object *) type->tp_alloc(type, 0); - self->cpp_obj = NULL; - return (PyObject *) self; - } +//--------------------------------------------------------- +// Python Bunch class definition +//--------------------------------------------------------- + +// constructor for python class wrapping Bunch instance +// It never will be called directly +static PyObject *Bunch_new(PyTypeObject *type, PyObject *args, PyObject *kwds) +{ + pyORBIT_Object *self; + self = (pyORBIT_Object *)type->tp_alloc(type, 0); + self->cpp_obj = NULL; + return (PyObject *)self; +} - //initializator for python Bunch class - //this is implementation of the __init__ method - static int Bunch_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds){ - //std::cerr<<"The Bunch __init__ has been called!"<cpp_obj = (void*) new Bunch(); - ((Bunch*) self->cpp_obj)->setPyWrapper((PyObject*) self); - //This is the way to create new class instance from the C-level - // Template: PyObject* PyObject_CallMethod( PyObject *o, char *method, char *format, ...) - //see Python/C API documentation - //It will create a SyncParticle object and set the reference to it from pyBunch - PyObject* mod = PyImport_ImportModule("orbit.core.bunch"); - PyObject* pySyncPart = PyObject_CallMethod(mod,const_cast("SyncParticle"),const_cast("O"),self); - - //the references should be decreased because they were created as "new reference" - Py_DECREF(pySyncPart); - Py_DECREF(mod); - return 0; - } - - //--------------------------------------------------------------- - // - // methods related to synchronous particle etc. - // - //---------------------------------------------------------------- - - //returns the SyncPart python class wrapper instance - static PyObject* Bunch_getSyncParticle(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pySyncPart = cpp_bunch->getSyncPart()->getPyWrapper(); - Py_INCREF(pySyncPart); - return pySyncPart; - } - - //--------------------------------------------------------------- - // - // set and get MPI Communicators - // - //---------------------------------------------------------------- - - //returns the local MPI Comm for this bunch - static PyObject* Bunch_getMPIComm(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyMPIComm = (PyObject*) cpp_bunch->getMPI_Comm_Local(); - Py_INCREF(pyMPIComm); - return pyMPIComm; - } - - //sets a new local MPI Comm for this bunch - static PyObject* Bunch_setMPIComm(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - cpp_bunch->setMPI_Comm_Local( (pyORBIT_MPI_Comm*) arg); - Py_INCREF(Py_None); - return Py_None; - } - - //--------------------------------------------------------------- - // - // add and remove particles, compress etc. - // - //---------------------------------------------------------------- - - //adds a particle to the Bunch object - //this is implementation of the addParticle(...) method - static PyObject* Bunch_addParticle(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - - double x = 0.; double xp = 0.; double y = 0.; - double yp = 0.; double z = 0.; double zp = 0.; - - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"dddddd:coordinates",&x,&xp,&y,&yp,&z,&zp)){ - error("PyBunch - addParticle - cannot parse arguments! It should be (x,xp,y,yp,z,zp)"); - } - int ind = cpp_bunch->addParticle(x,xp,y,yp,z,zp); - return Py_BuildValue("i",ind); +// initializator for python Bunch class +// this is implementation of the __init__ method +static int Bunch_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds) +{ + // std::cerr<<"The Bunch __init__ has been called!"<cpp_obj = (void *)new Bunch(); + ((Bunch *)self->cpp_obj)->setPyWrapper((PyObject *)self); + // This is the way to create new class instance from the C-level + // Template: PyObject* PyObject_CallMethod( PyObject *o, char *method, char *format, ...) + // see Python/C API documentation + // It will create a SyncParticle object and set the reference to it from pyBunch + PyObject *mod = PyImport_ImportModule("orbit.core.bunch"); + PyObject *pySyncPart = + PyObject_CallMethod(mod, const_cast("SyncParticle"), const_cast("O"), self); + + // the references should be decreased because they were created as "new reference" + Py_DECREF(pySyncPart); + Py_DECREF(mod); + return 0; +} + +//--------------------------------------------------------------- +// +// methods related to synchronous particle etc. +// +//---------------------------------------------------------------- + +// returns the SyncPart python class wrapper instance +static PyObject *Bunch_getSyncParticle(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + PyObject *pySyncPart = cpp_bunch->getSyncPart()->getPyWrapper(); + Py_INCREF(pySyncPart); + return pySyncPart; +} + +//--------------------------------------------------------------- +// +// set and get MPI Communicators +// +//---------------------------------------------------------------- + +// returns the local MPI Comm for this bunch +static PyObject *Bunch_getMPIComm(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + PyObject *pyMPIComm = (PyObject *)cpp_bunch->getMPI_Comm_Local(); + Py_INCREF(pyMPIComm); + return pyMPIComm; +} + +// sets a new local MPI Comm for this bunch +static PyObject *Bunch_setMPIComm(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + cpp_bunch->setMPI_Comm_Local((pyORBIT_MPI_Comm *)arg); + Py_RETURN_NONE; +} + +//--------------------------------------------------------------- +// +// add and remove particles, compress etc. +// +//---------------------------------------------------------------- + +// adds a particle to the Bunch object +// this is implementation of the addParticle(...) method +static PyObject *Bunch_addParticle(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + + double x = 0.; + double xp = 0.; + double y = 0.; + double yp = 0.; + double z = 0.; + double zp = 0.; + + // NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() + if (!PyArg_ParseTuple(args, "dddddd:coordinates", &x, &xp, &y, &yp, &z, &zp)) { + error("PyBunch - addParticle - cannot parse arguments! It should be (x,xp,y,yp,z,zp)"); + } + int ind = cpp_bunch->addParticle(x, xp, y, yp, z, zp); + return PyLong_FromLong(ind); +} + +// removes a particle to the Bunch object +// returns the number of particles in the bunch +// this is implementation of the deleteParticle(int index) method +static PyObject *Bunch_deleteParticle(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int ind; + + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "i:deleteParticle", &ind)) { + return NULL; } + cpp_bunch->deleteParticle(ind); + int size = cpp_bunch->getSize(); - //removes a particle to the Bunch object - //returns the number of particles in the bunch - //this is implementation of the deleteParticle(int index) method - static PyObject* Bunch_deleteParticle(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - int ind; + return PyLong_FromLong(size); +} - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"i:deleteParticle",&ind)){ - return NULL; - } +// removes a particle to the Bunch object +// returns the index of removed macro-particle +// this is implementation of the deleteParticleFast(int index) method +static PyObject *Bunch_deleteParticleFast(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int ind; + + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "i:deleteParticleFast", &ind)) { + return NULL; + } + + cpp_bunch->deleteParticleFast(ind); + return PyLong_FromLong(ind); +} - cpp_bunch->deleteParticle(ind); - int size = cpp_bunch->getSize(); +static PyObject *Bunch_recoverParticle(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int ind; - return Py_BuildValue("i",size); + if (!PyArg_Parse(arg, "i:recoverParticle", &ind)) { + return NULL; } - //removes a particle to the Bunch object - //returns the index of removed macro-particle - //this is implementation of the deleteParticleFast(int index) method - static PyObject* Bunch_deleteParticleFast(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - int ind; + cpp_bunch->recoverParticle(ind); + Py_RETURN_NONE; +} + +// removes all particles from the Bunch object +// this is implementation of the deleteAllParticles() method +static PyObject *Bunch_deleteAllParticles(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + cpp_bunch->deleteAllParticles(); + Py_RETURN_NONE; +} - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"i:deleteParticleFast",&ind)){ - return NULL; - } +// compress the bunch. This method should be called after deleting one +// or more macro-particles +// this is implementation of the compress() method +static PyObject *Bunch_compress(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + cpp_bunch->compress(); + Py_RETURN_NONE; +} - cpp_bunch->deleteParticleFast(ind); - return Py_BuildValue("i",ind); +//--------------------------------------------------------------- +// +// related to the macro-particles' coordinates +// +//---------------------------------------------------------------- + +// Sets or returns x coordinate of the macro-particle +// the action is depended on the number of arguments +// (index) - returns x-coordinate +// (index, value) - sets the new value to the x-coordinate +// this is implementation of the x(int index) method +static PyObject *Bunch_x(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int index; + double value; + + if (!PyArg_ParseTuple(args, "i|d:x", &index, &value)) { + return NULL; } - static PyObject* Bunch_recoverParticle(PyObject *self, PyObject *arg){ - Bunch *cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - int ind; + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->x(index); + } + else { + cpp_bunch->x(index) = value; + } - if(!PyArg_Parse(arg,"i:recoverParticle",&ind)){ - return NULL; - } + return PyFloat_FromDouble(value); +} - cpp_bunch->recoverParticle(ind); - Py_INCREF(Py_None); - return Py_None; - } - - //removes all particles from the Bunch object - //this is implementation of the deleteAllParticles() method - static PyObject* Bunch_deleteAllParticles(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - cpp_bunch->deleteAllParticles(); - Py_INCREF(Py_None); - return Py_None; - } - - //compress the bunch. This method should be called after deleting one - // or more macro-particles - //this is implementation of the compress() method - static PyObject* Bunch_compress(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - cpp_bunch->compress(); - Py_INCREF(Py_None); - return Py_None; - } - - //--------------------------------------------------------------- - // - // related to the macro-particles' coordinates - // - //---------------------------------------------------------------- - - //Sets or returns x coordinate of the macro-particle - // the action is depended on the number of arguments - // (index) - returns x-coordinate - // (index, value) - sets the new value to the x-coordinate - //this is implementation of the x(int index) method - static PyObject* Bunch_x(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 get coordinate - //if nVars == 2 set coordinate - int nVars = PyTuple_Size(args); - - int index = 0; - double val = 0.; - - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:x",&index)){ - error("PyBunch - x(index) - index is needed"); - } - val = cpp_bunch->x(index); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"id:x",&index,&val)){ - error("PyBunch - x(index, value) - index and value are needed"); - } - cpp_bunch->x(index) = val; - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call x(index) or x(index,value)"); - } +// Sets or returns y coordinate of the macro-particle +// the action is depended on the number of arguments +// (index) - returns y-coordinate +// (index, value) - sets the new value to the y-coordinate +// this is implementation of the y(int index) method +static PyObject *Bunch_y(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int index; + double value; + + if (!PyArg_ParseTuple(args, "i|d:y", &index, &value)) { + return NULL; + } - Py_INCREF(Py_None); - return Py_None; - } - - //Sets or returns y coordinate of the macro-particle - // the action is depended on the number of arguments - // (index) - returns y-coordinate - // (index, value) - sets the new value to the y-coordinate - //this is implementation of the y(int index) method - static PyObject* Bunch_y(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 get coordinate - //if nVars == 2 set coordinate - int nVars = PyTuple_Size(args); - - int index = 0; - double val = 0.; - - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:y",&index)){ - error("PyBunch - y(index) - index is needed"); - } - val = cpp_bunch->y(index); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"id:y",&index,&val)){ - error("PyBunch - y(index, value) - index and value are needed"); - } - cpp_bunch->y(index) = val; - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call y(index) or y(index,value)"); - } + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->y(index); + } + else { + cpp_bunch->y(index) = value; + } - Py_INCREF(Py_None); - return Py_None; - } - - //Sets or returns z coordinate of the macro-particle - // the action is depended on the number of arguments - // (index) - returns z(phi)-coordinate - // (index, value) - sets the new value to the z(phi)-coordinate - //this is implementation of the (z or phi)(int index) method - static PyObject* Bunch_z(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 get coordinate - //if nVars == 2 set coordinate - int nVars = PyTuple_Size(args); - - int index = 0; - double val = 0.; - - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:z",&index)){ - error("PyBunch - z(index) - index is needed"); - } - val = cpp_bunch->z(index); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"id:z",&index,&val)){ - error("PyBunch - z(index, value) - index and value are needed"); - } - cpp_bunch->z(index) = val; - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call z(index) or z(index,value)"); - } + return PyFloat_FromDouble(value); +} - Py_INCREF(Py_None); - return Py_None; +// Sets or returns z coordinate of the macro-particle +// the action is depended on the number of arguments +// (index) - returns z(phi)-coordinate +// (index, value) - sets the new value to the z(phi)-coordinate +// this is implementation of the (z or phi)(int index) method +static PyObject *Bunch_z(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int index; + double value; + + if (!PyArg_ParseTuple(args, "i|d:z", &index, &value)) { + return NULL; } + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->z(index); + } + else { + cpp_bunch->z(index) = value; + } - //Sets or returns px coordinate of the macro-particle - // the action is depended on the number of arguments - // (index) - returns px-coordinate - // (index, value) - sets the new value to the px-coordinate - //this is implementation of the px(int index) method - static PyObject* Bunch_px(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 get coordinate - //if nVars == 2 set coordinate - int nVars = PyTuple_Size(args); + return PyFloat_FromDouble(value); +} - int index = 0; - double val = 0.; +// Sets or returns px coordinate of the macro-particle +// the action is depended on the number of arguments +// (index) - returns px-coordinate +// (index, value) - sets the new value to the px-coordinate +// this is implementation of the px(int index) method +static PyObject *Bunch_px(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int index; + double value; + + if (!PyArg_ParseTuple(args, "i|d:px", &index, &value)) { + return NULL; + } - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:px",&index)){ - error("PyBunch - px(index) - index is needed"); - } - val = cpp_bunch->px(index); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"id:px",&index,&val)){ - error("PyBunch - px(index, value) - index and value are needed"); - } - cpp_bunch->px(index) = val; - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call px(index) or px(index,value)"); - } + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->px(index); + } + else { + cpp_bunch->px(index) = value; + } - Py_INCREF(Py_None); - return Py_None; - } - - //Sets or returns py coordinate of the macro-particle - // the action is depended on the number of arguments - // (index) - returns y-coordinate - // (index, value) - sets the new value to the py-coordinate - //this is implementation of the py(int index) method - static PyObject* Bunch_py(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 get coordinate - //if nVars == 2 set coordinate - int nVars = PyTuple_Size(args); - - int index = 0; - double val = 0.; - - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:py",&index)){ - error("PyBunch - py(index) - index is needed"); - } - val = cpp_bunch->py(index); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"id:py",&index,&val)){ - error("PyBunch - py(index, value) - index and value are needed"); - } - cpp_bunch->py(index) = val; - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call py(index) or py(index,value)"); - } + return PyFloat_FromDouble(value); +} - Py_INCREF(Py_None); - return Py_None; - } - - //Sets or returns pz or dE coordinate of the macro-particle - // the action is depended on the number of arguments - // (index) - returns pz(dE)-coordinate - // (index, value) - sets the new value to the pz(dE)-coordinate - //this is implementation of the (pz or dE)(int index) method - static PyObject* Bunch_pz(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 get coordinate - //if nVars == 2 set coordinate - int nVars = PyTuple_Size(args); - - int index = 0; - double val = 0.; - - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"i:pz",&index)){ - error("PyBunch - pz(index) - index is needed"); - } - val = cpp_bunch->pz(index); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"id:pz",&index,&val)){ - error("PyBunch - pz(index, value) - index and value are needed"); - } - cpp_bunch->pz(index) = val; - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call pz(index) or pz(index,value)"); - } +// Sets or returns py coordinate of the macro-particle +// the action is depended on the number of arguments +// (index) - returns y-coordinate +// (index, value) - sets the new value to the py-coordinate +// this is implementation of the py(int index) method +static PyObject *Bunch_py(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int index; + double value; + + if (!PyArg_ParseTuple(args, "i|d:py", &index, &value)) { + return NULL; + } - Py_INCREF(Py_None); - return Py_None; + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->py(index); + } + else { + cpp_bunch->py(index) = value; } - //Sets or returns flag of the macro-particle - // the action is depended on the number of arguments - // (index) - returns flag - //this is implementation of the flag(int index) method - static PyObject* Bunch_flag(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - int index = 0; - if(!PyArg_Parse(arg,"i:flag",&index)){ - return NULL; - } - int flag = cpp_bunch->flag(index); - return Py_BuildValue("i",flag); + return PyFloat_FromDouble(value); +} + +// Sets or returns pz or dE coordinate of the macro-particle +// the action is depended on the number of arguments +// (index) - returns pz(dE)-coordinate +// (index, value) - sets the new value to the pz(dE)-coordinate +// this is implementation of the (pz or dE)(int index) method +static PyObject *Bunch_pz(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int index; + double value; + + if (!PyArg_ParseTuple(args, "i|d:pz", &index, &value)) { + return NULL; } - //Wraps long. coords in the bunch - //ringwrap(ring_length) - static PyObject* Bunch_ringwrap(PyObject *self, PyObject *arg) { - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - double ring_length = 0.; + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->pz(index); + } + else { + cpp_bunch->pz(index) = value; + } - //NO NEW OBJECT CREATED BY PyArg_Parse! //NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"d:ringwrap",&ring_length)){ - return NULL; - } + return PyFloat_FromDouble(value); +} - cpp_bunch->ringwrap(ring_length); - Py_INCREF(Py_None); - return Py_None; +// Sets or returns flag of the macro-particle +// the action is depended on the number of arguments +// (index) - returns flag +// this is implementation of the flag(int index) method +static PyObject *Bunch_flag(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + int index = 0; + if (!PyArg_Parse(arg, "i:flag", &index)) { + return NULL; } + int flag = cpp_bunch->flag(index); + return PyLong_FromLong(flag); +} - //--------------------------------------------------------------- - // - // related to the bunch predefined attributes - // - //---------------------------------------------------------------- +// Wraps long. coords in the bunch +// ringwrap(ring_length) +static PyObject *Bunch_ringwrap(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + double ring_length = 0.; - //Sets or returns mass of the macro-particle in MeV - // the action is depended on the number of arguments - // mass() - returns mass - // mass(value) - sets the new value - //this is implementation of the getMass() and setMass methods of the Bunch class - static PyObject* Bunch_mass(PyObject *self, PyObject *args){ - Bunch* bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; + // NO NEW OBJECT CREATED BY PyArg_Parse! //NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "d:ringwrap", &ring_length)) { + return NULL; + } - if(0 == PyTuple_GET_SIZE(args)) { - return PyFloat_FromDouble(bunch->getMass()); - } + cpp_bunch->ringwrap(ring_length); + Py_RETURN_NONE; +} - double value; - if (!PyArg_ParseTuple(args, "d:mass", &value)) { - return NULL; - } +//--------------------------------------------------------------- +// +// related to the bunch predefined attributes +// +//---------------------------------------------------------------- + +// Sets or returns mass of the macro-particle in MeV +// the action is depended on the number of arguments +// mass() - returns mass +// mass(value) - sets the new value +// this is implementation of the getMass() and setMass methods of the Bunch class +static PyObject *Bunch_mass(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; - bunch->setMass(value); - return PyFloat_FromDouble(value); + if (PyTuple_GET_SIZE(args) == 0) { + return PyFloat_FromDouble(cpp_bunch->getMass()); } - //Returns classicalRadius of the particle in meters - static PyObject* Bunch_classicalRadius(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - double val = cpp_bunch->getClassicalRadius(); - return Py_BuildValue("d",val); + double value; + if (!PyArg_ParseTuple(args, "d:mass", &value)) { + return NULL; } - //Returns B_Rho of the particle in [Tesla*meter]. Parameter is used in TEAPOT - static PyObject* Bunch_B_Rho(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - double val = cpp_bunch->getB_Rho(); - return Py_BuildValue("d",val); + cpp_bunch->setMass(value); + return PyFloat_FromDouble(value); +} + +// Returns classicalRadius of the particle in meters +static PyObject *Bunch_classicalRadius(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + double val = cpp_bunch->getClassicalRadius(); + return PyFloat_FromDouble(val); +} + +// Returns B_Rho of the particle in [Tesla*meter]. Parameter is used in TEAPOT +static PyObject *Bunch_B_Rho(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + double val = cpp_bunch->getB_Rho(); + return PyFloat_FromDouble(val); +} + +// Sets or returns charge of the macro-particle in e-charge +// the action is depended on the number of arguments +// charge() - returns charge +// charge(value) - sets the new value +// this is implementation of the getCharge() and setCharge methods of the Bunch class +static PyObject *Bunch_charge(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + + if (PyTuple_GET_SIZE(args) == 0) { + return PyFloat_FromDouble(cpp_bunch->getCharge()); } - //Sets or returns charge of the macro-particle in e-charge - // the action is depended on the number of arguments - // charge() - returns charge - // charge(value) - sets the new value - //this is implementation of the getCharge() and setCharge methods of the Bunch class - static PyObject* Bunch_charge(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 0 get charge - //if nVars == 1 set charge - int nVars = PyTuple_Size(args); + double value; + if (!PyArg_ParseTuple(args, "d:charge", &value)) { + return NULL; + } - double val = 0.; + cpp_bunch->setCharge(value); + return PyFloat_FromDouble(value); +} - if(nVars == 0 || nVars == 1){ - if(nVars == 0){ - val = cpp_bunch->getCharge(); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"d:charge",&val)){ - error("PyBunch - charge(value) - value is needed"); - } - cpp_bunch->setCharge(val); - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call charge() or charge(value)"); - } +// Sets or returns macroSize of the macro-particle +// the action is depended on the number of arguments +// macroSize() - returns macroSize +// macroSize(value) - sets the new value +// this is implementation of the getMacroSize() and setMacroSize methods of the Bunch class +static PyObject *Bunch_macroSize(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; - Py_INCREF(Py_None); - return Py_None; - } + if (PyTuple_GET_SIZE(args) == 0) { + return PyFloat_FromDouble(cpp_bunch->getMacroSize()); + } - //Sets or returns macroSize of the macro-particle - // the action is depended on the number of arguments - // macroSize() - returns macroSize - // macroSize(value) - sets the new value - //this is implementation of the getMacroSize() and setMacroSize methods of the Bunch class - static PyObject* Bunch_macroSize(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 0 get macroSize - //if nVars == 1 set macroSize - int nVars = PyTuple_Size(args); - - double val = 0.; - - if(nVars == 0 || nVars == 1){ - if(nVars == 0){ - val = cpp_bunch->getMacroSize(); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"d:macroSize",&val)){ - error("PyBunch - macroSize(value) - value is needed"); - } - cpp_bunch->setMacroSize(val); - } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call macroSize() or macroSize(value)"); - } + double value; + if (!PyArg_ParseTuple(args, "d:macroSize", &value)) { + return NULL; + } + + cpp_bunch->setMacroSize(value); + return PyFloat_FromDouble(value); +} - Py_INCREF(Py_None); - return Py_None; +//--------------------------------------------------------------- +// +// related to the bunch attributes +// +//---------------------------------------------------------------- + +// initilizes bunch attributes from the bunch file +static PyObject *Bunch_initBunchAttr(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *file_name = NULL; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:initBunchAttr", &file_name)) { + return NULL; } + cpp_bunch->initBunchAttributes(file_name); + Py_RETURN_NONE; +} - //--------------------------------------------------------------- - // - // related to the bunch attributes - // - //---------------------------------------------------------------- +// Sets or returns a double bunch attribute +// the action is depended on the number of arguments +// (attr_name) - returns double-value +// (index, value) - sets the new value to the attribute +// this is implementation of +// getBunchAttributeDouble(name) +// setBunchAttributeDouble(name,value) Bunch methods +static PyObject *Bunch_bunchAttrDouble(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name; + double value; + + if (!PyArg_ParseTuple(args, "s|d:bunchAttrDouble", &attr_name, &value)) { + return NULL; + } - //initilizes bunch attributes from the bunch file - static PyObject* Bunch_initBunchAttr(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* file_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:initBunchAttr",&file_name)){ - return NULL; - } - cpp_bunch->initBunchAttributes(file_name); - Py_INCREF(Py_None); - return Py_None; - } - - //Sets or returns a double bunch attribute - // the action is depended on the number of arguments - // (attr_name) - returns double-value - // (index, value) - sets the new value to the attribute - //this is implementation of - // getBunchAttributeDouble(name) - // setBunchAttributeDouble(name,value) Bunch methods - static PyObject* Bunch_bunchAttrDouble(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 this is get attribute - //if nVars == 2 this is set attribute - int nVars = PyTuple_Size(args); - - const char* attr_name = NULL; - double val = 0.; - - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:bunchAttrDouble",&attr_name)){ - error("PyBunch - bunchAttrDouble(name) - name are needed"); - } - std::string attr_name_str(attr_name); - val = cpp_bunch->getBunchAttributeDouble(attr_name_str); - return Py_BuildValue("d",val); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"sd:bunchAttrDouble",&attr_name,&val)){ - error("PyBunch - bunchAttrDouble(name,value) - name and double value are needed"); - } + std::string attr_name_str(attr_name); + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->getBunchAttributeDouble(attr_name_str); + } + else { + cpp_bunch->setBunchAttribute(attr_name_str, value); + } - std::string attr_name_str(attr_name); - cpp_bunch->setBunchAttribute( attr_name_str, val); - return Py_BuildValue("d",val); - } - } - else{ - error("PyBunch. You should call bunchAttrDouble(name) or bunchAttrDouble(name,value)"); - } + return PyFloat_FromDouble(value); +} - Py_INCREF(Py_None); - return Py_None; - } - - //Sets or returns a integer bunch attribute - // the action is depended on the number of arguments - // (attr_name) - returns int-value - // (index, value) - sets the new value to the attribute - //this is implementation of - // getBunchAttributeInt(name) - // setBunchAttributeInt(name,value) Bunch methods - static PyObject* Bunch_bunchAttrInt(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 this is get attribute - //if nVars == 2 this is set attribute - int nVars = PyTuple_Size(args); - - const char* attr_name = NULL; - int val = 0; - - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple!NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:bunchAttrInt",&attr_name)){ - error("PyBunch - bunchAttrInt(name) - pyBunch object and name are needed"); - } - std::string attr_name_str(attr_name); - val = cpp_bunch->getBunchAttributeInt(attr_name_str); - return Py_BuildValue("i",val); - } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"si:bunchAttrInt",&attr_name,&val)){ - error("PyBunch - bunchAttrInt(name,value) - name, and double value are needed"); - } +// Sets or returns a integer bunch attribute +// the action is depended on the number of arguments +// (attr_name) - returns int-value +// (index, value) - sets the new value to the attribute +// this is implementation of +// getBunchAttributeInt(name) +// setBunchAttributeInt(name,value) Bunch methods +static PyObject *Bunch_bunchAttrInt(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name; + int value; + + if (!PyArg_ParseTuple(args, "s|i:bunchAttrInt", &attr_name, &value)) { + return NULL; + } - std::string attr_name_str(attr_name); - cpp_bunch->setBunchAttribute( attr_name_str, val); - return Py_BuildValue("i",val); - } + std::string attr_name_str(attr_name); + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->getBunchAttributeInt(attr_name_str); + } + else { + cpp_bunch->setBunchAttribute(attr_name_str, value); + } + + return PyLong_FromLong(value); +} + +// Returns a list (tuple) of ther double bunch attribute names +static PyObject *Bunch_bunchAttrDoubleNames(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + std::vector names; + cpp_bunch->getDoubleBunchAttributeNames(names); + // create tuple with names + PyObject *resTuple = PyTuple_New(names.size()); + for (int i = 0, n = names.size(); i < n; i++) { + PyObject *py_nm = PyUnicode_FromString(names[i].c_str()); + if (PyTuple_SetItem(resTuple, i, py_nm)) { + error("PyBunch - bunchAttrDoubleNames - cannot create tuple with bunch attr names"); } - else{ - error("PyBunch. You should call bunchAttrInt(name) or bunchAttrInt(name,value)"); + } + return resTuple; +} + +// Returns a list (tuple) of ther integer bunch attribute names +static PyObject *Bunch_bunchAttrIntNames(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + std::vector names; + cpp_bunch->getIntBunchAttributeNames(names); + // create tuple with names + PyObject *resTuple = PyTuple_New(names.size()); + for (int i = 0, n = names.size(); i < n; i++) { + PyObject *py_nm = PyUnicode_FromString(names[i].c_str()); + if (PyTuple_SetItem(resTuple, i, py_nm)) { + error("PyBunch - bunchAttrIntNames() - cannot create tuple with bunch attr names"); } + } + return resTuple; +} - Py_INCREF(Py_None); - return Py_None; - } - - //Returns a list (tuple) of ther double bunch attribute names - static PyObject* Bunch_bunchAttrDoubleNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - std::vector names; - cpp_bunch->getDoubleBunchAttributeNames(names); - //create tuple with names - PyObject* resTuple = PyTuple_New(names.size()); - for(int i = 0, n = names.size(); i < n; i++){ - PyObject* py_nm = PyUnicode_FromString(names[i].c_str()); - if(PyTuple_SetItem(resTuple,i,py_nm)){ - error("PyBunch - bunchAttrDoubleNames - cannot create tuple with bunch attr names"); - } - } - return resTuple; - } - - //Returns a list (tuple) of ther integer bunch attribute names - static PyObject* Bunch_bunchAttrIntNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - std::vector names; - cpp_bunch->getIntBunchAttributeNames(names); - //create tuple with names - PyObject* resTuple = PyTuple_New(names.size()); - for(int i = 0, n = names.size(); i < n; i++){ - PyObject* py_nm = PyUnicode_FromString(names[i].c_str()); - if(PyTuple_SetItem(resTuple,i,py_nm)){ - error("PyBunch - bunchAttrIntNames() - cannot create tuple with bunch attr names"); - } - } - return resTuple; +// Returns 0 or 1. The result is 1 if the bunch has an attribute with a particular name +static PyObject *Bunch_hasBunchAttrDouble(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name = NULL; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:hasBunchAttrDouble", &attr_name)) { + return NULL; + } + std::string attr_name_str(attr_name); + int res = cpp_bunch->getBunchAttributes()->hasDoubleAttribute(attr_name_str); + return PyLong_FromLong(res); +} + +// Returns 0 or 1. The result is 1 if the bunch has an attribute with a particular name +static PyObject *Bunch_hasBunchAttrInt(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name = NULL; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:hasBunchAttrInt", &attr_name)) { + return NULL; } + std::string attr_name_str(attr_name); + int res = cpp_bunch->getBunchAttributes()->hasIntAttribute(attr_name_str); + return PyLong_FromLong(res); +} - //Returns 0 or 1. The result is 1 if the bunch has an attribute with a particular name - static PyObject* Bunch_hasBunchAttrDouble(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:hasBunchAttrDouble",&attr_name)){ - return NULL; +//--------------------------------------------------------------- +// +// related to particles' attributes +// +//---------------------------------------------------------------- + +// Adds a particles' attributes with a particular name to the bunch +static PyObject *Bunch_addPartAttr(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name = NULL; + PyObject *py_attrParamsDict = NULL; + // NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() + if (!PyArg_ParseTuple(args, "s|O:addPartAttr", &attr_name, &py_attrParamsDict)) { + error("PyBunch - addPartAttr(name, [param_dict]) - a particle attr. name are needed"); + } + std::string attr_name_str(attr_name); + std::map part_attr_dict; + if (py_attrParamsDict != NULL) { + if (!PyDict_Check(py_attrParamsDict)) { + error("PyBunch - addPartAttr(name, [param_dict]) - param_dict is not a dictionary"); + } + PyObject *key, *value; + Py_ssize_t pos = 0; + while (PyDict_Next(py_attrParamsDict, &pos, &key, &value)) { + if (!PyUnicode_Check(key) || !PyNumber_Check(value)) { + error( + "PyBunch - addPartAttr(name, [param_dict]) - param_dict is not a {str:val} dictionary" + ); + } + std::string par_name((char *)PyUnicode_AsUTF8(key)); + double d_val = PyFloat_AsDouble(value); + part_attr_dict[par_name] = d_val; } - std::string attr_name_str(attr_name); - int res = cpp_bunch->getBunchAttributes()->hasDoubleAttribute(attr_name_str); - return Py_BuildValue("i",res); } + cpp_bunch->addParticleAttributes(attr_name_str, part_attr_dict); + Py_RETURN_NONE; +} - //Returns 0 or 1. The result is 1 if the bunch has an attribute with a particular name - static PyObject* Bunch_hasBunchAttrInt(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:hasBunchAttrInt",&attr_name)){ - return NULL; - } - std::string attr_name_str(attr_name); - int res = cpp_bunch->getBunchAttributes()->hasIntAttribute(attr_name_str); - return Py_BuildValue("i",res); - } - - //--------------------------------------------------------------- - // - // related to particles' attributes - // - //---------------------------------------------------------------- - - //Adds a particles' attributes with a particular name to the bunch - static PyObject* Bunch_addPartAttr(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* attr_name = NULL; - PyObject* py_attrParamsDict = NULL; - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s|O:addPartAttr",&attr_name,&py_attrParamsDict)){ - error("PyBunch - addPartAttr(name, [param_dict]) - a particle attr. name are needed"); +// Removes a particles' attributes with a particular name from the bunch +static PyObject *Bunch_removePartAttr(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name = NULL; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:removePartAttr", &attr_name)) { + return NULL; + } + std::string attr_name_str(attr_name); + cpp_bunch->removeParticleAttributes(attr_name_str); + Py_RETURN_NONE; +} + +// Removes all particles' attributes from the bunch +static PyObject *Bunch_removeAllPartAttr(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + cpp_bunch->removeAllParticleAttributes(); + Py_RETURN_NONE; +} + +// Returns a list (tuple) of the particles' attributes names +static PyObject *Bunch_getPartAttrNames(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + std::vector names; + cpp_bunch->getParticleAttributesNames(names); + // create tuple with names + PyObject *resTuple = PyTuple_New(names.size()); + for (int i = 0, n = names.size(); i < n; i++) { + PyObject *py_nm = PyUnicode_FromString(names[i].c_str()); + if (PyTuple_SetItem(resTuple, i, py_nm)) { + error("PyBunch - getPartAttrNames() - cannot create tuple with bunch attr names"); } - std::string attr_name_str(attr_name); - std::map part_attr_dict; - if(py_attrParamsDict != NULL){ - if(!PyDict_Check(py_attrParamsDict)){ - error("PyBunch - addPartAttr(name, [param_dict]) - param_dict is not a dictionary"); - } - PyObject *key, *value; - Py_ssize_t pos = 0; - while (PyDict_Next(py_attrParamsDict, &pos, &key, &value)) { - if(!PyUnicode_Check(key) || !PyNumber_Check(value)){ - error("PyBunch - addPartAttr(name, [param_dict]) - param_dict is not a {str:val} dictionary"); - } - std::string par_name((char *)PyUnicode_AsUTF8(key)); - double d_val = PyFloat_AsDouble(value); - part_attr_dict[par_name] = d_val; - } - } - cpp_bunch->addParticleAttributes(attr_name_str,part_attr_dict); - Py_INCREF(Py_None); - return Py_None; } + return resTuple; +} - //Removes a particles' attributes with a particular name from the bunch - static PyObject* Bunch_removePartAttr(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:removePartAttr",&attr_name)){ +// Returns a dict{"part. attribute name":dict{"key":val}} +static PyObject *Bunch_getPartAttrDicts(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + std::vector names; + cpp_bunch->getParticleAttributesNames(names); + // make dict{"part. attribute name":dict{"key":val}} + PyObject *resDict = PyDict_New(); + if (resDict == NULL) { + return NULL; + } + for (int i = 0, n = names.size(); i < n; i++) { + PyObject *py_param_dict = PyDict_New(); + if (py_param_dict == NULL) { + Py_DECREF(resDict); return NULL; } - std::string attr_name_str(attr_name); - cpp_bunch->removeParticleAttributes(attr_name_str); - Py_INCREF(Py_None); - return Py_None; - } - - //Removes all particles' attributes from the bunch - static PyObject* Bunch_removeAllPartAttr(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - cpp_bunch->removeAllParticleAttributes(); - Py_INCREF(Py_None); - return Py_None; - } - - //Returns a list (tuple) of the particles' attributes names - static PyObject* Bunch_getPartAttrNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - std::vector names; - cpp_bunch->getParticleAttributesNames(names); - //create tuple with names - PyObject* resTuple = PyTuple_New(names.size()); - for(int i = 0, n = names.size(); i < n; i++){ - PyObject* py_nm = PyUnicode_FromString(names[i].c_str()); - if(PyTuple_SetItem(resTuple,i,py_nm)){ - error("PyBunch - getPartAttrNames() - cannot create tuple with bunch attr names"); - } - } - return resTuple; - } - - //Returns a dict{"part. attribute name":dict{"key":val}} - static PyObject* Bunch_getPartAttrDicts(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - std::vector names; - cpp_bunch->getParticleAttributesNames(names); - //make dict{"part. attribute name":dict{"key":val}} - PyObject* resDict = PyDict_New(); - for(int i = 0, n = names.size(); i < n; i++){ - PyObject* py_param_dict = PyDict_New(); - PyDict_SetItemString(resDict,names[i].c_str(),py_param_dict); - std::map param_dict = cpp_bunch->getParticleAttributes(names[i])->parameterDict; - std::map::iterator pos; - for (pos = param_dict.begin(); pos != param_dict.end(); ++pos) { - std::string key = pos->first; - double val = pos->second; - PyDict_SetItemString(py_param_dict,key.c_str(),Py_BuildValue("d",val)); - } - } - return resDict; - } - - //Returns a list (tuple) of the possible particles' attributes names - static PyObject* Bunch_getPossiblePartAttrNames(PyObject *self, PyObject *Py_UNUSED(ignored)){ - std::vector names; - ParticleAttributesFactory::getParticleAttributesNames(names); - //create tuple with names - PyObject* resTuple = PyTuple_New(names.size()); - for(int i = 0, n = names.size(); i < n; i++){ - PyObject* py_nm = PyUnicode_FromString(names[i].c_str()); - if(PyTuple_SetItem(resTuple,i,py_nm)){ - error("PyBunch - getPossiblePartAttrNames - cannot create tuple with bunch attr names"); - } - } - return resTuple; - } - - //temporary removes and memorizes all particles' attributes names - static PyObject* Bunch_clearAllPartAttrAndMemorize(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - cpp_bunch->clearAllParticleAttributesAndMemorize(); - Py_INCREF(Py_None); - return Py_None; - } - - //restores all particles' attributes names from memory - static PyObject* Bunch_restoreAllPartAttrFromMemory(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - cpp_bunch->restoreAllParticleAttributesFromMemory(); - Py_INCREF(Py_None); - return Py_None; - } - - //Returns 0 or 1. The result is 1 if the bunch has a particles' attributes with a particular name - static PyObject* Bunch_hasPartAttr(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:hasPartAttr",&attr_name)){ + if (PyDict_SetItemString(resDict, names[i].c_str(), py_param_dict) < 0) { + Py_DECREF(py_param_dict); + Py_DECREF(resDict); return NULL; } - std::string attr_name_str(attr_name); - int res = cpp_bunch->hasParticleAttributes(attr_name_str); - return Py_BuildValue("i",res); - } - - //Returns a list (tuple) of their bunch particles attribute names specified in the bunch file - static PyObject* Bunch_readPartAttrNames(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* file_name = NULL; - std::vector names; - std::map > part_attr_dicts; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:readPartAttrNames",&file_name)){ - return NULL; + std::map param_dict = + cpp_bunch->getParticleAttributes(names[i])->parameterDict; + std::map::iterator pos; + for (pos = param_dict.begin(); pos != param_dict.end(); ++pos) { + std::string key = pos->first; + double val = pos->second; + PyObject *py_val = PyFloat_FromDouble(val); + if (py_val == NULL || PyDict_SetItemString(py_param_dict, key.c_str(), py_val) < 0) { + Py_XDECREF(py_val); + Py_DECREF(py_param_dict); + Py_DECREF(resDict); + return NULL; + } + Py_DECREF(py_val); } - cpp_bunch->readParticleAttributesNames(file_name,names,part_attr_dicts); - //create tuple with names - PyObject* resTuple = PyTuple_New(names.size()); - for(int i = 0, n = names.size(); i < n; i++){ - PyObject* py_nm = PyUnicode_FromString(names[i].c_str()); - if(PyTuple_SetItem(resTuple,i,py_nm)){ - error("PyBunch - readPartAttrNames(fileName) - cannot create tuple with particles attr. names"); - } - } - return resTuple; + Py_DECREF(py_param_dict); } + return resDict; +} - //Returns a dictionary with the bunch particles attribute names as keys and - //dictionaries with parameter:value for each attribute - static PyObject* Bunch_readPartAttrDicts(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* file_name = NULL; - std::vector names; - std::map > part_attr_dicts; - if(!PyArg_Parse(arg,"s:readPartAttrDicts",&file_name)){ - return NULL; +// Returns a list (tuple) of the possible particles' attributes names +static PyObject *Bunch_getPossiblePartAttrNames(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + std::vector names; + ParticleAttributesFactory::getParticleAttributesNames(names); + // create tuple with names + PyObject *resTuple = PyTuple_New(names.size()); + for (int i = 0, n = names.size(); i < n; i++) { + PyObject *py_nm = PyUnicode_FromString(names[i].c_str()); + if (PyTuple_SetItem(resTuple, i, py_nm)) { + error("PyBunch - getPossiblePartAttrNames - cannot create tuple with bunch attr names"); } - cpp_bunch->readParticleAttributesNames(file_name,names,part_attr_dicts); - PyObject* resDict = PyDict_New(); - for(int i = 0, n = names.size(); i < n; i++){ - if(part_attr_dicts.count(names[i]) > 0){ - PyObject* py_param_dict = PyDict_New(); - PyDict_SetItemString(resDict,names[i].c_str(),py_param_dict); - std::map param_dict = part_attr_dicts[names[i]]; - std::map::iterator pos; - for (pos = param_dict.begin(); pos != param_dict.end(); ++pos) { - std::string key = pos->first; - double val = pos->second; - PyDict_SetItemString(py_param_dict,key.c_str(),Py_BuildValue("d",val)); - } - } - } - return resDict; } + return resTuple; +} - //initilizes particles' attributes from the bunch file - static PyObject* Bunch_readPartAttr(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* file_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:readPartAttr",&file_name)){ - return NULL; - } - cpp_bunch->readParticleAttributes(file_name); - Py_INCREF(Py_None); - return Py_None; +// temporary removes and memorizes all particles' attributes names +static PyObject *Bunch_clearAllPartAttrAndMemorize(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + cpp_bunch->clearAllParticleAttributesAndMemorize(); + Py_RETURN_NONE; +} + +// restores all particles' attributes names from memory +static PyObject *Bunch_restoreAllPartAttrFromMemory(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + cpp_bunch->restoreAllParticleAttributesFromMemory(); + Py_RETURN_NONE; +} + +// Returns 0 or 1. The result is 1 if the bunch has a particles' attributes with a particular name +static PyObject *Bunch_hasPartAttr(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name = NULL; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:hasPartAttr", &attr_name)) { + return NULL; } + std::string attr_name_str(attr_name); + int res = cpp_bunch->hasParticleAttributes(attr_name_str); + return PyLong_FromLong(res); +} - //Returns the number of variables in the particles' attributes with a particular name - static PyObject* Bunch_getPartAttrSize(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - const char* attr_name = NULL; - //NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() - if(!PyArg_Parse(arg,"s:getPartAttrSize",&attr_name)){ - return NULL; - } - std::string attr_name_str(attr_name); - int size = cpp_bunch->getParticleAttributes(attr_name_str)->getAttSize(); - return Py_BuildValue("i",size); - } - - //Sets or returns a particles' attributes' value - // the action is depended on the number of arguments - // (attr_name,part_index, attr_index) - returns double-value - // (attr_name, part_index, attr_index, value) - sets the new value to the attribute - //This is slow. In the C++ code you have to get reference to - //particles' attributes object and operate through it - static PyObject* Bunch_partAttrValue(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 3 this is get attribute - //if nVars == 4 this is set attribute - int nVars = PyTuple_Size(args); - - const char* attr_name = NULL; - int part_index = 0; - int attr_index = 0; - double val = 0.; - - if(nVars == 3 || nVars == 4){ - if(!PyArg_ParseTuple(args,"sii|d:partAttrValue",&attr_name,&part_index ,&attr_index,&val)){ - error("PyBunch - partAttrValue(attr_name,part_index,atr_index,[val]) - params. are needed"); - } - std::string attr_name_str(attr_name); - int bunch_size = cpp_bunch->getSize(); - int attr_size = cpp_bunch->getParticleAttributes(attr_name_str)->getAttSize(); - if(part_index >= bunch_size || attr_size <= attr_index){ - error("PyBunch - partAttrValue(attr_name,part_index,atr_index,[val]) - indexes out of limits! Stop!"); - } - if(nVars == 3){ - val = cpp_bunch->getParticleAttributes(attr_name_str)->attValue(part_index,attr_index); - return Py_BuildValue("d",val); - } - else{ - cpp_bunch->getParticleAttributes(attr_name_str)->attValue(part_index,attr_index) = val; - return Py_BuildValue("d",val); - } - } - else{ - error("PyBunch. You should call partAttrValue(attr_name,part_ind,attr_ind) or partAttrValue(attr_name,part_ind,attr_ind,value)"); +// Returns a list (tuple) of their bunch particles attribute names specified in the bunch file +static PyObject *Bunch_readPartAttrNames(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *file_name = NULL; + std::vector names; + std::map> part_attr_dicts; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:readPartAttrNames", &file_name)) { + return NULL; + } + cpp_bunch->readParticleAttributesNames(file_name, names, part_attr_dicts); + // create tuple with names + PyObject *resTuple = PyTuple_New(names.size()); + for (int i = 0, n = names.size(); i < n; i++) { + PyObject *py_nm = PyUnicode_FromString(names[i].c_str()); + if (PyTuple_SetItem(resTuple, i, py_nm)) { + error( + "PyBunch - readPartAttrNames(fileName) - cannot create tuple with particles attr. names" + ); } + } + return resTuple; +} - Py_INCREF(Py_None); - return Py_None; - } - - //--------------------------------------------------------------- - // - // getSize, getSizeGlobal, getSizeGlobalFromMemory, setTotalCount - // getCapacity - // - //---------------------------------------------------------------- - - //returns the number of macro-particles in the bunch - //this is implementation of the "getSize()" method - static PyObject* Bunch_getSize(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - return Py_BuildValue("i",cpp_bunch->getSize()); - } - - //returns the number of macro-particles in the bunch in all CPUs - //this is implementation of the "getSizeGlobal()" method - static PyObject* Bunch_getSizeGlobal(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - return Py_BuildValue("i",cpp_bunch->getSizeGlobal()); - } - - //returns the number of macro-particles in the bunch in all CPUs - // that was calculated in the previous call of getSizeGlobal() - //this is implementation of the "getSizeGlobalFromMemory()" method - static PyObject* Bunch_getSizeGlobalFromMemory(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - return Py_BuildValue("i",cpp_bunch->getSizeGlobalFromMemory()); - } - - //returns the number of all macro-particles - alive, dead, new - //this is implementation of the "getTotalCount()" method - static PyObject* Bunch_getTotalCount(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - return Py_BuildValue("i",cpp_bunch->getTotalCount()); - } - - //returns the capacity of the bunch-container. It could be changed. - //this is implementation of the "getCapacity()" method - static PyObject* Bunch_getCapacity(PyObject *self, PyObject *Py_UNUSED(ignored)){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - return Py_BuildValue("i",cpp_bunch->getCapacity()); - } - - //--------------------------------------------------------------- - // - // write into file or print Bunch - // - //---------------------------------------------------------------- - - //Prints bunch into the std::cout stream - static PyObject* Bunch_dumpBunch(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 0 dumpBunchs into std::cout - //if nVars == 1 dumpBunchs into the file - int nVars = PyTuple_Size(args); - const char* file_name = NULL; - if(nVars == 0 || nVars == 1){ - if(nVars == 0){ - cpp_bunch->print(std::cout); +// Returns a dictionary with the bunch particles attribute names as keys and +// dictionaries with parameter:value for each attribute +static PyObject *Bunch_readPartAttrDicts(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *file_name = NULL; + std::vector names; + std::map> part_attr_dicts; + if (!PyArg_Parse(arg, "s:readPartAttrDicts", &file_name)) { + return NULL; + } + cpp_bunch->readParticleAttributesNames(file_name, names, part_attr_dicts); + PyObject *resDict = PyDict_New(); + if (resDict == NULL) { + return NULL; + } + for (int i = 0, n = names.size(); i < n; i++) { + if (part_attr_dicts.count(names[i]) > 0) { + PyObject *py_param_dict = PyDict_New(); + if (py_param_dict == NULL) { + Py_DECREF(resDict); + return NULL; } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:dumpBunch",&file_name)){ - error("PyBunch - dumpBunch(fileName) - a new value are needed"); + if (PyDict_SetItemString(resDict, names[i].c_str(), py_param_dict) < 0) { + Py_DECREF(py_param_dict); + Py_DECREF(resDict); + return NULL; + } + std::map param_dict = part_attr_dicts[names[i]]; + std::map::iterator pos; + for (pos = param_dict.begin(); pos != param_dict.end(); ++pos) { + std::string key = pos->first; + double val = pos->second; + PyObject *py_val = PyFloat_FromDouble(val); + if (py_val == NULL || PyDict_SetItemString(py_param_dict, key.c_str(), py_val) < 0) { + Py_XDECREF(py_val); + Py_DECREF(py_param_dict); + Py_DECREF(resDict); + return NULL; } - cpp_bunch->print(file_name); + Py_DECREF(py_val); } + Py_DECREF(py_param_dict); } - else{ - error("PyBunch. You should call dumpBunch() or dumpBunch(file_name)"); - } + } + return resDict; +} - Py_INCREF(Py_None); - return Py_None; - } - - //Reads bunch info from the file - static PyObject* Bunch_readBunch(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - //if nVars == 1 reads all macro-particles - //if nVars == 2 reads only specified number of macro-particles - int nVars = PyTuple_Size(args); - const char* file_name = NULL; - int nParts = 0; - if(nVars == 1 || nVars == 2){ - if(nVars == 1){ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"s:read",&file_name)){ - error("PyBunch - readBunch(fileName) - a file name are needed"); - } - cpp_bunch->initBunchAttributes(file_name); - cpp_bunch->readParticleAttributes(file_name); - cpp_bunch->readBunchCoords(file_name); +// initilizes particles' attributes from the bunch file +static PyObject *Bunch_readPartAttr(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *file_name = NULL; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:readPartAttr", &file_name)) { + return NULL; + } + cpp_bunch->readParticleAttributes(file_name); + Py_RETURN_NONE; +} + +// Returns the number of variables in the particles' attributes with a particular name +static PyObject *Bunch_getPartAttrSize(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name = NULL; + // NO NEW OBJECT CREATED BY PyArg_Parse! NO NEED OF Py_DECREF() + if (!PyArg_Parse(arg, "s:getPartAttrSize", &attr_name)) { + return NULL; + } + std::string attr_name_str(attr_name); + int size = cpp_bunch->getParticleAttributes(attr_name_str)->getAttSize(); + return PyLong_FromLong(size); +} + +// Sets or returns a particles' attributes' value +// the action is depended on the number of arguments +// (attr_name,part_index, attr_index) - returns double-value +// (attr_name, part_index, attr_index, value) - sets the new value to the attribute +// This is slow. In the C++ code you have to get reference to +// particles' attributes object and operate through it +static PyObject *Bunch_partAttrValue(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + const char *attr_name; + int part_index; + int attr_index; + double value; + + if ( + !PyArg_ParseTuple(args, "sii|d:partAttrValue", &attr_name, &part_index, &attr_index, &value) + ) { + return NULL; + } + + std::string attr_name_str(attr_name); + int bunch_size = cpp_bunch->getSize(); + int attr_size = cpp_bunch->getParticleAttributes(attr_name_str)->getAttSize(); + if (part_index >= bunch_size || attr_size <= attr_index) { + error( + "PyBunch - partAttrValue(attr_name,part_index,atr_index,[val]) - indexes out of limits! " + "Stop!" + ); + } + + if (PyTuple_GET_SIZE(args) == 3) { + value = cpp_bunch->getParticleAttributes(attr_name_str)->attValue(part_index, attr_index); + } + else { + cpp_bunch->getParticleAttributes(attr_name_str)->attValue(part_index, attr_index) = value; + } + + return PyFloat_FromDouble(value); +} + +//--------------------------------------------------------------- +// +// getSize, getSizeGlobal, getSizeGlobalFromMemory, setTotalCount +// getCapacity +// +//---------------------------------------------------------------- + +// returns the number of macro-particles in the bunch +// this is implementation of the "getSize()" method +static PyObject *Bunch_getSize(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + return PyLong_FromLong(cpp_bunch->getSize()); +} + +// returns the number of macro-particles in the bunch in all CPUs +// this is implementation of the "getSizeGlobal()" method +static PyObject *Bunch_getSizeGlobal(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + return PyLong_FromLong(cpp_bunch->getSizeGlobal()); +} + +// returns the number of macro-particles in the bunch in all CPUs +// that was calculated in the previous call of getSizeGlobal() +// this is implementation of the "getSizeGlobalFromMemory()" method +static PyObject *Bunch_getSizeGlobalFromMemory(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + return PyLong_FromLong(cpp_bunch->getSizeGlobalFromMemory()); +} + +// returns the number of all macro-particles - alive, dead, new +// this is implementation of the "getTotalCount()" method +static PyObject *Bunch_getTotalCount(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + return PyLong_FromLong(cpp_bunch->getTotalCount()); +} + +// returns the capacity of the bunch-container. It could be changed. +// this is implementation of the "getCapacity()" method +static PyObject *Bunch_getCapacity(PyObject *self, PyObject *Py_UNUSED(ignored)) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + return PyLong_FromLong(cpp_bunch->getCapacity()); +} + +//--------------------------------------------------------------- +// +// write into file or print Bunch +// +//---------------------------------------------------------------- + +// Prints bunch into the std::cout stream +static PyObject *Bunch_dumpBunch(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + // if nVars == 0 dumpBunchs into std::cout + // if nVars == 1 dumpBunchs into the file + int nVars = PyTuple_Size(args); + const char *file_name = NULL; + if (nVars == 0 || nVars == 1) { + if (nVars == 0) { + cpp_bunch->print(std::cout); + } + else { + // NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() + if (!PyArg_ParseTuple(args, "s:dumpBunch", &file_name)) { + error("PyBunch - dumpBunch(fileName) - a new value are needed"); } - else{ - //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() - if(!PyArg_ParseTuple(args,"si:read",&file_name,&nParts)){ - error("PyBunch - readBunch(fileName,nParts) - file name, and number of particles are needed"); - } - cpp_bunch->initBunchAttributes(file_name); - cpp_bunch->readParticleAttributes(file_name); - cpp_bunch->readBunchCoords(file_name,nParts); + cpp_bunch->print(file_name); + } + } + else { + error("PyBunch. You should call dumpBunch() or dumpBunch(file_name)"); + } + + Py_RETURN_NONE; +} + +// Reads bunch info from the file +static PyObject *Bunch_readBunch(PyObject *self, PyObject *args) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + // if nVars == 1 reads all macro-particles + // if nVars == 2 reads only specified number of macro-particles + int nVars = PyTuple_Size(args); + const char *file_name = NULL; + int nParts = 0; + if (nVars == 1 || nVars == 2) { + if (nVars == 1) { + // NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() + if (!PyArg_ParseTuple(args, "s:read", &file_name)) { + error("PyBunch - readBunch(fileName) - a file name are needed"); } + cpp_bunch->initBunchAttributes(file_name); + cpp_bunch->readParticleAttributes(file_name); + cpp_bunch->readBunchCoords(file_name); } - else{ - error("PyBunch. You should call readBunch(file_name) or readBunch(file_name,nParts)"); + else { + // NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() + if (!PyArg_ParseTuple(args, "si:read", &file_name, &nParts)) { + error( + "PyBunch - readBunch(fileName,nParts) - file name, and number of particles are needed" + ); + } + cpp_bunch->initBunchAttributes(file_name); + cpp_bunch->readParticleAttributes(file_name); + cpp_bunch->readBunchCoords(file_name, nParts); } - Py_INCREF(Py_None); - return Py_None; - } - - //Copy bunch attrubutes and structure to another bunch - static PyObject* Bunch_copyEmptyBunchTo(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyBunch_Target = arg; - Bunch* cpp_target_bunch = (Bunch*) ((pyORBIT_Object *) pyBunch_Target)->cpp_obj; - cpp_bunch->copyEmptyBunchTo(cpp_target_bunch); - Py_INCREF(Py_None); - return Py_None; - } - - //Copy bunch all info including particles coordinates and attributes to another bunch - static PyObject* Bunch_copyBunchTo(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyBunch_Target = arg; - Bunch* cpp_target_bunch = (Bunch*) ((pyORBIT_Object *) pyBunch_Target)->cpp_obj; - cpp_bunch->copyBunchTo(cpp_target_bunch); - Py_INCREF(Py_None); - return Py_None; - } - - //Copy particles coordinates from one bunch to another - static PyObject* Bunch_addParticlesTo(PyObject *self, PyObject *arg){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyBunch_Target = arg; - Bunch* cpp_target_bunch =(Bunch*) ((pyORBIT_Object *) pyBunch_Target)->cpp_obj ; - cpp_bunch->addParticlesTo(cpp_target_bunch); - Py_INCREF(Py_None); - return Py_None; - } - - //----------------------------------------------------- - //destructor for python Bunch class - //----------------------------------------------------- - //this is implementation of the __del__ method - static void Bunch_del(pyORBIT_Object* self){ - Bunch* cpp_bunch = (Bunch*) self->cpp_obj; - delete cpp_bunch; - self->ob_base.ob_type->tp_free((PyObject*)self); } + else { + error("PyBunch. You should call readBunch(file_name) or readBunch(file_name,nParts)"); + } + Py_RETURN_NONE; +} + +// Copy bunch attrubutes and structure to another bunch +static PyObject *Bunch_copyEmptyBunchTo(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + PyObject *pyBunch_Target = arg; + Bunch *cpp_target_bunch = (Bunch *)((pyORBIT_Object *)pyBunch_Target)->cpp_obj; + cpp_bunch->copyEmptyBunchTo(cpp_target_bunch); + Py_RETURN_NONE; +} + +// Copy bunch all info including particles coordinates and attributes to another bunch +static PyObject *Bunch_copyBunchTo(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + PyObject *pyBunch_Target = arg; + Bunch *cpp_target_bunch = (Bunch *)((pyORBIT_Object *)pyBunch_Target)->cpp_obj; + cpp_bunch->copyBunchTo(cpp_target_bunch); + Py_RETURN_NONE; +} + +// Copy particles coordinates from one bunch to another +static PyObject *Bunch_addParticlesTo(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + PyObject *pyBunch_Target = arg; + Bunch *cpp_target_bunch = (Bunch *)((pyORBIT_Object *)pyBunch_Target)->cpp_obj; + cpp_bunch->addParticlesTo(cpp_target_bunch); + Py_RETURN_NONE; +} + +//----------------------------------------------------- +// destructor for python Bunch class +//----------------------------------------------------- +// this is implementation of the __del__ method +static void Bunch_del(pyORBIT_Object *self) +{ + Bunch *cpp_bunch = (Bunch *)self->cpp_obj; + delete cpp_bunch; + self->ob_base.ob_type->tp_free((PyObject *)self); +} #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY -static PyArrayObject *parse_bunch_array(PyObject *input) { - PyArrayObject *array = (PyArrayObject *)PyArray_FROM_OTF( - input, - NPY_FLOAT64, - NPY_ARRAY_IN_ARRAY - ); +static PyArrayObject *parse_bunch_array(PyObject *input) +{ + PyArrayObject *array = (PyArrayObject *)PyArray_FROM_OTF(input, NPY_FLOAT64, NPY_ARRAY_IN_ARRAY); - if (array == NULL) { - return NULL; - } + if (array == NULL) { + return NULL; + } - if (PyArray_NDIM(array) != 2) { - PyErr_SetString(PyExc_ValueError, - "array must be 2-dimensional with shape (nparts, 6)" - ); - Py_DECREF(array); - return NULL; - } + if (PyArray_NDIM(array) != 2) { + PyErr_SetString(PyExc_ValueError, "array must be 2-dimensional with shape (nparts, 6)"); + Py_DECREF(array); + return NULL; + } - if (PyArray_DIM(array, 1) != 6) { - PyErr_SetString(PyExc_ValueError, - "expected shape (nparts, 6) for (x, xp, y, yp, z, dE)" - ); - Py_DECREF(array); - return NULL; - } + if (PyArray_DIM(array, 1) != 6) { + PyErr_SetString(PyExc_ValueError, "expected shape (nparts, 6) for (x, xp, y, yp, z, dE)"); + Py_DECREF(array); + return NULL; + } - return array; // new ref; caller MUST Py_DECREF(). + return array; // new ref; caller MUST Py_DECREF(). } -static void append_bunch_with_PyArray(Bunch *bunch, PyArrayObject *array) { - const npy_intp nparts = PyArray_DIM(array, 0); - const double *data = (const double *)PyArray_DATA(array); +static void append_bunch_with_PyArray(Bunch *cpp_bunch, PyArrayObject *array) +{ + const npy_intp nparts = PyArray_DIM(array, 0); + const double *data = (const double *)PyArray_DATA(array); - for (npy_intp i = 0; i < nparts; ++i) { - const double *coords = data + i*6; + for (npy_intp i = 0; i < nparts; ++i) { + const double *coords = data + i * 6; - bunch->addParticle(coords[0], coords[1], coords[2], coords[3], coords[4], coords[5]); - } + cpp_bunch->addParticle(coords[0], coords[1], coords[2], coords[3], coords[4], coords[5]); + } } PyDoc_STRVAR( - Bunch_to_numpy_doc, - "to_numpy($self, /, gather=False, root=0)\n" - "--\n" - "\n" - "Return the particle coordinates as a NumPy array.\n" - "\n" - "Parameters\n" - "----------\n" - "gather : bool, optional\n" - " Collect particles from all MPI ranks. The default is False.\n" - "root : int, optional\n" - " Rank that receives the collected array. The default is 0.\n" - "\n" - "Returns\n" - "-------\n" - "numpy.ndarray or None\n" - " Array with shape ``(n_particles, 6)`` and dtype ``float64``.\n" - " When gathering, non-root ranks return None.\n" - "\n" - "Raises\n" - "------\n" - "ValueError\n" - " If root is not a valid MPI rank.\n" - "OverflowError\n" - " If the bunch is too large for the MPI gather operation.\n" + Bunch_to_numpy_doc, + "to_numpy($self, /, gather=False, root=0)\n" + "--\n" + "\n" + "Return the particle coordinates as a NumPy array.\n" + "\n" + "Parameters\n" + "----------\n" + "gather : bool, optional\n" + " Collect particles from all MPI ranks. The default is False.\n" + "root : int, optional\n" + " Rank that receives the collected array. The default is 0.\n" + "\n" + "Returns\n" + "-------\n" + "numpy.ndarray or None\n" + " Array with shape ``(n_particles, 6)`` and dtype ``float64``.\n" + " When gathering, non-root ranks return None.\n" + "\n" + "Raises\n" + "------\n" + "ValueError\n" + " If root is not a valid MPI rank.\n" + "OverflowError\n" + " If the bunch is too large for the MPI gather operation.\n" ); -static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args, PyObject *kwargs) { - static const char* kwlist[] = {"gather", "root", NULL}; +static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args, PyObject *kwargs) +{ + static const char *kwlist[] = {"gather", "root", NULL}; int gather = 0; int root = 0; if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|pi:to_numpy", kwlist, &gather, &root)) { - return NULL; + return NULL; } - Bunch *bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; - const int rank = bunch->getMPI_Rank(); - const int size = bunch->getMPI_Size(); - const npy_intp nparts = (npy_intp)bunch->getSize(); + const int rank = cpp_bunch->getMPI_Rank(); + const int size = cpp_bunch->getMPI_Size(); + const npy_intp nparts = (npy_intp)cpp_bunch->getSize(); const npy_intp ncoords = 6; if (gather && (root < 0 || root >= size)) { - PyErr_Format(PyExc_ValueError, "root must be between 0 and %d, got %d", size - 1, rank); - return NULL; + PyErr_Format(PyExc_ValueError, "root must be between 0 and %d, got %d", size - 1, rank); + return NULL; } - npy_intp dims[2] = { nparts, ncoords }; + npy_intp dims[2] = {nparts, ncoords}; PyObject *local_array = PyArray_SimpleNew(2, dims, NPY_FLOAT64); @@ -1253,43 +1182,58 @@ static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args, PyObject *kwargs // or if we try to allocate a buffer to hold the bunch data and it fails on one rank // then propagate that result to the other ranks and raise. if (gather && size > 1) { - int local_counts_ok = nparts <= INT_MAX / ncoords; - int all_counts_ok = 0; - - ORBIT_MPI_Allreduce(&local_counts_ok, &all_counts_ok, 1, MPI_INT, MPI_MIN, bunch->getMPI_Comm_Local()->comm); - - if(!all_counts_ok) { - PyErr_SetString(PyExc_OverflowError, "local bunch is too large for MPI_Gatherv"); - return NULL; - } + int local_counts_ok = nparts <= INT_MAX / ncoords; + int all_counts_ok = 0; + + ORBIT_MPI_Allreduce( + &local_counts_ok, + &all_counts_ok, + 1, + MPI_INT, + MPI_MIN, + cpp_bunch->getMPI_Comm_Local()->comm + ); - int local_alloc_ok = local_array != NULL; - int all_alloc_ok = 0; + if (!all_counts_ok) { + PyErr_SetString(PyExc_OverflowError, "local bunch is too large for MPI_Gatherv"); + return NULL; + } - ORBIT_MPI_Allreduce(&local_alloc_ok, &all_alloc_ok, 1, MPI_INT, MPI_MIN, bunch->getMPI_Comm_Local()->comm); + int local_alloc_ok = local_array != NULL; + int all_alloc_ok = 0; - if(!all_alloc_ok) { - Py_XDECREF(local_array); + ORBIT_MPI_Allreduce( + &local_alloc_ok, + &all_alloc_ok, + 1, + MPI_INT, + MPI_MIN, + cpp_bunch->getMPI_Comm_Local()->comm + ); - if(!PyErr_Occurred()) { - PyErr_NoMemory(); - } + if (!all_alloc_ok) { + Py_XDECREF(local_array); - return NULL; - } - } else { - if (local_array == NULL) { - return NULL; + if (!PyErr_Occurred()) { + PyErr_NoMemory(); } + + return NULL; + } + } + else { + if (local_array == NULL) { + return NULL; + } } - PyArrayObject *arr_obj = (PyArrayObject*)local_array; - double *local_data = (double*)PyArray_DATA(arr_obj); - double **src = bunch->coordArr(); + PyArrayObject *arr_obj = (PyArrayObject *)local_array; + double *local_data = (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) { - local_data[j + i*ncoords] = src[i][j]; + local_data[j + i * ncoords] = src[i][j]; } } @@ -1298,7 +1242,7 @@ static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args, PyObject *kwargs } #if USE_MPI > 0 - MPI_Comm comm = bunch->getMPI_Comm_Local()->comm; + MPI_Comm comm = cpp_bunch->getMPI_Comm_Local()->comm; std::vector global_nparts(size); std::vector displacements(size); @@ -1306,69 +1250,87 @@ static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args, PyObject *kwargs const int local_nparts = (int)nparts; - if (MPI_SUCCESS != MPI_Gather(&local_nparts, 1, MPI_INT, global_nparts.data(), 1, MPI_INT, root, comm)) { - Py_DECREF(local_array); - PyErr_SetString(PyExc_RuntimeError, "MPI_Gather failed to collect the bunch sizes across ranks"); - return NULL; + if ( + MPI_Gather(&local_nparts, 1, MPI_INT, global_nparts.data(), 1, MPI_INT, root, comm) != + MPI_SUCCESS + ) { + Py_DECREF(local_array); + PyErr_SetString( + PyExc_RuntimeError, + "MPI_Gather failed to collect the bunch sizes across ranks" + ); + return NULL; } int total = 0; int layout_ok = 1; if (rank == root) { - for (int mpi_rank = 0; mpi_rank < size; ++mpi_rank) { - const int rank_nvalues = global_nparts[mpi_rank] * ncoords; - - if (total + rank_nvalues > INT_MAX) { - layout_ok = 0; - break; - } + for (int mpi_rank = 0; mpi_rank < size; ++mpi_rank) { + const int rank_nvalues = global_nparts[mpi_rank] * ncoords; - displacements[mpi_rank] = total; - recv_counts[mpi_rank] = rank_nvalues; - total += rank_nvalues; + if (total + rank_nvalues > INT_MAX) { + layout_ok = 0; + break; } + + displacements[mpi_rank] = total; + recv_counts[mpi_rank] = rank_nvalues; + total += rank_nvalues; + } } ORBIT_MPI_Bcast(&layout_ok, 1, MPI_INT, root, comm); - if(!layout_ok) { - Py_DECREF(local_array); - PyErr_SetString(PyExc_OverflowError, "global bunch is too large to collect"); - return NULL; + if (!layout_ok) { + Py_DECREF(local_array); + PyErr_SetString(PyExc_OverflowError, "global bunch is too large to collect"); + return NULL; } PyObject *global_array = NULL; if (rank == root) { - npy_intp global_dims[2] = {(npy_intp)(total / ncoords), ncoords}; - global_array = PyArray_SimpleNew(2, global_dims, NPY_FLOAT64); + npy_intp global_dims[2] = {(npy_intp)(total / ncoords), ncoords}; + global_array = PyArray_SimpleNew(2, global_dims, NPY_FLOAT64); } int global_alloc_ok = rank != root || global_array != NULL; ORBIT_MPI_Bcast(&global_alloc_ok, 1, MPI_INT, root, comm); - if(!global_alloc_ok) { - Py_DECREF(local_array); + if (!global_alloc_ok) { + Py_DECREF(local_array); - if(!PyErr_Occurred()) { - PyErr_NoMemory(); - } + if (!PyErr_Occurred()) { + PyErr_NoMemory(); + } - return NULL; + return NULL; } - double* global_data = rank == root ? (double *)PyArray_DATA((PyArrayObject*)global_array) : NULL; - - if(MPI_SUCCESS != MPI_Gatherv(local_data, nparts*ncoords, MPI_DOUBLE, global_data, recv_counts.data(), displacements.data(), MPI_DOUBLE, root, comm)) { - Py_XDECREF(global_array); - PyErr_SetString(PyExc_RuntimeError, "MPI_Gatherv failed to collect bunch."); - return NULL; + double *global_data = rank == root ? (double *)PyArray_DATA((PyArrayObject *)global_array) : NULL; + + if ( + MPI_Gatherv( + local_data, + nparts * ncoords, + MPI_DOUBLE, + global_data, + recv_counts.data(), + displacements.data(), + MPI_DOUBLE, + root, + comm + ) != MPI_SUCCESS + ) { + Py_XDECREF(global_array); + PyErr_SetString(PyExc_RuntimeError, "MPI_Gatherv failed to collect bunch."); + return NULL; } if (rank == root) { - return global_array; + return global_array; } #endif // USE_MPI > 0 @@ -1376,69 +1338,71 @@ static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args, PyObject *kwargs } PyDoc_STRVAR( - Bunch_update_from_numpy_doc, - "update_from_numpy($self, array, /)\n" - "--\n" - "\n" - "Replace the local particle coordinates from an array.\n" - "\n" - "Parameters\n" - "----------\n" - "array : array_like\n" - " Particle coordinates with shape ``(n_particles, 6)`` ordered as\n" - " ``(x, xp, y, yp, z, dE)``. Values are converted to ``float64``.\n" - "\n" - "Returns\n" - "-------\n" - "None\n" - "\n" - "Raises\n" - "------\n" - "ValueError\n" - " If the array does not have shape ``(n_particles, 6)``.\n" + Bunch_update_from_numpy_doc, + "update_from_numpy($self, array, /)\n" + "--\n" + "\n" + "Replace the local particle coordinates from an array.\n" + "\n" + "Parameters\n" + "----------\n" + "array : array_like\n" + " Particle coordinates with shape ``(n_particles, 6)`` ordered as\n" + " ``(x, xp, y, yp, z, dE)``. Values are converted to ``float64``.\n" + "\n" + "Returns\n" + "-------\n" + "None\n" + "\n" + "Raises\n" + "------\n" + "ValueError\n" + " If the array does not have shape ``(n_particles, 6)``.\n" ); -static PyObject *Bunch_update_from_numpy(PyObject *self, PyObject *arg) { - Bunch *bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; +static PyObject *Bunch_update_from_numpy(PyObject *self, PyObject *arg) +{ + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; PyArrayObject *array = parse_bunch_array(arg); if (array == NULL) { - return NULL; + return NULL; } - bunch->deleteAllParticles(); - append_bunch_with_PyArray(bunch, array); + cpp_bunch->deleteAllParticles(); + append_bunch_with_PyArray(cpp_bunch, array); Py_DECREF(array); Py_RETURN_NONE; } PyDoc_STRVAR( - Bunch_from_numpy_doc, - "from_numpy($type, array, /)\n" - "--\n" - "\n" - "Construct a bunch from particle coordinates.\n" - "\n" - "Parameters\n" - "----------\n" - "array : array_like\n" - " Particle coordinates with shape ``(n_particles, 6)`` ordered as\n" - " ``(x, xp, y, yp, z, dE)``. Values are converted to ``float64``.\n" - "\n" - "Returns\n" - "-------\n" - "Bunch\n" - " A new bunch containing the supplied particles.\n" - "\n" - "Raises\n" - "------\n" - "ValueError\n" - " If the array does not have shape ``(n_particles, 6)``.\n" + Bunch_from_numpy_doc, + "from_numpy($type, array, /)\n" + "--\n" + "\n" + "Construct a bunch from particle coordinates.\n" + "\n" + "Parameters\n" + "----------\n" + "array : array_like\n" + " Particle coordinates with shape ``(n_particles, 6)`` ordered as\n" + " ``(x, xp, y, yp, z, dE)``. Values are converted to ``float64``.\n" + "\n" + "Returns\n" + "-------\n" + "Bunch\n" + " A new bunch containing the supplied particles.\n" + "\n" + "Raises\n" + "------\n" + "ValueError\n" + " If the array does not have shape ``(n_particles, 6)``.\n" ); -static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *arg) { +static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *arg) +{ PyArrayObject *array = parse_bunch_array(arg); if (array == NULL) { @@ -1448,188 +1412,277 @@ static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *arg) { PyObject *py_bunch_obj = PyObject_CallNoArgs(cls); if (py_bunch_obj == NULL) { - Py_DECREF(array); - return NULL; + Py_DECREF(array); + return NULL; } - Bunch *bunch = (Bunch*)((pyORBIT_Object*)py_bunch_obj)->cpp_obj; + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)py_bunch_obj)->cpp_obj; - append_bunch_with_PyArray(bunch, array); + append_bunch_with_PyArray(cpp_bunch, array); Py_DECREF(array); return py_bunch_obj; } #endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY - static PyMethodDef BunchClassMethods[] = { - //-------------------------------------------------------- - // class Bunch wrapper START - //-------------------------------------------------------- - { "getMPIComm", Bunch_getMPIComm ,METH_NOARGS,"Returns MPI Comm of this bunch"}, - { "setMPIComm", Bunch_setMPIComm ,METH_O,"Sets a new MPI Comm for this bunch"}, - { "getSyncParticle", Bunch_getSyncParticle ,METH_NOARGS,"Returns syncParticle class instance"}, - { "addParticle", Bunch_addParticle ,METH_VARARGS,"Adds a macro-particle to the bunch"}, - { "deleteParticle", Bunch_deleteParticle ,METH_O,"Removes macro-particle from the bunch and call compress inside"}, - { "deleteParticleFast", Bunch_deleteParticleFast ,METH_O,"Removes macro-particle from the bunch very fast"}, - { "recoverParticle", Bunch_recoverParticle ,METH_O,"Recovers a particle marked for removal"}, - { "deleteAllParticles", Bunch_deleteAllParticles ,METH_NOARGS,"Removes all macro-particles from the bunch"}, - { "compress", Bunch_compress ,METH_NOARGS,"Compress the bunch"}, - { "x", Bunch_x ,METH_VARARGS,"Set x(index,value) or get x(index) coordinate"}, - { "y", Bunch_y ,METH_VARARGS,"Set y(index,value) or get y(index) coordinate"}, - { "z", Bunch_z ,METH_VARARGS,"Set z(index,value) or get z(index) coordinate"}, - { "px", Bunch_px ,METH_VARARGS,"Set px(index,value) or get px(index) coordinate"}, - { "py", Bunch_py ,METH_VARARGS,"Set py(index,value) or get py(index) coordinate"}, - { "pz", Bunch_pz ,METH_VARARGS,"Set pz(index,value) or get pz(index) coordinate"}, - { "dE", Bunch_pz ,METH_VARARGS,"Set dE(index,value) or get dE(index) coordinate"}, - { "xp", Bunch_px ,METH_VARARGS,"Set xp(index,value) or get xp(index) coordinate"}, - { "yp", Bunch_py ,METH_VARARGS,"Set yp(index,value) or get yp(index) coordinate"}, - { "flag", Bunch_flag ,METH_O,"Returns flag(index) for particle with index"}, - { "ringwrap", Bunch_ringwrap ,METH_O,"Perform the ring wrap. Usage: ringwrap(ring_length)"}, - { "mass", Bunch_mass ,METH_VARARGS,"Set mass(value) or get mass() the mass of particle in MeV"}, - { "classicalRadius", Bunch_classicalRadius ,METH_NOARGS,"Returns a classical radius of particle in [m]"}, - { "B_Rho", Bunch_B_Rho ,METH_NOARGS,"Returns B*Rho parameter of particle in [T*m]"}, - { "charge", Bunch_charge ,METH_VARARGS,"Set charge(value) or get charge() the charge of particle in e-charge"}, - { "macroSize", Bunch_macroSize ,METH_VARARGS,"Set macroSize(value) or get macroSize() the charge of particle in e-charge"}, - { "initBunchAttr", Bunch_initBunchAttr ,METH_O,"Reads and initilizes bunch attributes from a bunch file"}, - { "bunchAttrDouble", Bunch_bunchAttrDouble ,METH_VARARGS,"Returns and sets a double bunch attribute"}, - { "bunchAttrInt", Bunch_bunchAttrInt ,METH_VARARGS,"Returns and sets an integer bunch attribute"}, - { "bunchAttrDoubleNames", Bunch_bunchAttrDoubleNames ,METH_NOARGS,"Returns a list of double bunch attribute names"}, - { "bunchAttrIntNames", Bunch_bunchAttrIntNames ,METH_NOARGS,"Returns a list of integer bunch attribute names"}, - { "hasBunchAttrDouble", Bunch_hasBunchAttrDouble ,METH_O,"Returns 1 if there is a double bunch attr. with this name, 0 - otherwise"}, - { "hasBunchAttrInt", Bunch_hasBunchAttrInt ,METH_O,"Returns 1 if there is a int bunch attr. with this name, 0 - otherwise"}, - { "addPartAttr", Bunch_addPartAttr ,METH_VARARGS,"Adds a particles' attributes to the bunch"}, - { "removePartAttr", Bunch_removePartAttr ,METH_O,"Removes a particles' attributes from the bunch"}, - { "removeAllPartAttr", Bunch_removeAllPartAttr ,METH_NOARGS,"Removes all particles' attributes from the bunch"}, - { "getPartAttrNames", Bunch_getPartAttrNames ,METH_NOARGS,"Returns all particles' attributes names in the bunch at this moment"}, - { "getPartAttrDicts", Bunch_getPartAttrDicts ,METH_NOARGS,"Returns dict{part. attribute name:dict{key:val}}"}, - { "getPossiblePartAttrNames", Bunch_getPossiblePartAttrNames ,METH_NOARGS,"Returns all possible particles' attributes names"}, - { "clearAllPartAttrAndMemorize", Bunch_clearAllPartAttrAndMemorize ,METH_NOARGS,"Temporary removes and memorizes all particles' attributes names"}, - { "restoreAllPartAttrFromMemory", Bunch_restoreAllPartAttrFromMemory ,METH_NOARGS,"Restores all particles' attributes names from memory"}, - { "hasPartAttr", Bunch_hasPartAttr ,METH_O,"Returns 1 if there is a particles' attr. with this name, 0 - otherwis"}, - { "readPartAttrNames", Bunch_readPartAttrNames ,METH_O,"Returns a tuple with particles' attr. names in the bunch file"}, - { "readPartAttrDicts", Bunch_readPartAttrDicts ,METH_O,"Returns a dict{attr_name:dicts{param_name:val}} in the bunch file"}, - { "readPartAttr", Bunch_readPartAttr ,METH_O,"Initializes the particles' attr. from the bunch file"}, - { "getPartAttrSize", Bunch_getPartAttrSize ,METH_O,"Returns the number of variables in the particles' attributes with a particular name"}, - { "partAttrValue", Bunch_partAttrValue ,METH_VARARGS,"Sets or returns a particles' attribute value"}, - { "getSize", Bunch_getSize ,METH_NOARGS,"Returns number of macro-particles"}, - { "getSizeGlobal", Bunch_getSizeGlobal ,METH_NOARGS,"Returns number of macro-particles in all CPUs"}, - { "getSizeGlobalFromMemory", Bunch_getSizeGlobalFromMemory ,METH_NOARGS,"Returns number of macro-particles in all CPUs from memory"}, - { "getTotalCount", Bunch_getTotalCount ,METH_NOARGS,"Returns number of all particles - alive,dead,new"}, - { "getCapacity", Bunch_getCapacity ,METH_NOARGS,"Returns the capacity of the bunch-contaiter"}, - { "dumpBunch", Bunch_dumpBunch ,METH_VARARGS,"Prints the bunch info into a standart output stream or file"}, - { "readBunch", Bunch_readBunch ,METH_VARARGS,"Reads the bunch info from a file"}, - { "copyEmptyBunchTo", Bunch_copyEmptyBunchTo ,METH_O,"Copy bunch attrubutes and structure to another bunch"}, - { "copyBunchTo", Bunch_copyBunchTo ,METH_O,"Copy bunch all info including particles coordinates and attributes to another bunch"}, - { "addParticlesTo", Bunch_addParticlesTo ,METH_O,"Copy particles coordinates from one bunch to another"}, +static PyMethodDef BunchClassMethods[] = { + //-------------------------------------------------------- + // class Bunch wrapper START + //-------------------------------------------------------- + {"getMPIComm", Bunch_getMPIComm, METH_NOARGS, "Returns MPI Comm of this bunch"}, + {"setMPIComm", Bunch_setMPIComm, METH_O, "Sets a new MPI Comm for this bunch"}, + {"getSyncParticle", Bunch_getSyncParticle, METH_NOARGS, "Returns syncParticle class instance"}, + {"addParticle", Bunch_addParticle, METH_VARARGS, "Adds a macro-particle to the bunch"}, + {"deleteParticle", + Bunch_deleteParticle, + METH_O, + "Removes macro-particle from the bunch and call compress inside"}, + {"deleteParticleFast", + Bunch_deleteParticleFast, + METH_O, + "Removes macro-particle from the bunch very fast"}, + {"recoverParticle", Bunch_recoverParticle, METH_O, "Recovers a particle marked for removal"}, + {"deleteAllParticles", + Bunch_deleteAllParticles, + METH_NOARGS, + "Removes all macro-particles from the bunch"}, + {"compress", Bunch_compress, METH_NOARGS, "Compress the bunch"}, + {"x", Bunch_x, METH_VARARGS, "Set x(index,value) or get x(index) coordinate"}, + {"y", Bunch_y, METH_VARARGS, "Set y(index,value) or get y(index) coordinate"}, + {"z", Bunch_z, METH_VARARGS, "Set z(index,value) or get z(index) coordinate"}, + {"px", Bunch_px, METH_VARARGS, "Set px(index,value) or get px(index) coordinate"}, + {"py", Bunch_py, METH_VARARGS, "Set py(index,value) or get py(index) coordinate"}, + {"pz", Bunch_pz, METH_VARARGS, "Set pz(index,value) or get pz(index) coordinate"}, + {"dE", Bunch_pz, METH_VARARGS, "Set dE(index,value) or get dE(index) coordinate"}, + {"xp", Bunch_px, METH_VARARGS, "Set xp(index,value) or get xp(index) coordinate"}, + {"yp", Bunch_py, METH_VARARGS, "Set yp(index,value) or get yp(index) coordinate"}, + {"flag", Bunch_flag, METH_O, "Returns flag(index) for particle with index"}, + {"ringwrap", Bunch_ringwrap, METH_O, "Perform the ring wrap. Usage: ringwrap(ring_length)"}, + {"mass", Bunch_mass, METH_VARARGS, "Set mass(value) or get mass() the mass of particle in MeV"}, + {"classicalRadius", + Bunch_classicalRadius, + METH_NOARGS, + "Returns a classical radius of particle in [m]"}, + {"B_Rho", Bunch_B_Rho, METH_NOARGS, "Returns B*Rho parameter of particle in [T*m]"}, + {"charge", + Bunch_charge, + METH_VARARGS, + "Set charge(value) or get charge() the charge of particle in e-charge"}, + {"macroSize", + Bunch_macroSize, + METH_VARARGS, + "Set macroSize(value) or get macroSize() the charge of particle in e-charge"}, + {"initBunchAttr", + Bunch_initBunchAttr, + METH_O, + "Reads and initilizes bunch attributes from a bunch file"}, + {"bunchAttrDouble", + Bunch_bunchAttrDouble, + METH_VARARGS, + "Returns and sets a double bunch attribute"}, + {"bunchAttrInt", Bunch_bunchAttrInt, METH_VARARGS, "Returns and sets an integer bunch attribute"}, + {"bunchAttrDoubleNames", + Bunch_bunchAttrDoubleNames, + METH_NOARGS, + "Returns a list of double bunch attribute names"}, + {"bunchAttrIntNames", + Bunch_bunchAttrIntNames, + METH_NOARGS, + "Returns a list of integer bunch attribute names"}, + {"hasBunchAttrDouble", + Bunch_hasBunchAttrDouble, + METH_O, + "Returns 1 if there is a double bunch attr. with this name, 0 - otherwise"}, + {"hasBunchAttrInt", + Bunch_hasBunchAttrInt, + METH_O, + "Returns 1 if there is a int bunch attr. with this name, 0 - otherwise"}, + {"addPartAttr", Bunch_addPartAttr, METH_VARARGS, "Adds a particles' attributes to the bunch"}, + {"removePartAttr", + Bunch_removePartAttr, + METH_O, + "Removes a particles' attributes from the bunch"}, + {"removeAllPartAttr", + Bunch_removeAllPartAttr, + METH_NOARGS, + "Removes all particles' attributes from the bunch"}, + {"getPartAttrNames", + Bunch_getPartAttrNames, + METH_NOARGS, + "Returns all particles' attributes names in the bunch at this moment"}, + {"getPartAttrDicts", + Bunch_getPartAttrDicts, + METH_NOARGS, + "Returns dict{part. attribute name:dict{key:val}}"}, + {"getPossiblePartAttrNames", + Bunch_getPossiblePartAttrNames, + METH_NOARGS, + "Returns all possible particles' attributes names"}, + {"clearAllPartAttrAndMemorize", + Bunch_clearAllPartAttrAndMemorize, + METH_NOARGS, + "Temporary removes and memorizes all particles' attributes names"}, + {"restoreAllPartAttrFromMemory", + Bunch_restoreAllPartAttrFromMemory, + METH_NOARGS, + "Restores all particles' attributes names from memory"}, + {"hasPartAttr", + Bunch_hasPartAttr, + METH_O, + "Returns 1 if there is a particles' attr. with this name, 0 - otherwis"}, + {"readPartAttrNames", + Bunch_readPartAttrNames, + METH_O, + "Returns a tuple with particles' attr. names in the bunch file"}, + {"readPartAttrDicts", + Bunch_readPartAttrDicts, + METH_O, + "Returns a dict{attr_name:dicts{param_name:val}} in the bunch file"}, + {"readPartAttr", + Bunch_readPartAttr, + METH_O, + "Initializes the particles' attr. from the bunch file"}, + {"getPartAttrSize", + Bunch_getPartAttrSize, + METH_O, + "Returns the number of variables in the particles' attributes with a particular name"}, + {"partAttrValue", + Bunch_partAttrValue, + METH_VARARGS, + "Sets or returns a particles' attribute value"}, + {"getSize", Bunch_getSize, METH_NOARGS, "Returns number of macro-particles"}, + {"getSizeGlobal", + Bunch_getSizeGlobal, + METH_NOARGS, + "Returns number of macro-particles in all CPUs"}, + {"getSizeGlobalFromMemory", + Bunch_getSizeGlobalFromMemory, + METH_NOARGS, + "Returns number of macro-particles in all CPUs from memory"}, + {"getTotalCount", + Bunch_getTotalCount, + METH_NOARGS, + "Returns number of all particles - alive,dead,new"}, + {"getCapacity", Bunch_getCapacity, METH_NOARGS, "Returns the capacity of the bunch-contaiter"}, + {"dumpBunch", + Bunch_dumpBunch, + METH_VARARGS, + "Prints the bunch info into a standart output stream or file"}, + {"readBunch", Bunch_readBunch, METH_VARARGS, "Reads the bunch info from a file"}, + {"copyEmptyBunchTo", + Bunch_copyEmptyBunchTo, + METH_O, + "Copy bunch attrubutes and structure to another bunch"}, + {"copyBunchTo", + Bunch_copyBunchTo, + METH_O, + "Copy bunch all info including particles coordinates and attributes to another bunch"}, + {"addParticlesTo", + Bunch_addParticlesTo, + METH_O, + "Copy particles coordinates from one bunch to another"}, #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY - { "to_numpy", _PyCFunction_CAST(Bunch_to_numpy) ,METH_VARARGS | METH_KEYWORDS, Bunch_to_numpy_doc }, - { "update_from_numpy", Bunch_update_from_numpy ,METH_O, Bunch_update_from_numpy_doc }, - { "from_numpy", Bunch_from_numpy ,METH_O | METH_CLASS, Bunch_from_numpy_doc }, + {"to_numpy", _PyCFunction_CAST(Bunch_to_numpy), METH_VARARGS | METH_KEYWORDS, Bunch_to_numpy_doc}, + {"update_from_numpy", Bunch_update_from_numpy, METH_O, Bunch_update_from_numpy_doc}, + {"from_numpy", Bunch_from_numpy, METH_O | METH_CLASS, Bunch_from_numpy_doc}, #endif - {NULL,NULL} - //-------------------------------------------------------- - // class Bunch wrapper STOP - //-------------------------------------------------------- - }; - - // defenition of the memebers of the python Bunch wrapper class - // they will be vailable from python level - static PyMemberDef BunchClassMembers [] = { - {NULL} - }; - - //new python Bunch wrapper type definition - static PyTypeObject pyORBIT_Bunch_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - "Bunch", /*tp_name*/ - sizeof(pyORBIT_Object), /*tp_basicsize*/ - 0, /*tp_itemsize*/ - (destructor) Bunch_del , /*tp_dealloc*/ - 0, /*tp_print*/ - 0, /*tp_getattr*/ - 0, /*tp_setattr*/ - 0, /*tp_compare*/ - 0, /*tp_repr*/ - 0, /*tp_as_number*/ - 0, /*tp_as_sequence*/ - 0, /*tp_as_mapping*/ - 0, /*tp_hash */ - 0, /*tp_call*/ - 0, /*tp_str*/ - 0, /*tp_getattro*/ - 0, /*tp_setattro*/ - 0, /*tp_as_buffer*/ - Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ - "The Bunch python wrapper", /* tp_doc */ - 0, /* tp_traverse */ - 0, /* tp_clear */ - 0, /* tp_richcompare */ - 0, /* tp_weaklistoffset */ - 0, /* tp_iter */ - 0, /* tp_iternext */ - BunchClassMethods, /* tp_methods */ - BunchClassMembers, /* tp_members */ - 0, /* tp_getset */ - 0, /* tp_base */ - 0, /* tp_dict */ - 0, /* tp_descr_get */ - 0, /* tp_descr_set */ - 0, /* tp_dictoffset */ - (initproc) Bunch_init, /* tp_init */ - 0, /* tp_alloc */ - Bunch_new, /* tp_new */ - }; - - static PyMethodDef BunchModuleMethods[] = { {NULL,NULL} }; - + {NULL, NULL} + //-------------------------------------------------------- + // class Bunch wrapper STOP + //-------------------------------------------------------- +}; + +// defenition of the memebers of the python Bunch wrapper class +// they will be vailable from python level +static PyMemberDef BunchClassMembers[] = {{NULL}}; + +// new python Bunch wrapper type definition +static PyTypeObject pyORBIT_Bunch_Type = { + PyVarObject_HEAD_INIT(NULL, 0) "Bunch", /*tp_name*/ + sizeof(pyORBIT_Object), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor)Bunch_del, /*tp_dealloc*/ + 0, /*tp_print*/ + 0, /*tp_getattr*/ + 0, /*tp_setattr*/ + 0, /*tp_compare*/ + 0, /*tp_repr*/ + 0, /*tp_as_number*/ + 0, /*tp_as_sequence*/ + 0, /*tp_as_mapping*/ + 0, /*tp_hash */ + 0, /*tp_call*/ + 0, /*tp_str*/ + 0, /*tp_getattro*/ + 0, /*tp_setattro*/ + 0, /*tp_as_buffer*/ + Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/ + "The Bunch python wrapper", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + BunchClassMethods, /* tp_methods */ + BunchClassMembers, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc)Bunch_init, /* tp_init */ + 0, /* tp_alloc */ + Bunch_new, /* tp_new */ +}; + +static PyMethodDef BunchModuleMethods[] = {{NULL, NULL}}; #ifdef __cplusplus extern "C" { #endif - static struct PyModuleDef cModPyDem = - { - PyModuleDef_HEAD_INIT, - "bunch", "Bunch class", - -1, - BunchModuleMethods - }; - - /* 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) { - return NULL; - } - #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); - PyObject* module = PyModule_Create(&cModPyDem); - PyModule_AddObject(module, "Bunch", (PyObject *)&pyORBIT_Bunch_Type); - //add the SyncParticle python class - wrap_orbit_syncpart::initsyncpart(module); - wrap_bunch_twiss_analysis::initbunchtwissanalysis(module); - wrap_bunch_tune_analysis::initbunchtuneanalysis(module); - wrap_synch_part_redefinition::initsynchpartredefinition(module); - return module; - } - - PyObject* getBunchType(const char* name){ - PyObject* mod = PyImport_ImportModule("orbit.core.bunch"); - PyObject* pyType = PyObject_GetAttrString(mod,name); - Py_DECREF(mod); - Py_DECREF(pyType); - return pyType; - } +static struct PyModuleDef cModPyDem = + {PyModuleDef_HEAD_INIT, "bunch", "Bunch class", -1, BunchModuleMethods}; + +/* 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) { + return NULL; + } +#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); + PyObject *module = PyModule_Create(&cModPyDem); + PyModule_AddObject(module, "Bunch", (PyObject *)&pyORBIT_Bunch_Type); + // add the SyncParticle python class + wrap_orbit_syncpart::initsyncpart(module); + wrap_bunch_twiss_analysis::initbunchtwissanalysis(module); + wrap_bunch_tune_analysis::initbunchtuneanalysis(module); + wrap_synch_part_redefinition::initsynchpartredefinition(module); + return module; +} +PyObject *getBunchType(const char *name) +{ + PyObject *mod = PyImport_ImportModule("orbit.core.bunch"); + PyObject *pyType = PyObject_GetAttrString(mod, name); + Py_DECREF(mod); + Py_DECREF(pyType); + return pyType; +} #ifdef __cplusplus } #endif -//end of namespace wrap_orbit_bunch -} +// end of namespace wrap_orbit_bunch +} // namespace wrap_orbit_bunch /////////////////////////////////////////////////////////////////////////// //