#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""This module provides lattice picture from FLAME lattice file(.lat) or FLAME Machine object."""
from flame import Machine
from flame_utils.core import ModelFlame
from flame_utils.core import BeamState
from flame_utils.io import collect_data
import logging
import matplotlib.lines as lin
import matplotlib.patches as ptc
import matplotlib.pyplot as plt
import numpy as np
_LOGGER = logging.getLogger(__name__)
[docs]def hplot(*args, **kws):
"""Plot beam state history with input arguments.
Parameters
----------
machine : object
FLAME machine or ModelFlame object.
result :
List of beam state, generated by Machine.propagete() or ModelFlame.run() method.
args : str
Beam state attributes to plot, seprated by comma. (see Note)
colors : dict
Color definition of each ``args``.
lattice : bool
Plot lattice layout.
legend : bool
Plot legend of ``args``
cs : int
Index of the charge state to plot.
Notes
-----
List of supported :py:func:`BeamState <flame_utils.core.state.BeamState>` arguments
pos : float
Longitudinally propagating position, [m].
ref_beta : float
Speed in the unit of light velocity in vacuum of reference charge
state, Lorentz beta.
ref_bg : float
Multiplication of beta and gamma of reference charge state.
ref_gamma : float
Relativistic energy of reference charge state, Lorentz gamma.
ref_IonEk : float
Kinetic energy of reference charge state, [eV/u].
ref_IonEs : float
Rest energy of reference charge state, [eV/u].
ref_IonQ : int
Macro particle number of reference charge state.
ref_IonW : float
Total energy of reference charge state, [eV/u],
i.e. :math:`W = E_s + E_k`.
ref_IonZ : float
Reference charge to mass ratio, e.g.
:math:`^{33^{+}}_{238}U: Q[33]/A[238]`.
ref_phis : float
Absolute synchrotron phase of reference charge state, [rad].
ref_SampleIonK : float
Wave-vector in cavities with different beta values of reference charge
state.
beta : Array
Speed in the unit of light velocity in vacuum of all charge states,
Lorentz beta.
bg : Array
Multiplication of beta and gamma of all charge states.
gamma : Array
Relativistic energy of all charge states, Lorentz gamma.
IonEk : Array
Kinetic energy of all charge states, [eV/u].
IonEs : Array
Rest energy of all charge states, [eV/u].
IonQ : Array
Macro particle number of all charge states.
IonW : Array
Total energy of all charge states, [eV/u], i.e. :math:`W = E_s + E_k`.
IonZ : Array
All charge to mass ratios
phis : Array
Absolute synchrotron phase of all charge states, [rad]
SampleIonK : Array
Wave-vector in cavities with different beta values of all charge
states.
x0_env, xcen : float
Weight average of all charge states for x', [rad].
y0_env, ycen : float
Weight average of all charge states for y, [mm].
xp0_env, xpcen : float
Weight average of all charge states for x', [rad].
yp0_env, ypcen : float
Weight average of all charge states for y', [rad].
phi0_env, phicen, zcen: float
Weight average of all charge states for :math:`\phi`, [rad].
dEk0_env, dEkcen, zpcen : float
Weight average of all charge states for :math:`\delta E_k`, [MeV/u].
x0, xcen_all : Array
X centroid for all charge states, [mm].
y0, ycen_all : Array
Y centroid for all charge states, [mm].
xp0, xpcen_all : Array
X centroid divergence for all charge states, [rad].
yp0, ypcen_all : Array
Y centroid divergence for all charge states, [rad].
phi0, phicen_all, zcen_all : Array
Longitudinal beam length, measured in RF frequency for all charge
states, [rad].
dEk0, dEkcen_all, zpcen_all : Array
Kinetic energy deviation w.r.t. reference charge state, for all charge
states, [MeV/u].
x0_rms, xrms : float
General rms beam envelope for x, [mm].
y0_rms, yrms : float
General rms beam envelope for y, [mm].
xp0_rms, xprms : float
General rms beam envelope for x', [rad].
yp0_rms, yprms : float
General rms beam envelope for y', [rad].
phi0_rms, phirms, zrms : float
General rms beam envelope for :math:`\phi`, [rad].
dEk0_rms, dEkrms, zprms : float
General rms beam envelope for :math:`\delta E_k`, [MeV/u].
xrms_all : Array
General rms beam envelope for x of all charge states, [mm].
yrms_all : Array
General rms beam envelope for y of all charge states, [mm].
xprms_all : Array
General rms beam envelope for x' of all charge states, [rad].
yprms_all : Array
General rms beam envelope for y' of all charge states, [rad].
phirms_all : Array
General rms beam envelope for :math:`\phi` of all charge states, [rad].
dEkrms_all : Array
General rms beam envelope for :math:`\delta E_k` of all charge states, [MeV/u].
moment0_env, cenvector : Array
Weight average of centroid for all charge states, array of
``[x, x', y, y', phi, dEk, 1]``, with the units of
``[mm, rad, mm, rad, rad, MeV/u, 1]``.
moment0, cenvector_all : Array
Centroid for all charge states, array of ``[x, x', y, y', phi, dEk, 1]``.
moment0_rms, rmsvector : Array
RMS beam envelope, part of statistical results from ``moment1``.
moment1, beammatrix_all : Array
Covariance matrices of all charge states, for each charge state.
moment1_env, beammatrix : Array
Covariance matrices of all charge states, average over all charge states.
xemittance, xeps : float
Weight average of geometrical x emittance, [mm-mrad].
yemittance, yeps : float
Weight average of geometrical y emittance, [mm-mrad].
zemittance, zeps : float
Weight average of geometrical z emittance, [rad-MeV/u].
xnemittance, xepsn : float
Weight average of normalized x emittance, [mm-mrad].
ynemittance, yepsn : float
Weight average of normalized y emittance, [mm-mrad].
znemittance, zepsn : float
Weight average of normalized z emittance, [rad-MeV/u].
xemittance_all, xeps_all : Array
Geometrical x emittance of all charge states, [mm-mrad].
yemittance_all, yeps_all : Array
Geometrical y emittance of all charge states, [mm-mrad].
zemittance_all, zeps_all : Array
Geometrical z emittance of all charge states, [rad-MeV/u].
xnemittance_all, xepsn_all : Array
Normalized x emittance of all charge states, [mm-mrad].
ynemittance_all, yepsn_all : Array
Normalized y emittance of all charge states, [mm-mrad].
znemittance_all, zepsn_all : Array
Normalized z emittance of all charge states, [rad-MeV/u].
xtwiss_beta, xtwsb : float
Weight average of twiss beta x, [m/rad].
ytwiss_beta, ytwsb : float
Weight average of twiss beta y, [m/rad].
ztwiss_beta, ztwsb : float
Weight average of twiss beta z, [rad/MeV/u].
xtwiss_alpha, xtwsa : float
Weight average of twiss alpha x, [1].
ytwiss_alpha, ytwsa : float
Weight average of twiss alpha y, [1].
ztwiss_alpha, ztwsa : float
Weight average of twiss alpha z, [1].
xtwiss_beta_all, xtwsb_all : Array
Twiss beta x of all charge states, [m/rad].
ytwiss_beta_all, ytwsb_all : Array
Twiss beta y of all charge states, [m/rad].
ztwiss_beta_all, ztwsb_all : Array
Twiss beta z of all charge states, [rad/MeV/u].
xtwiss_alpha_all, xtwsa_all : Array
Twiss alpha x of all charge states, [1].
ytwiss_alpha_all, ytwsa_all : Array
Twiss alpha y of all charge states, [1].
ztwiss_alpha_all, ztwsa_all : Array
Twiss alpha z of all charge states, [1].
couple_xy, cxy : float
Weigth average of normalized x-y coupling term, [1].
couple_xpy, cxpy : float
Weigth average of normalized xp-y coupling term, [1].
couple_xyp, cxyp : float
Weigth average of normalized x-yp coupling term, [1].
couple_xpyp, cxpyp : float
Weigth average of normalized xp-yp coupling term, [1].
couple_xy_all, cxy_all : Array
Normalized x-y coupling term of all charge states, [1].
couple_xpy_all, cxpy_all : Array
Normalized xp-y coupling term of all charge states, [1].
couple_xyp_all, cxyp_all : Array
Normalized x-yp coupling term of all charge states, [1].
couple_xpyp_all, cxpyp_all : Array
Normalized xp-yp coupling term of all charge states, [1].
Examples
--------
>>> fm = ModelFlame(lat_file = 'userfile.lat')
>>> hplot('xrms', 'yrms', machine=fm)
"""
ldct = {'rms': 'RMS',
'cen': 'centroid',
'eps': 'geom. emittance',
'epsn': 'norm. emittance',
'twsb': 'twiss beta',
'twsa': 'twiss alpha'}
cdct = {'x': 'b', 'y': 'r', 'z': 'g'}
if 'colors' in kws:
for key in kws['colors']:
cdct[key] = kws['colors'][key]
lsls = ['-', '--', ':', '-.']
lttc = bool(kws['lattice']) if 'lattice' in kws else True
lgnd = bool(kws['legend']) if 'legend' in kws else True
cs = int(kws['cs']) if 'cs' in kws else 0
m = kws['machine'] if 'machine' in kws else None
r = kws['result'] if 'result' in kws else None
if isinstance(m, str):
m = ModelFlame(lat_file=m)
if m is None and r is None:
_LOGGER.error('Nothing to plot.')
return None
elif r is None:
if isinstance(m, Machine):
s = m.allocState({})
r = m.propagate(s, observe = range(len(m)))
elif isinstance(m, ModelFlame):
r, s = m.run(monitor = 'all')
if r[0][0] != 0:
r = [(0, m.bmstate)] + r
m = m.machine
else:
_LOGGER.error('Unsuported type of machine.')
return None
else:
if isinstance(m, Machine):
pass
elif isinstance(m, ModelFlame):
m = m.machine
else:
_LOGGER.error('Unsuported type of machine.')
return None
keys = ['pos'] + list(args)
d = collect_data(r, *keys)
for key in keys:
if len(d[key].shape) == 2:
d[key] = d[key][:, cs]
if lttc:
if len(args) != 0:
ymax = -1e256
ymin = 1e256
for key in args:
ymax = max([ymax, max(d[key])])
ymin = min([ymin, min(d[key])])
ydif = ymax - ymin
yscl = ydif
if ydif == 0.0:
ydif = ymax*0.1 if ymax != 0.0 else 0.1
yscl = ydif*0.2
else:
ymin = 0
ydif = 0
yscl = 10
l = PlotLat(m, auto_scaling=False)
l.generate(ycen=ymin-0.2*ydif, yscl=0.1*yscl, legend=False, option=False)
cntx = -1
cnty = -1
cntz = -1
cnto = -1
for key in args:
if key[0] == 'x':
color = cdct[key] if key in cdct else cdct['x']
cnt = cntx = (cntx + 1)%4
label = 'x' if cnt == 0 else key
elif key[0] == 'y':
color = cdct[key] if key in cdct else cdct['y']
cnt = cnty = (cnty + 1)%4
label = 'y' if cnt == 0 else key
elif key[0] == 'z':
color = cdct[key] if key in cdct else cdct['z']
cnt = cntz = (cntz + 1)%4
label = 'z' if cnt == 0 else key
else:
color = cdct[key] if key in cdct else None
cnt = cnto = (cnto + 1)%4
label = key
plt.plot(d['pos'], d[key], c=color, ls=lsls[cnt], label = label)
if lgnd and len(args) != 0:
plt.legend(loc='best')
plt.xlabel('z [m]')
plt.xlim([min(d['pos']), max(d['pos'])])
if len(args) != 0:
ykey = args[0]
if ykey[0] in ['x', 'y', 'z']:
ylabel = ykey[1:]
elif ykey[0:4] == 'ref_':
ylabel = ykey[4:]
else:
ylabel = ykey
ylabel = ldct[ylabel] if ylabel in ldct else ylabel
try:
txt = getattr(BeamState, ykey).__doc__
u0 = txt.index('[')
u1 = txt.index(']')
unit = ' ' + txt[u0:u1+1]
except:
unit = ''
plt.ylabel(ylabel + unit)
else:
plt.ylim([-2, 2])
[docs]class PlotLat:
"""Lattice picture class from FLAME lattice file or FLAME Machine object.
Parameters
----------
source : str or callable
File path of the lattic file (str) or FLAME Machine object (callable)
output : str (None), optional
Output file name. If defined, the lattice plot is generated automatically.
auto_scaling : bool (True), optional
Flag for y-axis scaling by strength of the optical elements
starting_offset : float (0.0), optional
Position offset of starting point in the lattice file
Attributes
----------
types : dict
Element type list of the lattice. Each element type contains
on-off 'flag', plotting 'color', and y-axis 'scale'.
"""
def __init__(self, source, output=None, auto_scaling=False , starting_offset=0.0, **kws):
self._source = source
self._auto_scaling = auto_scaling
self._starting_offset = starting_offset
if type(self._source) == str:
with open(self._source, 'rb') as lat:
self.M = Machine(lat)
elif type(self._source) == Machine:
self.M = self._source
else:
raise ValueError('source must be a file path of .lat or flame.Machine object')
self.types = {'rfcavity': {'flag':True, 'name':'rfcavity', 'color':'orange', 'scale':0.0},
'solenoid': {'flag':True, 'name':'solenoid', 'color':'red', 'scale':0.0},
'quadrupole': {'flag':True, 'name':'quad', 'color':'purple', 'scale':0.0},
'sextupole': {'flag':True, 'name':'sext', 'color':'navy', 'scale':0.0},
'sbend': {'flag':True, 'name':'bend', 'color':'green', 'scale':0.0},
'equad': {'flag':True, 'name':'e-quad', 'color':'blue', 'scale':0.0},
'edipole': {'flag':True, 'name':'e-dipole', 'color':'lime', 'scale':0.0},
'bpm': {'flag':True, 'name':'bpm', 'color':'m', 'scale':0.0},
'orbtrim': {'flag':True, 'name':'corr', 'color':'black', 'scale':0.0},
'stripper': {'flag':True, 'name':'stripper', 'color':'y', 'scale':0.0},
'marker': {'flag':True, 'name':'pm', 'color':'c', 'scale':0.0}
}
if self._auto_scaling:
for i in range(len(self.M)):
elem = self.M.conf(i)
if elem['type'] in self.types.keys():
prv_scl = self.types[elem['type']]['scale']
tmp_scl = np.abs(self._get_scl(elem))
self.types[elem['type']]['scale'] = max(prv_scl,tmp_scl)
if isinstance(output, str):
self.generate(legend=True, option=True)
self.output(window=True, fname=output, **kws)
def _get_scl(self, elem):
"""Get arbital strength of the optical element.
"""
scl = 0.0
if elem['type'] == 'rfcavity':
scl = elem['scl_fac']*np.cos(2.0*np.pi*elem['phi']/360.0)
elif elem['type'] == 'solenoid':
scl = elem['B']
elif elem['type'] == 'quadrupole':
scl = elem['B2'] if 'B2' in elem else 1.0
elif elem['type'] == 'sextupole':
scl = elem['B3']
elif elem['type'] == 'sbend':
scl = elem['phi']
elif elem['type'] == 'equad':
scl = elem['V']/elem['radius']**2.0
elif elem['type'] == 'edipole':
scl = elem['phi']
return scl
[docs] def generate(self, start=None, end=None, xlim=None, ycen=0.0, yscl=1.0, aspect=5.0,
legend=True, option=True, axes = None):
"""Generate matplotlib Axes class object from lattice file.
Parameters
----------
start : int
Index of the lattice start.
end : int
Index of the lattice end.
xlim : list[2], optinal
Plot range of the lattice.
ycen : float (0.0)
Vertical offset of the plot.
yscl : float (1.0)
Plot scale of the element size
aspect : float (5.0), optional
Aspect ratio of the picture.
legend : bool
Add legend for the elements.
option : bool
Add optional setting for the plot.
Attributes
----------
axes : callable
Axes class object of matplotlib.
total_length : float
Total length of the lattice.
"""
if axes is None:
self._fig = plt.figure()
self.axes = self._fig.add_subplot(111)
else:
self.axes = axes
pos = self._starting_offset
bp = ycen
indexes = range(len(self.M))[start:end]
foundelm = []
for i in indexes:
elem = self.M.conf(i)
try:
dL = elem['L']
except:
dL = 0.0
if elem['type'] in self.types.keys():
info = self.types[elem['type']]
if foundelm.count(elem['type']) == 0:
foundelm.append(elem['type'])
if legend and info['flag']:
self.axes.fill_between([0,0],[0,0],[0,0], color=info['color'], label=info['name'])
if info['flag']:
if dL != 0.0:
bpp = bp
if info['scale'] != 0.0:
ht = yscl*self._get_scl(elem)/info['scale'] + 0.05
else:
ht = yscl*np.sign(self._get_scl(elem))
if elem['type'] == 'rfcavity' or elem['type'] == 'solenoid':
bpp = bp-yscl*0.7
ht = yscl*2.0*0.7
self.axes.add_patch(ptc.Rectangle((pos, bpp), dL, ht,
edgecolor='none',facecolor=info['color']))
else:
self.axes.add_line(lin.Line2D([pos,pos],[-yscl*0.3+bp, yscl*0.3+bp],color=info['color']))
pos += dL
self.total_length = pos
self.axes.add_line(lin.Line2D([0.0, pos], [bp,bp], color='gray', zorder=-5))
if len(foundelm) <= 4 :
ncol = 4
elif len(foundelm) <= 6:
ncol = 3
elif len(foundelm) <= 8 :
ncol = 4
elif len(foundelm) <= 9:
ncol = 3
else :
ncol = 4
if option:
if xlim != None:
self.axes.set_xlim(xlim)
ancscl = xlim[1]-xlim[0]
else :
self.axes.set_xlim((0.0,pos))
ancscl = pos
if legend:
self.axes.legend(ncol=ncol, loc=10, bbox_to_anchor=(0.5, -0.2*ancscl))
self.axes.set_aspect(aspect)
self.axes.set_ylim((-1.0,1.0))
self.axes.set_yticks(())
self.axes.grid()
self.axes.spines['right'].set_visible(False)
self.axes.spines['left'].set_visible(False)
self.axes.spines['top'].set_visible(False)
self.axes.spines['bottom'].set_visible(False)
plt.tight_layout()
[docs] def output(self,window=True,fname=None,**kws):
"""Output the lattice picture to window and/or file.
Parameters
----------
window : bool (True)
Output flag to the window.
fname : str, optional
File path of the output picutre.
**kwargs :
The same kwargs as pyplot.savefig() are available.
"""
if type(fname) == str :
plt.savefig(fname, **kws)
if window: plt.show()
if __name__ == "__main__":
import sys
try:
lat = sys.argv[1]
except:
raise ValueError('First argument must be a lattice file.')
try:
out = sys.argv[2]
except:
raise ValueError('Second argument must be a output file name.')
pl = PlotLat(lat, output=out)