Skip to content

Commit

Permalink
Merge pull request #45 from ChristopherMayes/astra_fixes
Browse files Browse the repository at this point in the history
Add probe particle option
  • Loading branch information
ChristopherMayes authored Aug 2, 2023
2 parents 8a209eb + 811352c commit 0b589b7
Showing 4 changed files with 243 additions and 626 deletions.
764 changes: 143 additions & 621 deletions docs/examples/fields/field_examples.ipynb

Large diffs are not rendered by default.

15 changes: 13 additions & 2 deletions pmd_beamphysics/fields/fieldmesh.py
Original file line number Diff line number Diff line change
@@ -9,7 +9,7 @@
from pmd_beamphysics.plot import plot_fieldmesh_cylindrical_2d, plot_fieldmesh_cylindrical_1d

from pmd_beamphysics.interfaces.ansys import read_ansys_ascii_3d_fields
from pmd_beamphysics.interfaces.astra import read_astra_3d_fieldmaps, write_astra_3d_fieldmaps
from pmd_beamphysics.interfaces.astra import write_astra_1d_fieldmap, read_astra_3d_fieldmaps, write_astra_3d_fieldmaps, astra_1d_fieldmap_data
from pmd_beamphysics.interfaces.gpt import write_gpt_fieldmesh
from pmd_beamphysics.interfaces.impact import create_impact_solrf_fieldmap_fourier, create_impact_solrf_ele
from pmd_beamphysics.interfaces.superfish import write_superfish_t7, read_superfish_t7
@@ -102,8 +102,10 @@ class FieldMesh:
Writers
- `.write`
- `.write_astra_1d`
- `.write_astra_3d`
- `.to_cylindrical`
- `.to_astra_1d`
- `.to_impact_solrf`
- `.write_gpt`
- `.write_superfish`
@@ -345,7 +347,16 @@ def write(self, h5, name=None):
g = h5

write_pmd_field(g, self.data, name=name)


@functools.wraps(write_astra_1d_fieldmap)
def write_astra_1d(self, filePath):
return write_astra_1d_fieldmap(self, filePath)

def to_astra_1d(self):
z, fz = astra_1d_fieldmap_data(self)
dat = np.array([z, fz]).T
return {'attrs': {'type': 'astra_1d'}, 'data': dat}

def write_astra_3d(self, common_filePath, verbose=False):
return write_astra_3d_fieldmaps(self, common_filePath)

85 changes: 84 additions & 1 deletion pmd_beamphysics/interfaces/astra.py
Original file line number Diff line number Diff line change
@@ -153,7 +153,10 @@ def vprint(*a, **k):
sigma = {}
for k in ['x', 'y', 'z', 'px', 'py', 'pz', 't']:
ref_particle[k] = particle_group.avg(k)
sigma[k] = particle_group.std(k)
std = particle_group.std(k)
if std == 0:
std = 1e-12 # Give some size
sigma[k] = std
ref_particle['t'] *= 1e9 # s -> nS

# Make structured array
@@ -203,13 +206,93 @@ def vprint(*a, **k):
data[5]['x'] = 1.5*sigma['x'];data[5]['t'] = 1.5*sigma['t']
data[6]['y'] = 1.5*sigma['y'];data[6]['t'] = -1.5*sigma['t']
data[1:7]['status'] = -3
data[1:7]['q'] = 0.5e-5 # ? Seems to be required
data[1:7]['pz'] = 0 #? This is what the Astra Generator does

# Save in the 'high_res = T' format
np.savetxt(outfile, data, fmt = ' '.join(8*['%20.12e']+2*['%4i']))




def write_astra_1d_fieldmap(fm, filePath):
"""
Writes an Astra fieldmap file from a FieldMesh object.
Requires cylindrical geometry for now.
Parameters
----------
filePath: str
Filename to write to
"""

z, fz = astra_1d_fieldmap_data(fm)

# Flatten dat
dat = np.array([z, fz]).T

np.savetxt(filePath, dat, comments='')



def astra_1d_fieldmap_data(fm):
"""
Astra fieldmap data.
Requires cylindrical geometry for now.
Returns
-------
z: array-like
z coordinate in meters
fz: array-like
field amplitude corresponding to z
"""


assert fm.geometry == 'cylindrical', f'Geometry: {fm.geometry} not implemented'

assert fm.shape[1] == 1, 'Cylindrical symmetry required'

dat = {}
z = fm.coord_vec('z')

if fm.is_static:
if fm.is_pure_magnetic:
fz = np.real(fm['Bz'][0,0,:])
elif fm.is_pure_electric:
fz = np.real(fm['Ez'][0,0,:])
else:
raise ValueError('Mixed static field TODO')

else:
# Assume RF cavity, rotate onto the real axis
# Get complex fields
fz = fm.Ez[0,0,:]

# Get imaginary argument (angle)
iangle = np.unique(np.mod(np.angle(fz), np.pi))

# Make sure there is only one
assert len(iangle) == 1, print(iangle)
iangle = iangle[0]

if iangle != 0:
rot = np.exp(-1j*iangle)
# print(f'Rotating complex field by {-iangle}')
else:
rot = 1
fz = np.real(fz*rot)

return z, fz




def vec_spacing(vec):
"""
Returns the spacing and minimum of a coordinate.
5 changes: 3 additions & 2 deletions pmd_beamphysics/particles.py
Original file line number Diff line number Diff line change
@@ -759,8 +759,9 @@ def drift_to_t(self, t=None):
self.t = np.full(self.n_particle, t)

# Writers
def write_astra(self, filePath, verbose=False):
write_astra(self, filePath, verbose=verbose)
@functools.wraps(write_astra)
def write_astra(self, filePath, verbose=False, probe=False):
write_astra(self, filePath, verbose=verbose, probe=probe)

def write_bmad(self, filePath, p0c=None, t_ref=0, verbose=False):
write_bmad(self, filePath, p0c=p0c, t_ref=t_ref, verbose=verbose)

0 comments on commit 0b589b7

Please sign in to comment.