PyTao beam tracking with progress bar¶

In [1]:

import numpy as np

from pytao import Tao
from pmd_beamphysics import ParticleGroup

import numpy as np from pytao import Tao from pmd_beamphysics import ParticleGroup

In [2]:

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

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

In [3]:

print("\n".join(tao.cmd("set global n_threads = 2")))

print("\n".join(tao.cmd("set global n_threads = 2")))

[MESSAGE | 2025-APR-03 01:23:12] tao_set_openmp_n_threads:
    OpenMP active with number of threads: 2

In [4]:

def set_gaussian(
    tao,
    n_particle,
    a_norm_emit=1.0e-6,
    b_norm_emit=1.0e-6,
    bunch_charge=1e-9,
    sig_pz0=2e-6,
    sig_z=200e-6,
    center_pz=0,
    chirp=0,  # 1/m
):
    sig_pz = np.hypot(sig_pz0, chirp * sig_z)

    cmds = [
        f"set beam_init n_particle = {n_particle}",
        "set beam_init random_engine = quasi",
        "set beam_init saved_at = MARKER::*, BEGINNING, END",
        f"set beam_init a_norm_emit = {a_norm_emit}",
        f"set beam_init b_norm_emit = {b_norm_emit}",
        f"set beam_init bunch_charge = {bunch_charge}",
        f"set beam_init sig_pz = {sig_pz}",
        f"set beam_init sig_z = {sig_z}",
        f"set beam_init dpz_dz = {chirp}",
        f"set beam_init center(6) = {center_pz}",
    ]
    tao.cmds(cmds)
    tao.cmd("set global lattice_calc_on = T")


def get_particles(tao, ele_id):
    return ParticleGroup(data=tao.bunch_data(ele_id))

set_gaussian(tao, n_particle=10_000)

def set_gaussian( tao, n_particle, a_norm_emit=1.0e-6, b_norm_emit=1.0e-6, bunch_charge=1e-9, sig_pz0=2e-6, sig_z=200e-6, center_pz=0, chirp=0, # 1/m ): sig_pz = np.hypot(sig_pz0, chirp * sig_z) cmds = [ f"set beam_init n_particle = {n_particle}", "set beam_init random_engine = quasi", "set beam_init saved_at = MARKER::*, BEGINNING, END", f"set beam_init a_norm_emit = {a_norm_emit}", f"set beam_init b_norm_emit = {b_norm_emit}", f"set beam_init bunch_charge = {bunch_charge}", f"set beam_init sig_pz = {sig_pz}", f"set beam_init sig_z = {sig_z}", f"set beam_init dpz_dz = {chirp}", f"set beam_init center(6) = {center_pz}", ] tao.cmds(cmds) tao.cmd("set global lattice_calc_on = T") def get_particles(tao, ele_id): return ParticleGroup(data=tao.bunch_data(ele_id)) set_gaussian(tao, n_particle=10_000)

In [5]:

set_gaussian(tao, n_particle=10_000)

set_gaussian(tao, n_particle=10_000)

In [6]:

tao.track_beam(use_progress_bar=True, jupyter=True)

tao.track_beam(use_progress_bar=True, jupyter=True)

Out[6]:

['',
 'Tao: set global track_type = beam',
 '',
 'Tao: set global track_type = single']