diff --git a/examples/write_examples.ipynb b/examples/write_examples.ipynb index 4269fd3..8d01617 100644 --- a/examples/write_examples.ipynb +++ b/examples/write_examples.ipynb @@ -128,16 +128,16 @@ "name": "stdout", "output_type": "stream", "text": [ - " -1.289814080985e-05 1.712192978919e-05 0.000000000000e+00 -9.245284702413e-01 -3.316650265292e+00 2.210558337183e+02 1.819664001274e-05 0.000000000000e+00 1 -1\n", - " -1.184861337727e-03 -2.101371437059e-03 0.000000000000e+00 -3.047044656318e+02 -3.039342419008e+02 -1.005645891013e+02 -9.745607002167e-04 1.000000000000e-06 1 -1\n", - " -5.181307340245e-04 -2.178353405029e-03 0.000000000000e+00 5.525456229648e+02 2.416723877028e+02 -6.554342847563e+01 1.280843753434e-03 1.000000000000e-06 1 -1\n", - " -1.773501610902e-03 2.864979597813e-03 0.000000000000e+00 -2.226004747820e+02 9.450238076106e+00 -1.055085411491e+02 3.835366744569e-04 1.000000000000e-06 1 -1\n", - " 1.686555815999e-03 -2.401048305081e-04 0.000000000000e+00 -1.891692499417e+02 4.859547751754e+01 3.339263495319e+02 1.902998338336e-03 1.000000000000e-06 1 -1\n", - " -7.779454935491e-04 -6.800063114796e-04 0.000000000000e+00 6.716138938638e+01 -2.064173000222e+02 -1.405963302134e+02 1.779005092730e-04 1.000000000000e-06 1 -1\n", - " -2.593702199590e-03 -2.301030494125e-03 0.000000000000e+00 -1.455653402031e+01 2.074634953296e+02 -1.397453142110e+02 1.368567098305e-03 1.000000000000e-06 1 -1\n", - " 1.997801509161e-03 2.648416193086e-03 0.000000000000e+00 -2.124047726665e+00 -6.792723569247e+01 6.931081537770e+01 2.616497721112e-04 1.000000000000e-06 1 -1\n", - " 1.999741847023e-03 -6.945690451493e-04 0.000000000000e+00 -9.991142908925e+01 -9.189412445573e+01 2.259539675809e+02 -8.681109991004e-04 1.000000000000e-06 1 -1\n", - " -7.033822974359e-04 -5.677746866954e-04 0.000000000000e+00 7.520962264129e+02 3.125940718167e+02 -1.451032665210e+02 4.315191990002e-04 1.000000000000e-06 1 -1\n" + " -1.289814080985e-05 1.712192978919e-05 0.000000000000e+00 -9.245284702413e-01 -3.316650265292e+00 2.210558337183e+02 1.819664001274e-05 0.000000000000e+00 1 5\n", + " -1.184861337727e-03 -2.101371437059e-03 0.000000000000e+00 -3.047044656318e+02 -3.039342419008e+02 -1.005645891013e+02 -9.745607002167e-04 1.000000000000e-06 1 5\n", + " -5.181307340245e-04 -2.178353405029e-03 0.000000000000e+00 5.525456229648e+02 2.416723877028e+02 -6.554342847563e+01 1.280843753434e-03 1.000000000000e-06 1 5\n", + " -1.773501610902e-03 2.864979597813e-03 0.000000000000e+00 -2.226004747820e+02 9.450238076106e+00 -1.055085411491e+02 3.835366744569e-04 1.000000000000e-06 1 5\n", + " 1.686555815999e-03 -2.401048305081e-04 0.000000000000e+00 -1.891692499417e+02 4.859547751754e+01 3.339263495319e+02 1.902998338336e-03 1.000000000000e-06 1 5\n", + " -7.779454935491e-04 -6.800063114796e-04 0.000000000000e+00 6.716138938638e+01 -2.064173000222e+02 -1.405963302134e+02 1.779005092730e-04 1.000000000000e-06 1 5\n", + " -2.593702199590e-03 -2.301030494125e-03 0.000000000000e+00 -1.455653402031e+01 2.074634953296e+02 -1.397453142110e+02 1.368567098305e-03 1.000000000000e-06 1 5\n", + " 1.997801509161e-03 2.648416193086e-03 0.000000000000e+00 -2.124047726665e+00 -6.792723569247e+01 6.931081537770e+01 2.616497721112e-04 1.000000000000e-06 1 5\n", + " 1.999741847023e-03 -6.945690451493e-04 0.000000000000e+00 -9.991142908925e+01 -9.189412445573e+01 2.259539675809e+02 -8.681109991004e-04 1.000000000000e-06 1 5\n", + " -7.033822974359e-04 -5.677746866954e-04 0.000000000000e+00 7.520962264129e+02 3.125940718167e+02 -1.451032665210e+02 4.315191990002e-04 1.000000000000e-06 1 5\n" ] } ], @@ -145,6 +145,20 @@ "!head astra_particles.txt" ] }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [], + "source": [ + "# Check the readback\n", + "from pmd_beamphysics.interfaces.astra import parse_astra_phase_file\n", + "import numpy as np\n", + "P1 = ParticleGroup(data=parse_astra_phase_file('astra_particles.txt'))\n", + "for k in ['x', 'px', 'y', 'py', 'z', 'pz']:\n", + " assert np.allclose(P[k], P1[k])" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -154,7 +168,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 11, "metadata": {}, "outputs": [], "source": [ @@ -163,7 +177,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 12, "metadata": {}, "outputs": [ { @@ -196,7 +210,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 13, "metadata": {}, "outputs": [ { @@ -213,7 +227,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 14, "metadata": {}, "outputs": [ { @@ -256,7 +270,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 15, "metadata": {}, "outputs": [ { @@ -273,7 +287,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 16, "metadata": {}, "outputs": [ { @@ -306,7 +320,7 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 17, "metadata": {}, "outputs": [ { @@ -323,7 +337,7 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 18, "metadata": {}, "outputs": [ { @@ -355,7 +369,7 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": 19, "metadata": {}, "outputs": [ { @@ -363,7 +377,7 @@ "output_type": "stream", "text": [ "writing 10000 particles to gpt_particles.txt\n", - "ASCII particles written. Convert to GDF using: asci2df -o particles.gdf gpt_particles.txt\n" + "ASCII particles written. Convert to GDF using: asci2gdf -o particles.gdf gpt_particles.txt\n" ] } ], @@ -374,7 +388,7 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": 20, "metadata": {}, "outputs": [ { @@ -392,7 +406,7 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 21, "metadata": {}, "outputs": [ { @@ -425,7 +439,7 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": 22, "metadata": {}, "outputs": [], "source": [ @@ -435,7 +449,7 @@ }, { "cell_type": "code", - "execution_count": 22, + "execution_count": 23, "metadata": {}, "outputs": [ { @@ -447,7 +461,7 @@ " 'Flagimg': 0}" ] }, - "execution_count": 22, + "execution_count": 23, "metadata": {}, "output_type": "execute_result" } @@ -459,7 +473,7 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": 24, "metadata": {}, "outputs": [ { @@ -492,7 +506,7 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": 25, "metadata": {}, "outputs": [], "source": [ @@ -502,7 +516,7 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": 26, "metadata": {}, "outputs": [ { @@ -519,7 +533,7 @@ "'litrack.zd'" ] }, - "execution_count": 25, + "execution_count": 26, "metadata": {}, "output_type": "execute_result" } @@ -531,7 +545,7 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": 27, "metadata": {}, "outputs": [ { @@ -574,7 +588,7 @@ }, { "cell_type": "code", - "execution_count": 27, + "execution_count": 28, "metadata": {}, "outputs": [ { @@ -591,7 +605,7 @@ }, { "cell_type": "code", - "execution_count": 28, + "execution_count": 29, "metadata": {}, "outputs": [ { @@ -605,10 +619,10 @@ { "data": { "text/plain": [ - "" + "" ] }, - "execution_count": 28, + "execution_count": 29, "metadata": {}, "output_type": "execute_result" } @@ -622,7 +636,7 @@ }, { "cell_type": "code", - "execution_count": 29, + "execution_count": 30, "metadata": {}, "outputs": [ { @@ -631,7 +645,7 @@ "['BEGINNING']" ] }, - "execution_count": 29, + "execution_count": 30, "metadata": {}, "output_type": "execute_result" } @@ -650,7 +664,7 @@ }, { "cell_type": "code", - "execution_count": 30, + "execution_count": 31, "metadata": {}, "outputs": [], "source": [ @@ -663,7 +677,7 @@ }, { "cell_type": "code", - "execution_count": 31, + "execution_count": 32, "metadata": {}, "outputs": [ { @@ -689,7 +703,7 @@ }, { "cell_type": "code", - "execution_count": 32, + "execution_count": 33, "metadata": {}, "outputs": [], "source": [ @@ -700,7 +714,7 @@ }, { "cell_type": "code", - "execution_count": 33, + "execution_count": 34, "metadata": {}, "outputs": [ { @@ -733,7 +747,7 @@ }, { "cell_type": "code", - "execution_count": 34, + "execution_count": 35, "metadata": {}, "outputs": [], "source": [ @@ -742,7 +756,7 @@ }, { "cell_type": "code", - "execution_count": 35, + "execution_count": 36, "metadata": {}, "outputs": [], "source": [ @@ -751,7 +765,7 @@ }, { "cell_type": "code", - "execution_count": 36, + "execution_count": 37, "metadata": {}, "outputs": [], "source": [ @@ -760,7 +774,7 @@ }, { "cell_type": "code", - "execution_count": 37, + "execution_count": 38, "metadata": {}, "outputs": [], "source": [ diff --git a/pmd_beamphysics/interfaces/astra.py b/pmd_beamphysics/interfaces/astra.py index c5cc378..5b1b89c 100644 --- a/pmd_beamphysics/interfaces/astra.py +++ b/pmd_beamphysics/interfaces/astra.py @@ -1,10 +1,122 @@ import numpy as np +from pmd_beamphysics.readers import component_alias +import os + + + +astra_species_name = {1:'electron', 2:'positron', 3:'proton', 4:'hydrogen'} +astra_species_index = {v:k for k, v in astra_species_name.items()} # Inverse mapping + +astra_particle_status_names = {-1:'standard particle, at the cathode', + 3:'trajectory probe particle', + 5:'standard particle'} + + + + +def parse_astra_phase_file(filePath): + """ + + Parses astra particle dumps to data dict, that corresponds to the + openpmd-beamphysics ParticeGroup data= input. + + Units are in m, s, eV/c + + Live particles (status==5) are relabeled as status = 1. + Original status == 2 are relabeled to status = 2 (previously unused by Astra) + + """ + + # + # Internal Astra Columns + # x y z px py pz t macho_charge astra_index status_flag + # m m m eV/c eV/c eV/c ns nC 1 1 + + # The first line is the reference particle in absolute corrdinate. Subsequent particles have: + # z pz t + # relative to the reference. + # + # + # astra_index represents the species: 1:electrons, 2:positrons, 3:protons, 4:hydroger, ... + # There is a large table of status. Status_flag = 5 is a standard particle. + + assert os.path.exists(filePath), f'particle file does not exist: {filePath}' + + data = np.loadtxt(filePath) + ref = data[0,:] # Reference particle. + + # position in m + x = data[1:,0] + y = data[1:,1] + + z_rel = data[1:,2] + z_ref = ref[2] + #z = z_rel + z_ref + + # momenta in eV/c + px = data[1:,3] + py = data[1:,4] + pz_rel = data[1:,5] + + pz_ref = ref[5] + #pz = pz_rel + pz_ref + + # Time in seconds + t_ref = ref[6]*1e-9 + t_rel = data[1:,6]*1e-9 + #t = t_rel + t_ref + + # macro charge in Coulomb. The sign doesn't matter, so make positive + qmacro = np.abs(data[1:,7]*1e-9) + + species_index = data[1:,8].astype(np.int) + status = data[1:,9].astype(np.int) + + # Select particle by status + #probe_particles = np.where(status == 3) + #good_particles = np.where(status == 5) + + data = {} + + n_particle = len(x) + + data['x'] = x + data['y'] = y + data['z'] = z_rel + z_ref + data['px'] = px + data['py'] = py + data['pz'] = pz_rel + pz_ref + data['t_clock'] = t_rel + t_ref #np.full(n_particle, t_ref) # full array + data['t'] = t_ref + + # Status + # The standard defines 1 as a live particle, but astra uses 1 as a 'passive' particle + # and 5 as a 'standard' particle. 2 is not used. + # To preserve this information, make 1->2 and then 5->1 + where_1 = np.where(status==1) + where_5 = np.where(status == 5) + status[where_1] = 2 + status[where_5] = 1 + data['status'] = status + + data['weight'] = qmacro + + unique_species = set(species_index) + assert len(unique_species) == 1, 'All species must be the same' + + # Scalars + data['species'] = astra_species_name[list(unique_species)[0]] + data['n_particle'] = n_particle + + return data + + + def write_astra(particle_group, outfile, verbose=False, - species='electron', probe=False): """ Writes Astra style particles from particle_group type data. @@ -19,8 +131,6 @@ def vprint(*a, **k): print(*a, **k) vprint(f'writing {particle_group.n_particle} particles to {outfile}') - - assert species == 'electron' # TODO: add more species # number of lines in file size = particle_group.n_particle + 1 # Allow one for reference particle @@ -54,8 +164,23 @@ def vprint(*a, **k): data['q'][i_start:] = particle_group.weight*1e9 # C -> nC # Set these to be the same - data['index'] = 1 # electron - data['status'] = -1 # Particle at cathode + data['index'] = astra_species_index[particle_group.species] + + # Status + # The standard defines 1 as a live particle, but astra uses 1 as a 'passive' particle + # and 5 as a 'standard' particle. 2 is not used. + # On parsing 1->2 and then 5->1 + # Revese: 1->5, 2->1 + status = particle_group.status + astra_status = status.copy() + astra_status[ np.where(status==1) ] = 5 + astra_status[ np.where(status==2) ] = 1 + data['status'][i_start:] = astra_status + # Handle reference particle. If any -1 are found, assume we are starting at the cathode + if -1 in astra_status: + ref_particle['status']= -1 # At cathode + else: + ref_particle['status']= 5 # standard particle # Subtract off reference z, pz, t for k in ['z', 'pz', 't']: @@ -79,3 +204,214 @@ def vprint(*a, **k): # Save in the 'high_res = T' format np.savetxt(outfile, data, fmt = ' '.join(8*['%20.12e']+2*['%4i'])) + + +def vec_spacing(vec): + """ + Returns the spacing and minimum of a coordinate. + Asserts that the spacing is uniform + """ + vmin = vec.min() + vmax = vec.max() + n = len(vec) + assert np.allclose(np.linspace(vmin, vmax, n), vec), 'Non-uniform spacing detected!' + #return (vmax-vmin)/(n-1) + return np.mean(np.diff(vec)) + + +def parse_astra_fieldmap_3d(filePath, frequency=0): + """ + Parses a single Astra 3D fieldmap TXT file, described in the Astra manual. + + The format is: + Nx x[1] x[2] ....... x[Nx-1] x[Nx] + Ny y[1] y[2] ....... y[Ny-1] y[Ny] + Nz z[1] z[2] ....... z[Nz-1] z[Nz] + F[ 1, 1, 1] F[ 2, 1, 1] ... F[Nx, 1, 1] F[ 1, 2, 1] F[ 2, 2, 1]... F[Nx, 2, 1] + F[ 1,Ny,Nz] F[ 2,Ny,Nz]................... F[Nx,Ny,Nz] + where the items can be written in a free format, ignoring line breaks. + + This routine should be used by: + read_astra_3d_fieldmaps + + + Parameters + ---------- + filePath : str + A single 3d fieldmap file, with extension in : + 'ex', 'ey', 'ez', 'bx', 'by', 'bz' + + Returns + ------- + attrs : dict of attributes + + components: dict of one component, named from the extension. + + + Notes + ----- + + Data to be written back to the file should be done as: + components['Ex'].reshape(nx, ny*nz, order='F').T + + + """ + # Get as a flat array. + txt = open(filePath).read().split() + dat = np.asarray(txt, dtype=float) + + # Pick out Nx, Ny, Nz + nx = int(dat[0]) + ny = int(dat[nx+1]) + nz = int(dat[nx+ny+2]) + + # Pick out coordinate vectors. These can be irregular! + xvec = dat[1:nx+1] + yvec = dat[nx+2:nx+ny+2] + zvec = dat[nx+ny+3:nx+ny+nz+3] + + # Get the grid data in 3D. + # To write back to the file: grid.reshape(nx, ny*nz, order='F').T + grid = dat[nx+ny+nz+3:].reshape(nx, ny, nz, order='F') + + + # Debug + #if raw: + # out = {} + # out['xvec'] = xvec + # out['yvec'] = yvec + # out['zvec'] = zvec + # out['grid'] = grid + # + # return out + + # Form proper attrs + dx = vec_spacing(xvec) + dy = vec_spacing(yvec) + dz = vec_spacing(zvec) + + attrs = {} + attrs['eleAnchorPt'] = 'beginning' + attrs['gridGeometry'] = 'rectangular' + attrs['axisLabels'] = ('x', 'y', 'z') + attrs['gridLowerBound'] = (0, 0, 0) + attrs['gridOriginOffset'] = (xvec.min(), yvec.min(), zvec.min()) + attrs['gridSpacing'] = (dx, dy, dz) + attrs['gridSize'] = (nx, ny, nz) + attrs['harmonic'] = 1 + attrs['fundamentalFrequency'] = frequency + + + # Get a name (actually an alias Ex, By, ...) + component_name = filePath.split('.')[-1].title() + + # There is only one component + components = {component_name:grid} + + return attrs, components + + +def read_astra_3d_fieldmaps(common_filePath, frequency=0): + """ + Reads multiple files from the common_filePath without the ex, + and returns a data dict to instantiate a FieldMesh object: + + Examples + -------- + + data = read_astra_3d_fieldmaps('3D_file_from_astra') + # will parse 3D_file_from_astra.bx, .ey, etc. + FieldMesh(data=data) + + + + Parameters + ---------- + common_filePath: str + File path of the common fieldmap without the extension. + The actual field files should have extensions in: + 'ex', 'ey', 'ez', 'bx', 'by', 'bz' + + Returns + ------- + data : dict + dict of 'attrs' and 'components' to instantiate a FieldMesh object + + + + Notes + ----- + This can only accept regular grids. Irregular grids are not in the openPMD standard. + + """ + attrs = {} + components = {} + for ext in ['ex', 'ey', 'ez', 'bx', 'by', 'bz']: + file = f'{common_filePath}.{ext}' + + if not os.path.exists(file): + continue + + # Actually parse + attrs1, components1 = parse_astra_fieldmap_3d(file, frequency=frequency) + + # Assert that the attrs are the same + if not attrs: + attrs = attrs1 + else: + assert attrs1 == attrs, 'attrs do not match' + + components.update(components1) + + data= dict(attrs=attrs, components=components) + + return data + + + +def write_astra_3d_fieldmaps(fieldmesh_object, common_filePath): + """ + + Parameters + ---------- + fieldmesh_object : FieldMesh + .geometry must be 'rectangular' + + common_filePath : str + common filename to write the .ex, .by etc. files to. + + + Returns + ------- + flist : list of files written (str) + """ + + base = os.path.split(common_filePath)[1] + assert base[0:2]=='3D' or base[3:5]=='3D', 'The base filename must begin with 3D or have 3D starting at the fourth character, according to Astra' + + assert fieldmesh_object.geometry == 'rectangular' + + nx, ny, nz = fieldmesh_object.shape + + flist = [] + for comp in fieldmesh_object.components: + # This makes the correct extension + ext = component_alias[comp].lower() + fname = f'{common_filePath}.{ext}' + flist.append(fname) + + # Form header + header = '' + for key in ['x', 'y', 'z']: + vec = fieldmesh_object.coord_vec(key) + vec = np.around(vec, 15) # remove tiny irregularities + vstr = ' '.join(vec.astype(str)) + header += f'{len(vec)} {vstr}\n' + + # 2D data to write + dat = fieldmesh_object.components[comp].reshape(nx, ny*nz, order='F').T + + np.savetxt(fname, dat, header=header, comments='') + + return flist +