diff --git a/.github/workflows/compilation.yml b/.github/workflows/compilation.yml index c40130af..76102b1a 100644 --- a/.github/workflows/compilation.yml +++ b/.github/workflows/compilation.yml @@ -147,5 +147,81 @@ jobs: run: | .github/workflows/pip-versions.sh + python-tests: + name: Python tests (${{ matrix.mpi }}) + runs-on: ubuntu-latest + timeout-minutes: 20 + + strategy: + fail-fast: false + matrix: + include: + - mpi: none + mpi_packages: "" + mpiexec_preflags: "" + - mpi: mpich + mpi_packages: >- + mpich + libfabric-devel + mpiexec_preflags: "" + - mpi: ompi + mpi_packages: openmpi + mpiexec_preflags: --oversubscribe + + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: mamba-org/setup-micromamba@v3 + with: + environment-name: pyorbit-test + create-args: >- + python=3.13 + numpy + pytest + fftw + meson + ninja + pkg-config + setuptools + setuptools-scm + ${{ matrix.mpi_packages }} + cache-downloads: true + + - name: Configure + shell: bash -el {0} + run: | + meson setup build-${{ matrix.mpi }} \ + --prefix="$CONDA_PREFIX" \ + --libdir=lib \ + --buildtype=release \ + -DUSE_MPI=${{ matrix.mpi }} \ + -DPyORBIT_EXPERIMENTAL_WITH_NUMPY=true + + - name: Compile + shell: bash -el {0} + run: | + meson compile -C build-${{ matrix.mpi }} + + - name: Install + shell: bash -el {0} + run: | + meson install --no-rebuild -C build-${{ matrix.mpi }} + + - name: Show MPI configuration + if: matrix.mpi != 'none' + shell: bash -el {0} + run: | + mpirun --version + + - name: Run Python tests + shell: bash -el {0} + env: + PYORBIT_MPIEXEC_PREFLAGS: ${{ matrix.mpiexec_preflags }} + run: | + export LD_LIBRARY_PATH="$CONDA_PREFIX/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + python -m pytest tests/py/ -v + build-docs: uses: ./.github/workflows/docs-build.yml 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/.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/ diff --git a/meson_options.txt b/meson_options.txt index 3ae5f6ab..90e09823 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('PyORBIT_EXPERIMENTAL_WITH_NUMPY', type: 'boolean', value: false, description: 'Use numpy') diff --git a/py/orbit/bunch_utils/__init__.py b/py/orbit/bunch_utils/__init__.py index 42ea363e..3096e0d9 100644 --- a/py/orbit/bunch_utils/__init__.py +++ b/py/orbit/bunch_utils/__init__.py @@ -7,23 +7,20 @@ from .particleidnumber import ParticleIdNumber -# This guards against missing numpy. -# Should be imporved with some meaningful (and MPI friendly?) warning printed out. -try: - from .serialize import collect_bunch, save_bunch, load_bunch - from .serialize import BunchDict, SyncPartDict - from .serialize import FileHandler, NumPyHandler -except ImportError: - pass +__all__ = ["ParticleIdNumber"] -__all__ = [] -# __all__.append("addParticleIdNumbers") # doesn't exist -__all__.append("ParticleIdNumber") -__all__.append("collect_bunch") -__all__.append("save_bunch") -__all__.append("load_bunch") -__all__.append("BunchDict") -__all__.append("SyncPartDict") -__all__.append("FileHandler") -__all__.append("NumPyHandler") +from .numpy_utils import bunch_from_shared_numpy +from .serialize import collect_bunch, save_bunch, load_bunch +from .serialize import BunchDict, SyncPartDict +from .serialize import FileHandler, NumPyHandler +__all__ += [ + "bunch_from_shared_numpy", + "collect_bunch", + "save_bunch", + "load_bunch", + "BunchDict", + "SyncPartDict", + "FileHandler", + "NumPyHandler", +] 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 new file mode 100644 index 00000000..76ff52c2 --- /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/pyproject.toml b/pyproject.toml index d601642b..2a8c6a06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,16 +1,15 @@ [build-system] build-backend = 'mesonpy' -requires = ['meson-python', "setuptools>=45", "wheel", "setuptools_scm"] +requires = ['meson-python', "setuptools>=45", "wheel", "setuptools_scm", "numpy>=2.0"] [project] name = 'PyORBIT' dynamic = ["version"] description = 'Use meson-python to build c++ anf python modules.' requires-python = '>=3.9' -#dependencies = [ -#'setuptools', -#'setuptools-scm' -#] +dependencies = [ + 'numpy>=2.0', +] authors = [ {name = 'Alexander Zhukov', email = 'zhukovap@ornl.gov'}, ] @@ -18,9 +17,3 @@ authors = [ [tool.setuptools_scm] -[project.optional-dependencies] -numpy = [ - "numpy", -] - - diff --git a/src/meson.build b/src/meson.build index c223553d..a3e86e34 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('PyORBIT_EXPERIMENTAL_WITH_NUMPY') + +if with_numpy + message('Compiling with Numpy support') + python = pymod.find_installation('python3', pure: false, modules: ['numpy']) + numpy_inc = run_command( + python, '-c', 'import numpy; print(numpy.get_include())', + check: true + ).stdout().strip() + add_project_arguments('-DPyORBIT_EXPERIMENTAL_WITH_NUMPY=1', language: 'cpp') + add_project_arguments('-I' + numpy_inc, language: 'cpp') + +else + python = pymod.find_installation('python3', pure: false) +endif + # Add C++ compiler details cpp = meson.get_compiler('cpp') @@ -17,39 +33,37 @@ dependencies += dependency('fftw3', version: '>= 3.0.0', required: true) # Detecting if MPICH or OPENMPI are installed and enabling support if present - mpi_use = get_option('USE_MPI') -# message('MPI_USE is set to', mpi_use) + +cpp_args = ['-fPIC', '-std=c++11', '-O3', '-march=native'] if mpi_use == 'mpich' message('Requested to use MPICH as the MPI implementation.') dependencies += dependency('mpich', version: '>= 4.0.0', required: true) - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] + cpp_args += ['-DUSE_MPI=1'] # Configure dependencies or settings specific to MPICH elif mpi_use == 'ompi' message('Requested to use OpenMPI as the MPI implementation.') dependencies += dependency('ompi', version: '>= 4.0.0', required: true) - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] + cpp_args += ['-DUSE_MPI=1'] elif mpi_use == 'none' message('Requested to not use MPI.') - cpp_args = ['-fPIC', '-std=c++11'] else mpich_dependency = dependency('mpich', version: '>= 4.0.0', required: false) openmpi_dependency = dependency('ompi', version: '>= 4.0.0', required: false) if mpich_dependency.found() - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] + cpp_args = ['-DUSE_MPI=1'] dependencies += mpich_dependency message('Using MPICH as the MPI implementation.') elif openmpi_dependency.found() - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] + cpp_args = ['-DUSE_MPI=1'] dependencies += openmpi_dependency message('Using OpenMPI as the MPI implementation.') else - cpp_args = ['-fPIC', '-std=c++11'] message('MPI will not be used.') endif @@ -280,6 +294,7 @@ inc = include_directories([ ]) + core_lib = library('core', sources: sources, include_directories: inc, @@ -304,7 +319,7 @@ python.extension_module('orbit_mpi', python.extension_module('bunch', sources: [base + '/bunch_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, #['-fPIC', '-std=c++11'], dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -313,7 +328,7 @@ python.extension_module('bunch', python.extension_module('spacecharge', sources: [base + '/spacecharge_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, #['-fPIC', '-std=c++11'], dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -322,7 +337,7 @@ python.extension_module('spacecharge', python.extension_module('trackerrk4', sources: [base + '/trackerrk4_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, #['-fPIC', '-std=c++11'], dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -331,7 +346,7 @@ python.extension_module('trackerrk4', python.extension_module('teapot_base', sources: [base + '/teapot_base_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, #['-fPIC', '-std=c++11'], dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -340,7 +355,7 @@ python.extension_module('teapot_base', python.extension_module('linac', sources: [base + '/linac_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -349,7 +364,7 @@ python.extension_module('linac', python.extension_module('orbit_utils', sources: [base + '/utils_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -358,7 +373,7 @@ python.extension_module('orbit_utils', python.extension_module('aperture', sources: [base + '/aperture_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -367,7 +382,7 @@ python.extension_module('aperture', python.extension_module('foil', sources: [base + '/foil_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -376,7 +391,7 @@ python.extension_module('foil', python.extension_module('field_sources', sources: [base + '/field_sources_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -385,7 +400,7 @@ python.extension_module('field_sources', python.extension_module('rfcavities', sources: [base + '/rfcavities_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -394,7 +409,7 @@ python.extension_module('rfcavities', python.extension_module('impedances', sources: [base + '/impedances_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -403,7 +418,7 @@ python.extension_module('impedances', python.extension_module('fieldtracker', sources: [base + '/fieldtracker_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -412,7 +427,7 @@ python.extension_module('fieldtracker', python.extension_module('collimator', sources: [base + '/collimator_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', @@ -421,7 +436,7 @@ python.extension_module('collimator', python.extension_module('error_base', sources: [base + '/error_base_init.cc'], include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], + cpp_args: cpp_args, dependencies: [core_dep], install: true, subdir: 'orbit/core', 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/src/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index 8bda038d..cdeb9563 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -9,1340 +9,1680 @@ // /////////////////////////////////////////////////////////////////////////// #include "wrap_bunch.hh" -#include "wrap_syncpart.hh" -#include "wrap_bunch_twiss_analysis.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" -#include "pyORBIT_Object.hh" +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY +#include + +static int ensure_numpy() +{ + static int numpy_initialized = 0; + if (!numpy_initialized) { + import_array1(-1); + numpy_initialized = 1; + } + return 0; +} +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY #include "Bunch.hh" #include "ParticleAttributesFactory.hh" -namespace wrap_orbit_bunch{ +namespace wrap_orbit_bunch +{ + +void error(const char *msg) +{ + ORBIT_MPI_Finalize(msg); +} - void error(const char* msg){ ORBIT_MPI_Finalize(msg); } - //--------------------------------------------------------- - //Python Bunch class definition - //--------------------------------------------------------- +//--------------------------------------------------------- +// 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; +} - //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; +} - //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 *args){ - 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 *args){ - 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 *args){ - 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)."); - } - Py_INCREF(Py_None); - return Py_None; +//--------------------------------------------------------------- +// +// 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); +} - //--------------------------------------------------------------- - // - // add and remove particles, compress etc. - // - //---------------------------------------------------------------- +// 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; + } - //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; + cpp_bunch->deleteParticle(ind); + int size = cpp_bunch->getSize(); - double x = 0.; double xp = 0.; double y = 0.; - double yp = 0.; double z = 0.; double zp = 0.; + return PyLong_FromLong(size); +} - //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); +// 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); +} - //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){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - int ind; +static PyObject *Bunch_recoverParticle(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"); - } + if (!PyArg_Parse(arg, "i:recoverParticle", &ind)) { + return NULL; + } + + cpp_bunch->recoverParticle(ind); + Py_RETURN_NONE; +} - cpp_bunch->deleteParticle(ind); - int size = cpp_bunch->getSize(); +// 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; +} + +// 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; +} - return Py_BuildValue("i",size); +//--------------------------------------------------------------- +// +// 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; } - //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){ - 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; + } - //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"); - } + return PyFloat_FromDouble(value); +} - cpp_bunch->deleteParticleFast(ind); - return Py_BuildValue("i",ind); +// 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; } - static PyObject* Bunch_recoverParticle(PyObject *self, PyObject *args){ - Bunch *cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - int ind; + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->y(index); + } + else { + cpp_bunch->y(index) = value; + } - if(!PyArg_ParseTuple(args,"i:recoverParticle",&ind)){ - error("PyBunch - recoverParticle - needs index of particle for recovering"); - } + 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 *args){ - 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 *args){ - 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 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; + } - 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->z(index); + } + else { + cpp_bunch->z(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 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 (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->px(index); + } + else { + cpp_bunch->px(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 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; + } - 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); + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->py(index); + } + else { + cpp_bunch->py(index) = value; + } + + 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; + } + + if (PyTuple_GET_SIZE(args) == 1) { + value = cpp_bunch->pz(index); + } + else { + cpp_bunch->pz(index) = value; + } + + return PyFloat_FromDouble(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 PyLong_FromLong(flag); +} + +// 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.; + + // 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); + Py_RETURN_NONE; +} + +//--------------------------------------------------------------- +// +// 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; + + if (PyTuple_GET_SIZE(args) == 0) { + return PyFloat_FromDouble(cpp_bunch->getMass()); + } + + double value; + if (!PyArg_ParseTuple(args, "d:mass", &value)) { + return NULL; + } + + 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()); + } + + double value; + if (!PyArg_ParseTuple(args, "d:charge", &value)) { + return NULL; + } + + cpp_bunch->setCharge(value); + return PyFloat_FromDouble(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; + + if (PyTuple_GET_SIZE(args) == 0) { + return PyFloat_FromDouble(cpp_bunch->getMacroSize()); + } + + double value; + if (!PyArg_ParseTuple(args, "d:macroSize", &value)) { + return NULL; + } + + cpp_bunch->setMacroSize(value); + return PyFloat_FromDouble(value); +} + +//--------------------------------------------------------------- +// +// 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; +} + +// 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; + } + + 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); + } + + return PyFloat_FromDouble(value); +} + +// 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); + 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 px(index) or px(index,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; - } - - //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; +// 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); +} + +//--------------------------------------------------------------- +// +// 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" + ); } - return Py_BuildValue("d",val); + std::string par_name((char *)PyUnicode_AsUTF8(key)); + double d_val = PyFloat_AsDouble(value); + part_attr_dict[par_name] = d_val; } - else{ - error("PyBunch. You should call py(index) or py(index,value)"); + } + cpp_bunch->addParticleAttributes(attr_name_str, part_attr_dict); + Py_RETURN_NONE; +} + +// 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"); } + } + return resTuple; +} - 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); +// 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; } - else{ - error("PyBunch. You should call pz(index) or pz(index,value)"); + if (PyDict_SetItemString(resDict, names[i].c_str(), py_param_dict) < 0) { + Py_DECREF(py_param_dict); + Py_DECREF(resDict); + return NULL; } - - 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 *args){ - 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); + 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); } - else{ - error("PyBunch. You should call bunch.flag(index)"); + Py_DECREF(py_param_dict); + } + 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"); } - Py_INCREF(Py_None); - return Py_None; } + return resTuple; +} - //Wraps long. coords in the bunch - //ringwrap(ring_length) - static PyObject* Bunch_ringwrap(PyObject *self, PyObject *args) { - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - double ring_length = 0.; +// 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; +} - //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"); - } +// 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; +} - cpp_bunch->ringwrap(ring_length); - Py_INCREF(Py_None); - return Py_None; - } - - //--------------------------------------------------------------- - // - // 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; - //if nVars == 0 get mass - //if nVars == 1 set mass - int nVars = PyTuple_Size(args); - - double val = 0.; - - if(nVars == 0 || nVars == 1){ - if(nVars == 0){ - val = cpp_bunch->getMass(); +// 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 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; +} + +// 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; + } + if (PyDict_SetItemString(resDict, names[i].c_str(), py_param_dict) < 0) { + Py_DECREF(py_param_dict); + Py_DECREF(resDict); + return NULL; } - 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"); + 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->setMass(val); + Py_DECREF(py_val); } - return Py_BuildValue("d",val); - } - else{ - error("PyBunch. You should call mass() or mass(value)"); + Py_DECREF(py_param_dict); } + } + return resDict; +} - Py_INCREF(Py_None); - return Py_None; +// 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 classicalRadius of the particle in meters - static PyObject* Bunch_classicalRadius(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - double val = cpp_bunch->getClassicalRadius(); - return Py_BuildValue("d",val); +// 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); +} - //Returns B_Rho of the particle in [Tesla*meter]. Parameter is used in TEAPOT - static PyObject* Bunch_B_Rho(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - double val = cpp_bunch->getB_Rho(); - return Py_BuildValue("d",val); +// 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; } - //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); + 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!" + ); + } - double val = 0.; + 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; + } - 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); + 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{ - error("PyBunch. You should call charge() or charge(value)"); + 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"); + } + cpp_bunch->print(file_name); } + } + else { + error("PyBunch. You should call dumpBunch() or dumpBunch(file_name)"); + } - Py_INCREF(Py_None); - return Py_None; - } + Py_RETURN_NONE; +} - //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(); +// 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"); } - 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); + cpp_bunch->initBunchAttributes(file_name); + cpp_bunch->readParticleAttributes(file_name); + cpp_bunch->readBunchCoords(file_name); } - else{ - error("PyBunch. You should call macroSize() or macroSize(value)"); + 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); } + } + 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; +} - 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_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); + + if (array == NULL) { + return NULL; } - //--------------------------------------------------------------- - // - // related to the bunch attributes - // - //---------------------------------------------------------------- + if (PyArray_NDIM(array) != 2) { + PyErr_SetString(PyExc_ValueError, "array must be 2-dimensional with shape (nparts, 6)"); + Py_DECREF(array); + return NULL; + } - //initilizes bunch attributes from the bunch file - static PyObject* Bunch_initBunchAttr(PyObject *self, PyObject *args){ - 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"); - } - 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"); - } + 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; + } - 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 array; // new ref; caller MUST Py_DECREF(). +} - 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"); - } +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); - std::string attr_name_str(attr_name); - cpp_bunch->setBunchAttribute( attr_name_str, val); - return Py_BuildValue("i",val); - } - } - else{ - error("PyBunch. You should call bunchAttrInt(name) or bunchAttrInt(name,value)"); - } + for (npy_intp i = 0; i < nparts; ++i) { + const double *coords = data + i * 6; - Py_INCREF(Py_None); - return Py_None; - } - - //Returns a list (tuple) of ther double bunch attribute names - static PyObject* Bunch_bunchAttrDoubleNames(PyObject *self, PyObject *args){ - 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 *args){ - 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; + cpp_bunch->addParticle(coords[0], coords[1], coords[2], coords[3], coords[4], coords[5]); } +} - //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){ - 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"); - } - std::string attr_name_str(attr_name); - int res = cpp_bunch->getBunchAttributes()->hasDoubleAttribute(attr_name_str); - return Py_BuildValue("i",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 *args){ - 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"); - } - 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"); - } - 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; - } - - //Removes a particles' attributes with a particular name from the bunch - static PyObject* Bunch_removePartAttr(PyObject *self, PyObject *args){ - 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"); - } - 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 *args){ - 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 *args){ - 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 *args){ - 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 *args){ - 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 *args){ - 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 *args){ - 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 *args){ - 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"); - } - 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 *args){ - 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"); - } - 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; - } - - //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){ - 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"); - } - 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; +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; } - //initilizes particles' attributes from the bunch file - static PyObject* Bunch_readPartAttr(PyObject *self, PyObject *args){ - 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"); - } - cpp_bunch->readParticleAttributes(file_name); - Py_INCREF(Py_None); - return Py_None; - } - - //Returns the number of variables in the particles' attributes with a particular name - static PyObject* Bunch_getPartAttrSize(PyObject *self, PyObject *args){ - 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"); + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + + 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; + } + + npy_intp dims[2] = {nparts, ncoords}; + + 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, + cpp_bunch->getMPI_Comm_Local()->comm + ); + + if (!all_counts_ok) { + PyErr_SetString(PyExc_OverflowError, "local bunch is too large for MPI_Gatherv"); + 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); + + 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, + cpp_bunch->getMPI_Comm_Local()->comm + ); + + if (!all_alloc_ok) { + Py_XDECREF(local_array); + + if (!PyErr_Occurred()) { + PyErr_NoMemory(); } + + return NULL; } - else{ - error("PyBunch. You should call partAttrValue(attr_name,part_ind,attr_ind) or partAttrValue(attr_name,part_ind,attr_ind,value)"); + } + else { + if (local_array == NULL) { + return NULL; } + } - 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 *args){ - 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){ - 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 *args){ - 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){ - 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){ - 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); - } - 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"); - } - cpp_bunch->print(file_name); - } - } - else{ - error("PyBunch. You should call dumpBunch() or dumpBunch(file_name)"); + 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]; } + } - 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); - } - 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); + if (!gather || size == 1) { + return local_array; + } + +#if USE_MPI > 0 + MPI_Comm comm = cpp_bunch->getMPI_Comm_Local()->comm; + + std::vector global_nparts(size); + std::vector displacements(size); + std::vector recv_counts(size); + + const int local_nparts = (int)nparts; + + 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; } + + displacements[mpi_rank] = total; + recv_counts[mpi_rank] = rank_nvalues; + total += rank_nvalues; } - else{ - error("PyBunch. You should call readBunch(file_name) or readBunch(file_name,nParts)"); + } + + 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; + } + + 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(); } - Py_INCREF(Py_None); - return Py_None; + + return NULL; } - //Copy bunch attrubutes and structure to another bunch - static PyObject* Bunch_copyEmptyBunchTo(PyObject *self, PyObject *args){ - 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"); - } - 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 *args){ - 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"); - } - Bunch* cpp_target_bunch = (Bunch*) ((pyORBIT_Object *) pyBunch_Target)->cpp_obj; - cpp_bunch->copyBunchTo(cpp_target_bunch); - Py_INCREF(Py_None); - return Py_None; + 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; } - //Copy particles coordinates from one bunch to another - static PyObject* Bunch_addParticlesTo(PyObject *self, PyObject *args){ - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object *) self)->cpp_obj; - PyObject* pyBunch_Target; +#endif // USE_MPI > 0 + Py_RETURN_NONE; +} - //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"); - } - 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); - } - - static PyMethodDef BunchClassMethods[] = { - //-------------------------------------------------------- - // 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"}, - { "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"}, - { "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_VARARGS,"Returns flag(index) for particle with index"}, - { "ringwrap", Bunch_ringwrap ,METH_VARARGS,"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]"}, - { "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"}, - { "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"}, - { "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"}, - { "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"}, - { "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"}, - {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} }; +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 *cpp_bunch = (Bunch *)((pyORBIT_Object *)self)->cpp_obj; + + PyArrayObject *array = parse_bunch_array(arg); + + if (array == NULL) { + return NULL; + } + + 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" +); + +static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *arg) +{ + PyArrayObject *array = parse_bunch_array(arg); + + if (array == NULL) { + return NULL; + } + + PyObject *py_bunch_obj = PyObject_CallNoArgs(cls); + + if (py_bunch_obj == NULL) { + Py_DECREF(array); + return NULL; + } + + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)py_bunch_obj)->cpp_obj; + + 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"}, +#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}, +#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}}; #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) { - //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 /////////////////////////////////////////////////////////////////////////// // diff --git a/src/spacecharge/wrap_grid3D.cc b/src/spacecharge/wrap_grid3D.cc index 5dcc32b0..956cfecc 100644 --- a/src/spacecharge/wrap_grid3D.cc +++ b/src/spacecharge/wrap_grid3D.cc @@ -1,414 +1,538 @@ +#include "wrap_grid3D.hh" + +#include + #include "orbit_mpi.hh" #include "pyORBIT_Object.hh" - -#include "wrap_grid3D.hh" -#include "wrap_spacecharge.hh" #include "wrap_bunch.hh" +#include "wrap_spacecharge.hh" -#include +#ifdef PyORBIT_EXPERIMENTAL_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 // PyORBIT_EXPERIMENTAL_WITH_NUMPY #include "Grid3D.hh" 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; +} - //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; +// 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)); +} - //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); +// 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)); +} - // 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"}, - {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 - //-------------------------------------------------- - 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); - } +// 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; +} + +// 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); +} + +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY +static PyObject *Grid3D_to_numpy(PyObject *self, + PyObject *Py_UNUSED(ignored)) { + Grid3D *cpp_Grid3D = (Grid3D *)((pyORBIT_Object *)self)->cpp_obj; + + 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}; + + 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 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]; + } + } + } + + return arr_obj; +} + +static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *arg) { + Grid3D *cpp_Grid3D = (Grid3D *)((pyORBIT_Object *)self)->cpp_obj; + + PyArrayObject *arr = (PyArrayObject *)PyArray_FROM_OTF(arg, 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 (nx,ny,nz)"); + return NULL; + } + + 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 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 (nx_in != nx_grid || ny_in != ny_grid || nz_in != nz_grid) { + Py_DECREF(arr); + PyErr_SetString( + PyExc_ValueError, + "from_numpy: shape mismatch; expected (xSize, ySize, zSize)"); + 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[(ix * ny_grid + iy) * nz_grid + iz]; + } + } + } + + Py_DECREF(arr); + Py_RETURN_NONE; +} +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY + +// 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"}, +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY + {"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}}; + +// 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 +//-------------------------------------------------- +int initGrid3D(PyObject *module) { +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY + if (ensure_numpy() != 0) { + return -1; + } +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY + if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) { + return -1; + } + Py_INCREF(&pyORBIT_Grid3D_Type); + if (PyModule_AddObject(module, "Grid3D", + (PyObject *)&pyORBIT_Grid3D_Type) < 0) { + Py_DECREF(&pyORBIT_Grid3D_Type); + return -1; + } + return 0; +} #ifdef __cplusplus } #endif -//end of namespace wrap_spacecharge -} +// end of namespace wrap_spacecharge +} // namespace wrap_spacecharge 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); 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..cb47a5d3 --- /dev/null +++ b/tests/py/orbit/core/test_numpy_interop.py @@ -0,0 +1,308 @@ +import os +import shlex +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 _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( + 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) + + +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) + + 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_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) + + np.testing.assert_array_equal(bunch.to_numpy(gather=True), expected) + + def test_missing_numpy_raises_import_error_instead_of_aborting(self): + _assert_missing_numpy_raises_import_error("bunch") + + 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 + + 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) + ) + """ + ) + ) + + result = subprocess.run( + _mpi_test_command(mpirun, 2, script_path), + capture_output=True, + text=True, + timeout=60, + check=False, + env=_mpi_test_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_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, + ) + + 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) + """ + ) + ) + + result = subprocess.run( + _mpi_test_command(mpirun, 3, script_path), + capture_output=True, + text=True, + timeout=60, + check=False, + env=_mpi_test_env(), + ) + + if "PYORBIT_MPI_DISABLED" in result.stdout: + pytest.skip("PyORBIT was built without MPI support") + + assert result.returncode == 0, _process_failure(result) + + +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 + + 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_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") 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