Bunch Commands¶

In [1]:

from pytao import Tao

from pytao import Tao

In [2]:

import numpy as np
import matplotlib.pyplot as plt

import numpy as np import matplotlib.pyplot as plt

Initialize Tao on the CSR beam tracking example¶

In [3]:

tao = Tao(
    "-init $ACC_ROOT_DIR/bmad-doc/tao_examples/csr_beam_tracking/tao.init -noplot"
)

tao = Tao( "-init $ACC_ROOT_DIR/bmad-doc/tao_examples/csr_beam_tracking/tao.init -noplot" )

bunch_params¶

Bunch statistics can be retrieved from any element as a dict.

In [4]:

stats = tao.bunch_params("end")
stats

stats = tao.bunch_params("end") stats

Out[4]:

{'twiss_beta_x': 0.30134457126613,
 'twiss_alpha_x': -2.15210497361386,
 'twiss_gamma_x': 18.6880944753442,
 'twiss_phi_x': 0.0,
 'twiss_eta_x': -0.0481939319203434,
 'twiss_etap_x': -0.454973392151944,
 'twiss_sigma_x': 6.05472920750717e-05,
 'twiss_sigma_p_x': 0.000476810212736377,
 'twiss_emit_x': 1.21653911408494e-08,
 'twiss_norm_emit_x': 9.99817439089443e-07,
 'twiss_beta_y': 0.407832041608165,
 'twiss_alpha_y': 1.94408133190086,
 'twiss_gamma_y': 11.7191680334852,
 'twiss_phi_y': 0.0,
 'twiss_eta_y': -0.0457317680153808,
 'twiss_etap_y': 0.0263984769451972,
 'twiss_sigma_y': 7.04370887189114e-05,
 'twiss_sigma_p_y': 0.000377580116845047,
 'twiss_emit_y': 1.21652615808019e-08,
 'twiss_norm_emit_y': 9.99806791146156e-07,
 'twiss_beta_z': 95.8242541877309,
 'twiss_alpha_z': -1.24059179274334,
 'twiss_gamma_z': 0.0264971328787783,
 'twiss_phi_z': 0.0,
 'twiss_eta_z': 0.0,
 'twiss_etap_z': 0.0,
 'twiss_sigma_z': 0.0008994584510898,
 'twiss_sigma_p_z': 1.49569425187929e-05,
 'twiss_emit_z': 8.44280513419792e-09,
 'twiss_norm_emit_z': 6.9387524907938e-07,
 'twiss_beta_a': 0.248530085331005,
 'twiss_alpha_a': -1.77499320489309,
 'twiss_gamma_a': 15.4136110252092,
 'twiss_phi_a': 0.0,
 'twiss_eta_a': 0.00124844704286057,
 'twiss_etap_a': 0.0111101590168622,
 'twiss_sigma_a': 0.0,
 'twiss_sigma_p_a': 0.0,
 'twiss_emit_a': 1.21664545105205e-08,
 'twiss_norm_emit_a': 9.99904832542641e-07,
 'twiss_beta_b': 0.335544868515057,
 'twiss_alpha_b': 1.59897669921574,
 'twiss_gamma_b': 9.63845045654778,
 'twiss_phi_b': 0.0,
 'twiss_eta_b': 0.00454529158526127,
 'twiss_etap_b': -0.0218933442291171,
 'twiss_sigma_b': 0.0,
 'twiss_sigma_p_b': 0.0,
 'twiss_emit_b': 1.21741461356351e-08,
 'twiss_norm_emit_b': 1.00053697176739e-06,
 'twiss_beta_c': 95.7148956246458,
 'twiss_alpha_c': -1.2389734965013,
 'twiss_gamma_c': 0.0264430499032274,
 'twiss_phi_c': 0.0,
 'twiss_eta_c': 1.2389734965013,
 'twiss_etap_c': 0.0264430499032274,
 'twiss_sigma_c': 0.0,
 'twiss_sigma_p_c': 0.0,
 'twiss_emit_c': 8.43349923229653e-09,
 'twiss_norm_emit_c': 6.931104398842e-07,
 'sigma_11': 3.66649417909144e-09,
 'sigma_12': 2.61861040625002e-08,
 'sigma_13': -1.48398138750058e-13,
 'sigma_14': -4.74391717648667e-14,
 'sigma_15': -4.67916776621505e-10,
 'sigma_16': -1.0781470751519e-11,
 'sigma_21': 2.61861040625002e-08,
 'sigma_22': 2.27394287142704e-07,
 'sigma_23': -1.40343913296626e-12,
 'sigma_24': 4.92122676913874e-14,
 'sigma_25': -4.44903635737791e-09,
 'sigma_26': -1.01782156482131e-10,
 'sigma_31': -1.48398138750058e-13,
 'sigma_32': -1.40343913296626e-12,
 'sigma_33': 4.96185133335392e-09,
 'sigma_34': -2.36505280107631e-08,
 'sigma_35': 1.51286568416869e-13,
 'sigma_36': -1.02306597454638e-11,
 'sigma_41': -4.74391717648667e-14,
 'sigma_42': 4.92122676913874e-14,
 'sigma_43': -2.36505280107631e-08,
 'sigma_44': 1.42566900535742e-07,
 'sigma_45': -3.27450023923243e-13,
 'sigma_46': 5.90560669628933e-12,
 'sigma_51': -4.67916776621505e-10,
 'sigma_52': -4.44903635737791e-09,
 'sigma_53': 1.51286568416869e-13,
 'sigma_54': -3.27450023923243e-13,
 'sigma_55': 8.09025505236861e-07,
 'sigma_56': 1.04740747572173e-08,
 'sigma_61': -1.0781470751519e-11,
 'sigma_62': -1.01782156482131e-10,
 'sigma_63': -1.02306597454638e-11,
 'sigma_64': 5.90560669628933e-12,
 'sigma_65': 1.04740747572173e-08,
 'sigma_66': 2.23710129510474e-10,
 'rel_min_1': -0.000190703788737588,
 'rel_max_1': 0.00017509947414712,
 'centroid_vec_1': -1.0077897456587e-07,
 'rel_min_2': -0.00145920753512278,
 'rel_max_2': 0.00133407063334133,
 'centroid_vec_2': -9.21157518145357e-07,
 'rel_min_3': -0.000205532005074851,
 'rel_max_3': 0.000204275481451832,
 'centroid_vec_3': 5.35418309934568e-13,
 'rel_min_4': -0.00119689779899842,
 'rel_max_4': 0.00125449039844888,
 'centroid_vec_4': -2.53681610205117e-11,
 'rel_min_5': -0.00261715445721329,
 'rel_max_5': 0.00278367096446691,
 'centroid_vec_5': -7.53655670444979e-08,
 'rel_min_6': -1.67983140050979e-05,
 'rel_max_6': 2.80117072986177e-05,
 'centroid_vec_6': -5.77005927578576e-06,
 'centroid_t': 1.48447035006252e-09,
 'centroid_p0c': 41996891.3143949,
 'centroid_beta': 0.999925982822,
 'ix_ele': 8,
 'direction': 1,
 'species': 'Electron',
 'location': 'Downstream_End',
 's': 0.444999999999986,
 't': 1.48447035006252e-09,
 'sigma_t': 3.00049253442181e-12,
 'charge_live': 7.70000000000011e-11,
 'n_particle_tot': 1000,
 'n_particle_live': 1000,
 'n_particle_lost_in_ele': 0,
 'beam_saved': True}

This says that the full beam is saved at this element

In [5]:

stats["beam_saved"]

stats["beam_saved"]

Out[5]:

True

bunch1¶

Array data from a bunch can be retrieved. Available coordinates are:

x, px, y, py, z, pz, s, t, charge, p0c, state, ix_ele

Appropriate data types are returned

In [6]:

tao.bunch1("end", "x")[0:10]

tao.bunch1("end", "x")[0:10]

Out[6]:

array([-1.69327762e-07, -4.58088151e-06,  4.46517206e-06, -1.72815751e-06,
        9.38761727e-06,  5.93297413e-05, -6.70659004e-05,  4.15474873e-05,
       -7.45202256e-05,  5.06597769e-05])

In [7]:

tao.bunch1("end", "ix_ele")[0:10]

tao.bunch1("end", "ix_ele")[0:10]

Out[7]:

array([8, 8, 8, 8, 8, 8, 8, 8, 8, 8], dtype=int32)

Plot in matplotlib¶

This can be used to plot particles.

In [8]:

import matplotlib.pyplot as plt

import matplotlib.pyplot as plt

In [9]:

%config InlineBackend.figure_format = 'retina' # Nicer plotting
%matplotlib inline

%config InlineBackend.figure_format = 'retina' # Nicer plotting %matplotlib inline

In [10]:

xdat = tao.bunch1("end", "x")
pxdat = tao.bunch1("end", "px")
chargedat = tao.bunch1("end", "charge")

xdata = 1000 * xdat
ydata = 1000 * pxdat
weights = chargedat

xdat = tao.bunch1("end", "x") pxdat = tao.bunch1("end", "px") chargedat = tao.bunch1("end", "charge") xdata = 1000 * xdat ydata = 1000 * pxdat weights = chargedat

In [11]:

# hist2d

mycmap = plt.get_cmap("plasma")  # viridis plasma inferno magma and _r versions
mycmap.set_under(color="white")  # map 0 to this color
myvmin = 1e-30  # something tiny
# Bin particles
plt.hist2d(
    x=1000 * xdata, y=ydata, bins=2 * [40], weights=weights, cmap=mycmap, vmin=myvmin
)
plt.xlabel("x (mm)")
plt.ylabel("px (mrad)")
plt.show()

# hist2d mycmap = plt.get_cmap("plasma") # viridis plasma inferno magma and _r versions mycmap.set_under(color="white") # map 0 to this color myvmin = 1e-30 # something tiny # Bin particles plt.hist2d( x=1000 * xdata, y=ydata, bins=2 * [40], weights=weights, cmap=mycmap, vmin=myvmin ) plt.xlabel("x (mm)") plt.ylabel("px (mrad)") plt.show()

Numpy histogram 2d, with custom color map

In [12]:

import matplotlib.colors as colors

mycmap = plt.get_cmap("viridis")  # viridis plasma inferno magma and _r versions
mycmap.set_under(color="white")  # map 0 to this color
H, xedges, yedges = np.histogram2d(xdata, ydata, weights=chargedat, bins=40)

xmin, xmax = min(xedges), max(xedges)
ymin, ymax = min(yedges), max(yedges)

image = np.flip(H.T, axis=0)  #
imax = np.max(image)
norm = colors.Normalize(vmin=1e-12 * imax, vmax=imax)
plt.xlabel("x (mm)")
plt.ylabel("px (mrad)")
plt.imshow(
    image, cmap=mycmap, norm=norm, extent=[xmin, xmax, ymin, ymax], aspect="auto"
);

import matplotlib.colors as colors mycmap = plt.get_cmap("viridis") # viridis plasma inferno magma and _r versions mycmap.set_under(color="white") # map 0 to this color H, xedges, yedges = np.histogram2d(xdata, ydata, weights=chargedat, bins=40) xmin, xmax = min(xedges), max(xedges) ymin, ymax = min(yedges), max(yedges) image = np.flip(H.T, axis=0) # imax = np.max(image) norm = colors.Normalize(vmin=1e-12 * imax, vmax=imax) plt.xlabel("x (mm)") plt.ylabel("px (mrad)") plt.imshow( image, cmap=mycmap, norm=norm, extent=[xmin, xmax, ymin, ymax], aspect="auto" );

In [13]:

np.min(image), np.max(image)

np.min(image), np.max(image)

Out[13]:

(np.float64(0.0), np.float64(8.469999999999999e-13))

Plot in Bokeh¶

In [14]:

from bokeh.plotting import figure, show, output_notebook
from bokeh import palettes, colors
from bokeh.models import ColumnDataSource, HoverTool

output_notebook(verbose=False, hide_banner=True)

pal = palettes.Viridis[256]
# white=colors.named.white
# pal[0] = white # replace 0 with white

from bokeh.plotting import figure, show, output_notebook from bokeh import palettes, colors from bokeh.models import ColumnDataSource, HoverTool output_notebook(verbose=False, hide_banner=True) pal = palettes.Viridis[256] # white=colors.named.white # pal[0] = white # replace 0 with white

In [15]:

H, xedges, yedges = np.histogram2d(xdata, ydata, weights=chargedat, bins=40)
xmin, xmax = min(xedges), max(xedges)
ymin, ymax = min(yedges), max(yedges)

H, xedges, yedges = np.histogram2d(xdata, ydata, weights=chargedat, bins=40) xmin, xmax = min(xedges), max(xedges) ymin, ymax = min(yedges), max(yedges)

In [16]:

ds = ColumnDataSource(data=dict(image=[H.transpose()]))
p = figure(
    x_range=[xmin, xmax],
    y_range=[ymin, ymax],
    title="Bunch at end",
    x_axis_label="x (mm)",
    y_axis_label="px (mrad)",
    width=500,
    height=500,
)
p.image(
    image="image",
    source=ds,
    x=xmin,
    y=ymin,
    dw=xmax - xmin,
    dh=ymax - ymin,
    palette=pal,
)
show(p)

ds = ColumnDataSource(data=dict(image=[H.transpose()])) p = figure( x_range=[xmin, xmax], y_range=[ymin, ymax], title="Bunch at end", x_axis_label="x (mm)", y_axis_label="px (mrad)", width=500, height=500, ) p.image( image="image", source=ds, x=xmin, y=ymin, dw=xmax - xmin, dh=ymax - ymin, palette=pal, ) show(p)

Data for ParticleGroup¶

The above commands have been packaged into two functions for easier use, and to easily create ParticleGroup objects

In [17]:

data = tao.bunch_data("end")
data.keys()

data = tao.bunch_data("end") data.keys()

Out[17]:

dict_keys(['x', 'px', 'y', 'py', 't', 'pz', 'status', 'weight', 'z', 'species'])

In [18]:

from pmd_beamphysics import ParticleGroup

P = ParticleGroup(data=data)
P

from pmd_beamphysics import ParticleGroup P = ParticleGroup(data=data) P

Out[18]:

<ParticleGroup with 1000 particles at 0x7f77c6a18bf0>

In [19]:

P.plot("x", "px")

P.plot("x", "px")

In [20]:

P.twiss("xy")

P.twiss("xy")

Out[20]:

{'alpha_x': np.float64(-2.1521049736138536),
 'beta_x': np.float64(0.3013428324900903),
 'gamma_x': np.float64(18.688202307379193),
 'emit_x': np.float64(1.2177638975257874e-08),
 'eta_x': np.float64(-0.04819365383894791),
 'etap_x': np.float64(-0.45497339215410726),
 'norm_emit_x': np.float64(1.0008240308093852e-06),
 'alpha_y': np.float64(1.9440813319008754),
 'beta_y': np.float64(0.40782968839311173),
 'gamma_y': np.float64(11.719235654169719),
 'emit_y': np.float64(1.2177509284772165e-08),
 'eta_y': np.float64(-0.045731504141568825),
 'etap_y': np.float64(0.02639847694869729),
 'norm_emit_y': np.float64(1.0008133721459997e-06)}

bunch_comb¶

In [21]:

tao.bunch_comb("x")

tao.bunch_comb("x")

Out[21]:

array([ 7.40357485e-22,  1.52112707e-14,  3.04225419e-14,  4.56338120e-14,
        6.08450826e-14,  7.60563532e-14,  9.12676229e-14, -1.16282940e-09,
       -1.14870454e-09, -1.12526206e-09, -1.09261268e-09, -1.05084624e-09,
       -9.99970453e-10, -9.39813419e-10, -8.69809925e-10, -7.88320099e-10,
       -6.91745919e-10, -5.74407015e-10, -4.27876634e-10,  1.42339073e-09,
        1.63091681e-09,  1.83850625e-09,  2.04615922e-09,  2.25387404e-09,
        2.46164780e-09,  2.66947691e-09,  2.87735726e-09,  9.86158250e-09,
        9.32199246e-09,  8.20091852e-09,  6.43516375e-09,  3.96715881e-09,
        7.44494068e-10, -3.28054979e-09, -8.15630332e-09, -1.39318757e-08,
       -2.06607116e-08, -2.83988809e-08, -4.55143879e-08, -5.47252817e-08,
       -6.39361217e-08, -7.31469081e-08, -8.23576435e-08, -9.15683309e-08,
       -1.00778975e-07])

Make a nice plot with the beam envelope

In [22]:

s = tao.bunch_comb("s")
mean_x = tao.bunch_comb("x")
max_x = mean_x + tao.bunch_comb("rel_max.1")
min_x = mean_x + tao.bunch_comb("rel_min.1")
sigma_x = np.sqrt(tao.bunch_comb("sigma.11"))
fig, ax = plt.subplots()

ax.fill_between(s, min_x, max_x, alpha=0.2)
ax.plot(s, sigma_x, label=r"$+\sigma_x$")
ax.plot(s, mean_x, label=r"$<x>$", marker=".")
ax.plot(s, -sigma_x, label=r"$-\sigma_x$")
ax.set_xlabel("s (m)")
ax.set_ylabel("beam sizes (m)")
plt.legend();

s = tao.bunch_comb("s") mean_x = tao.bunch_comb("x") max_x = mean_x + tao.bunch_comb("rel_max.1") min_x = mean_x + tao.bunch_comb("rel_min.1") sigma_x = np.sqrt(tao.bunch_comb("sigma.11")) fig, ax = plt.subplots() ax.fill_between(s, min_x, max_x, alpha=0.2) ax.plot(s, sigma_x, label=r"$+\sigma_x$") ax.plot(s, mean_x, label=r"$$", marker=".") ax.plot(s, -sigma_x, label=r"$-\sigma_x$") ax.set_xlabel("s (m)") ax.set_ylabel("beam sizes (m)") plt.legend();

Beam betas

In [23]:

plt.plot(tao.bunch_comb("s"), 1000 * tao.bunch_comb("x.beta"), label="beam beta_x")
plt.plot(tao.bunch_comb("s"), 1000 * tao.bunch_comb("y.beta"), label="beam beta_y")
plt.xlabel("s (m)")
plt.ylabel("beam Twiss beta (m)")
plt.legend();

plt.plot(tao.bunch_comb("s"), 1000 * tao.bunch_comb("x.beta"), label="beam beta_x") plt.plot(tao.bunch_comb("s"), 1000 * tao.bunch_comb("y.beta"), label="beam beta_y") plt.xlabel("s (m)") plt.ylabel("beam Twiss beta (m)") plt.legend();