def read_lucretia(filename, ele_name='', t_ref=0, kill_dead_particles=True, verbose=False): """ Lucretia's format is described in: https://www.slac.stanford.edu/accel/ilc/codes/Lucretia/web/beam.html One Lucretia ".mat" file can include beams at multiple lattice elements. To find the beam at one element, one has to follow down this order of "fields": bstore >> ele_name >> Bunch >> x, in which x is a 6-to-Np array with: Lucretia x = x in m Lucretia px = px/p in radian Lucretia y = y in m Lucretia py = py/p in radian Lucretia z = (t - t_ref)*c in m Lucretia p = p in GeV/c Note that p is the total, not reference, momentum. t_ref is zero by default. To access valid element names in a Lucretia beam file, do: dat = sio.loadmat('filename.mat') print(dat['bstore'].dtype) """ mdat = sio.loadmat(filename) coords = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['x'][0, 0] charges = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['Q'][0, 0] Np = coords.shape[1] x = coords[0] px_luc = coords[1] # normalized by total momentum y = coords[2] py_luc = coords[3] # normalized by total momentum z_luc = coords[4] ptot = coords[5] # total momentum in GeV/c px = px_luc * ptot * 1E9 # in eV/c py = py_luc * ptot * 1E9 pz = np.sqrt((ptot * 1E9)**2 - px**2 - py**2) t = z_luc / 299792458 + t_ref status = np.ones(Np) ix = np.where(ptot == 0) status[ix] = 0 n_dead = len(ix[0]) if verbose: print(Np, 'particles detected,', n_dead, 'found dead!') data = { 'x': x, 'px': px, 'y': y, 'py': py, 'z': np.zeros(Np), 'pz': pz, 't': t, 'status': status, 'weight': charges, 'species': 'electron' } P = ParticleGroup(data=data) if (kill_dead_particles): if verbose: print('Excluding dead particles (if any)...') P = P.where(P.p > 0) return P
def read_lucretia(filename, ele_name='BEGINNING', t_ref=0, exclude_dead_particles=True, verbose=False): """ Load one beam in a Lucretia beam file into a ParticleGroup Parameters: ---------- filename : str Lucretia '.mat' file name. ele_name : str name of the element at which the beam is located. An invalid name results in an error. If the beam file has one element, this only one beam is read. Default: 'BEGINNING' t_ref : float, optional reference time of the beam in seconds. Default: 0. exclude_dead_particles : bool, optional if True, excludes dead particles. Default: True. ---------- Returns: ---------- A ParticleGroup object ---------- Lucretia's format is described in: https://www.slac.stanford.edu/accel/ilc/codes/Lucretia/web/beam.html One Lucretia ".mat" file can include beams at multiple lattice elements. To find the beam at one element, one has to follow down this order of "fields": bstore >> ele_name >> Bunch >> x, in which x is a 6-to-Np array with: Lucretia x = x in m Lucretia px = px/p in radian Lucretia y = y in m Lucretia py = py/p in radian Lucretia z = (t - t_ref)*c in m Lucretia p = p in GeV/c Note that p is the total, not reference, momentum. To access valid element names in a Lucretia beam file, use the helper function list_element_names(filename). dat = sio.loadmat('filename.mat') print(dat['bstore'].dtype) """ ele_list = list_element_names(filename) if verbose: print(len(ele_list), 'elements found in the file!') # Check if the element exists if (ele_name not in ele_list): raise ValueError('The provided element name ' + str(ele_name) + ' does not exist in the file!') elif (len(ele_list) == 1): ele_name = ele_list[0] mdat = sio.loadmat(filename) coords = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['x'][0, 0] charges = mdat['bstore'][ele_name][0, 0]['Bunch'][0, 0]['Q'][0, 0] Np = coords.shape[1] x = coords[0] px_luc = coords[1] # normalized by total momentum y = coords[2] py_luc = coords[3] # normalized by total momentum z_luc = coords[4] ptot = coords[5] # total momentum in GeV/c px = px_luc * ptot * 1E9 # in eV/c py = py_luc * ptot * 1E9 pz = np.sqrt((ptot * 1E9)**2 - px**2 - py**2) t = z_luc / 299792458 + t_ref status = np.ones(Np) ix = np.where(ptot == 0) status[ix] = 0 n_dead = len(ix[0]) if verbose: print(Np, 'particles detected,', n_dead, 'found dead!') data = { 'x': x, 'px': px, 'y': y, 'py': py, 'z': np.zeros(Np), 'pz': pz, 't': t, 'status': status, 'weight': charges, 'species': 'electron' } P = ParticleGroup(data=data) if (exclude_dead_particles): if verbose: print('Excluding dead particles (if any)...') P = P.where(P.p > 0) return P