Merge branch 'multibunch_dev'

Author: giadarol

communication_helpers.py

@@ -1,11 +1,13 @@
 import numpy as np
 from PyHEADTAIL.particles.particles import Particles
+import json
+import pickle
 
 
 def combine_float_buffers(list_of_buffers):
 	N_buffers = len(list_of_buffers)
 	len_buffers = np.array(map(lambda seq: float(len(seq)), list_of_buffers))
-	return np.concatenate([np.array([float(N_buffers)]), len_buffers]+list_of_buffers)
+	return np.array(np.concatenate([np.array([float(N_buffers)]), len_buffers]+list_of_buffers),dtype=np.float64)
 
 def split_float_buffers(megabuffer):
 	i_mbuf = 0
@@ -39,7 +41,7 @@ def buffer_2_list_of_strings(buf):
 	return strlist
 
 
-def beam_2_buffer(beam):
+def beam_2_buffer(beam, mode='pickle', verbose=False):
 
 	#print beam
 
@@ -50,32 +52,62 @@ def beam_2_buffer(beam):
 			raise ValueError('particlenumber_per_mp is a vector! Not implemented!')
 
 
+		# if not hasattr(beam, 'slice_info'):
+		# 	sl_info_buf = np.array([-1., 0., 0., 0.])
+		# elif beam.slice_info is None:
+		# 	sl_info_buf = np.array([-1., 0., 0., 0.])
+		# elif beam.slice_info == 'unsliced':
+		# 	sl_info_buf = np.array([0., 0., 0., 0.])
+		# else:
+		# 	sl_info_buf = np.array([1.,
+		# 					beam.slice_info['z_bin_center'],
+		# 					beam.slice_info['z_bin_right'],
+		# 					beam.slice_info['z_bin_left']])
+
+		# Prepare slice info buffer
 		if not hasattr(beam, 'slice_info'):
-			sl_info_buf = np.array([-1., 0., 0., 0.])
-		elif beam.slice_info is None:
-			sl_info_buf = np.array([-1., 0., 0., 0.])
-		elif beam.slice_info == 'unsliced':
-			sl_info_buf = np.array([0., 0., 0., 0.])
+			sinfo = None
+		else:
+			sinfo = beam.slice_info
+
+		# Beam data buffer
+		if mode=='json':
+			sinfo_str = json.dumps(sinfo)
+			sinfo_int = np.array(map(ord, sinfo_str), dtype=np.int)
+			sinfo_float_buf = sinfo_int.astype(np.float, casting='safe')
+		elif mode=='pickle':
+			pss = pickle.dumps(sinfo, protocol=2)
+			# Pad to have a multiple of 8 bytes
+			s1arr = np.frombuffer(pss, dtype='S1')
+			ll = len(s1arr)
+			s1arr_padded = np.concatenate((s1arr, np.zeros(8-ll%8, dtype='S1')))
+			# Cast to array of floats
+			f8arr = np.frombuffer(s1arr_padded, dtype=np.float64)
+			sinfo_float_buf = np.concatenate((np.array([ll], dtype=np.float64),f8arr))
 		else:
-			sl_info_buf = np.array([1.,
-							beam.slice_info['z_bin_center'],
-							beam.slice_info['z_bin_right'],
-							beam.slice_info['z_bin_left']])
-		
-		buf = np.concatenate((
-			np.array([float(beam.macroparticlenumber)]),
-			np.array([float(beam.particlenumber_per_mp)]), 
+			raise ValueError('Unknown mode!')
+
+		
+		buf = np.array(np.concatenate((
+			np.array([np.float64(beam.macroparticlenumber)]),
+			np.array([np.float64(beam.particlenumber_per_mp)]), 
 			np.array([beam.charge]),
 			np.array([beam.mass]),
 			np.array([beam.circumference]),
 			np.array([beam.gamma]),
 			np.atleast_1d(np.float_(beam.id)),
 			beam.x, beam.xp, beam.y, beam.yp, beam.z, beam.dp,
-			sl_info_buf))
+			np.array([float(len(sinfo_float_buf))]),sinfo_float_buf)), dtype=np.float64)
+
+		if verbose:
+			print('mode=%s'%mode)
+			print('beam.macroparticlenumber:%d'%beam.macroparticlenumber)
+			print('len(buf):%d'%len(buf))
+			print('len(sinfo_float_buf):%d'%len(sinfo_float_buf))
 
 	return buf
 
-def buffer_2_beam(buf):
+def buffer_2_beam(buf, mode='pickle'):
 
 	if buf[0]<0:
 		beam=None
@@ -121,8 +153,11 @@ def buffer_2_beam(buf):
 		dp = buf[i_buf:i_buf+macroparticlenumber]
 		i_buf += macroparticlenumber
 
-		slice_info_buf = buf[i_buf:i_buf+4]
-		i_buf += 4
+		slice_info_buf_size = int(buf[i_buf])
+		i_buf += 1
+		
+		slice_info_buf = buf[i_buf:i_buf+slice_info_buf_size]
+		i_buf += slice_info_buf_size
 
 		beam = Particles(macroparticlenumber=macroparticlenumber,
 						particlenumber_per_mp=particlenumber_per_mp, charge=charge,
@@ -136,15 +171,35 @@ def buffer_2_beam(buf):
 								'dp':np.atleast_1d(dp)})
 
 		beam.id = np.atleast_1d(id_)
+
 
-		if slice_info_buf[0] < 0.:
-			beam.slice_info = None
-		elif slice_info_buf[0] == 0.:
-			beam.slice_info = 'unsliced'
+		if mode=='json':
+			si_int = slice_info_buf.astype(np.int)
+			si_str = ''.join(map(unichr, list(si_int)))
+			beam.slice_info = json.loads(si_str)
+		elif mode=='pickle':
+			# Get length in bytes
+			llrec = int(slice_info_buf[0])
+			s1back_padded = np.frombuffer(slice_info_buf[1:].tobytes(), dtype='S1') 
+			s1back = s1back_padded[:llrec]
+			pss_rec = s1back.tobytes()
+			beam.slice_info = pickle.loads(pss_rec)
 		else:
-			beam.slice_info = {\
-                    'z_bin_center':slice_info_buf[1] ,
-                    'z_bin_right':slice_info_buf[2],
-                    'z_bin_left':slice_info_buf[3]}
-	
+			raise ValueError('Unknown mode!')
+
+
+		
+		# if slice_info_buf[0] < 0.:
+		# 	beam.slice_info = None
+		# elif slice_info_buf[0] == 0.:
+		# 	beam.slice_info = 'unsliced'
+		# else:
+		# 	beam.slice_info = {\
+  #                   'z_bin_center':slice_info_buf[1] ,
+  #                   'z_bin_right':slice_info_buf[2],
+  #                   'z_bin_left':slice_info_buf[3]}
+
+
+
+
 	return beam

gen_multibunch_beam.py

@@ -0,0 +1,59 @@
+from scipy.constants import c as clight, e as qe
+import numpy as np
+
+from PyHEADTAIL.particles.slicing import UniformBinSlicer
+
+def gen_matched_multibunch_beam(machine, n_macroparticles_per_bunch, filling_pattern, 
+        b_spac_s, bunch_intensity, epsn_x, epsn_y, sigma_z, non_linear_long_matching, min_inten_slice4EC):
+
+    bucket_length_m = machine.circumference/(machine.longitudinal_map.harmonics[0])
+    b_spac_m =  b_spac_s*machine.beta*clight
+    b_spac_buckets = np.round(b_spac_m/bucket_length_m)
+
+    if np.abs(b_spac_buckets*bucket_length_m-b_spac_m)/b_spac_m > 0.03:
+        raise ValueError('Bunch spacing is not a multiple of the bucket length!')
+
+    if non_linear_long_matching:
+        generate_bunch = machine.generate_6D_Gaussian_bunch_matched
+    else:
+        generate_bunch = machine.generate_6D_Gaussian_bunch
+
+    list_genbunches = []
+    for i_slot, inten_slot in enumerate(filling_pattern):
+        if inten_slot>0:
+            bunch = generate_bunch(n_macroparticles_per_bunch, inten_slot*bunch_intensity, epsn_x, epsn_y, sigma_z=sigma_z)
+            bunch.z -= b_spac_buckets*bucket_length_m*i_slot
+            list_genbunches.append(bunch)
+
+    beam = sum(list_genbunches)
+
+    bucket = machine.longitudinal_map.get_bucket(gamma=machine.gamma, mass=machine.mass, charge=machine.charge)
+    # z_beam_center = bucket.z_ufp_separatrix + bucket_length - self.circumference/2.
+
+    # Here the center of the bucket
+    bucket.z_sfp
+
+    # I want to re-separate the bunches
+    buncher = UniformBinSlicer(n_slices = 0, z_sample_points = np.arange(bucket.z_sfp-len(filling_pattern)*bucket_length_m*b_spac_buckets, 
+                                            bucket.z_sfp+bucket_length_m, bucket_length_m*b_spac_buckets))
+    buncher_slice_set = beam.get_slices(buncher, statistics=True)
+    list_bunches = beam.extract_slices(buncher, include_non_sliced='never')
+    # The bunch at position 0 is the tail
+    
+    # If last bunch is empty remove it
+    if (list_bunches[0].intensity<min_inten_slice4EC):
+        list_bunches = list_bunches[1:]
+    
+    # Add further information to bunches
+    for i_bb, bb in enumerate(list_bunches[::-1]): # I want bunch 0 at the head of the train
+        slice4EC = bb.intensity>min_inten_slice4EC
+        bb.slice_info['slice_4_EC'] = slice4EC
+        bb.slice_info['interact_with_EC'] = slice4EC
+        bb.slice_info['N_bunches_tot_beam'] = len(list_bunches)
+        bb.slice_info['i_bunch'] = i_bb
+	bb.slice_info['i_turn'] = 0
+        
+
+       
+    
+    return list_bunches