From d56d5fc063cc407f928c0ebf845556fd8f2ede1d Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Mon, 15 Jun 2026 16:58:02 -0400 Subject: [PATCH 1/6] Clean up meson files. - tabs -> 2 spaces per style guide - foreach loop for the extension modules - clean up cpp_args - migrate c++ standard specification to project() - global project argument for `-DUSE_MPI` - Removed `-fPIC`. Meson handles this for us. - set default buildtype to 'release' in project() - can be overridden with `meson setup ... -Dbuildtype=debug` - Add `-march=native` global project argument, conditionally. - Since we don't distribute binaries and expect users to compile pyo3 for themselves, I don't see a reason not to include this flag as it can only help performance. At worst, it does nothing. --- meson.build | 10 +- src/meson.build | 658 ++++++++++++++++++++---------------------------- 2 files changed, 283 insertions(+), 385 deletions(-) diff --git a/meson.build b/meson.build index b2bb9c79..428193ff 100644 --- a/meson.build +++ b/meson.build @@ -1,4 +1,12 @@ -project('PyORBIT', ['cpp'], version: run_command('python', '-m', 'setuptools_scm').stdout().strip()) +project('PyORBIT', + ['cpp'], + version: run_command('python', '-m', 'setuptools_scm').stdout().strip(), + default_options: [ + 'buildtype=release', + 'cpp_std=c++11', + 'warning_level=3', + ], +) subdir('src') subdir('py/orbit') diff --git a/src/meson.build b/src/meson.build index c223553d..c7474e58 100644 --- a/src/meson.build +++ b/src/meson.build @@ -1,14 +1,13 @@ - -#project('orbit', ['cpp'], version: '0.0.1') - - # Add Python installation details python = import('python').find_installation('python3', pure: false) # Add C++ compiler details cpp = meson.get_compiler('cpp') -base = meson.current_source_dir() + '/core' +# Add -march=native if the compiler supports it +if cpp.has_argument('-march=native') + add_project_arguments('-march=native', language: 'cpp') +endif dependencies = [] @@ -17,412 +16,303 @@ 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) 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'] + message('Requested to use MPICH as the MPI implementation.') + dependencies += dependency('mpich', version: '>= 4.0.0', required: true) + add_project_arguments('-DUSE_MPI=1', language: 'cpp') - # Configure dependencies or settings specific to MPICH + # 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'] + message('Requested to use OpenMPI as the MPI implementation.') + dependencies += dependency('ompi', version: '>= 4.0.0', required: true) + add_project_arguments('-DUSE_MPI=1', language: 'cpp') elif mpi_use == 'none' - message('Requested to not use MPI.') - cpp_args = ['-fPIC', '-std=c++11'] + message('Requested to not use MPI.') 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'] - dependencies += mpich_dependency - message('Using MPICH as the MPI implementation.') + add_project_arguments('-DUSE_MPI=1', language: 'cpp') + dependencies += mpich_dependency + message('Using MPICH as the MPI implementation.') elif openmpi_dependency.found() - cpp_args = ['-fPIC', '-std=c++11', '-DUSE_MPI=1'] - dependencies += openmpi_dependency - message('Using OpenMPI as the MPI implementation.') + add_project_arguments('-DUSE_MPI=1', language: 'cpp') + dependencies += openmpi_dependency + message('Using OpenMPI as the MPI implementation.') else - cpp_args = ['-fPIC', '-std=c++11'] - message('MPI will not be used.') + message('MPI will not be used.') endif endif sources = files([ - 'linac/wrap_linacmodule.cc', - 'linac/tracking/linac_tracking.cc', - 'linac/tracking/wrap_linac_tracking.cc', - 'linac/rfgap/wrap_BaseRfGap.cc', - 'linac/rfgap/BaseRfGap_slow.cc', - 'linac/rfgap/wrap_RfGapTTF_slow.cc', - 'linac/rfgap/wrap_RfGapThreePointTTF.cc', - 'linac/rfgap/BaseRfGap.cc', - 'linac/rfgap/RfGapThreePointTTF.cc', - 'linac/rfgap/MatrixRfGap.cc', - 'linac/rfgap/wrap_BaseRfGap_slow.cc', - 'linac/rfgap/wrap_RfGapThreePointTTF_slow.cc', - 'linac/rfgap/wrap_SuperFishFieldSource.cc', - 'linac/rfgap/RfGapTTF.cc', - 'linac/rfgap/SuperFishFieldSource.cc', - 'linac/rfgap/RfGapTTF_slow.cc', - 'linac/rfgap/RfGapThreePointTTF_slow.cc', - 'linac/rfgap/wrap_RfGapTTF.cc', - 'linac/rfgap/wrap_MatrixRfGap.cc', - 'orbit/OrbitConst.cc', - 'orbit/wrap_bunch.cc', - 'orbit/wrap_syncpart.cc', - 'orbit/Bunch.cc', - 'orbit/SyncPart.cc', - 'orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc', - 'orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc', - 'orbit/BunchDiagnostics/BunchTwissAnalysis.cc', - 'orbit/BunchDiagnostics/BunchTuneAnalysis.cc', - 'orbit/FieldTracker/wrap_fieldtracker.cc', - 'orbit/FieldTracker/FieldTracker.cc', - 'orbit/MaterialInteractions/wrap_foil.cc', - 'orbit/MaterialInteractions/Foil.cc', - 'orbit/MaterialInteractions/Collimator.cc', - 'orbit/MaterialInteractions/numrecipes.cc', - 'orbit/MaterialInteractions/cross_sections.cc', - 'orbit/MaterialInteractions/wrap_collimator.cc', - 'orbit/MaterialInteractions/MaterialInteractions.cc', - 'orbit/RFCavities/Dual_Harmonic_Cav.cc', - 'orbit/RFCavities/Frequency_Cav.cc', - 'orbit/RFCavities/Barrier_Cav.cc', - 'orbit/RFCavities/wrap_Harmonic_Cav.cc', - 'orbit/RFCavities/Harmonic_Cav.cc', - 'orbit/RFCavities/wrap_Dual_Harmonic_Cav.cc', - 'orbit/RFCavities/wrap_Barrier_Cav.cc', - 'orbit/RFCavities/wrap_Frequency_Cav.cc', - 'orbit/RFCavities/wrap_rfcavities.cc', - 'orbit/Impedances/wrap_TImpedance.cc', - 'orbit/Impedances/TImpedance.cc', - 'orbit/Impedances/wrap_impedances.cc', - 'orbit/Impedances/LImpedance.cc', - 'orbit/Impedances/wrap_LImpedance.cc', - 'orbit/SynchPartRedefinition/SynchPartRedefinitionZdE.cc', - 'orbit/SynchPartRedefinition/wrap_synch_part_redefinition_z_de.cc', - 'orbit/Errors/errorbase.cc', - 'orbit/Errors/wrap_errorbase.cc', - 'orbit/ParticlesAttributes/ParticleInitialCoordinates.cc', - 'orbit/ParticlesAttributes/ParticleIdNumber.cc', - 'orbit/ParticlesAttributes/ParticleMacroSize.cc', - 'orbit/ParticlesAttributes/ParticleAttributesFactory.cc', - 'orbit/ParticlesAttributes/WaveFunctionAmplitudes.cc', - 'orbit/ParticlesAttributes/pq_coordinates.cc', - 'orbit/ParticlesAttributes/Evolution.cc', - 'orbit/ParticlesAttributes/ParticleAttributes.cc', - 'orbit/ParticlesAttributes/TurnNumberAttributes.cc', - 'orbit/ParticlesAttributes/LostParticleAttributes.cc', - 'orbit/ParticlesAttributes/ParticlePhaseAttributes.cc', - 'orbit/ParticlesAttributes/part_time.cc', - 'orbit/ParticlesAttributes/AtomPopulations.cc', - 'orbit/Apertures/CircleApertureShape.cc', - 'orbit/Apertures/PhaseAperture.cc', - 'orbit/Apertures/Aperture.cc', - 'orbit/Apertures/wrap_aperture.cc', - 'orbit/Apertures/wrap_CompositeApertureShape.cc', - 'orbit/Apertures/RectangularApertureShape.cc', - 'orbit/Apertures/PyBaseApertureShape.cc', - 'orbit/Apertures/ConvexApertureShape.cc', - 'orbit/Apertures/wrap_ConvexApertureShape.cc', - 'orbit/Apertures/wrap_PyBaseApertureShape.cc', - 'orbit/Apertures/EnergyAperture.cc', - 'orbit/Apertures/BaseAperture.cc', - 'orbit/Apertures/EllipseApertureShape.cc', - 'orbit/Apertures/wrap_BaseAperture.cc', - 'orbit/Apertures/wrap_PrimitiveApertureShape.cc', - 'orbit/Apertures/CompositeApertureShape.cc', - 'orbit/Apertures/wrap_EnergyAperture.cc', - 'orbit/Apertures/wrap_TAperture.cc', - 'orbit/Apertures/wrap_PhaseAperture.cc', - 'orbit/Apertures/BaseApertureShape.cc', - 'trackerrk4/wrap_runge_kutta_tracker.cc', - 'trackerrk4/wrap_py_external_effects.cc', - 'trackerrk4/ExternalEffects.cc', - 'trackerrk4/ExtEffectsContainer.cc', - 'trackerrk4/wrap_trackerrk4.cc', - 'trackerrk4/PyExternalEffects.cc', - 'trackerrk4/RungeKuttaTracker.cc', - 'trackerrk4/wrap_ext_effects_container.cc', - 'mpi/wrap_orbit_mpi.cc', - 'mpi/wrap_mpi_comm.cc', - 'mpi/orbit_mpi.cc', - 'mpi/wrap_mpi_datatype.cc', - 'mpi/wrap_mpi_request.cc', - 'mpi/wrap_mpi_op.cc', - 'mpi/wrap_mpi_status.cc', - 'mpi/wrap_mpi_group.cc', - 'spacecharge/PoissonSolverFFT3D.cc', - 'spacecharge/SpaceChargeCalcUnifEllipse.cc', - 'spacecharge/ForceSolver2D.cc', - 'spacecharge/PoissonSolverFFT2D.cc', - 'spacecharge/wrap_poissonsolverfft2d.cc', - 'spacecharge/wrap_poissonsolverfft3d.cc', - 'spacecharge/PoissonSolver2D.cc', - 'spacecharge/SpaceChargeForceCalc2p5D.cc', - 'spacecharge/PoissonSolver3D.cc', - 'spacecharge/SpaceChargeCalc3D.cc', - 'spacecharge/wrap_boundary2d.cc', - 'spacecharge/wrap_forcesolverfft2d.cc', - 'spacecharge/ForceSolverFFT2D.cc', - 'spacecharge/SpaceChargeCalcSliceBySlice2D.cc', - 'spacecharge/Grid3D.cc', - 'spacecharge/wrap_spacecharge.cc', - 'spacecharge/wrap_spacechargecalc2p5d_rb.cc', - 'spacecharge/Grid2D.cc', - 'spacecharge/SpaceChargeCalc2p5Drb.cc', - 'spacecharge/Grid1D.cc', - 'spacecharge/LSpaceChargeCalc.cc', - 'spacecharge/wrap_grid1D.cc', - 'spacecharge/wrap_uniform_ellipsoid_field_calculator.cc', - 'spacecharge/ShapedBoundary2D.cc', - 'spacecharge/wrap_grid3D.cc', - 'spacecharge/BaseBoundary2D.cc', - 'spacecharge/wrap_grid2D.cc', - 'spacecharge/wrap_spacechargecalc_uniform_ellipse.cc', - 'spacecharge/wrap_spacechargecalc3d.cc', - 'spacecharge/UniformEllipsoidFieldCalculator.cc', - 'spacecharge/wrap_spacechargeforcecalc2p5d.cc', - 'spacecharge/SpaceChargeCalc2p5D.cc', - 'spacecharge/wrap_spacechargecalc_slicebyslice_2D.cc', - 'spacecharge/wrap_spacechargecalc2p5d.cc', - 'spacecharge/wrap_lspacechargecalc.cc', - 'utils/wrap_function.cc', - 'utils/FieldSourceContainer.cc', - 'utils/AttributesBucket.cc', - 'utils/bessel.cc', - 'utils/PyBaseFieldSource.cc', - 'utils/wrap_field_source_container.cc', - 'utils/OU_SplineCH.cc', - 'utils/CppPyWrapper.cc', - 'utils/wrap_random.cc', - 'utils/wrap_utils.cc', - 'utils/StringUtils.cc', - 'utils/wrap_py_base_field_source.cc', - 'utils/wrap_splinech.cc', - 'utils/Random.cc', - 'utils/wrap_numrecipes.cc', - 'utils/BaseFieldSource.cc', - 'utils/OU_Function.cc', - 'utils/MathPolynomial.cc', - 'utils/BufferStore.cc', - 'utils/statistics/wrap_statmoments2d.cc', - 'utils/statistics/StatMoments2D.cc', - 'utils/field_sources/wrap_loop_field_source.cc', - 'utils/field_sources/LoopFieldSource.cc', - 'utils/field_sources/QuadFieldSource.cc', - 'utils/field_sources/ShiftedFieldSource.cc', - 'utils/field_sources/wrap_dipole_field_source.cc', - 'utils/field_sources/wrap_quad_field_source.cc', - 'utils/field_sources/wrap_magnetic_field_source_grid3d.cc', - 'utils/field_sources/DipoleFieldSource.cc', - 'utils/field_sources/wrap_field_sources_module.cc', - 'utils/field_sources/MagnetFieldSourceGrid3D.cc', - 'utils/integration/gauss_legendre_points.cc', - 'utils/integration/wrap_gauss_legendre_integrator.cc', - 'utils/integration/GaussLegendreIntegrator.cc', - 'utils/bunch/ParticlesWithIdFunctions.cc', - 'utils/bunch/wrap_bunch_extrema_calculator.cc', - 'utils/bunch/TwissFilteringFunctions.cc', - 'utils/bunch/InitialCoordsAttrFunctions.cc', - 'utils/bunch/BunchExtremaCalculator.cc', - 'utils/bunch/wrap_bunch_utils_functions.cc', - 'utils/harmonic_analysis/HarmonicData.cc', - 'utils/harmonic_analysis/wrap_harmonic_data.cc', - 'utils/polynomial/wrap_polynomial.cc', - 'utils/polynomial/OU_Polynomial.cc', - 'utils/matrix/wrap_matrix.cc', - 'utils/matrix/PhaseVector.cc', - 'utils/matrix/MatrixOperations.cc', - 'utils/matrix/wrap_phase_vector.cc', - 'utils/matrix/Matrix.cc', - 'teapot/wrap_teapotbase.cc', - 'teapot/wrap_matrix_generator.cc', - 'teapot/teapotbase.cc', - 'teapot/MatrixGenerator.cc' + 'linac/wrap_linacmodule.cc', + 'linac/tracking/linac_tracking.cc', + 'linac/tracking/wrap_linac_tracking.cc', + 'linac/rfgap/wrap_BaseRfGap.cc', + 'linac/rfgap/BaseRfGap_slow.cc', + 'linac/rfgap/wrap_RfGapTTF_slow.cc', + 'linac/rfgap/wrap_RfGapThreePointTTF.cc', + 'linac/rfgap/BaseRfGap.cc', + 'linac/rfgap/RfGapThreePointTTF.cc', + 'linac/rfgap/MatrixRfGap.cc', + 'linac/rfgap/wrap_BaseRfGap_slow.cc', + 'linac/rfgap/wrap_RfGapThreePointTTF_slow.cc', + 'linac/rfgap/wrap_SuperFishFieldSource.cc', + 'linac/rfgap/RfGapTTF.cc', + 'linac/rfgap/SuperFishFieldSource.cc', + 'linac/rfgap/RfGapTTF_slow.cc', + 'linac/rfgap/RfGapThreePointTTF_slow.cc', + 'linac/rfgap/wrap_RfGapTTF.cc', + 'linac/rfgap/wrap_MatrixRfGap.cc', + 'orbit/OrbitConst.cc', + 'orbit/wrap_bunch.cc', + 'orbit/wrap_syncpart.cc', + 'orbit/Bunch.cc', + 'orbit/SyncPart.cc', + 'orbit/BunchDiagnostics/wrap_bunch_tune_analysis.cc', + 'orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc', + 'orbit/BunchDiagnostics/BunchTwissAnalysis.cc', + 'orbit/BunchDiagnostics/BunchTuneAnalysis.cc', + 'orbit/FieldTracker/wrap_fieldtracker.cc', + 'orbit/FieldTracker/FieldTracker.cc', + 'orbit/MaterialInteractions/wrap_foil.cc', + 'orbit/MaterialInteractions/Foil.cc', + 'orbit/MaterialInteractions/Collimator.cc', + 'orbit/MaterialInteractions/numrecipes.cc', + 'orbit/MaterialInteractions/cross_sections.cc', + 'orbit/MaterialInteractions/wrap_collimator.cc', + 'orbit/MaterialInteractions/MaterialInteractions.cc', + 'orbit/RFCavities/Dual_Harmonic_Cav.cc', + 'orbit/RFCavities/Frequency_Cav.cc', + 'orbit/RFCavities/Barrier_Cav.cc', + 'orbit/RFCavities/wrap_Harmonic_Cav.cc', + 'orbit/RFCavities/Harmonic_Cav.cc', + 'orbit/RFCavities/wrap_Dual_Harmonic_Cav.cc', + 'orbit/RFCavities/wrap_Barrier_Cav.cc', + 'orbit/RFCavities/wrap_Frequency_Cav.cc', + 'orbit/RFCavities/wrap_rfcavities.cc', + 'orbit/Impedances/wrap_TImpedance.cc', + 'orbit/Impedances/TImpedance.cc', + 'orbit/Impedances/wrap_impedances.cc', + 'orbit/Impedances/LImpedance.cc', + 'orbit/Impedances/wrap_LImpedance.cc', + 'orbit/SynchPartRedefinition/SynchPartRedefinitionZdE.cc', + 'orbit/SynchPartRedefinition/wrap_synch_part_redefinition_z_de.cc', + 'orbit/Errors/errorbase.cc', + 'orbit/Errors/wrap_errorbase.cc', + 'orbit/ParticlesAttributes/ParticleInitialCoordinates.cc', + 'orbit/ParticlesAttributes/ParticleIdNumber.cc', + 'orbit/ParticlesAttributes/ParticleMacroSize.cc', + 'orbit/ParticlesAttributes/ParticleAttributesFactory.cc', + 'orbit/ParticlesAttributes/WaveFunctionAmplitudes.cc', + 'orbit/ParticlesAttributes/pq_coordinates.cc', + 'orbit/ParticlesAttributes/Evolution.cc', + 'orbit/ParticlesAttributes/ParticleAttributes.cc', + 'orbit/ParticlesAttributes/TurnNumberAttributes.cc', + 'orbit/ParticlesAttributes/LostParticleAttributes.cc', + 'orbit/ParticlesAttributes/ParticlePhaseAttributes.cc', + 'orbit/ParticlesAttributes/part_time.cc', + 'orbit/ParticlesAttributes/AtomPopulations.cc', + 'orbit/Apertures/CircleApertureShape.cc', + 'orbit/Apertures/PhaseAperture.cc', + 'orbit/Apertures/Aperture.cc', + 'orbit/Apertures/wrap_aperture.cc', + 'orbit/Apertures/wrap_CompositeApertureShape.cc', + 'orbit/Apertures/RectangularApertureShape.cc', + 'orbit/Apertures/PyBaseApertureShape.cc', + 'orbit/Apertures/ConvexApertureShape.cc', + 'orbit/Apertures/wrap_ConvexApertureShape.cc', + 'orbit/Apertures/wrap_PyBaseApertureShape.cc', + 'orbit/Apertures/EnergyAperture.cc', + 'orbit/Apertures/BaseAperture.cc', + 'orbit/Apertures/EllipseApertureShape.cc', + 'orbit/Apertures/wrap_BaseAperture.cc', + 'orbit/Apertures/wrap_PrimitiveApertureShape.cc', + 'orbit/Apertures/CompositeApertureShape.cc', + 'orbit/Apertures/wrap_EnergyAperture.cc', + 'orbit/Apertures/wrap_TAperture.cc', + 'orbit/Apertures/wrap_PhaseAperture.cc', + 'orbit/Apertures/BaseApertureShape.cc', + 'trackerrk4/wrap_runge_kutta_tracker.cc', + 'trackerrk4/wrap_py_external_effects.cc', + 'trackerrk4/ExternalEffects.cc', + 'trackerrk4/ExtEffectsContainer.cc', + 'trackerrk4/wrap_trackerrk4.cc', + 'trackerrk4/PyExternalEffects.cc', + 'trackerrk4/RungeKuttaTracker.cc', + 'trackerrk4/wrap_ext_effects_container.cc', + 'mpi/wrap_orbit_mpi.cc', + 'mpi/wrap_mpi_comm.cc', + 'mpi/orbit_mpi.cc', + 'mpi/wrap_mpi_datatype.cc', + 'mpi/wrap_mpi_request.cc', + 'mpi/wrap_mpi_op.cc', + 'mpi/wrap_mpi_status.cc', + 'mpi/wrap_mpi_group.cc', + 'spacecharge/PoissonSolverFFT3D.cc', + 'spacecharge/SpaceChargeCalcUnifEllipse.cc', + 'spacecharge/ForceSolver2D.cc', + 'spacecharge/PoissonSolverFFT2D.cc', + 'spacecharge/wrap_poissonsolverfft2d.cc', + 'spacecharge/wrap_poissonsolverfft3d.cc', + 'spacecharge/PoissonSolver2D.cc', + 'spacecharge/SpaceChargeForceCalc2p5D.cc', + 'spacecharge/PoissonSolver3D.cc', + 'spacecharge/SpaceChargeCalc3D.cc', + 'spacecharge/wrap_boundary2d.cc', + 'spacecharge/wrap_forcesolverfft2d.cc', + 'spacecharge/ForceSolverFFT2D.cc', + 'spacecharge/SpaceChargeCalcSliceBySlice2D.cc', + 'spacecharge/Grid3D.cc', + 'spacecharge/wrap_spacecharge.cc', + 'spacecharge/wrap_spacechargecalc2p5d_rb.cc', + 'spacecharge/Grid2D.cc', + 'spacecharge/SpaceChargeCalc2p5Drb.cc', + 'spacecharge/Grid1D.cc', + 'spacecharge/LSpaceChargeCalc.cc', + 'spacecharge/wrap_grid1D.cc', + 'spacecharge/wrap_uniform_ellipsoid_field_calculator.cc', + 'spacecharge/ShapedBoundary2D.cc', + 'spacecharge/wrap_grid3D.cc', + 'spacecharge/BaseBoundary2D.cc', + 'spacecharge/wrap_grid2D.cc', + 'spacecharge/wrap_spacechargecalc_uniform_ellipse.cc', + 'spacecharge/wrap_spacechargecalc3d.cc', + 'spacecharge/UniformEllipsoidFieldCalculator.cc', + 'spacecharge/wrap_spacechargeforcecalc2p5d.cc', + 'spacecharge/SpaceChargeCalc2p5D.cc', + 'spacecharge/wrap_spacechargecalc_slicebyslice_2D.cc', + 'spacecharge/wrap_spacechargecalc2p5d.cc', + 'spacecharge/wrap_lspacechargecalc.cc', + 'utils/wrap_function.cc', + 'utils/FieldSourceContainer.cc', + 'utils/AttributesBucket.cc', + 'utils/bessel.cc', + 'utils/PyBaseFieldSource.cc', + 'utils/wrap_field_source_container.cc', + 'utils/OU_SplineCH.cc', + 'utils/CppPyWrapper.cc', + 'utils/wrap_random.cc', + 'utils/wrap_utils.cc', + 'utils/StringUtils.cc', + 'utils/wrap_py_base_field_source.cc', + 'utils/wrap_splinech.cc', + 'utils/Random.cc', + 'utils/wrap_numrecipes.cc', + 'utils/BaseFieldSource.cc', + 'utils/OU_Function.cc', + 'utils/MathPolynomial.cc', + 'utils/BufferStore.cc', + 'utils/statistics/wrap_statmoments2d.cc', + 'utils/statistics/StatMoments2D.cc', + 'utils/field_sources/wrap_loop_field_source.cc', + 'utils/field_sources/LoopFieldSource.cc', + 'utils/field_sources/QuadFieldSource.cc', + 'utils/field_sources/ShiftedFieldSource.cc', + 'utils/field_sources/wrap_dipole_field_source.cc', + 'utils/field_sources/wrap_quad_field_source.cc', + 'utils/field_sources/wrap_magnetic_field_source_grid3d.cc', + 'utils/field_sources/DipoleFieldSource.cc', + 'utils/field_sources/wrap_field_sources_module.cc', + 'utils/field_sources/MagnetFieldSourceGrid3D.cc', + 'utils/integration/gauss_legendre_points.cc', + 'utils/integration/wrap_gauss_legendre_integrator.cc', + 'utils/integration/GaussLegendreIntegrator.cc', + 'utils/bunch/ParticlesWithIdFunctions.cc', + 'utils/bunch/wrap_bunch_extrema_calculator.cc', + 'utils/bunch/TwissFilteringFunctions.cc', + 'utils/bunch/InitialCoordsAttrFunctions.cc', + 'utils/bunch/BunchExtremaCalculator.cc', + 'utils/bunch/wrap_bunch_utils_functions.cc', + 'utils/harmonic_analysis/HarmonicData.cc', + 'utils/harmonic_analysis/wrap_harmonic_data.cc', + 'utils/polynomial/wrap_polynomial.cc', + 'utils/polynomial/OU_Polynomial.cc', + 'utils/matrix/wrap_matrix.cc', + 'utils/matrix/PhaseVector.cc', + 'utils/matrix/MatrixOperations.cc', + 'utils/matrix/wrap_phase_vector.cc', + 'utils/matrix/Matrix.cc', + 'teapot/wrap_teapotbase.cc', + 'teapot/wrap_matrix_generator.cc', + 'teapot/teapotbase.cc', + 'teapot/MatrixGenerator.cc' ]) + inc = include_directories([ - 'main', - 'orbit/SynchPartRedefinition', - 'trackerrk4', - 'utils/ellipticalint', - 'teapot', - 'mpi', - 'orbit/ParticlesAttributes', - 'orbit/Errors', - 'orbit/FieldTracker', - 'utils/statistics', - 'utils', - 'linac/tracking', - 'spacecharge', - 'orbit/MaterialInteractions', - 'utils/field_sources', - 'utils/bunch', - 'orbit/Impedances', - 'utils/harmonic_analysis', - 'linac', - 'utils/polynomial', - 'linac/rfgap', - 'orbit/RFCavities', - 'utils/matrix', - 'orbit/BunchDiagnostics', - 'orbit', - 'utils/integration', - 'orbit/Apertures' + 'main', + 'orbit/SynchPartRedefinition', + 'trackerrk4', + 'utils/ellipticalint', + 'teapot', + 'mpi', + 'orbit/ParticlesAttributes', + 'orbit/Errors', + 'orbit/FieldTracker', + 'utils/statistics', + 'utils', + 'linac/tracking', + 'spacecharge', + 'orbit/MaterialInteractions', + 'utils/field_sources', + 'utils/bunch', + 'orbit/Impedances', + 'utils/harmonic_analysis', + 'linac', + 'utils/polynomial', + 'linac/rfgap', + 'orbit/RFCavities', + 'utils/matrix', + 'orbit/BunchDiagnostics', + 'orbit', + 'utils/integration', + 'orbit/Apertures' ]) - -core_lib = library('core', - sources: sources, - include_directories: inc, - cpp_args: cpp_args, - override_options: ['b_lundef=false'], - dependencies: dependencies, - install: true, +core_lib = library( + 'core', + sources: sources, + include_directories: inc, + override_options: ['b_lundef=false'], + dependencies: dependencies, + install: true, ) core_dep = declare_dependency(link_with : core_lib) -python.extension_module('orbit_mpi', - sources: [base + '/mpi_init.cc'], - include_directories: inc, - cpp_args: cpp_args, - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - +ext_modules = { + 'orbit_mpi': 'mpi_init.cc', + 'bunch': 'bunch_init.cc', + 'spacecharge': 'spacecharge_init.cc', + 'trackerrk4': 'trackerrk4_init.cc', + 'teapot_base': 'teapot_base_init.cc', + 'linac': 'linac_init.cc', + 'orbit_utils': 'utils_init.cc', + 'aperture': 'aperture_init.cc', + 'foil': 'foil_init.cc', + 'field_sources': 'field_sources_init.cc', + 'rfcavities': 'rfcavities_init.cc', + 'impedances': 'impedances_init.cc', + 'fieldtracker': 'fieldtracker_init.cc', + 'collimator': 'collimator_init.cc', + 'error_base': 'error_base_init.cc', +} -python.extension_module('bunch', - sources: [base + '/bunch_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('spacecharge', - sources: [base + '/spacecharge_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('trackerrk4', - sources: [base + '/trackerrk4_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('teapot_base', - sources: [base + '/teapot_base_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('linac', - sources: [base + '/linac_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('orbit_utils', - sources: [base + '/utils_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('aperture', - sources: [base + '/aperture_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('foil', - sources: [base + '/foil_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('field_sources', - sources: [base + '/field_sources_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('rfcavities', - sources: [base + '/rfcavities_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('impedances', - sources: [base + '/impedances_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('fieldtracker', - sources: [base + '/fieldtracker_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) - -python.extension_module('collimator', - sources: [base + '/collimator_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) +base = meson.current_source_dir() + '/core' -python.extension_module('error_base', - sources: [base + '/error_base_init.cc'], - include_directories: inc, - cpp_args: ['-fPIC', '-std=c++11'], - dependencies: [core_dep], - install: true, - subdir: 'orbit/core', -) +foreach name, src : ext_modules + python.extension_module( + name, + sources: base / src, + include_directories: inc, + dependencies: [core_dep], + install: true, + subdir: 'orbit/core', + ) +endforeach From eeadd5ff732036bf5a2a65c61be8313a63353208 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Tue, 16 Jun 2026 10:30:45 -0400 Subject: [PATCH 2/6] Also enable link time optimization --- meson.build | 1 + 1 file changed, 1 insertion(+) diff --git a/meson.build b/meson.build index 428193ff..d37025ea 100644 --- a/meson.build +++ b/meson.build @@ -3,6 +3,7 @@ project('PyORBIT', version: run_command('python', '-m', 'setuptools_scm').stdout().strip(), default_options: [ 'buildtype=release', + 'b_lto=true', 'cpp_std=c++11', 'warning_level=3', ], From 541af43ff2e4da58810ee79598029d4e5232c9bb Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 23 Jul 2026 10:42:35 -0400 Subject: [PATCH 3/6] Fix moment calculation bugs. Refactor BunchTwissAnalysis. Deprecate BunchTwissAnalysis::analyzeBunch. Bump to c++17. The moment calculation bugs are corrected, and an additional `normalization` flag is added to computeBunchMoments to prevent normalization of x-y moments by sqrt(beta) or sqrt(beta*emitt) unless explicitly requested. The class has been reworked to expedite calculation of moments up to 2nd order, which is the most common usecase. Previously, calling computeBunchMoments would scale as O(N^2) due to a double loop over the Bunch. Now, the double loop is triggered only if moment orders >=3 are requested. Bumping the c++ std version allows us to more easily compose specialized implementations of the moments calculation based on whether the dispersion correction is enabled or the 'macrosize' bunch attribute is defined. When these two conditions are `false`, the compiler will be able to optimize-out the extra instructions within in the Bunch loop(s) that these scenarios require. --- meson.build | 2 +- .../BunchDiagnostics/BunchTwissAnalysis.cc | 803 ++++++++---------- .../BunchDiagnostics/BunchTwissAnalysis.hh | 233 ++--- .../wrap_bunch_twiss_analysis.cc | 701 ++++++++------- .../wrap_bunch_twiss_analysis.hh | 9 +- src/utils/bunch/TwissFilteringFunctions.cc | 7 +- src/utils/bunch/TwissFilteringFunctions.hh | 5 +- .../core/bunch/test_BunchTwissAnalysis.py | 54 ++ 8 files changed, 924 insertions(+), 890 deletions(-) create mode 100644 tests/py/orbit/core/bunch/test_BunchTwissAnalysis.py diff --git a/meson.build b/meson.build index d37025ea..609e22ef 100644 --- a/meson.build +++ b/meson.build @@ -4,7 +4,7 @@ project('PyORBIT', default_options: [ 'buildtype=release', 'b_lto=true', - 'cpp_std=c++11', + 'cpp_std=c++17', 'warning_level=3', ], ) diff --git a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc index 0a361326..bdcf5ae9 100644 --- a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc +++ b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc @@ -1,525 +1,384 @@ -#include "BunchTwissAnalysis.hh" +#include "orbit/BunchDiagnostics/BunchTwissAnalysis.hh" -#include -#include -#include -#include +#include "orbit/Bunch.hh" +#include "orbit/ParticlesAttributes/ParticleMacroSize.hh" +#include "orbit/SyncPart.hh" -#include "ParticleMacroSize.hh" -#include "SyncPart.hh" +#include +#include -/** Constructor */ -BunchTwissAnalysis::BunchTwissAnalysis(): CppPyWrapper(NULL) +/** Performs the Twiss analysis of the bunch */ +void BunchTwissAnalysis::analyzeBunch(Bunch* bunch) { - - avg_arr = (double* ) malloc (6*sizeof(double)); - corr_arr = (double* ) malloc (36*sizeof(double)); - - avg_arr_MPI = (double* ) malloc (6*sizeof(double)); - corr_arr_MPI = (double* ) malloc (36*sizeof(double)); - - for(int i = 0; i < 6; i++){ - avg_arr[i] = 0.; - avg_arr_MPI[i] = 0.; - } - - for(int i = 0; i < 6; i++){ - for(int j = 0; j < 6; j++){ - corr_arr[i+6*j] = 0.; - corr_arr_MPI[i+6*j] = 0.; - } - } - - count = 0; - _order = 0; - - total_macrosize = 0.; + static std::atomic_flag warned = ATOMIC_FLAG_INIT; + if (!warned.test_and_set(std::memory_order_relaxed)) { + std::cerr << "[WARNING] BunchTwissAnalysis::analyzeBunch is deprecated. Use " + "computeBunchMoments() instead.\n"; + } + computeBunchMoments(bunch); } -/** Destructor */ -BunchTwissAnalysis::~BunchTwissAnalysis() +/// @brief Templated implementation of computeBunchMoments specialized on macrosize and dispersion +/// flags. +/// +/// Computes low order (order <= 2) moments in the first loop over the bunch followed by an early +/// return. If higher order moments (order >= 3) are requested, then a second loop over the bunch +/// will be triggered to compute the remaining moments. +template +void BunchTwissAnalysis::computeBunchMomentsImpl(Bunch* bunch, bool normalize, bool emitnormflag) { - free(avg_arr); - free(corr_arr); - free(avg_arr_MPI); - free(corr_arr_MPI); + const auto size = bunch->getSize(); + + if (size == 0) { + return; + } + + auto* macrosize_attr = + HasMacrosizeAttr ? static_cast(bunch->getParticleAttributes("macrosize")) + : nullptr; + auto** coords = bunch->coordArr(); + + for (int ip = 0; ip < size; ++ip) { + const double w = HasMacrosizeAttr ? macrosize_attr->macrosize(ip) : 1.0; + + if constexpr (HasMacrosizeAttr) { + total_macrosize_ += w; + } + + for (int i = 0; i < N; ++i) { + avg_arr[i] += w * coords[ip][i]; + for (int j = 0; j <= i; ++j) { + cov_arr[covIdx(i, j)] += w * coords[ip][i] * coords[ip][j]; + } + } + } + + auto* comm = bunch->getMPI_Comm_Local()->comm; + + if constexpr (HasMacrosizeAttr) { + ORBIT_MPI_Allreduce(MPI_IN_PLACE, &total_macrosize_, 1, MPI_DOUBLE, MPI_SUM, comm); + } + else { + total_macrosize_ = bunch->getSizeGlobal(); + } + + ORBIT_MPI_Allreduce(MPI_IN_PLACE, avg_arr.data(), N, MPI_DOUBLE, MPI_SUM, comm); + ORBIT_MPI_Allreduce(MPI_IN_PLACE, cov_arr.data(), NN, MPI_DOUBLE, MPI_SUM, comm); + + // = - uhat*vhat + for (int i = 0; i < N; ++i) { + avg_arr[i] /= total_macrosize_; + for (int j = 0; j < i + 1; ++j) { + const int idx = covIdx(i, j); + cov_arr[idx] = cov_arr[idx] / total_macrosize_ - avg_arr[i] * avg_arr[j]; + cov_arr[covIdx(j, i)] = cov_arr[covIdx(i, j)]; // covariance matrix is symmetric. + } + } + + double inv_xbt = 1.0; + double inv_ybt = 1.0; + + if (normalize) { + inv_xbt /= std::sqrt(getBeta(0) * (emitnormflag ? getEmittance(0) : 1.0)); + inv_ybt /= std::sqrt(getBeta(1) * (emitnormflag ? getEmittance(1) : 1.0)); + } + + const double disp_scale = + Dispersion ? getDispersion(0) / (bunch_kinenergy_ + bunch_mass_) / (bunch_beta_ * bunch_beta_) + : 0.0; + const double xAvg = avg_arr[0] - disp_scale * avg_arr[5]; + const double yAvg = avg_arr[2]; + + const int nMoments = order_ + 1; + momentXY_[momentIdx(0, 0, nMoments)] = 1.0; + momentXY_[momentIdx(1, 0, nMoments)] = xAvg; + momentXY_[momentIdx(0, 1, nMoments)] = yAvg; + + if (order_ < 2) { + return; + } + + const double x2 = cov_arr[covIdx(0, 0)] - 2.0 * disp_scale * cov_arr[covIdx(0, 5)] + + disp_scale * disp_scale * cov_arr[covIdx(5, 5)]; + const double xy = cov_arr[covIdx(0, 2)] - disp_scale * cov_arr[covIdx(2, 5)]; + const double y2 = cov_arr[covIdx(2, 2)]; + momentXY_[momentIdx(2, 0, nMoments)] = x2 * inv_xbt * inv_xbt; + momentXY_[momentIdx(1, 1, nMoments)] = xy * inv_xbt * inv_ybt; + momentXY_[momentIdx(0, 2, nMoments)] = y2 * inv_ybt * inv_ybt; + + if (order_ == 2) { + return; + } + + for (int ip = 0; ip < size; ++ip) { + const double w = HasMacrosizeAttr ? macrosize_attr->macrosize(ip) : 1.0; + + const double dx = + Dispersion ? coords[ip][0] - xAvg - disp_scale * coords[ip][5] : coords[ip][0] - xAvg; + const double dy = coords[ip][2] - yAvg; + + const double normX = dx * inv_xbt; + const double normY = dy * inv_ybt; + + double ny = 1.0; + for (int j = 0; j < order_; ++j) { + int i_start = (j < 3) ? 3 - j : 0; + double nx = 1.0; + for (int ii = 0; ii < i_start; ++ii) { + nx *= normX; + } + for (int i = i_start; i < nMoments - j; ++i) { + momentXY_[momentIdx(i, j, nMoments)] += w * nx * ny; + nx *= normX; + } + ny *= normY; + } + } + + ORBIT_MPI_Allreduce( + MPI_IN_PLACE, + momentXY_.data(), + nMoments * nMoments, + MPI_DOUBLE, + MPI_SUM, + comm + ); + + for (int i = 0; i < nMoments; ++i) { + for (int j = 0; j < nMoments - i; ++j) { + if (i + j <= 2) { + continue; + } + momentXY_[momentIdx(i, j, nMoments)] /= total_macrosize_; + } + } } -/** Performs the Twiss analysis of the bunch */ -void BunchTwissAnalysis::analyzeBunch(Bunch* bunch){ - - //initialization - for(int i = 0; i < 6; i++){ - avg_arr[i] = 0.; - } - - for(int i = 0; i < 6; i++){ - for(int j = 0; j < 6; j++){ - corr_arr[i+6*j] = 0.; - } - } - count = 0; - total_macrosize = 0.; - - bunch->compress(); - double m_size = 0.; - int nParts = bunch->getSize(); - count += nParts; - double** part_coord_arr = bunch->coordArr(); - int has_msize = bunch->hasParticleAttributes("macrosize"); - if(has_msize > 0){ - ParticleMacroSize* macroSizeAttr = (ParticleMacroSize*) bunch->getParticleAttributes("macrosize"); - double m_size = 0.; - for(int ip = 0; ip < nParts; ip++){ - m_size = macroSizeAttr->macrosize(ip); - total_macrosize += m_size; - for(int i = 0; i < 6; i++){ - avg_arr[i] += m_size*part_coord_arr[ip][i]; - } - - for(int i = 0; i < 6; i++){ - for(int j = 0; j < i+1; j++){ - corr_arr[i+6*j] += m_size*part_coord_arr[ip][i]*part_coord_arr[ip][j]; - } - } - } - } else { - m_size = 1.0; - for(int ip = 0; ip < nParts; ip++){ - for(int i = 0; i < 6; i++){ - avg_arr[i] += part_coord_arr[ip][i]; - } - - for(int i = 0; i < 6; i++){ - for(int j = 0; j < i+1; j++){ - corr_arr[i+6*j] += part_coord_arr[ip][i]*part_coord_arr[ip][j]; - } - } - } - total_macrosize += nParts*m_size; - for(int i = 0; i < 6; i++){ - avg_arr[i] *= m_size; - } - for(int i = 0; i < 6; i++){ - for(int j = 0; j < i+1; j++){ - corr_arr[i+6*j] *= m_size; - } - } - } - - int count_MPI = 0; - ORBIT_MPI_Allreduce(&count,&count_MPI,1,MPI_INT,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - count = count_MPI; - - double total_macrosize_MPI = 0.; - ORBIT_MPI_Allreduce(&total_macrosize,&total_macrosize_MPI,1,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - total_macrosize = total_macrosize_MPI; - - ORBIT_MPI_Allreduce(avg_arr,avg_arr_MPI,6,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - ORBIT_MPI_Allreduce(corr_arr,corr_arr_MPI,36,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - - if(fabs(total_macrosize) > 0.){ - for(int i = 0; i < 6; i++){ - avg_arr[i] = avg_arr_MPI[i]/total_macrosize; - } - - for(int i = 0; i < 6; i++){ - for(int j = 0; j < i+1; j++){ - corr_arr[i+6*j] = corr_arr_MPI[i+6*j]/total_macrosize; - corr_arr[j+6*i] = corr_arr_MPI[i+6*j]/total_macrosize; - } - } - } - - SyncPart* syncPart = bunch->getSyncPart(); - - bunch_momentum = syncPart->getMomentum(); - bunch_beta = syncPart->getBeta(); - bunch_gamma = syncPart->getGamma(); - bunch_kinenergy = syncPart->getEnergy(); - bunch_mass = syncPart->getMass(); - -} - - -/** Performs the bunch moments computations */ -void BunchTwissAnalysis::computeBunchMoments(Bunch* bunch, int order, int dispersionflag, int emitnormflag){ - _order = order; - int i = 0; - int j = 0; - - bunch->compress(); - - double dispterm = 0.; - double xbetaterm = 0.; - double ybetaterm = 0.; - double m_size = 0.; - double xAvg = 0.; - double yAvg = 0.; - total_macrosize = 0; //Total macrosize (can different than number of macroparticles if m_size is specified) - int nParts = bunch->getSize(); - double total_macrosize_MPI = 0.; - double** part_coord_arr = bunch->coordArr(); - int has_msize = bunch->hasParticleAttributes("macrosize"); - - analyzeBunch(bunch); - - if(dispterm > 0){ - if(has_msize > 0){ - ParticleMacroSize* macroSizeAttr = (ParticleMacroSize*) bunch->getParticleAttributes("macrosize"); - double m_size = 0.; - for(int ip = 0; ip < nParts; ip++){ - m_size = macroSizeAttr->macrosize(ip); - total_macrosize += m_size; - if (dispersionflag > 0) { - dispterm = getDispersion(0) * part_coord_arr[ip][5] / (bunch_kinenergy + bunch_mass) / (bunch_beta*bunch_beta); - } - xAvg += m_size*(part_coord_arr[ip][0] - dispterm); - } - } else { - m_size = 1.0; - for(int ip = 0; ip < nParts; ip++){ - if (dispersionflag > 0) { - dispterm = getDispersion(0) * part_coord_arr[ip][5] / (bunch_kinenergy + bunch_mass) / (bunch_beta*bunch_beta); - } - xAvg += part_coord_arr[ip][0] - dispterm; - } - total_macrosize += nParts*m_size; - xAvg *= m_size; - } - - - ORBIT_MPI_Allreduce(&total_macrosize,&total_macrosize_MPI,1,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - total_macrosize = total_macrosize_MPI; - - double xAvg_MPI = 0; - ORBIT_MPI_Allreduce(&xAvg,&xAvg_MPI,1,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - if(fabs(total_macrosize) > 0.){ - xAvg = xAvg_MPI/total_macrosize; - } - } - else{ - xAvg = getAverage(0); - } - yAvg = getAverage(2); - - double momX [order+1]; - double momY [order+1]; - - momentXY = new double*[order+1]; - - for(i=0; i < order+1; i++) - momentXY[i] = new double[order+1]; - - //initialization - for (int n=0; n < _order+1;n++){ - momX[n] = 0.; - momY[n] = 0.; - } - - for (int n=0; n<_order+1;n++) - for (int m=0; m<_order+1;m++) - momentXY[n][m]=0.; - - momX[0]=1.0; - momY[0]=1.0; - - bunch->compress(); - total_macrosize = 0.; - if(has_msize > 0){ - ParticleMacroSize* macroSizeAttr = (ParticleMacroSize*) bunch->getParticleAttributes("macrosize"); - double m_size = 0.; - for(int ip = 0; ip < nParts; ip++){ - - m_size = macroSizeAttr->macrosize(ip); - total_macrosize += m_size; - - if (dispersionflag > 0) { - dispterm = getDispersion(0) * part_coord_arr[ip][5] / (bunch_kinenergy + bunch_mass) / (bunch_beta*bunch_beta); - } - - if(emitnormflag > 0){ - xbetaterm = sqrt(getBeta(0)*getEmittance(0)); - ybetaterm = sqrt(getBeta(1)*getEmittance(1)); - } - else{ - xbetaterm = sqrt(getBeta(0)); - ybetaterm = sqrt(getBeta(1)); - } - - for(i = 0; i < _order; i++) - momX[i+1] = momX[i]*m_size*((part_coord_arr[ip][0] - dispterm) - xAvg); - - for(i = 0; i< _order; i++) - momY[i+1] = momY[i]*m_size*(part_coord_arr[ip][2] - yAvg); - - for(j = 0; j<_order; j++) - for(i=0 ; i< _order+1-j; i++){ - momentXY[i][j] += momX[i]/pow(xbetaterm, double(i)) * momY[j]/pow(ybetaterm, double(j)); - momentXY[i][j] += momX[i] * momY[j]; - } - } - - } else { - m_size = 1.0; - for(int ip = 0; ip < nParts; ip++){ - - if (dispersionflag > 0) { - dispterm = getDispersion(0) * part_coord_arr[ip][5] / (bunch_kinenergy + bunch_mass) / (bunch_beta*bunch_beta); - } - - if(emitnormflag > 0){ - xbetaterm = sqrt(getBeta(0)*getEmittance(0)); - ybetaterm = sqrt(getBeta(1)*getEmittance(1)); - } - else{ - xbetaterm = sqrt(getBeta(0)); - ybetaterm = sqrt(getBeta(1)); - } - - for(i = 0; i < _order; i++) - momX[i+1] = momX[i]*((part_coord_arr[ip][0] - dispterm) - xAvg); - - for(i = 0; i< _order; i++) - momY[i+1] = momY[i]*(part_coord_arr[ip][2] - yAvg); - - for(j = 0; j<_order; j++) - for(i=0 ; i< _order+1-j; i++) - momentXY[i][j] += momX[i]/pow(xbetaterm, double(i)) * momY[j]/pow(ybetaterm, double(j)); - momentXY[i][j] += momX[i] * momY[j]; - - } - - total_macrosize += nParts*m_size; - - } - - ORBIT_MPI_Allreduce(&total_macrosize,&total_macrosize_MPI,1,MPI_DOUBLE,MPI_SUM,bunch->getMPI_Comm_Local()->comm); - total_macrosize = total_macrosize_MPI; - - //if( nMPIsize_ > 1){ - double* buff_0 = (double *) malloc (sizeof(double)*(_order+1)*(_order+1)); - double* buff_1 = (double *) malloc (sizeof(double)*(_order+1)*(_order+1)); - int count = 0; - for(j=0; j<_order+1; j++){ - for(i=0 ; i< _order+1-j; i++){ - buff_0[count]= momentXY[i][j]; - count++; - } - } - - //MPI_Allreduce(buff_0, buff_1, count, MPI_DOUBLE,MPI_SUM,MPI_COMM_WORLD); - - ORBIT_MPI_Allreduce(buff_0, buff_1, count, MPI_DOUBLE, MPI_SUM, bunch->getMPI_Comm_Local()->comm); - - count = 0; - for(j=0; j<_order+1; j++){ - for(i=0 ; i< _order+1-j; i++){ - momentXY[i][j] = buff_1[count]; - count++; - } - } - - free(buff_0); - free(buff_1); - - if(fabs(total_macrosize) > 0.){ - for(i=0; i< _order+1; i++) - for(j=0; j< _order+1-i ; j++) - momentXY[i][j] /= total_macrosize; - } - - momentXY[0][0] = 1.; // 0th moment - momentXY[1][0] = xAvg; - momentXY[0][1] = yAvg; - -//return momentXY; - +void BunchTwissAnalysis::computeBunchMoments( + Bunch* bunch, + int order, + bool normalize, + bool emitnormflag, + bool dispersionflag +) +{ + avg_arr.fill(0.0); + cov_arr.fill(0.0); + total_macrosize_ = 0.0; + + // the number of unique pairs i, j for moments up to a max order, n: + // (n+1)(n+2)/2 + momentXY_.assign((order + 1)*(order + 2)/2, 0.0); + + bunch->compress(); + count_ = bunch->getSizeGlobal(); + order_ = order; + + SyncPart* syncPart = bunch->getSyncPart(); + bunch_momentum_ = syncPart->getMomentum(); + bunch_beta_ = syncPart->getBeta(); + bunch_gamma_ = syncPart->getGamma(); + bunch_kinenergy_ = syncPart->getEnergy(); + bunch_mass_ = syncPart->getMass(); + + if (bunch->hasParticleAttributes("macrosize")) { + if (dispersionflag) { + computeBunchMomentsImpl(bunch, normalize, emitnormflag); + } + else { + computeBunchMomentsImpl(bunch, normalize, emitnormflag); + } + } + else { + if (dispersionflag) { + computeBunchMomentsImpl(bunch, normalize, emitnormflag); + } + else { + computeBunchMomentsImpl(bunch, normalize, emitnormflag); + } + } } -double BunchTwissAnalysis::getCorrelation(int ic, int jc){ - if(ic < 0 || ic > 5 || jc < 0 || jc >5) return 0.; - return (corr_arr[ic+6*jc] - avg_arr[ic]* avg_arr[jc]); +double BunchTwissAnalysis::getCorrelation(int ic, int jc) const +{ + if (ic < 0 || ic > 5 || jc < 0 || jc > 5) { + return 0.; + } + return cov_arr[covIdx(ic, jc)]; } -/** Returns the XY moment of the beam */ -double BunchTwissAnalysis::getBunchMoment(int i, int j){ - if(i < 0 || i > _order || j < 0 || j > _order) return 0.; - return momentXY[i][j]; +double BunchTwissAnalysis::getBunchMoment(int i, int j) const +{ + if (i < 0 || i > order_ || j < 0 || j > order_) { + return 0.; + } + return momentXY_[momentIdx(i, j, order_ + 1)]; } -/** Returns the average value for coordinate with index ic */ -double BunchTwissAnalysis::getAverage(int ic){ - if(ic < 0 || ic > 5 ) return 0.; - return avg_arr[ic]; +double BunchTwissAnalysis::getAverage(int ic) const +{ + if (ic < 0 || ic > 5) { + return 0.; + } + return avg_arr[ic]; } -/** Returns the total number of analysed macroparticles */ -int BunchTwissAnalysis::getGlobalCount(){ - return count; +int BunchTwissAnalysis::getGlobalCount() const +{ + return count_; } -/** Returns the total macrosize */ -double BunchTwissAnalysis::getGlobalMacrosize(){ - return total_macrosize; +double BunchTwissAnalysis::getGlobalMacrosize() const +{ + return total_macrosize_; } -/** Returns the emittance for index 0,1,2 - x,y,z planes. */ -double BunchTwissAnalysis::getEmittance(int ic) +double BunchTwissAnalysis::getEmittance(int ic) const { - // for x and y the pure betatron emittance is computed (subtracting the dispersive contribution) - if(ic < 0 || ic > 2 ) return 0.; - double x2_avg = fabs(this->getCorrelation(2*ic,2*ic)); - double xp2_avg = fabs(this->getCorrelation(2*ic+1,2*ic+1)); - double x_xp_avg = this->getCorrelation(2*ic,2*ic+1); - double x_dE_avg = this->getCorrelation(2*ic,5); - double xp_dE_avg = this->getCorrelation(2*ic+1,5); - double dE2_avg = fabs(this->getCorrelation(5,5)); - double emitt_rms; - if(ic==2 or dE2_avg==0){ - emitt_rms = sqrt(fabs(x2_avg*xp2_avg - x_xp_avg*x_xp_avg)); - } else { - emitt_rms = sqrt(fabs( (x2_avg - x_dE_avg*x_dE_avg/dE2_avg) * (xp2_avg - xp_dE_avg*xp_dE_avg/dE2_avg) - - (x_xp_avg - x_dE_avg*xp_dE_avg/dE2_avg) * (x_xp_avg - x_dE_avg*xp_dE_avg/dE2_avg) )); - } - return emitt_rms; + if (ic < 0 || ic > 2) { + return 0.; + } + double x2_avg = std::abs(getCorrelation(2 * ic, 2 * ic)); + double xp2_avg = std::abs(getCorrelation(2 * ic + 1, 2 * ic + 1)); + double x_xp_avg = getCorrelation(2 * ic, 2 * ic + 1); + double x_dE_avg = getCorrelation(2 * ic, 5); + double xp_dE_avg = getCorrelation(2 * ic + 1, 5); + double dE2_avg = std::abs(getCorrelation(5, 5)); + if (ic == 2 || dE2_avg == 0) { + return std::sqrt(std::abs(x2_avg * xp2_avg - x_xp_avg * x_xp_avg)); + } + double x1 = x2_avg - x_dE_avg * x_dE_avg / dE2_avg; + double x2 = xp2_avg - xp_dE_avg * xp_dE_avg / dE2_avg; + double x3 = x_xp_avg - x_dE_avg * xp_dE_avg / dE2_avg; + return std::sqrt(std::abs(x1 * x2 - x3 * x3)); } -/** Returns the normalized betatron emittance for index 0,1 - x,y planes. */ -double BunchTwissAnalysis::getEmittanceNormalized(int ic) +double BunchTwissAnalysis::getEmittanceNormalized(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - if(ic == 2) return this->getEmittance(ic); - return this->getEmittance(ic) * bunch_gamma * bunch_beta; + if (ic < 0 || ic > 2) { + return 0.; + } + if (ic == 2) { + return getEmittance(ic); + } + return getEmittance(ic) * bunch_gamma_ * bunch_beta_; } -/** Returns Twiss alpha (without dispersive part for x,y) for index 0,1,2 - x,y,z planes.*/ -double BunchTwissAnalysis::getAlpha(int ic) +double BunchTwissAnalysis::getAlpha(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - double x_xp_avg = this->getCorrelation(2*ic,2*ic+1); - double alpha; - if(ic == 2){ - alpha = - x_xp_avg/this->getEmittance(ic); - } else { - double x_dE_avg = this->getCorrelation(2*ic, 5); - double xp_dE_avg = this->getCorrelation(2*ic+1, 5); - double dE2_avg = fabs(this->getCorrelation(5, 5)); - alpha = -(x_xp_avg - x_dE_avg * xp_dE_avg / dE2_avg) / this->getEmittance(ic); - } - return alpha; + if (ic < 0 || ic > 2) { + return 0.; + } + double x_xp_avg = getCorrelation(2 * ic, 2 * ic + 1); + if (ic == 2) { + return -x_xp_avg / getEmittance(ic); + } + double x_dE_avg = getCorrelation(2 * ic, 5); + double xp_dE_avg = getCorrelation(2 * ic + 1, 5); + double dE2_avg = std::abs(getCorrelation(5, 5)); + return -(x_xp_avg - x_dE_avg * xp_dE_avg / dE2_avg) / getEmittance(ic); } -/** Returns Twiss beta (without dispersive part for x,y) for index 0,1,2 - x,y,z planes.*/ -double BunchTwissAnalysis::getBeta(int ic) +double BunchTwissAnalysis::getBeta(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - double x2_avg = fabs(this->getCorrelation(2*ic,2*ic)); - - double beta; - if(ic == 2){ - beta = x2_avg/this->getEmittance(ic); - } else { - double x_dE_avg = this->getCorrelation(2*ic, 5); - double dE2_avg = fabs(this->getCorrelation(5, 5)); - beta = (x2_avg - x_dE_avg * x_dE_avg / dE2_avg) / this->getEmittance(ic); - } - return beta; + if (ic < 0 || ic > 2) { + return 0.; + } + double x2_avg = std::abs(getCorrelation(2 * ic, 2 * ic)); + if (ic == 2) { + return x2_avg / getEmittance(ic); + } + double x_dE_avg = getCorrelation(2 * ic, 5); + double dE2_avg = std::abs(getCorrelation(5, 5)); + return (x2_avg - x_dE_avg * x_dE_avg / dE2_avg) / getEmittance(ic); } -/** Returns Twiss gamma (without dispersive part for x,y) for index 0,1,2 - x,y,z planes.*/ -double BunchTwissAnalysis::getGamma(int ic) +double BunchTwissAnalysis::getGamma(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - double xp2_avg = fabs(this->getCorrelation(2*ic+1,2*ic+1)); - double gamma; - if(ic == 2){ - gamma = xp2_avg/this->getEmittance(ic); - } else { - double xp_dE_avg = this->getCorrelation(2*ic+1, 5); - double dE2_avg = fabs(this->getCorrelation(5, 5)); - gamma = (xp2_avg - xp_dE_avg * xp_dE_avg / dE2_avg) / this->getEmittance(ic); - } - return gamma; + if (ic < 0 || ic > 2) { + return 0.; + } + double xp2_avg = std::abs(getCorrelation(2 * ic + 1, 2 * ic + 1)); + if (ic == 2) { + return xp2_avg / getEmittance(ic); + } + double xp_dE_avg = getCorrelation(2 * ic + 1, 5); + double dE2_avg = std::abs(getCorrelation(5, 5)); + return (xp2_avg - xp_dE_avg * xp_dE_avg / dE2_avg) / getEmittance(ic); } -/** Returns Twiss dispersion function for index 0,1 - x,y planes.*/ -double BunchTwissAnalysis::getDispersion(int ic) +double BunchTwissAnalysis::getDispersion(int ic) const { - if(ic < 0 || ic > 1 ) return 0.; - double x_dE_avg = this->getCorrelation(2*ic, 5); - double dE2_avg = fabs(this->getCorrelation(5, 5)); - double dispersion = x_dE_avg/dE2_avg * bunch_momentum * bunch_beta; - return dispersion; + if (ic < 0 || ic > 1) { + return 0.; + } + double x_dE_avg = getCorrelation(2 * ic, 5); + double dE2_avg = std::abs(getCorrelation(5, 5)); + return x_dE_avg / dE2_avg * bunch_momentum_ * bunch_beta_; } -/** Returns Twiss dispersion_prime function for index 0,1 - x,y planes.*/ -double BunchTwissAnalysis::getDispersionDerivative(int ic) +double BunchTwissAnalysis::getDispersionDerivative(int ic) const { - if(ic < 0 || ic > 1 ) return 0.; - double xp_dE_avg = this->getCorrelation(2*ic+1, 5); - double dE2_avg = fabs(this->getCorrelation(5, 5)); - double dispersion_prime = xp_dE_avg/dE2_avg * bunch_momentum * bunch_beta; - return dispersion_prime; + if (ic < 0 || ic > 1) { + return 0.; + } + double xp_dE_avg = getCorrelation(2 * ic + 1, 5); + double dE2_avg = std::abs(getCorrelation(5, 5)); + return xp_dE_avg / dE2_avg * bunch_momentum_ * bunch_beta_; } -/** Returns the effective emittance for index 0,1 - x,y planes. */ -double BunchTwissAnalysis::getEffectiveEmittance(int ic) +double BunchTwissAnalysis::getEffectiveEmittance(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - double x_avg = this->getAverage(2*ic); - double xp_avg = this->getAverage(2*ic+1); - double x2_avg = fabs(this->getCorrelation(2*ic,2*ic)); - double xp2_avg = fabs(this->getCorrelation(2*ic+1,2*ic+1)); - double x_xp_avg = this->getCorrelation(2*ic,2*ic+1); - double emitt_rms = sqrt(fabs(x2_avg*xp2_avg - x_xp_avg*x_xp_avg)); - return emitt_rms; + if (ic < 0 || ic > 2) { + return 0.; + } + double x2_avg = std::abs(getCorrelation(2 * ic, 2 * ic)); + double xp2_avg = std::abs(getCorrelation(2 * ic + 1, 2 * ic + 1)); + double x_xp_avg = getCorrelation(2 * ic, 2 * ic + 1); + return std::sqrt(std::abs(x2_avg * xp2_avg - x_xp_avg * x_xp_avg)); } -/** Returns effective Twiss alpha for index 0,1 - x,y planes.*/ -double BunchTwissAnalysis::getEffectiveAlpha(int ic) +double BunchTwissAnalysis::getEffectiveAlpha(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - double x_avg = this->getAverage(2*ic); - double xp_avg = this->getAverage(2*ic+1); - double x2_avg = fabs(this->getCorrelation(2*ic,2*ic)); - double xp2_avg = fabs(this->getCorrelation(2*ic+1,2*ic+1)); - double x_xp_avg = this->getCorrelation(2*ic,2*ic+1); - double emitt2_rms = x2_avg*xp2_avg - x_xp_avg*x_xp_avg; - if( emitt2_rms <= 0.) return 0.; - double emitt_rms = sqrt(emitt2_rms); - double alpha = - x_xp_avg/emitt_rms; - return alpha; + if (ic < 0 || ic > 2) { + return 0.; + } + double x2_avg = std::abs(getCorrelation(2 * ic, 2 * ic)); + double xp2_avg = std::abs(getCorrelation(2 * ic + 1, 2 * ic + 1)); + double x_xp_avg = getCorrelation(2 * ic, 2 * ic + 1); + double emitt2_rms = x2_avg * xp2_avg - x_xp_avg * x_xp_avg; + if (emitt2_rms <= 0.) { + return 0.; + } + return -x_xp_avg / std::sqrt(emitt2_rms); } -/** Returns effective Twiss beta for index 0,1 - x,y planes.*/ -double BunchTwissAnalysis::getEffectiveBeta(int ic) +double BunchTwissAnalysis::getEffectiveBeta(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - double x_avg = this->getAverage(2*ic); - double xp_avg = this->getAverage(2*ic+1); - double x2_avg = fabs(this->getCorrelation(2*ic,2*ic)); - double xp2_avg = fabs(this->getCorrelation(2*ic+1,2*ic+1)); - double x_xp_avg = this->getCorrelation(2*ic,2*ic+1); - double emitt2_rms = x2_avg*xp2_avg - x_xp_avg*x_xp_avg; - if( emitt2_rms <= 0.) return 0.; - double emitt_rms = sqrt(emitt2_rms); - double beta = x2_avg/emitt_rms; - return beta; + if (ic < 0 || ic > 2) { + return 0.; + } + double x2_avg = std::abs(getCorrelation(2 * ic, 2 * ic)); + double xp2_avg = std::abs(getCorrelation(2 * ic + 1, 2 * ic + 1)); + double x_xp_avg = getCorrelation(2 * ic, 2 * ic + 1); + double emitt2_rms = x2_avg * xp2_avg - x_xp_avg * x_xp_avg; + if (emitt2_rms <= 0.) { + return 0.; + } + return x2_avg / std::sqrt(emitt2_rms); } -/** Returns effective Twiss gamma for index 0,1 - x,y planes.*/ -double BunchTwissAnalysis::getEffectiveGamma(int ic) +double BunchTwissAnalysis::getEffectiveGamma(int ic) const { - if(ic < 0 || ic > 2 ) return 0.; - double x_avg = this->getAverage(2*ic); - double xp_avg = this->getAverage(2*ic+1); - double x2_avg = fabs(this->getCorrelation(2*ic,2*ic)); - double xp2_avg = fabs(this->getCorrelation(2*ic+1,2*ic+1)); - double x_xp_avg = this->getCorrelation(2*ic,2*ic+1); - double emitt2_rms = x2_avg*xp2_avg - x_xp_avg*x_xp_avg; - if( emitt2_rms <= 0.) return DBL_MAX; - double emitt_rms = sqrt(emitt2_rms); - double gamma = xp2_avg/emitt_rms; - return gamma; + if (ic < 0 || ic > 2) { + return 0.; + } + double xp2_avg = std::abs(getCorrelation(2 * ic + 1, 2 * ic + 1)); + double x2_avg = std::abs(getCorrelation(2 * ic, 2 * ic)); + double x_xp_avg = getCorrelation(2 * ic, 2 * ic + 1); + double emitt2_rms = x2_avg * xp2_avg - x_xp_avg * x_xp_avg; + if (emitt2_rms <= 0.) { + return std::numeric_limits::max(); + } + return xp2_avg / std::sqrt(emitt2_rms); } diff --git a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.hh b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.hh index deaa20ee..b24c962b 100644 --- a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.hh +++ b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.hh @@ -1,124 +1,137 @@ #ifndef BUNCH_TWISS_ANALYSIS_H #define BUNCH_TWISS_ANALYSIS_H -//pyORBIT utils -#include "CppPyWrapper.hh" +#include "utils/CppPyWrapper.hh" -#include "Bunch.hh" +#include +#include -using namespace std; +class Bunch; +class ParticleMacroSize; /** - The BunchTwissAnalysis class calculates the average of 6D coordinates and they correlations. - As results it returns the Twiss parameters for each plane. -*/ - -class BunchTwissAnalysis: public OrbitUtils::CppPyWrapper + * Calculates the average of 6D coordinates and their correlations. + * Returns the Twiss parameters for each plane. + */ +class BunchTwissAnalysis : public OrbitUtils::CppPyWrapper { - public: - - /** Constructor*/ - BunchTwissAnalysis(); - - /** Destructor */ - virtual ~BunchTwissAnalysis(); - - /** Performs the Twiss analysis of the bunch */ - void analyzeBunch(Bunch* bunch); - - /** - * @brief Returns the centered correlation between two components. - * - * \f$\langle (x-\langle x\rangle)(y-\langle y\rangle)\rangle - * = \langle x y\rangle - \langle x\rangle\langle y\rangle\f$ - * - * @param ic Index of the first component (x). - * @param jc Index of the second component (y). - * @return The centered correlation value as a double. - */ - double getCorrelation(int ic, int jc); - - /** Returns the average value for coordinate with index ic */ - double getAverage(int ic); - - /** Returns the total number of analysed macroparticles */ - int getGlobalCount(); - - /** Returns the total macrosize */ - double getGlobalMacrosize(); - - /** Returns the emittance for index 0,1,2 - x,y,z planes. */ - double getEffectiveEmittance(int ic); - - /** Returns Twiss alpha for index 0,1,2 - x,y,z planes.*/ - double getEffectiveAlpha(int ic); - - /** Returns Twiss beta for index 0,1,2 - x,y,z planes.*/ - double getEffectiveBeta(int ic); - - /** Returns Twiss gamma for index 0,1,2 - x,y,z planes.*/ - double getEffectiveGamma(int ic); - - /** Returns the pure betatron emittance for index 0,1 - x,y planes. */ - double getEmittance(int ic); - - /** Returns the normalized pure betatron emittance for index 0,1 - x,y planes. */ - double getEmittanceNormalized(int ic); - - /** Returns Twiss alpha without dispersive part for index 0,1 - x,y planes.*/ - double getAlpha(int ic); - - /** Returns Twiss beta without dispersive part for index 0,1 - x,y planes.*/ - double getBeta(int ic); - - /** Returns Twiss gamma without dispersive part for index 0,1 - x,y planes.*/ - double getGamma(int ic); - - /** Returns Twiss dispersion function for index 0,1 - x,y planes.*/ - double getDispersion(int ic); - - /** Returns Twiss dispersion_prime function for index 0,1 - x,y planes.*/ - double getDispersionDerivative(int ic); - - /** Computes the XY moments of the bunch up to a prescribed order */ - void computeBunchMoments(Bunch* bunch, int order, int dispersionflag, int emitnormflag); - - /** Returns the XY moment of the beam */ - double getBunchMoment(int i, int j); - - private: - - /** Number of points accounted */ - int count; - - /** total macrosize accounted */ - double total_macrosize; - - /** array with average values for 6D coordinates */ - double* avg_arr; - double* avg_arr_MPI; - - /** array with correlations between 6D coordinates. - It is a packed 1D array 6*6 = 36. - To reach i,j = 0-5 coerrelation use [i+6*j] - It is excessive because = etc.*/ - double* corr_arr; - double* corr_arr_MPI; - - /** array with XY Moments */ - double** momentXY; +public: + BunchTwissAnalysis() = default; + virtual ~BunchTwissAnalysis() = default; + + BunchTwissAnalysis(const BunchTwissAnalysis&) = delete; + BunchTwissAnalysis& operator=(const BunchTwissAnalysis&) = delete; + BunchTwissAnalysis(BunchTwissAnalysis&&) = delete; + BunchTwissAnalysis& operator=(BunchTwissAnalysis&&) = delete; + + [[deprecated("Use computeBunchMoments instead.")]] + void analyzeBunch(Bunch* bunch); - /** order for the moments **/ - double _order; + /** + * @brief Computes central moments of the bunch up to the given order. + * + * @param bunch The bunch to analyze (compressed before use). + * @param order Maximum moment order (default 2). + * @param normalize Whether to normalize the moments to \f$\sqrt{\beta_{x,y}}\f$. + * @param emitnormflag If true, moments are normalized to \f$\sqrt{\beta_{x,y} + * \varepsilon{x,y}}\f$. Implies normalize=true. (default false). + * @param dispersionflag If true, x-coordinates are corrected for dispersion (default false). + */ + void computeBunchMoments( + Bunch* bunch, + int order = 2, + bool normalize = false, + bool emitnormflag = false, + bool dispersionflag = false + ); + + /** + * @brief Returns the centered covariance for a pair of coordinates. + * + * \f$\langle \left(u - \langle u \rangle \right) \left( v - \langle v \rangle \right) \rangle \f$ + * + * @param ic Integer corresponding to the first coordinate, \f$u\f$. + * @param jc Integer corresponding to the second coordinate, \f$u\f$. + **/ + double getCorrelation(int ic, int jc) const; + + /** Returns the average value for coordinate with index ic. */ + double getAverage(int ic) const; + + /** Returns the total number of analysed macroparticles. */ + int getGlobalCount() const; + + /** Returns the total macrosize. */ + double getGlobalMacrosize() const; + + /** Returns the XY moment of the beam. */ + double getBunchMoment(int i, int j) const; + + /** Returns the pure betatron emittance for index 0,1,2 - x,y,z planes. */ + double getEmittance(int ic) const; + + /** Returns the normalized pure betatron emittance for index 0,1 - x,y planes. */ + double getEmittanceNormalized(int ic) const; + + /** Returns Twiss alpha (without dispersive part for x,y) for index 0,1,2 - x,y,z planes. */ + double getAlpha(int ic) const; + + /** Returns Twiss beta (without dispersive part for x,y) for index 0,1,2 - x,y,z planes. */ + double getBeta(int ic) const; + + /** Returns Twiss gamma (without dispersive part for x,y) for index 0,1,2 - x,y,z planes. */ + double getGamma(int ic) const; + + /** Returns Twiss dispersion function for index 0,1 - x,y planes. */ + double getDispersion(int ic) const; + + /** Returns Twiss dispersion_prime function for index 0,1 - x,y planes. */ + double getDispersionDerivative(int ic) const; + + /** Returns the effective emittance for index 0,1 - x,y planes. */ + double getEffectiveEmittance(int ic) const; + + /** Returns effective Twiss alpha for index 0,1 - x,y planes. */ + double getEffectiveAlpha(int ic) const; + + /** Returns effective Twiss beta for index 0,1 - x,y planes. */ + double getEffectiveBeta(int ic) const; + + /** Returns effective Twiss gamma for index 0,1 - x,y planes. */ + double getEffectiveGamma(int ic) const; + +private: + template + void computeBunchMomentsImpl(Bunch* bunch, bool normalize, bool emitnormflag); + + static constexpr int momentIdx(int i, int j, int maxOrder) + { + return i * maxOrder - (i * (i - 1)) / 2 + j; + } + + static constexpr int covIdx(int i, int j) + { + return i * N + j; + } + + static constexpr int N = 6; // {x, x', y, y', z, dE} + static constexpr int NN = N * N; // 36 + + int count_{}; + double total_macrosize_{}; + + std::array avg_arr{}; + std::array cov_arr{}; - /** bunch parameters */ - double bunch_momentum; - double bunch_gamma; - double bunch_beta; - double bunch_kinenergy; - double bunch_mass; + std::vector momentXY_; + + int order_{}; + double bunch_momentum_{}; + double bunch_gamma_{}; + double bunch_beta_{}; + double bunch_kinenergy_{}; + double bunch_mass_{}; }; #endif -//endif for BUNCH_TWISS_ANALYSIS_H diff --git a/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc b/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc index ca3efe39..bb4af851 100644 --- a/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc +++ b/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.cc @@ -1,344 +1,453 @@ -#include "orbit_mpi.hh" -#include "pyORBIT_Object.hh" +#include "orbit/Bunchdiagnostics/wrap_bunch_twiss_analysis.hh" -#include "wrap_bunch_twiss_analysis.hh" -#include "wrap_bunch.hh" +#include "main/pyORBIT_Object.hh" +#include "mpi/orbit_mpi.hh" +#include "orbit/BunchDiagnostics/BunchTwissAnalysis.hh" +#include "orbit/wrap_bunch.hh" -#include +namespace wrap_bunch_twiss_analysis +{ -#include "BunchTwissAnalysis.hh" - -namespace wrap_bunch_twiss_analysis{ - - void error(const char* msg){ ORBIT_MPI_Finalize(msg); } +void error(const char* msg) +{ + ORBIT_MPI_Finalize(msg); +} #ifdef __cplusplus extern "C" { #endif - /** - Constructor for python class wrapping c++ BunchTwissAnalysis instance. - It never will be called directly. - */ - static PyObject* BunchTwissAnalysis_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; - } - - /** This is implementation of the __init__ method */ - static int BunchTwissAnalysis_init(pyORBIT_Object *self, PyObject *args, PyObject *kwds){ - self->cpp_obj = new BunchTwissAnalysis(); - ((BunchTwissAnalysis*) self->cpp_obj)->setPyWrapper((PyObject*) self); - return 0; - } +/** + Constructor for python class wrapping c++ BunchTwissAnalysis instance. + It never will be called directly. +*/ +static PyObject* BunchTwissAnalysis_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; +} - /** Performs the Twiss analysis of the bunch */ - static PyObject* BunchTwissAnalysis_analyzeBunch(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - PyObject* pyBunch; - if(!PyArg_ParseTuple(args,"O:analyzeBunch",&pyBunch)){ - ORBIT_MPI_Finalize("BunchTwissAnalysis - analyzeBunch(Bunch* bunch) - parameter are needed."); - } - PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); - if(!PyObject_IsInstance(pyBunch,pyORBIT_Bunch_Type)){ - ORBIT_MPI_Finalize("BunchTwissAnalysis - analyzeBunch(Bunch* bunch) - method needs a Bunch."); - } - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object*)pyBunch)->cpp_obj; - cpp_BunchTwissAnalysis->analyzeBunch(cpp_bunch); - Py_INCREF(Py_None); - return Py_None; - } +/** This is implementation of the __init__ method */ +static int BunchTwissAnalysis_init(pyORBIT_Object* self, PyObject* args, PyObject* kwds) +{ + self->cpp_obj = new BunchTwissAnalysis(); + ((BunchTwissAnalysis*)self->cpp_obj)->setPyWrapper((PyObject*)self); + return 0; +} - /** Returns the XY moments of the beam up to a prescribed order */ - static PyObject* BunchTwissAnalysis_computeBunchMoments(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - PyObject* pyBunch; - int order; - int dispterm; - int emitnormterm; - if(!PyArg_ParseTuple(args,"Oiii:computeBunchMoments",&pyBunch,&order,&dispterm,&emitnormterm)){ - ORBIT_MPI_Finalize("BunchTwissAnalysis - computeBunchMoments(Bunch* bunch, int order, int dispterm, int emitnormterm) - parameters are needed."); - } - PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); - if(!PyObject_IsInstance(pyBunch,pyORBIT_Bunch_Type)){ - ORBIT_MPI_Finalize("BunchTwissAnalysis - computeBunchMoments(Bunch* bunch, int order, int dispterm, int emitnormterm) - method needs a Bunch."); - } - - Bunch* cpp_bunch = (Bunch*) ((pyORBIT_Object*)pyBunch)->cpp_obj; - cpp_BunchTwissAnalysis->computeBunchMoments(cpp_bunch, order, dispterm, emitnormterm); - Py_INCREF(Py_None); - return Py_None; - - } - - /** It will return the centered correlation <(x-)*(y-)> = - * for coordinates with indeces (ic,jc) */ - static PyObject* BunchTwissAnalysis_getCorrelation(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic,jc; - if(!PyArg_ParseTuple( args,"ii:getCorrelation",&ic,&jc)){ - error("pyBunchTwissAnalysis.getCorrelation(ic,jc) - parameters are needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getCorrelation(ic,jc)); - } - - /** It will return the (i,j) XY moment of the beam */ - static PyObject* BunchTwissAnalysis_getBunchMoment(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int i,j; - if(!PyArg_ParseTuple( args,"ii:getBunchMoment",&i,&j)){ - error("pyBunchTwissAnalysis.getBunchMoment(i,j) - parameters are needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getBunchMoment(i,j)); - } - - /** Returns the average value for coordinate with index ic */ - static PyObject* BunchTwissAnalysis_getAverage(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:",&ic)){ - error("pyBunchTwissAnalysis.getAverage(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getAverage(ic)); +/** Performs the Twiss analysis of the bunch */ +static PyObject* BunchTwissAnalysis_analyzeBunch(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + PyObject* pyBunch; + if (!PyArg_ParseTuple(args, "O:analyzeBunch", &pyBunch)) { + ORBIT_MPI_Finalize("BunchTwissAnalysis - analyzeBunch(Bunch* bunch) - parameter are needed."); + } + PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); + if (!PyObject_IsInstance(pyBunch, pyORBIT_Bunch_Type)) { + ORBIT_MPI_Finalize("BunchTwissAnalysis - analyzeBunch(Bunch* bunch) - method needs a Bunch."); } + Bunch* cpp_bunch = (Bunch*)((pyORBIT_Object*)pyBunch)->cpp_obj; + cpp_BunchTwissAnalysis->analyzeBunch(cpp_bunch); + Py_INCREF(Py_None); + return Py_None; +} - /** Returns the total number of analysed macroparticles */ - static PyObject* BunchTwissAnalysis_getGlobalCount(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - return Py_BuildValue("i",cpp_BunchTwissAnalysis->getGlobalCount()); +/** Returns the XY moments of the beam up to a prescribed order */ +static PyObject* BunchTwissAnalysis_computeBunchMoments(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + PyObject* pyBunch; + int order = 2; + int normalize = 0; + int emitnormterm = 0; + int dispterm = 0; + if (!PyArg_ParseTuple( + args, + "O|ippp:computeBunchMoments", + &pyBunch, + &order, + &normalize, + &emitnormterm, + &dispterm + )) { + ORBIT_MPI_Finalize( + "BunchTwissAnalysis - computeBunchMoments(Bunch* bunch, int order, int dispterm, int " + "emitnormterm) - parameters are needed." + ); + } + PyObject* pyORBIT_Bunch_Type = wrap_orbit_bunch::getBunchType("Bunch"); + if (!PyObject_IsInstance(pyBunch, pyORBIT_Bunch_Type)) { + ORBIT_MPI_Finalize( + "BunchTwissAnalysis - computeBunchMoments(Bunch* bunch, int order, int dispterm, int " + "emitnormterm) - method needs a Bunch." + ); } - /** Returns the total macrosize */ - static PyObject* BunchTwissAnalysis_getGlobalMacrosize(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getGlobalMacrosize()); + Bunch* cpp_bunch = (Bunch*)((pyORBIT_Object*)pyBunch)->cpp_obj; + cpp_BunchTwissAnalysis->computeBunchMoments(cpp_bunch, order, normalize, emitnormterm, dispterm); + Py_INCREF(Py_None); + return Py_None; +} + +/** It will return the centered correlation <(x-)*(y-)> = - * for coordinates with + * indeces (ic,jc) */ +static PyObject* BunchTwissAnalysis_getCorrelation(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic, jc; + if (!PyArg_ParseTuple(args, "ii:getCorrelation", &ic, &jc)) { + error("pyBunchTwissAnalysis.getCorrelation(ic,jc) - parameters are needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getCorrelation(ic, jc)); +} - //------------------------------------------------------------ - //Twiss functions - //------------------------------------------------------------ - /** It returns the emittance for index 0,1,2 - x,y,z planes*/ - static PyObject* BunchTwissAnalysis_getEmittance(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getEmittance",&ic)){ - error("pyBunchTwissAnalysis.getEmittance(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getEmittance(ic)); +/** It will return the (i,j) XY moment of the beam */ +static PyObject* BunchTwissAnalysis_getBunchMoment(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int i, j; + if (!PyArg_ParseTuple(args, "ii:getBunchMoment", &i, &j)) { + error("pyBunchTwissAnalysis.getBunchMoment(i,j) - parameters are needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getBunchMoment(i, j)); +} - /** It returns the normalized emittance for index 0,1 - x,y planes*/ - static PyObject* BunchTwissAnalysis_getEmittanceNormalized(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getEmittanceNormalized",&ic)){ - error("pyBunchTwissAnalysis.getEmittanceNormalized(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getEmittanceNormalized(ic)); +/** Returns the average value for coordinate with index ic */ +static PyObject* BunchTwissAnalysis_getAverage(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:", &ic)) { + error("pyBunchTwissAnalysis.getAverage(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getAverage(ic)); +} + +/** Returns the total number of analysed macroparticles */ +static PyObject* BunchTwissAnalysis_getGlobalCount(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + return Py_BuildValue("i", cpp_BunchTwissAnalysis->getGlobalCount()); +} + +/** Returns the total macrosize */ +static PyObject* BunchTwissAnalysis_getGlobalMacrosize(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getGlobalMacrosize()); +} - /** It returns the Twiss alpha for index 0,1,2 - x,y,z planes*/ - static PyObject* BunchTwissAnalysis_getAlpha(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getAlpha",&ic)){ - error("pyBunchTwissAnalysis.getAlpha(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getAlpha(ic)); +//------------------------------------------------------------ +// Twiss functions +//------------------------------------------------------------ +/** It returns the emittance for index 0,1,2 - x,y,z planes*/ +static PyObject* BunchTwissAnalysis_getEmittance(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getEmittance", &ic)) { + error("pyBunchTwissAnalysis.getEmittance(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getEmittance(ic)); +} - /** It returns the Twiss beta for index 0,1,2 - x,y,z planes*/ - static PyObject* BunchTwissAnalysis_getBeta(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getBeta",&ic)){ - error("pyBunchTwissAnalysis.getBeta(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getBeta(ic)); +/** It returns the normalized emittance for index 0,1 - x,y planes*/ +static PyObject* BunchTwissAnalysis_getEmittanceNormalized(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getEmittanceNormalized", &ic)) { + error("pyBunchTwissAnalysis.getEmittanceNormalized(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getEmittanceNormalized(ic)); +} - /** It returns the Twiss gamma for index 0,1,2 - x,y,z planes*/ - static PyObject* BunchTwissAnalysis_getGamma(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getGamma",&ic)){ - error("pyBunchTwissAnalysis.getGamma(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getGamma(ic)); +/** It returns the Twiss alpha for index 0,1,2 - x,y,z planes*/ +static PyObject* BunchTwissAnalysis_getAlpha(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getAlpha", &ic)) { + error("pyBunchTwissAnalysis.getAlpha(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getAlpha(ic)); +} - /** It returns Twiss dispersion function for index 0,1 - x,y planes*/ - static PyObject* BunchTwissAnalysis_getDispersion(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getDispersion",&ic)){ - error("pyBunchTwissAnalysis.getDispersion(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getDispersion(ic)); +/** It returns the Twiss beta for index 0,1,2 - x,y,z planes*/ +static PyObject* BunchTwissAnalysis_getBeta(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getBeta", &ic)) { + error("pyBunchTwissAnalysis.getBeta(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getBeta(ic)); +} - /** It returns Twiss dispersion_prime function for index 0,1 - x,y planes*/ - static PyObject* BunchTwissAnalysis_getDispersionDerivative(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getDispersionDerivative",&ic)){ - error("pyBunchTwissAnalysis.getDispersionDerivative(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getDispersionDerivative(ic)); +/** It returns the Twiss gamma for index 0,1,2 - x,y,z planes*/ +static PyObject* BunchTwissAnalysis_getGamma(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getGamma", &ic)) { + error("pyBunchTwissAnalysis.getGamma(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getGamma(ic)); +} - /** It returns the Twiss array (alpha,beta,gamma,emittance) for index 0,1,2 - x,y,z planes*/ - static PyObject* BunchTwissAnalysis_getTwiss(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getTwiss",&ic)){ - error("pyBunchTwissAnalysis.getTwiss(ic) - parameter is needed"); - } - double alpha = cpp_BunchTwissAnalysis->getAlpha(ic); - double beta = cpp_BunchTwissAnalysis->getBeta(ic); - double gamma = cpp_BunchTwissAnalysis->getGamma(ic); - double emitt = cpp_BunchTwissAnalysis->getEmittance(ic); - return Py_BuildValue("(dddd)",alpha,beta,gamma,emitt); +/** It returns Twiss dispersion function for index 0,1 - x,y planes*/ +static PyObject* BunchTwissAnalysis_getDispersion(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getDispersion", &ic)) { + error("pyBunchTwissAnalysis.getDispersion(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getDispersion(ic)); +} +/** It returns Twiss dispersion_prime function for index 0,1 - x,y planes*/ +static PyObject* BunchTwissAnalysis_getDispersionDerivative(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getDispersionDerivative", &ic)) { + error("pyBunchTwissAnalysis.getDispersionDerivative(ic) - parameter is needed"); + } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getDispersionDerivative(ic)); +} - /** It returns the effective emittance for index 0,1 - x,y planes*/ - static PyObject* BunchTwissAnalysis_getEffectiveEmittance(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getEffectiveEmittance",&ic)){ - error("pyBunchTwissAnalysis.getEffectiveEmittance(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getEffectiveEmittance(ic)); +/** It returns the Twiss array (alpha,beta,gamma,emittance) for index 0,1,2 - x,y,z planes*/ +static PyObject* BunchTwissAnalysis_getTwiss(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getTwiss", &ic)) { + error("pyBunchTwissAnalysis.getTwiss(ic) - parameter is needed"); } + double alpha = cpp_BunchTwissAnalysis->getAlpha(ic); + double beta = cpp_BunchTwissAnalysis->getBeta(ic); + double gamma = cpp_BunchTwissAnalysis->getGamma(ic); + double emitt = cpp_BunchTwissAnalysis->getEmittance(ic); + return Py_BuildValue("(dddd)", alpha, beta, gamma, emitt); +} - /** It returns the effective Twiss alpha for index 0,1 - x,y planes*/ - static PyObject* BunchTwissAnalysis_getEffectiveAlpha(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getEffectiveAlpha",&ic)){ - error("pyBunchTwissAnalysis.getEffectiveAlpha(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getEffectiveAlpha(ic)); +/** It returns the effective emittance for index 0,1 - x,y planes*/ +static PyObject* BunchTwissAnalysis_getEffectiveEmittance(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getEffectiveEmittance", &ic)) { + error("pyBunchTwissAnalysis.getEffectiveEmittance(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getEffectiveEmittance(ic)); +} - /** It returns the effective Twiss beta for index 0,1 - x,y planes*/ - static PyObject* BunchTwissAnalysis_getEffectiveBeta(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getEffectiveBeta",&ic)){ - error("pyBunchTwissAnalysis.getEffectiveBeta(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getEffectiveBeta(ic)); +/** It returns the effective Twiss alpha for index 0,1 - x,y planes*/ +static PyObject* BunchTwissAnalysis_getEffectiveAlpha(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getEffectiveAlpha", &ic)) { + error("pyBunchTwissAnalysis.getEffectiveAlpha(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getEffectiveAlpha(ic)); +} - /** It returns the effective Twiss gamma for index 0,1 - x,y planes*/ - static PyObject* BunchTwissAnalysis_getEffectiveGamma(PyObject *self, PyObject *args){ - BunchTwissAnalysis* cpp_BunchTwissAnalysis = (BunchTwissAnalysis*)((pyORBIT_Object*) self)->cpp_obj; - int ic; - if(!PyArg_ParseTuple( args,"i:getEffectiveGamma",&ic)){ - error("pyBunchTwissAnalysis.getEffectiveGamma(ic) - parameter is needed"); - } - return Py_BuildValue("d",cpp_BunchTwissAnalysis->getEffectiveGamma(ic)); +/** It returns the effective Twiss beta for index 0,1 - x,y planes*/ +static PyObject* BunchTwissAnalysis_getEffectiveBeta(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getEffectiveBeta", &ic)) { + error("pyBunchTwissAnalysis.getEffectiveBeta(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getEffectiveBeta(ic)); +} - //----------------------------------------------------- - //destructor for python BunchTwissAnalysis class (__del__ method). - //----------------------------------------------------- - static void BunchTwissAnalysis_del(pyORBIT_Object* self){ - //std::cerr<<"The BunchTwissAnalysis __del__ has been called!"<cpp_obj); - self->ob_base.ob_type->tp_free((PyObject*)self); +/** It returns the effective Twiss gamma for index 0,1 - x,y planes*/ +static PyObject* BunchTwissAnalysis_getEffectiveGamma(PyObject* self, PyObject* args) +{ + BunchTwissAnalysis* cpp_BunchTwissAnalysis = + (BunchTwissAnalysis*)((pyORBIT_Object*)self)->cpp_obj; + int ic; + if (!PyArg_ParseTuple(args, "i:getEffectiveGamma", &ic)) { + error("pyBunchTwissAnalysis.getEffectiveGamma(ic) - parameter is needed"); } + return Py_BuildValue("d", cpp_BunchTwissAnalysis->getEffectiveGamma(ic)); +} - // defenition of the methods of the python BunchTwissAnalysis wrapper class - // they will be vailable from python level - static PyMethodDef BunchTwissAnalysisClassMethods[] = { - { "analyzeBunch", BunchTwissAnalysis_analyzeBunch, METH_VARARGS,"Performs the Twiss analysis of the bunch."}, - { "computeBunchMoments", BunchTwissAnalysis_computeBunchMoments, METH_VARARGS,"Returns the XY moments of the beam up to a prescribed order"}, - { "getCorrelation", BunchTwissAnalysis_getCorrelation, METH_VARARGS,"Returns the centered correlation <(x-)*(y-)> = - * for coordinates with indeces (ic,jc)"}, - { "getBunchMoment", BunchTwissAnalysis_getBunchMoment, METH_VARARGS,"Returns the (i,j) xy moment of the beam"}, - { "getAverage", BunchTwissAnalysis_getAverage, METH_VARARGS,"Returns the average value for coordinate with index ic"}, - { "getGlobalCount", BunchTwissAnalysis_getGlobalCount, METH_VARARGS,"Returns the total number of analysed macroparticles"}, - { "getGlobalMacrosize", BunchTwissAnalysis_getGlobalMacrosize, METH_VARARGS,"Returns the total macrosize"}, - { "getEmittance", BunchTwissAnalysis_getEmittance, METH_VARARGS,"Returns the emittance for index 0,1,2 - x,y,z planes"}, - { "getEmittanceNormalized", BunchTwissAnalysis_getEmittanceNormalized, METH_VARARGS,"Returns the normalized emittance for index 0,1 - x,y planes"}, - { "getAlpha", BunchTwissAnalysis_getAlpha, METH_VARARGS,"Returns Twiss alpha for index 0,1,2 - x,y,z planes"}, - { "getBeta", BunchTwissAnalysis_getBeta, METH_VARARGS,"Returns Twiss beta for index 0,1,2 - x,y,z planes"}, - { "getGamma", BunchTwissAnalysis_getGamma, METH_VARARGS,"Returns Twiss gamma for index 0,1,2 - x,y,z planes"}, - { "getTwiss", BunchTwissAnalysis_getTwiss, METH_VARARGS,"Returns Twiss tuple (alpha,beta,gamma,emitt) for index 0,1,2 - x,y,z planes"}, - { "getDispersion", BunchTwissAnalysis_getDispersion, METH_VARARGS,"Returns Twiss dispersion function for index 0,1 - x,y planes"}, - { "getDispersionDerivative",BunchTwissAnalysis_getDispersionDerivative, METH_VARARGS,"Returns Twiss dispersion' function for index 0,1 - x,y planes"}, - { "getEffectiveEmittance", BunchTwissAnalysis_getEffectiveEmittance, METH_VARARGS,"Returns the effective emittance for index 0,1 - x,y planes"}, - { "getEffectiveAlpha", BunchTwissAnalysis_getEffectiveAlpha, METH_VARARGS,"Returns effective Twiss alpha for index 0,1 - x,y planes"}, - { "getEffectiveBeta", BunchTwissAnalysis_getEffectiveBeta, METH_VARARGS,"Returns effective Twiss beta for index 0,1 - x,y planes"}, - { "getEffectiveGamma", BunchTwissAnalysis_getEffectiveGamma, METH_VARARGS,"Returns effective Twiss gamma for index 0,1 - x,y planes"}, - {NULL} - }; - - // defenition of the memebers of the python BunchTwissAnalysis wrapper class - // they will be vailable from python level - static PyMemberDef BunchTwissAnalysisClassMembers [] = { - {NULL} - }; - - //new python BunchTwissAnalysis wrapper type definition - static PyTypeObject pyORBIT_BunchTwissAnalysis_Type = { - PyVarObject_HEAD_INIT(NULL, 0) - "BunchTwissAnalysis", /*tp_name*/ - sizeof(pyORBIT_Object), /*tp_basicsize*/ - 0, /*tp_itemsize*/ - (destructor) BunchTwissAnalysis_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 BunchTwissAnalysis python wrapper", /* tp_doc */ - 0, /* tp_traverse */ - 0, /* tp_clear */ - 0, /* tp_richcompare */ - 0, /* tp_weaklistoffset */ - 0, /* tp_iter */ - 0, /* tp_iternext */ - BunchTwissAnalysisClassMethods, /* tp_methods */ - BunchTwissAnalysisClassMembers, /* tp_members */ - 0, /* tp_getset */ - 0, /* tp_base */ - 0, /* tp_dict */ - 0, /* tp_descr_get */ - 0, /* tp_descr_set */ - 0, /* tp_dictoffset */ - (initproc) BunchTwissAnalysis_init, /* tp_init */ - 0, /* tp_alloc */ - BunchTwissAnalysis_new, /* tp_new */ - }; - - - - //-------------------------------------------------- - //Initialization BunchTwissAnalysis of the pyBunchTwissAnalysis class - //-------------------------------------------------- - void initbunchtwissanalysis(PyObject* module){ - if (PyType_Ready(&pyORBIT_BunchTwissAnalysis_Type) < 0) return; - Py_INCREF(&pyORBIT_BunchTwissAnalysis_Type); - PyModule_AddObject(module, "BunchTwissAnalysis", (PyObject *)&pyORBIT_BunchTwissAnalysis_Type); - } +//----------------------------------------------------- +// destructor for python BunchTwissAnalysis class (__del__ method). +//----------------------------------------------------- +static void BunchTwissAnalysis_del(pyORBIT_Object* self) +{ + // std::cerr<<"The BunchTwissAnalysis __del__ has been called!"<cpp_obj); + self->ob_base.ob_type->tp_free((PyObject*)self); +} + +// defenition of the methods of the python BunchTwissAnalysis wrapper class +// they will be vailable from python level +static PyMethodDef BunchTwissAnalysisClassMethods[] = + {{"analyzeBunch", + BunchTwissAnalysis_analyzeBunch, + METH_VARARGS, + "Performs the Twiss analysis of the bunch."}, + {"computeBunchMoments", + BunchTwissAnalysis_computeBunchMoments, + METH_VARARGS, + "Returns the XY moments of the beam up to a prescribed order"}, + {"getCorrelation", + BunchTwissAnalysis_getCorrelation, + METH_VARARGS, + "Returns the centered correlation <(x-)*(y-)> = - * for coordinates with " + "indeces (ic,jc)"}, + {"getBunchMoment", + BunchTwissAnalysis_getBunchMoment, + METH_VARARGS, + "Returns the (i,j) xy moment of the beam"}, + {"getAverage", + BunchTwissAnalysis_getAverage, + METH_VARARGS, + "Returns the average value for coordinate with index ic"}, + {"getGlobalCount", + BunchTwissAnalysis_getGlobalCount, + METH_VARARGS, + "Returns the total number of analysed macroparticles"}, + {"getGlobalMacrosize", + BunchTwissAnalysis_getGlobalMacrosize, + METH_VARARGS, + "Returns the total macrosize"}, + {"getEmittance", + BunchTwissAnalysis_getEmittance, + METH_VARARGS, + "Returns the emittance for index 0,1,2 - x,y,z planes"}, + {"getEmittanceNormalized", + BunchTwissAnalysis_getEmittanceNormalized, + METH_VARARGS, + "Returns the normalized emittance for index 0,1 - x,y planes"}, + {"getAlpha", + BunchTwissAnalysis_getAlpha, + METH_VARARGS, + "Returns Twiss alpha for index 0,1,2 - x,y,z planes"}, + {"getBeta", + BunchTwissAnalysis_getBeta, + METH_VARARGS, + "Returns Twiss beta for index 0,1,2 - x,y,z planes"}, + {"getGamma", + BunchTwissAnalysis_getGamma, + METH_VARARGS, + "Returns Twiss gamma for index 0,1,2 - x,y,z planes"}, + {"getTwiss", + BunchTwissAnalysis_getTwiss, + METH_VARARGS, + "Returns Twiss tuple (alpha,beta,gamma,emitt) for index 0,1,2 - x,y,z planes"}, + {"getDispersion", + BunchTwissAnalysis_getDispersion, + METH_VARARGS, + "Returns Twiss dispersion function for index 0,1 - x,y planes"}, + {"getDispersionDerivative", + BunchTwissAnalysis_getDispersionDerivative, + METH_VARARGS, + "Returns Twiss dispersion' function for index 0,1 - x,y planes"}, + {"getEffectiveEmittance", + BunchTwissAnalysis_getEffectiveEmittance, + METH_VARARGS, + "Returns the effective emittance for index 0,1 - x,y planes"}, + {"getEffectiveAlpha", + BunchTwissAnalysis_getEffectiveAlpha, + METH_VARARGS, + "Returns effective Twiss alpha for index 0,1 - x,y planes"}, + {"getEffectiveBeta", + BunchTwissAnalysis_getEffectiveBeta, + METH_VARARGS, + "Returns effective Twiss beta for index 0,1 - x,y planes"}, + {"getEffectiveGamma", + BunchTwissAnalysis_getEffectiveGamma, + METH_VARARGS, + "Returns effective Twiss gamma for index 0,1 - x,y planes"}, + {NULL}}; + +// defenition of the memebers of the python BunchTwissAnalysis wrapper class +// they will be vailable from python level +static PyMemberDef BunchTwissAnalysisClassMembers[] = {{NULL}}; + +// new python BunchTwissAnalysis wrapper type definition +static PyTypeObject pyORBIT_BunchTwissAnalysis_Type = { + PyVarObject_HEAD_INIT(NULL, 0) "BunchTwissAnalysis", /*tp_name*/ + sizeof(pyORBIT_Object), /*tp_basicsize*/ + 0, /*tp_itemsize*/ + (destructor)BunchTwissAnalysis_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 BunchTwissAnalysis python wrapper", /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + BunchTwissAnalysisClassMethods, /* tp_methods */ + BunchTwissAnalysisClassMembers, /* tp_members */ + 0, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + 0, /* tp_dictoffset */ + (initproc)BunchTwissAnalysis_init, /* tp_init */ + 0, /* tp_alloc */ + BunchTwissAnalysis_new, /* tp_new */ +}; + +//-------------------------------------------------- +// Initialization BunchTwissAnalysis of the pyBunchTwissAnalysis class +//-------------------------------------------------- +void initbunchtwissanalysis(PyObject* module) +{ + if (PyType_Ready(&pyORBIT_BunchTwissAnalysis_Type) < 0) + return; + Py_INCREF(&pyORBIT_BunchTwissAnalysis_Type); + PyModule_AddObject(module, "BunchTwissAnalysis", (PyObject*)&pyORBIT_BunchTwissAnalysis_Type); +} #ifdef __cplusplus } #endif - -} +} // namespace wrap_bunch_twiss_analysis diff --git a/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.hh b/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.hh index 7f01e87d..2fe8d9eb 100644 --- a/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.hh +++ b/src/orbit/BunchDiagnostics/wrap_bunch_twiss_analysis.hh @@ -1,15 +1,16 @@ #ifndef WRAP_BUNCH_TWISS_ANALYSIS_HH_ #define WRAP_BUNCH_TWISS_ANALYSIS_HH_ -#include "Python.h" +#include #ifdef __cplusplus extern "C" { #endif - namespace wrap_bunch_twiss_analysis{ - void initbunchtwissanalysis(PyObject* module); - } +namespace wrap_bunch_twiss_analysis +{ +void initbunchtwissanalysis(PyObject* module); +} #ifdef __cplusplus } diff --git a/src/utils/bunch/TwissFilteringFunctions.cc b/src/utils/bunch/TwissFilteringFunctions.cc index 56902267..5bf3aaeb 100644 --- a/src/utils/bunch/TwissFilteringFunctions.cc +++ b/src/utils/bunch/TwissFilteringFunctions.cc @@ -14,10 +14,11 @@ // according to their positions relative to the emittance phase space center // /////////////////////////////////////////////////////////////////////////// +#include "utils/bunch/TwissFilteringFunctions.hh" -#include "BunchTwissAnalysis.hh" - -#include "orbit_mpi.hh" +#include "mpi/orbit_mpi.hh" +#include "orbit/Bunch.hh" +#include "orbit/BunchDiagnostics/BunchTwissAnalysis.hh" namespace OrbitUtils{ diff --git a/src/utils/bunch/TwissFilteringFunctions.hh b/src/utils/bunch/TwissFilteringFunctions.hh index 63d99bc4..91c7fd9d 100644 --- a/src/utils/bunch/TwissFilteringFunctions.hh +++ b/src/utils/bunch/TwissFilteringFunctions.hh @@ -18,10 +18,7 @@ #ifndef BUNCH_TWISS_FILTERING_FUNCTIONS_H #define BUNCH_TWISS_FILTERING_FUNCTIONS_H -#include - -//ORBIT bunch -#include "Bunch.hh" +class Bunch; namespace OrbitUtils{ diff --git a/tests/py/orbit/core/bunch/test_BunchTwissAnalysis.py b/tests/py/orbit/core/bunch/test_BunchTwissAnalysis.py new file mode 100644 index 00000000..c5b30dfc --- /dev/null +++ b/tests/py/orbit/core/bunch/test_BunchTwissAnalysis.py @@ -0,0 +1,54 @@ +import pytest + +from orbit.core.bunch import Bunch, BunchTwissAnalysis + +import numpy as np + +@pytest.fixture +def gaus_dist(): + return np.random.normal(loc=0, scale=1e-3, size=(100, 6)) + +@pytest.fixture +def gaus_bunch(gaus_dist): + b = Bunch() + b.getSyncParticle().kinEnergy(2.5e-3) + b.mass(0.939294) + b.charge(-1.0) + for coords in gaus_dist: + b.addParticle(*coords) + return b + +@pytest.fixture +def twiss_analysis(): + return BunchTwissAnalysis() + +def test_computeBunchMoments(gaus_bunch, gaus_dist, twiss_analysis): + twiss_analysis.computeBunchMoments(gaus_bunch) + + mean = gaus_dist.mean(axis=0) + cov = np.cov(gaus_dist.T, bias=True) + + for i in range(6): + assert twiss_analysis.getAverage(i) == pytest.approx(mean[i]) + for j in range(6): + assert twiss_analysis.getCorrelation(i, j) == pytest.approx(cov[i, j]) + +def test_computeBunchMoments_FirstOrder(gaus_bunch, gaus_dist, twiss_analysis): + twiss_analysis.computeBunchMoments(gaus_bunch, 3) + + mean = gaus_dist.mean(axis=0) + cov = np.cov(gaus_dist.T, bias=True) + + assert twiss_analysis.getBunchMoment(1, 0) == pytest.approx(mean[0]) + assert twiss_analysis.getBunchMoment(0, 1) == pytest.approx(mean[2]) + assert twiss_analysis.getBunchMoment(1, 1) == pytest.approx(cov[0, 2]) + +def test_computeBunchMoments_SecondOrder(gaus_bunch, gaus_dist, twiss_analysis): + twiss_analysis.computeBunchMoments(gaus_bunch, 3) + + var_x = np.var(gaus_dist[:, 0]) + var_y = np.var(gaus_dist[:, 2]) + + assert twiss_analysis.getBunchMoment(2, 0) == pytest.approx(var_x) + assert twiss_analysis.getBunchMoment(0, 2) == pytest.approx(var_y) + From efd80882e9a76820365d98796c1a46f596a996b3 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 23 Jul 2026 11:55:50 -0400 Subject: [PATCH 4/6] emitnormflag should imply normalize --- src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc index bdcf5ae9..507d4fc7 100644 --- a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc +++ b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc @@ -78,7 +78,7 @@ void BunchTwissAnalysis::computeBunchMomentsImpl(Bunch* bunch, bool normalize, b double inv_xbt = 1.0; double inv_ybt = 1.0; - if (normalize) { + if (normalize || emitnormflag) { inv_xbt /= std::sqrt(getBeta(0) * (emitnormflag ? getEmittance(0) : 1.0)); inv_ybt /= std::sqrt(getBeta(1) * (emitnormflag ? getEmittance(1) : 1.0)); } From 3343381ad2f602065f8b5a61897b292a8ba4b7c7 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 23 Jul 2026 12:05:07 -0400 Subject: [PATCH 5/6] Use ORBIT_MPI_IN_PLACE --- src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc index 507d4fc7..d0ce840f 100644 --- a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc +++ b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc @@ -56,14 +56,14 @@ void BunchTwissAnalysis::computeBunchMomentsImpl(Bunch* bunch, bool normalize, b auto* comm = bunch->getMPI_Comm_Local()->comm; if constexpr (HasMacrosizeAttr) { - ORBIT_MPI_Allreduce(MPI_IN_PLACE, &total_macrosize_, 1, MPI_DOUBLE, MPI_SUM, comm); + ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, &total_macrosize_, 1, MPI_DOUBLE, MPI_SUM, comm); } else { total_macrosize_ = bunch->getSizeGlobal(); } - ORBIT_MPI_Allreduce(MPI_IN_PLACE, avg_arr.data(), N, MPI_DOUBLE, MPI_SUM, comm); - ORBIT_MPI_Allreduce(MPI_IN_PLACE, cov_arr.data(), NN, MPI_DOUBLE, MPI_SUM, comm); + ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, avg_arr.data(), N, MPI_DOUBLE, MPI_SUM, comm); + ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, cov_arr.data(), NN, MPI_DOUBLE, MPI_SUM, comm); // = - uhat*vhat for (int i = 0; i < N; ++i) { @@ -136,7 +136,7 @@ void BunchTwissAnalysis::computeBunchMomentsImpl(Bunch* bunch, bool normalize, b } ORBIT_MPI_Allreduce( - MPI_IN_PLACE, + ORBIT_MPI_IN_PLACE, momentXY_.data(), nMoments * nMoments, MPI_DOUBLE, From fb5d6279a4820f16a8b2405e257619c4e1085349 Mon Sep 17 00:00:00 2001 From: "Wood, Tony" Date: Thu, 23 Jul 2026 12:12:38 -0400 Subject: [PATCH 6/6] Interestingly this compiles in my version of clang++, but it's not supposed to be a ptr. --- src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc index d0ce840f..b6a498f0 100644 --- a/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc +++ b/src/orbit/BunchDiagnostics/BunchTwissAnalysis.cc @@ -53,7 +53,7 @@ void BunchTwissAnalysis::computeBunchMomentsImpl(Bunch* bunch, bool normalize, b } } - auto* comm = bunch->getMPI_Comm_Local()->comm; + auto comm = bunch->getMPI_Comm_Local()->comm; if constexpr (HasMacrosizeAttr) { ORBIT_MPI_Allreduce(ORBIT_MPI_IN_PLACE, &total_macrosize_, 1, MPI_DOUBLE, MPI_SUM, comm);