diff --git a/tests/test_monitor.py b/tests/test_monitor.py index 7c2746898..6ef5c48c5 100644 --- a/tests/test_monitor.py +++ b/tests/test_monitor.py @@ -52,6 +52,7 @@ def test_constructor(test_context): #xt.ParticlesMonitor(_context=test_context, start_at_turn=0, stop_at_turn=10, num_particles=42), #xt.LastTurnsMonitor(_context=test_context, n_last_turns=5, num_particles=42), xt.BeamPositionMonitor(_context=test_context, stop_at_turn=10), + xt.BunchMonitor(_context=test_context, stop_at_turn=10), ] # test to_dict / from_dict @@ -492,3 +493,69 @@ def test_beam_position_monitor(test_context): expected_x_centroid[18:20] = np.nan assert_allclose(monitor.x_cen, expected_x_centroid, err_msg="Monitor x centroid does not match expected values") assert_allclose(monitor.y_cen, -expected_x_centroid, err_msg="Monitor y centroid does not match expected values") + +@for_all_test_contexts +def test_bunch_monitor(test_context): + + npart = 512 # must be even and >= 512 + zeta = np.empty(npart+4) # generate a few more than we record to test "num_particles" + zeta[0:npart:2] = zeta1 = np.linspace(-0.2, 0.2, npart//2) # first bunch centered around 0 + zeta[1:npart:2] = zeta2 = -0.5 + zeta1 # second bunch centered around -1/harmonic + zeta[npart:] = np.linspace(2,4,4) + + particles = xp.Particles( + p0c=6.5e12, + zeta=zeta, + delta=1e-3*zeta, + _context=test_context, + ) + + monitor = xt.BunchMonitor( + num_particles=npart, + start_at_turn=0, + stop_at_turn=10, + frev=2.99792458e8, + harmonic=2, + _context=test_context, + ) + + line = xt.Line([monitor]) + line.build_tracker(_context=test_context) + + for turn in range(11): # track a few more than we record to test "stop_at_turn" + # Note that indicees are re-ordered upon particle loss on CPU contexts, + # so sort before manipulation + if isinstance(test_context, xo.ContextCpu): + particles.sort(interleave_lost_particles=True) + + # manipulate particles for testing + particles.zeta[0] += 0.005 + particles.delta[0] += 1e-3*0.005 + if turn == 8: + particles.state[256:] = 0 + if turn == 9: + particles.state[:] = 0 + + if isinstance(test_context, xo.ContextCpu): + particles.reorganize() + + # track and monitor + line.track(particles, num_turns=1) + + # Check against expected values + expected_count = np.tile([npart//2, npart//2], 10) + expected_count[16:18] = 128 + expected_count[18:20] = 0 + assert_equal(monitor.count, expected_count, err_msg="Monitor count does not match expected particle count") + + expected_zeta_sum = np.zeros(20) + expected_zeta_sum[0::2] = np.sum(zeta1) + 0.005*np.arange(1, 11) + expected_zeta_sum[1::2] = np.sum(zeta2) + expected_zeta_sum[16] -= np.sum(zeta1[128:]) + expected_zeta_sum[17] -= np.sum(zeta2[128:]) + expected_zeta_sum[18:] = 0 + assert_allclose(monitor.zeta_sum, expected_zeta_sum, err_msg="Monitor zeta sum does not match expected values") + assert_allclose(monitor.delta_sum, 1e-3*expected_zeta_sum, err_msg="Monitor delta sum does not match expected values") + + # TODO: test zeta_mean, delta_mean, zeta2_sum, delta2_sum, zeta_std, delta_std, zeta_var, delta_var + diff --git a/xtrack/monitors/__init__.py b/xtrack/monitors/__init__.py index d8bf6d53d..acc4219fd 100644 --- a/xtrack/monitors/__init__.py +++ b/xtrack/monitors/__init__.py @@ -4,5 +4,6 @@ from .beam_position_monitor import * from .beam_size_monitor import * from .beam_profile_monitor import * +from .bunch_monitor import * monitor_classes = tuple(v for v in globals().values() if isinstance(v, type) and issubclass(v, BeamElement)) diff --git a/xtrack/monitors/bunch_monitor.h b/xtrack/monitors/bunch_monitor.h index bf83b6769..3a892c75a 100644 --- a/xtrack/monitors/bunch_monitor.h +++ b/xtrack/monitors/bunch_monitor.h @@ -1,33 +1,32 @@ // ################################## // Bunch Monitor // -// Author: Philipp Niedermayer +// Author: Philipp Niedermayer, Cristopher Cortes // Date: 2023-08-14 // Edit: 2024-06-12 // ################################## -#ifndef XTRACK_BEAM_SIZE_MONITOR_H -#define XTRACK_BEAM_SIZE_MONITOR_H +#ifndef XTRACK_BUNCH_MONITOR_H +#define XTRACK_BUNCH_MONITOR_H #if !defined( C_LIGHT ) #define C_LIGHT ( 299792458.0 ) #endif /* !defined( C_LIGHT ) */ /*gpufun*/ -void BeamSizeMonitor_track_local_particle(BeamSizeMonitorData el, LocalParticle* part0){ +void BunchMonitor_track_local_particle(BunchMonitorData el, LocalParticle* part0){ // get parameters - int64_t const start_at_turn = BeamSizeMonitorData_get_start_at_turn(el); - int64_t particle_id_start = BeamSizeMonitorData_get_particle_id_start(el); - int64_t particle_id_stop = particle_id_start + BeamSizeMonitorData_get_num_particles(el); - // Should we do explicit casting to float (?) - int64_t const harmonic = BeamSizeMonitorData_get_harmonic(el); - double const frev = BeamSizeMonitorData_get_frev(el); + int64_t const start_at_turn = BunchMonitorData_get_start_at_turn(el); + int64_t particle_id_start = BunchMonitorData_get_particle_id_start(el); + int64_t particle_id_stop = particle_id_start + BunchMonitorData_get_num_particles(el); + int64_t const harmonic = BunchMonitorData_get_harmonic(el); + double const frev = BunchMonitorData_get_frev(el); - BeamSizeMonitorRecord record = BeamSizeMonitorData_getp_data(el); + BunchMonitorRecord record = BunchMonitorData_getp_data(el); - int64_t max_slot = BeamSizeMonitorRecord_len_count(record); + int64_t max_slot = BunchMonitorRecord_len_count(record); //start_per_particle_block(part0->part) @@ -46,24 +45,25 @@ void BeamSizeMonitor_track_local_particle(BeamSizeMonitorData el, LocalParticle* if (slot >= 0 && slot < max_slot){ - double delta = LocalParticle_get_delta(part); + double const delta = LocalParticle_get_delta(part); - /*gpuglmem*/ double* count = BeamSizeMonitorRecord_getp1_count(record, slot); + /*gpuglmem*/ double* count = BunchMonitorRecord_getp1_count(record, slot); atomicAdd(count, 1); - /*gpuglmem*/ double * zeta_sum = BeamSizeMonitorRecord_getp1_zeta_sum(record, slot); + /*gpuglmem*/ double * zeta_sum = BunchMonitorRecord_getp1_zeta_sum(record, slot); atomicAdd(zeta_sum, zeta); - /*gpuglmem*/ double * delta_sum = BeamSizeMonitorRecord_getp1_delta_sum(record, slot); + /*gpuglmem*/ double * delta_sum = BunchMonitorRecord_getp1_delta_sum(record, slot); atomicAdd(delta_sum, delta); - /*gpuglmem*/ double * zeta2_sum = BeamSizeMonitorRecord_getp1_zeta2_sum(record, slot); + /*gpuglmem*/ double * zeta2_sum = BunchMonitorRecord_getp1_zeta2_sum(record, slot); atomicAdd(zeta2_sum, zeta*zeta); - /*gpuglmem*/ double * delta2_sum = BeamSizeMonitorRecord_getp1_delta2_sum(record, slot); + /*gpuglmem*/ double * delta2_sum = BunchMonitorRecord_getp1_delta2_sum(record, slot); atomicAdd(delta2_sum, delta*delta); } } + //end_per_particle_block } diff --git a/xtrack/monitors/bunch_monitor.py b/xtrack/monitors/bunch_monitor.py index a461e1b08..c018fa0fd 100644 --- a/xtrack/monitors/bunch_monitor.py +++ b/xtrack/monitors/bunch_monitor.py @@ -1,7 +1,7 @@ """ Beam Size Monitor -Author: Philipp Niedermayer +Author: Philipp Niedermayer, Cristopher Cortes Date: 2023-08-14 Edit: 2024-06-12 """ @@ -15,14 +15,14 @@ from ..internal_record import RecordIndex from ..general import _pkg_root -class BeamSizeMonitorRecord(xo.Struct): +class BunchMonitorRecord(xo.Struct): count = xo.Float64[:] zeta_sum = xo.Float64[:] zeta2_sum = xo.Float64[:] delta_sum = xo.Float64[:] delta2_sum = xo.Float64[:] -class BeamSizeMonitor(BeamElement): +class BunchMonitor(BeamElement): _xofields={ 'particle_id_start': xo.Int64, @@ -32,13 +32,13 @@ class BeamSizeMonitor(BeamElement): 'frev': xo.Float64, 'harmonic': xo.Int64, '_index': RecordIndex, - 'data': BeamSizeMonitorRecord, + 'data': BunchMonitorRecord, } behaves_like_drift = True allow_loss_refinement = True - properties = [field.name for field in BeamSizeMonitorRecord._fields] + properties = [field.name for field in BunchMonitorRecord._fields] _extra_c_sources = [ _pkg_root.joinpath('headers/atomicadd.h'), @@ -49,7 +49,7 @@ def __init__(self, *, particle_id_range=None, particle_id_start=None, num_partic start_at_turn=None, stop_at_turn=None, frev=None, harmonic=None, _xobject=None, **kwargs): """ - Monitor to save the longitudinal beam size (standard deviation of the tracked particle positions) + Monitor to save the longitudinal bunch position and size (mean and std of zeta) as well as mean and std of momentum spread (delta) The monitor allows for arbitrary sampling rate and can thus not only be used to monitor @@ -117,18 +117,18 @@ def __init__(self, *, particle_id_range=None, particle_id_start=None, num_partic if "data" not in kwargs: # explicitely init with zeros (instead of size only) to have consistent initial values - size = int(round(( stop_at_turn - start_at_turn ) * sampling_frequency / frev)) + size = int(round(( stop_at_turn - start_at_turn ) * harmonic)) kwargs["data"] = {prop: np.zeros(size) for prop in self.properties} super().__init__(particle_id_start=particle_id_start, num_particles=num_particles, start_at_turn=start_at_turn, stop_at_turn=stop_at_turn, frev=frev, - sampling_frequency=sampling_frequency, **kwargs) + harmonic=harmonic, **kwargs) def __repr__(self): return ( f"{type(self).__qualname__}(start_at_turn={self.start_at_turn}, stop_at_turn={self.stop_at_turn}, " f"particle_id_start={self.particle_id_start}, num_particles={self.num_particles}, frev={self.frev}, " - f"harmonic={self.sampling_frequency}) at {hex(id(self))}" + f"harmonic={self.harmonic}) at {hex(id(self))}" ) def __getattr__(self, attr):