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/pyproject.toml b/pyproject.toml index a65921e6..a53ad656 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [build-system] build-backend = 'mesonpy' -requires = ['meson-python', "setuptools>=45", "wheel", "setuptools_scm"] +requires = ['meson-python', "setuptools>=45", "wheel", "setuptools_scm", "numpy>=2.0"] [project] name = 'PyORBIT' diff --git a/src/meson.build b/src/meson.build index c223553d..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/orbit/wrap_bunch.cc b/src/orbit/wrap_bunch.cc index 68824b4e..4293dd78 100644 --- a/src/orbit/wrap_bunch.cc +++ b/src/orbit/wrap_bunch.cc @@ -16,6 +16,19 @@ #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" @@ -623,7 +636,7 @@ namespace wrap_orbit_bunch{ } std::string attr_name_str(attr_name); val = cpp_bunch->getBunchAttributeDouble(attr_name_str); - return Py_BuildValue("d",val); + return Py_BuildValue("d",val); } else{ //NO NEW OBJECT CREATED BY PyArg_ParseTuple! NO NEED OF Py_DECREF() @@ -1171,6 +1184,114 @@ namespace wrap_orbit_bunch{ self->ob_base.ob_type->tp_free((PyObject*)self); } +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY +static PyObject *Bunch_to_numpy(PyObject *self, PyObject *args) { + pyORBIT_Object *pyBunch = (pyORBIT_Object*)self; + Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj; + + if (!PyArg_ParseTuple(args, ":to_numpy")) { + ORBIT_MPI_Finalize("PyBunch - to_numpy() - no parameters are needed."); + } + + const npy_intp nparts = (npy_intp)cpp_Bunch->getSize(); + const npy_intp ncoords = 6; + + npy_intp dims[2] = { nparts, ncoords }; + + PyObject *py_array = PyArray_SimpleNew(2, dims, NPY_FLOAT64); + PyArrayObject *arr_obj = (PyArrayObject*)py_array; + double *data_buffer = (double*)PyArray_DATA(arr_obj); + double **src = cpp_Bunch->coordArr(); + + for (npy_intp i = 0; i < nparts; ++i) { + for (npy_intp j = 0; j < ncoords; ++j) { + data_buffer[j + i*ncoords] = src[i][j]; + } + } + + return py_array; +} + + static int bunch_fill_from_numpy_args(Bunch *cpp_Bunch, PyObject *args) { + PyObject *arr_in = NULL; + + if (!PyArg_ParseTuple(args, "O:from_numpy", &arr_in)) { + return -1; + } + + PyArrayObject *arr = + (PyArrayObject*)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, NPY_ARRAY_IN_ARRAY); + if (!arr) return -1; + + if (PyArray_NDIM(arr) != 2) { + Py_DECREF(arr); + PyErr_SetString(PyExc_ValueError, + "from_numpy: array must be 2-dimensional with shape (nparts, 6)"); + return -1; + } + + const npy_intp nparts = PyArray_DIM(arr, 0); + const npy_intp ncoords = PyArray_DIM(arr, 1); + + if (ncoords != 6) { + Py_DECREF(arr); + PyErr_SetString(PyExc_ValueError, + "from_numpy: expected coordinate dimension with shape 6 (x, px, y, py, z, dE)"); + return -1; + } + + const double *data = (const double*)PyArray_DATA(arr); + + for (npy_intp i = 0; i < nparts; ++i) { + const npy_intp stride = i * ncoords; + cpp_Bunch->addParticle( + data[stride+0], data[stride+1], data[stride+2], + data[stride+3], data[stride+4], data[stride+5] + ); + } + + Py_DECREF(arr); + return 0; +} + +static PyObject *Bunch_update_from_numpy(PyObject *self, PyObject *args) { + pyORBIT_Object *pyBunch = (pyORBIT_Object*)self; + Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj; + + if (cpp_Bunch->getSizeGlobal() > 0) { + for (int i = 0; i < cpp_Bunch->getSizeGlobal(); ++i) { + cpp_Bunch->deleteParticleFast(i); + } + cpp_Bunch->compress(); + } + + if (bunch_fill_from_numpy_args(cpp_Bunch, args) < 0) return NULL; + + Py_RETURN_NONE; +} + +static PyObject *Bunch_from_numpy(PyObject *cls, PyObject *args) { + PyObject *py_bunch_obj = PyObject_CallNoArgs(cls); + if (!py_bunch_obj) return NULL; + + pyORBIT_Object *pyBunch = (pyORBIT_Object*)py_bunch_obj; + Bunch *cpp_Bunch = (Bunch*)pyBunch->cpp_obj; + + if (!cpp_Bunch) { + Py_DECREF(py_bunch_obj); + PyErr_SetString(PyExc_RuntimeError, "from_numpy: Constructed bunch has NULL cpp_obj"); + return NULL; + } + + if (bunch_fill_from_numpy_args(cpp_Bunch, args) < 0) { + Py_DECREF(py_bunch_obj); + return NULL; + } + + return py_bunch_obj; +} +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY + static PyMethodDef BunchClassMethods[] = { //-------------------------------------------------------- // class Bunch wrapper START @@ -1230,6 +1351,11 @@ namespace wrap_orbit_bunch{ { "copyEmptyBunchTo", Bunch_copyEmptyBunchTo ,METH_VARARGS,"Copy bunch attrubutes and structure to another bunch"}, { "copyBunchTo", Bunch_copyBunchTo ,METH_VARARGS,"Copy bunch all info including particles coordinates and attributes to another bunch"}, { "addParticlesTo", Bunch_addParticlesTo ,METH_VARARGS,"Copy particles coordinates from one bunch to another"}, +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY + { "to_numpy", Bunch_to_numpy ,METH_VARARGS, "Convert bunch coordinates to a numpy array" }, + { "update_from_numpy", Bunch_update_from_numpy ,METH_VARARGS, "Update bunch coordinates from a numpy array" }, + { "from_numpy", Bunch_from_numpy ,METH_VARARGS | METH_CLASS, "Construct a new Bunch from a numpy array" }, +#endif {NULL,NULL} //-------------------------------------------------------- // class Bunch wrapper STOP @@ -1301,6 +1427,11 @@ extern "C" { /* The name of the function was changed to avoid collision with PyImport magic naming */ PyMODINIT_FUNC initbunch(void) { + #ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY + if (ensure_numpy() != 0) { + throw std::runtime_error("NumPy C-API init failed"); + } + #endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY //check that the Bunch wrapper is ready if(PyType_Ready(&pyORBIT_Bunch_Type) < 0) return NULL; Py_INCREF(&pyORBIT_Bunch_Type); diff --git a/src/spacecharge/wrap_grid3D.cc b/src/spacecharge/wrap_grid3D.cc index 5dcc32b0..3100109a 100644 --- a/src/spacecharge/wrap_grid3D.cc +++ b/src/spacecharge/wrap_grid3D.cc @@ -1,414 +1,541 @@ +#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; +} + +// getGridX(ix) +static PyObject *Grid3D_getGridX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int ind = -1; + if (!PyArg_ParseTuple(args, "i:getGridX", &ind) || ind < 0 || + ind >= cpp_Grid3D->getSizeX()) { + ORBIT_MPI_Finalize( + "PyGrid3D - getGridX(ix) - parameter is needed. [0 - sizeX["); + } + return Py_BuildValue("d", cpp_Grid3D->getGridX(ind)); +} + +// getGridY(iy) +static PyObject *Grid3D_getGridY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int ind = -1; + if (!PyArg_ParseTuple(args, "i:getGridY", &ind) || ind < 0 || + ind >= cpp_Grid3D->getSizeY()) { + ORBIT_MPI_Finalize( + "PyGrid3D - getGridY(iy) - parameter is needed. [0 - sizeY["); + } + return Py_BuildValue("d", cpp_Grid3D->getGridY(ind)); +} + +// getGridZ(iy) +static PyObject *Grid3D_getGridZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int ind = -1; + if (!PyArg_ParseTuple(args, "i:getGridZ", &ind) || ind < 0 || + ind >= cpp_Grid3D->getSizeZ()) { + ORBIT_MPI_Finalize( + "PyGrid3D - getGridZ(iz) - parameter is needed. [0 - sizeZ["); + } + return Py_BuildValue("d", cpp_Grid3D->getGridZ(ind)); +} + +// getSizeX() +static PyObject *Grid3D_getSizeX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("i", cpp_Grid3D->getSizeX()); +} + +// getSizeY() +static PyObject *Grid3D_getSizeY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("i", cpp_Grid3D->getSizeY()); +} + +// getSizeZ() +static PyObject *Grid3D_getSizeZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("i", cpp_Grid3D->getSizeZ()); +} + +// It will synchronize through the MPI communicator +static PyObject *Grid3D_synchronizeMPI(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + int nVars = PyTuple_Size(args); + if (nVars == 0) { + cpp_Grid3D->synchronizeMPI(NULL); + } else { + PyObject *py_mpi_comm_type = + wrap_orbit_mpi_comm::getMPI_CommType("MPI_Comm"); + PyObject *pyMPIComm = PyTuple_GetItem(args, 0); + if ((!PyObject_IsInstance(pyMPIComm, py_mpi_comm_type))) { + ORBIT_MPI_Finalize("Grid3D.synchronizeMPI(MPI_Comm) - input " + "parameter is not MPI_Comm"); + } + cpp_Grid3D->synchronizeMPI((pyORBIT_MPI_Comm *)pyMPIComm); + } + Py_INCREF(Py_None); + return Py_None; +} + +// getMinX() +static PyObject *Grid3D_getMinX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMinX()); +} + +// getMaxX() +static PyObject *Grid3D_getMaxX(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMaxX()); +} + +// getMinY() +static PyObject *Grid3D_getMinY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMinY()); +} + +// getMaxY() +static PyObject *Grid3D_getMaxY(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMaxY()); +} + +// getMinZ() +static PyObject *Grid3D_getMinZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMinZ()); +} + +// getMaxZ() +static PyObject *Grid3D_getMaxZ(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + return Py_BuildValue("d", cpp_Grid3D->getMaxZ()); +} + +// longWrapping([isWrapped]) set or return the longitudinal wrapping policy +static PyObject *Grid3D_longWrapping(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + PyObject *pyIsWrapped = NULL; + if (!PyArg_ParseTuple(args, "|O:longWrapping", &pyIsWrapped)) { + ORBIT_MPI_Finalize("PyGrid3D - longWrapping([True/False]) - " + "parameter may be needed."); + } + if (pyIsWrapped != NULL) { + int isWrapped = PyObject_IsTrue(pyIsWrapped); + if (isWrapped != 1) + isWrapped = 0; + cpp_Grid3D->setLongWrapping(isWrapped); + } + return Py_BuildValue("i", cpp_Grid3D->getLongWrapping()); +} + +// binBunch(Bunch* bunch) +static PyObject *Grid3D_binBunch(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + PyObject *pyBunch; + double lambda = -1.0; + if (!PyArg_ParseTuple(args, "O|d:binBunch", &pyBunch, &lambda)) { + ORBIT_MPI_Finalize("PyGrid3D - binBunch(Bunch* bunch [,lambda]) - " + "parameters are needed."); + } + PyObject *pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); + if (!PyObject_IsInstance(pyBunch, pyORBIT_Bunch_Type)) { + ORBIT_MPI_Finalize("PyGrid3D - binBunch(Bunch* bunch [,lambda]) - " + "method needs a Bunch."); + } + Bunch *cpp_bunch = (Bunch *)((pyORBIT_Object *)pyBunch)->cpp_obj; + if (lambda > 0.) { + cpp_Grid3D->binBunch(cpp_bunch, lambda); + } else { + cpp_Grid3D->binBunch(cpp_bunch); + } + Py_INCREF(Py_None); + return Py_None; +} + +// 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 *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + + if (!PyArg_ParseTuple(args, ":to_numpy")) { + ORBIT_MPI_Finalize("PyGrid3D - to_numpy() - no parameters are needed."); + } + + const npy_intp nz = (npy_intp)cpp_Grid3D->getSizeZ(); + const npy_intp nx = (npy_intp)cpp_Grid3D->getSizeX(); + const npy_intp ny = (npy_intp)cpp_Grid3D->getSizeY(); + + npy_intp dims[3] = {nz, nx, ny}; + PyObject *arr_obj = PyArray_SimpleNew(3, dims, NPY_FLOAT64); + if (NULL == arr_obj) { + return NULL; } - //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; + PyArrayObject *arr = (PyArrayObject *)arr_obj; + double *out_buffer = (double *)PyArray_DATA(arr); + double ***src = cpp_Grid3D->getArr3D(); + + for (npy_intp iz = 0; iz < nz; ++iz) { + for (npy_intp ix = 0; ix < ny; ++ix) { + for (npy_intp iy = 0; iy < nx; ++iy) { + out_buffer[iy + ix * ny + iz * nx * ny] = src[iz][ix][iy]; + } + } } - //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); + return arr_obj; +} + +static PyObject *Grid3D_from_numpy(PyObject *self, PyObject *args) { + pyORBIT_Object *pyGrid3D = (pyORBIT_Object *)self; + Grid3D *cpp_Grid3D = (Grid3D *)pyGrid3D->cpp_obj; + + PyObject *arr_in = NULL; + const char* order = "zxy"; + + if (!PyArg_ParseTuple(args, "O|s:from_numpy", &arr_in, &order)) { + ORBIT_MPI_Finalize("PyGrid3D - from_numpy() - ndarray is needed."); } - // 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); - } + PyArrayObject *arr = (PyArrayObject *)PyArray_FROM_OTF(arr_in, NPY_FLOAT64, + NPY_ARRAY_IN_ARRAY); + if (NULL == arr) { + return NULL; + } + + if (PyArray_NDIM(arr) != 3) { + Py_DECREF(arr); + PyErr_SetString( + PyExc_ValueError, + "from_numpy: array must be 3-dimensional with shape (nz,nx,ny)"); + return NULL; + } + + const npy_intp nz_in = PyArray_DIM(arr, 0); + const npy_intp nx_in = PyArray_DIM(arr, 1); + const npy_intp ny_in = PyArray_DIM(arr, 2); + + const npy_intp nz_grid = (npy_intp)cpp_Grid3D->getSizeZ(); + const npy_intp nx_grid = (npy_intp)cpp_Grid3D->getSizeX(); + const npy_intp ny_grid = (npy_intp)cpp_Grid3D->getSizeY(); + + if (nz_in != nz_grid || nx_in != nx_grid || ny_in != ny_grid) { + Py_DECREF(arr); + PyErr_SetString( + PyExc_ValueError, + "from_numpy: shape mismatch; expected (zSize, xSize, ySize)"); + return NULL; + } + + const double *in_buffer = (double *)PyArray_DATA(arr); + double ***dst = cpp_Grid3D->getArr3D(); + + for (npy_intp iz = 0; iz < nz_grid; ++iz) { + for (npy_intp ix = 0; ix < nx_grid; ++ix) { + for (npy_intp iy = 0; iy < ny_grid; ++iy) { + dst[iz][ix][iy] = in_buffer[iy + ix * ny_grid + iz * nx_grid * ny_grid]; + } + } + } + + Py_DECREF(arr); + Py_RETURN_NONE; +} +#endif // 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_VARARGS, "converts the 3D grid to a numpy array"}, + {"from_numpy", Grid3D_from_numpy, METH_VARARGS, "converts the numpy array to a 3D grid"}, +#endif // 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 +//-------------------------------------------------- +void initGrid3D(PyObject *module) { +#ifdef PyORBIT_EXPERIMENTAL_WITH_NUMPY + if (ensure_numpy() != 0) { + throw std::runtime_error("NumPy C-API init failed"); + } +#endif // PyORBIT_EXPERIMENTAL_WITH_NUMPY + if (PyType_Ready(&pyORBIT_Grid3D_Type) < 0) + return; + Py_INCREF(&pyORBIT_Grid3D_Type); + PyModule_AddObject(module, "Grid3D", (PyObject *)&pyORBIT_Grid3D_Type); +} #ifdef __cplusplus } #endif -//end of namespace wrap_spacecharge -} +// end of namespace wrap_spacecharge +} // namespace wrap_spacecharge