Bunch Commands¶

In [1]:

import matplotlib.pyplot as plt
import numpy as np

from pytao import Tao, SubprocessTao

import matplotlib.pyplot as plt import numpy as np from pytao import Tao, SubprocessTao

Initialize Tao on the CSR beam tracking example¶

In [2]:

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 through tao.ele().bunch_params as a Pydantic model, or through tao.bunch_params as a dictionary.

In [3]:

end_ele = tao.ele("end")

end_ele = tao.ele("end")

In [4]:

end_ele.bunch_params

end_ele.bunch_params

Out[4]:

ElementBunchParams(
    beam_saved=True,
    centroid_beta=0.9999259828218,
    centroid_p0c=41996891.3143865,
    centroid_t=1.48447035006282e-09,
    centroid_vec_1=-1.00778974569331e-07,
    centroid_vec_2=-9.21157518197918e-07,
    centroid_vec_3=5.35418307193245e-13,
    centroid_vec_4=-2.53681610157894e-11,
    centroid_vec_5=-7.53655671508842e-08,
    centroid_vec_6=-5.77005927629709e-06,
    charge_live=7.70000000000011e-11,
    direction=1,
    ix_ele=8,
    location='Downstream_End',
    n_particle_live=1000,
    n_particle_lost_in_ele=0,
    n_particle_tot=1000,
    rel_max_1=0.000175099474265468,
    rel_max_2=0.00133407063424454,
    rel_max_3=0.000204275481589786,
    rel_max_4=0.00125449039929661,
    rel_max_5=0.00278367096445625,
    rel_max_6=2.80117072980844e-05,
    rel_min_1=-0.000190703788866682,
    rel_min_2=-0.00145920753611091,
    rel_min_3=-0.000205532005213596,
    rel_min_4=-0.00119689779980729,
    rel_min_5=-0.00261715445721058,
    rel_min_6=-1.67983140045068e-05,
    s=0.444999999999986,
    sigma_11=3.66649418404824e-09,
    sigma_12=2.61861040979124e-08,
    sigma_13=-1.48398138843564e-13,
    sigma_14=-4.74391719458135e-14,
    sigma_15=-4.67916777463431e-10,
    sigma_16=-1.07814707650408e-11,
    sigma_21=2.61861040979124e-08,
    sigma_22=2.27394287450161e-07,
    sigma_23=-1.40343913387095e-12,
    sigma_24=4.92122667757304e-14,
    sigma_25=-4.44903636535732e-09,
    sigma_26=-1.01782156616636e-10,
    sigma_31=-1.48398138843564e-13,
    sigma_32=-1.40343913387095e-12,
    sigma_33=4.96185134005342e-09,
    sigma_34=-2.36505280427116e-08,
    sigma_35=1.51286568676907e-13,
    sigma_36=-1.02306597520076e-11,
    sigma_41=-4.74391719458135e-14,
    sigma_42=4.92122667757304e-14,
    sigma_43=-2.36505280427116e-08,
    sigma_44=1.42566900728402e-07,
    sigma_45=-3.27450024554485e-13,
    sigma_46=5.9056066984631e-12,
    sigma_51=-4.67916777463431e-10,
    sigma_52=-4.44903636535732e-09,
    sigma_53=1.51286568676907e-13,
    sigma_54=-3.27450024554485e-13,
    sigma_55=8.09025505237071e-07,
    sigma_56=1.04740747567993e-08,
    sigma_61=-1.07814707650408e-11,
    sigma_62=-1.01782156616636e-10,
    sigma_63=-1.02306597520076e-11,
    sigma_64=5.9056066984631e-12,
    sigma_65=1.04740747567993e-08,
    sigma_66=2.23710129496356e-10,
    sigma_t=3.00049253467856e-12,
    species='Electron',
    t=1.48447035006282e-09,
    twiss_alpha_a=-1.77499320489309,
    twiss_alpha_b=1.59897669921574,
    twiss_alpha_c=-1.2389734965013,
    twiss_alpha_x=-2.15210497361757,
    twiss_alpha_y=1.94408133190159,
    twiss_alpha_z=-1.24059179271663,
    twiss_beta_a=0.248530085331005,
    twiss_beta_b=0.335544868515057,
    twiss_beta_c=95.7148956246458,
    twiss_beta_x=0.301344571266547,
    twiss_beta_y=0.407832041608049,
    twiss_beta_z=95.8242541895164,
    twiss_dalpha_dpz_a=0.0,
    twiss_dalpha_dpz_b=0.0,
    twiss_dalpha_dpz_c=0.0,
    twiss_dalpha_dpz_x=0.0,
    twiss_dalpha_dpz_y=0.0,
    twiss_dalpha_dpz_z=0.0,
    twiss_dbeta_dpz_a=0.0,
    twiss_dbeta_dpz_b=0.0,
    twiss_dbeta_dpz_c=0.0,
    twiss_dbeta_dpz_x=0.0,
    twiss_dbeta_dpz_y=0.0,
    twiss_dbeta_dpz_z=0.0,
    twiss_deta_dpz_a=0.0,
    twiss_deta_dpz_b=0.0,
    twiss_deta_dpz_c=0.0,
    twiss_deta_dpz_x=0.0,
    twiss_deta_dpz_y=0.0,
    twiss_deta_dpz_z=0.0,
    twiss_deta_ds_a=0.0,
    twiss_deta_ds_b=0.0,
    twiss_deta_ds_c=0.0,
    twiss_deta_ds_x=0.0,
    twiss_deta_ds_y=0.0,
    twiss_deta_ds_z=0.0,
    twiss_detap_dpz_a=0.0,
    twiss_detap_dpz_b=0.0,
    twiss_detap_dpz_c=0.0,
    twiss_detap_dpz_x=0.0,
    twiss_detap_dpz_y=0.0,
    twiss_detap_dpz_z=0.0,
    twiss_emit_a=1.21664545105205e-08,
    twiss_emit_b=1.21741461356351e-08,
    twiss_emit_c=8.43349923229653e-09,
    twiss_emit_x=1.21653911572771e-08,
    twiss_emit_y=1.21652615972309e-08,
    twiss_emit_z=8.4428051340428e-09,
    twiss_eta_a=0.00124844704286057,
    twiss_eta_b=0.00454529158526127,
    twiss_eta_c=1.2389734965013,
    twiss_eta_x=-0.0481939319838284,
    twiss_eta_y=-0.0457317680475184,
    twiss_eta_z=0.0,
    twiss_etap_a=0.0111101590168622,
    twiss_etap_b=-0.0218933442291171,
    twiss_etap_c=0.0264430499032274,
    twiss_etap_x=-0.454973392781903,
    twiss_etap_y=0.0263984769565801,
    twiss_etap_z=0.0,
    twiss_gamma_a=15.4136110252092,
    twiss_gamma_b=9.63845045654778,
    twiss_gamma_c=0.0264430499032274,
    twiss_gamma_x=18.6880944753713,
    twiss_gamma_y=11.7191680334955,
    twiss_gamma_z=0.0264971328775929,
    twiss_norm_emit_a=9.99904831192272e-07,
    twiss_norm_emit_b=1.00053697041617e-06,
    twiss_norm_emit_c=6.93110438948157e-07,
    twiss_norm_emit_x=9.99817439089306e-07,
    twiss_norm_emit_y=9.99806791146145e-07,
    twiss_norm_emit_z=6.93875248129554e-07,
    twiss_phi_a=0.0,
    twiss_phi_b=0.0,
    twiss_phi_c=0.0,
    twiss_phi_x=0.0,
    twiss_phi_y=0.0,
    twiss_phi_z=0.0,
    twiss_sigma_a=0.0,
    twiss_sigma_b=0.0,
    twiss_sigma_c=0.0,
    twiss_sigma_p_a=0.0,
    twiss_sigma_p_b=0.0,
    twiss_sigma_p_c=0.0,
    twiss_sigma_p_x=0.000476810213058656,
    twiss_sigma_p_y=0.000377580117100171,
    twiss_sigma_p_z=1.49569425183209e-05,
    twiss_sigma_x=6.0547292115994e-05,
    twiss_sigma_y=7.04370887664636e-05,
    twiss_sigma_z=0.000899458451089916
)

The following indicates that the full beam is saved at this element:

In [5]:

end_ele.bunch_params.beam_saved

end_ele.bunch_params.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.93297414e-05, -6.70659005e-05,  4.15474873e-05,
       -7.45202256e-05,  5.06597770e-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]:

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 [9]:

# 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 [10]:

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 [11]:

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

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

Out[11]:

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

Plot in Bokeh¶

In [12]:

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

output_notebook(verbose=False, hide_banner=True)

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

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

In [13]:

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 [14]:

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 '.particles' helper¶

In [15]:

P_end = tao.particles("end")
P_end.plot("x", "px")

P_end = tao.particles("end") P_end.plot("x", "px")

To reproduce the above manually, we can use tao.bunch_data and pass it to ParticleGroup(data=...):

In [16]:

from beamphysics import ParticleGroup
data = tao.bunch_data("end")
P = ParticleGroup(data=data)

data.keys(), P

from beamphysics import ParticleGroup data = tao.bunch_data("end") P = ParticleGroup(data=data) data.keys(), P

Out[16]:

(dict_keys(['x', 'px', 'y', 'py', 't', 'pz', 'status', 'weight', 'z', 'species']),
 <ParticleGroup with 1000 particles at 0x7f9be33dee40>)

Twiss parameters¶

In [17]:

P.twiss("xy")

P.twiss("xy")

Out[17]:

{'alpha_x': np.float64(-2.152104973617563),
 'beta_x': np.float64(0.301342832490508),
 'gamma_x': np.float64(18.688202307406268),
 'emit_x': np.float64(1.2177638991702104e-08),
 'eta_x': np.float64(-0.04819365390231965),
 'etap_x': np.float64(-0.4549733927846398),
 'norm_emit_x': np.float64(1.0008240321606599e-06),
 'alpha_y': np.float64(1.944081331901604),
 'beta_y': np.float64(0.40782968839299794),
 'gamma_y': np.float64(11.719235654179935),
 'emit_y': np.float64(1.2177509301217845e-08),
 'eta_y': np.float64(-0.04573150417022795),
 'etap_y': np.float64(0.02639847694381584),
 'norm_emit_y': np.float64(1.0008133734973938e-06)}

bunch_comb¶

In [18]:

conf = tao.get_config()
conf.beam.comb_ds_save = 0.01
print(conf.get_set_commands(tao))
conf.set(tao)

conf = tao.get_config() conf.beam.comb_ds_save = 0.01 print(conf.get_set_commands(tao)) conf.set(tao)

['set beam comb_ds_save = 0.01']

Out[18]:

True

In [19]:

tao.bunch_comb("x")

tao.bunch_comb("x")

Out[19]:

array([-7.60367147e-22,  7.83142510e-15,  1.56628500e-14,  2.34942744e-14,
        3.13256996e-14,  3.91571245e-14,  4.69885488e-14, -1.16287350e-09,
       -1.14874763e-09, -1.12530081e-09, -1.09263690e-09, -1.05083213e-09,
       -9.99872190e-10, -9.39557033e-10, -8.69310335e-10, -7.87597111e-10,
       -6.90732323e-10, -5.72648666e-10, -4.24514002e-10,  1.42964050e-09,
        1.64071157e-09,  1.85184445e-09,  2.27429080e-09,  2.69695616e-09,
        2.90836029e-09,  9.88977741e-09,  9.33396264e-09,  8.18206124e-09,
        6.36976778e-09,  3.83843952e-09,  5.34575968e-10, -3.59072952e-09,
       -8.58045823e-09, -1.44793707e-08, -2.13345175e-08, -2.91972195e-08,
       -4.64445773e-08, -6.51384965e-08, -7.44853986e-08, -8.38322676e-08,
       -1.02525931e-07])

Make a nice plot with the beam envelope

In [20]:

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 [21]:

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();

Initial particles¶

Tao.set_initial_particles and Tao.get_initial_particles make it easy to change or inspect the initial particles in use.

First, let's generate a beam in Tao using beam_init. Only n_particle and bunch_charge need to be set. Then, track the beam and read it back from track_start (typically BEGINNING).

In [22]:

conf = tao.get_config()
conf.beam_init.n_particle = 1000
conf.beam_init.bunch_charge = 1e-9
conf.beam_init.set(tao, only_changed=True)

tao.track_beam(use_progress_bar=False)  # tracks the beam through the lattice and saves the bunch data at each element
P0 = tao.particles("BEGINNING")
P0

conf = tao.get_config() conf.beam_init.n_particle = 1000 conf.beam_init.bunch_charge = 1e-9 conf.beam_init.set(tao, only_changed=True) tao.track_beam(use_progress_bar=False) # tracks the beam through the lattice and saves the bunch data at each element P0 = tao.particles("BEGINNING") P0

Out[22]:

<ParticleGroup with 1000 particles at 0x7f9be3e7a2d0>

get_initial_particles wraps that read-back. It returns the beam saved at track_start, or None if none is saved there.

In [23]:

tao.get_initial_particles()

tao.get_initial_particles()

Out[23]:

<ParticleGroup with 1000 particles at 0x7f9be3154680>

Setting initial particles with ParticleGroup¶

Pass a ParticleGroup (or an openPMD-beamphysics HDF5 filename) to set_initial_particles. It writes the particles to a file that lives for the duration of the Tao session, points beam_init at it, reinitializes the beam, and returns the beam read back from track_start.

For example:

• tao.set_initial_particles(ParticleGroup(...)) uses a temporary directory
• tao.set_initial_particles(ParticleGroup(...), write_to="foo.h5") writes to foo.h5 instead of the temporary directory
• tao.set_initial_particles("/path/to/particlegroup.h5") uses the existing file on disk

In [24]:

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

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

In [25]:

P1 = tao2.set_initial_particles(P0)
P1

P1 = tao2.set_initial_particles(P0) P1

Out[25]:

<ParticleGroup with 1000 particles at 0x7f9be31572f0>

The beam in the second instance now matches the one we generated:

In [26]:

P0 == tao2.get_initial_particles()

P0 == tao2.get_initial_particles()

Out[26]:

True

Use n_particle to take a subset of the incoming particles, or pass write_to to keep the particle file at a location of your choosing instead of the temporary session directory.

In [27]:

P_small = tao2.set_initial_particles(P0, n_particle=100)
len(P_small)

P_small = tao2.set_initial_particles(P0, n_particle=100) len(P_small)

Out[27]:

100

In [28]:

tao2.close_subprocess()

tao2.close_subprocess()