Examples¶
Example scripts are in python/examples/.
Lattice Basics¶
simple.py¶
Port of the Fortran simple_bmad_program to pybmad. Reads a lattice, modifies
a solenoid's strength, recalculates transfer matrices, propagates Twiss
parameters, and prints element info.
#!/usr/bin/env python3
"""
Port of the Fortran 'simple_bmad_program' program to pybmad.
Read in a lattice, modify CLEO_SOL solenoid KS strength, recalculate
transfer matrices, propagate Twiss, and print element info.
"""
from __future__ import annotations
import pybmad
from pybmad import EleKey
# Read in a lattice
res = pybmad.bmad_parser("$ACC_ROOT_DIR/code_examples/simple_bmad_program/lat.bmad")
if res.err_flag:
raise RuntimeError("Failed to parse lattice")
lat = res.lat
# Find the CLEO_SOL element
(cleo,) = [ele for ele in lat.ele if ele.name == "CLEO_SOL"]
# See also: lat_ele_locator.py
# Modify the KS solenoid strength
ks_val = cleo.value[pybmad.KS]
cleo.value[pybmad.KS] = ks_val + 0.001
# Set flags and re-bookkeep
pybmad.set_flags_for_changed_attribute(cleo, cleo.value[pybmad.KS])
pybmad.lattice_bookkeeper(lat)
# Remake transfer matrix for the modified element
pybmad.lat_make_mat6(lat, ix_ele=cleo.ix_ele)
# Calculate starting Twiss params if the lattice is closed,
# and then propagate the Twiss parameters through the lattice.
assert lat.param is not None
if lat.param.geometry == pybmad.CLOSED:
pybmad.twiss_at_start(lat)
pybmad.twiss_propagate_all(lat)
# Print info on the first 11 elements
print(f" {'Ix':>3} {'Name':<16} {'Ele_type':<16} {'S':>12} {'Beta_a':>12}")
for i in range(11):
ele = lat.ele[i]
key_name = EleKey(ele.key).name
print(f"{i:4d} {ele.name:<16} {key_name:<16} {ele.s:12.4f} {ele.a.beta:12.4f}")
info = pybmad.type_ele(cleo)
print()
print("!---------------------------------------------------------")
print("! Information on element: CLEO_SOL")
print()
print("\n".join(info.lines))
# Alternatively, all of the information is accessible directly through the element:
key_name = EleKey(cleo.key).name
print(f" Name: {cleo.name}")
print(f" Key: {key_name}")
print(f" ix_ele: {cleo.ix_ele}")
print(f" s: {cleo.s:.6f}")
print(f" KS: {cleo.value[pybmad.KS]:.6f}")
print(f" L: {cleo.value[pybmad.L]:.6f}")
print(f" beta_a: {cleo.a.beta:.6f}")
print(f" beta_b: {cleo.b.beta:.6f}")
print(f" alpha_a: {cleo.a.alpha:.6f}")
print(f" alpha_b: {cleo.b.alpha:.6f}")
parse.py¶
Parse a Bmad lattice file and iterate over elements, printing names and properties.
#!/usr/bin/env python3
from __future__ import annotations
import logging
import time
import pybmad
logger = logging.getLogger("pybmad-test")
logger.setLevel("DEBUG")
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s")
t0 = time.monotonic()
res = pybmad.bmad_parser("${ACC_ROOT_DIR}/bmad-doc/tao_examples/fodo/fodo.bmad")
t1 = time.monotonic()
lat = res.lat
print("Parsed lattice {lat.use_name} in {t1 - t0} s (as seen from Python)", lat.use_name, t1 - t0)
print("Branch 0 elements:", lat.branch[0].ele)
for ele in lat.branch[0].ele:
print(ele, ele.name, ele.lord, ele.n_lord)
lat_ele_locator.py¶
Select elements within a lattice by name or wildcard pattern using lat_ele_locator.
from __future__ import annotations
import logging
import time
import pybmad
logger = logging.getLogger("pybmad-test")
logger.setLevel("DEBUG")
logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s")
t0 = time.monotonic()
res = pybmad.bmad_parser("${ACC_ROOT_DIR}/bmad-doc/tao_examples/fodo/fodo.bmad")
t1 = time.monotonic()
lat = res.lat
ele = pybmad.pointer_to_ele(lat, 0)
print(ele)
ele_ptrs = pybmad.ElePointerStruct.new_array1d(0)
res = pybmad.lat_ele_locator(loc_str="BEGINNING", lat=lat, eles=ele_ptrs, n_loc=0)
assert len(ele_ptrs) == 1
res = pybmad.lat_ele_locator(loc_str="*", lat=lat, eles=ele_ptrs, n_loc=0)
assert len(ele_ptrs) == lat.n_ele_track + 2 # + beginning/end
Beam Tracking¶
csr.py¶
Track a particle distribution through a lattice with coherent synchrotron radiation
(CSR) and space charge effects. Shows how to configure bmad_com,
space_charge_com, initialize a beam, and call track_beam.
#!/usr/bin/env python3
from __future__ import annotations
import pathlib
import time
import pybmad
t0 = time.monotonic()
parsed = pybmad.bmad_parser("data/csr_example/lat.bmad")
if parsed.err_flag:
raise RuntimeError("bmad_parser failed")
lat = parsed.lat
pybmad.ran_seed_put(123456)
beam_init = pybmad.BeamInitStruct(
a_norm_emit=4e-12,
b_norm_emit=4e-12,
dPz_dz=0.0,
sig_z=0.3e-3,
sig_pz=0e-20,
bunch_charge=0.01e-10,
n_particle=1000,
n_bunch=1,
)
bmad_com = pybmad.get_bmad_com()
bmad_com.csr_and_space_charge_on = True
space_charge_com = pybmad.get_space_charge_com()
space_charge_com.ds_track_step = 0.1
space_charge_com.n_bin = 400
space_charge_com.beam_chamber_height = 0.02
space_charge_com.n_shield_images = 0
space_charge_com.particle_bin_span = 8
assert lat.param is not None # just for the type checker
init = pybmad.init_beam_distribution(lat.ele[0], lat.param, beam_init)
if init.err_flag:
raise RuntimeError("init_beam_distribution failed (1)")
# First bunch and its particles
beam = init.beam
bunch = beam.bunch[0]
particles = bunch.particle
n_particles = len(particles)
# Calculate the average (centroid)
ave = [0.0] * 6
for i in range(6):
total = sum(p.vec[i] for p in particles)
ave[i] = (total / n_particles) if n_particles > 0 else 0.0
centroid = pybmad.CoordStruct.new_array1d(0)
pybmad.reallocate_coord(centroid, lat, 0)
pybmad.init_coord(centroid[0], ave, lat.ele[0], pybmad.DOWNSTREAM_END)
pybmad.track_all(lat, centroid)
init = pybmad.init_beam_distribution(lat.ele[0], lat.param, beam_init)
if init.err_flag:
raise RuntimeError("init_beam_distribution failed (2)")
beam1 = init.beam
pybmad.track_beam(lat, beam1, ele1=None, ele2=None, centroid=centroid.view())
print("First particle coords at end of lattice:")
print(beam1.bunch[0].particle[0].vec.to_list())
with pathlib.Path("csr.dat").open("w") as out:
idx = 1
for part in beam1.bunch[0].particle:
vec_str = " ".join([f"{val:.8f}" for val in part.vec])
out.write(f"{idx:8d} ({vec_str})\n")
idx += 1
t1 = time.monotonic()
print(f"Elapsed: {t1 - t0:0.3} s")
bbu.py¶
Beam breakup (BBU) instability analysis: parse a lattice, compute Twiss
parameters, hybridize the lattice, and track with bbu_track_all.
from __future__ import annotations
import copy
import pathlib
import pybmad as pb
from pybmad import LatStruct
def find_elements_by_name(lat, name: str) -> list:
"""Replicates lat_ele_locator logic in Python."""
return [ele for ele in lat.ele if ele.name == name]
def set_hom_order_cutoff(lat: LatStruct, cutoff: float):
# Iterate over all elements in the lattice
for ele in lat.ele:
# Check availability of wakefields
if not ele.wake:
continue
# Check availability of Long Range (LR) modes
if not ele.wake.lr or not ele.wake.lr.mode:
continue
# Filter modes based on 'm' attribute
high_order_modes = [m for m in ele.wake.lr.mode if m.m > cutoff]
# If no modes exceed cutoff, continue
if not high_order_modes:
continue
# If all modes exceed cutoff, clear the list
if len(high_order_modes) == len(ele.wake.lr.mode):
# TODO this might not work in pybmad
ele.wake.lr.mode = []
continue
# Partial removal: keep modes <= cutoff
kept_modes = [m for m in ele.wake.lr.mode if m.m <= cutoff]
ele.wake.lr.mode = kept_modes
def hybridize(
lat: LatStruct,
ele_track_end: str = "",
use_taylor_for_hybrids: bool = False,
keep_all_lcavities: bool = False,
):
for ele in lat.ele:
ele.select = False # Default: Hybridize this element
if ele.name == ele_track_end:
ele.select = True
continue
if ele.key == pb.TAYLOR:
ele.select = True
continue
if ele.key == pb.LCAVITY:
if keep_all_lcavities:
pass
elif not ele.wake or not ele.wake.lr.mode:
# If no wake or no modes, default (false) applies?
# Fortran logic: if NOT keep_all_lcavities
# if not associated(wake) cycle
# if size(mode) == 0 cycle
# ele%select = true
continue
ele.select = True
return pb.make_hybrid_lat(lat, use_taylor_for_hybrids)
def main():
bbu_param = pb.BbuParamStruct(
lat_filename="$ACC_ROOT_DIR/regression_tests/bbu_test/oneturn_lat.bmad",
keep_overlays_and_groups=False,
simulation_turns_max=500,
elname="T1",
hybridize=True,
nrep=5,
limit_factor=3,
keep_all_lcavities=False,
ran_gauss_sigma_cut=3,
nstep=50,
current=1.000,
ran_seed=100,
rel_tol=0.001,
lat2_filename="",
bunch_freq=1300000000.0,
)
# ---------------------------------------------------------------
beam_init = pb.BeamInitStruct(
n_particle=1,
# Define distance between bunches (1 / freq)
dt_bunch=1.0 / bbu_param.bunch_freq if bbu_param.bunch_freq != 0 else 0.0,
)
bbu_param.n_ramp_pattern = len(bbu_param.ramp_pattern)
if bbu_param.ramp_on and bbu_param.n_ramp_pattern < 1:
raise RuntimeError("RAMP_ON = TRUE BUT THERE IS NO RAMP_PATTERN!")
pb.ran_seed_put(bbu_param.ran_seed)
if bbu_param.ran_gauss_sigma_cut > 0:
pb.ran_gauss_converter(set_sigma_cut=bbu_param.ran_gauss_sigma_cut)
print(f"Lattice file: {bbu_param.lat_filename}")
res = pb.bmad_parser(bbu_param.lat_filename)
if res.err_flag:
raise RuntimeError("bmad_parser failed")
lat_in: LatStruct = res.lat
# Set Bmad Com
bmad_com = pb.get_bmad_com()
bmad_com.auto_bookkeeper = False
if bbu_param.lat2_filename:
print(f"DR-scan or Phase-scan, parsing: {bbu_param.lat2_filename}")
# Note: bmad_parser2 updates the existing lat_in
pb.bmad_parser2(bbu_param.lat2_filename, lat_in)
# --------------------------------------------------------------------------
# Twiss and Track (Closed Orbit)
# --------------------------------------------------------------------------
orb = pb.CoordStruct.new_array1d(0)
pb.twiss_and_track(lat_in, orb)
if bbu_param.hom_order_cutoff > 0:
set_hom_order_cutoff(lat_in, cutoff=bbu_param.hom_order_cutoff)
# --------------------------------------------------------------------------
# Hybridization Logic
# --------------------------------------------------------------------------
# In Fortran, this replaces drift/magnets with Taylor maps
lat = lat_in # Default if not hybridizing
if bbu_param.hybridize:
print("Hybridizing lattice...")
hybrid_lat = hybridize(
lat,
ele_track_end=bbu_param.ele_track_end,
use_taylor_for_hybrids=bbu_param.use_taylor_for_hybrids,
keep_all_lcavities=bbu_param.keep_all_lcavities,
)
print("Hybridization complete !!!")
if bbu_param.write_digested_hybrid_lat:
# pybmad.write_digested_bmad_file("hybrid.digested", lat)
pb.write_bmad_lattice_file("hybrid.lat", hybrid_lat)
# --------------------------------------------------------------------------
# Setup Tracking End point
# --------------------------------------------------------------------------
# Keep copy of current state before restoration
lat0 = copy.copy(lat)
if bbu_param.ele_track_end:
locs = find_elements_by_name(lat, bbu_param.ele_track_end)
if not locs:
raise ValueError(f"No matching element found for {bbu_param.ele_track_end}")
if len(locs) > 1:
print(f"Multiple elements found for {bbu_param.ele_track_end}")
print("Will use the first instance.")
ele_end = locs[0]
# Handle Lord/Slave logic (slaves contain the tracking physics)
if ele_end.lord_status == pb.SUPER_LORD:
ele_end = ele_end.slave[-1]
ix = ele_end.ix_ele
if ix > lat.n_ele_track:
raise ValueError(f"STOPPING ELEMENT IS A LORD! {bbu_param.ele_track_end}")
bbu_param.ix_ele_track_end = ix
if bbu_param.write_hom_info:
pb.rf_cav_names(lat)
# Check RF Freq
pb.check_rf_freq(lat, bbu_param.bunch_freq)
# Prepare custom BBU beam structure (assumed binding)
bbu_beam = pb.BbuBeamStruct()
# BBU Setup call
pb.bbu_setup(lat, beam_init.dt_bunch, bbu_param, bbu_beam)
print("bbu_setup complete !!!")
# Recalculate bunch charge based on current
beam_init.bunch_charge = bbu_param.current * beam_init.dt_bunch
n_stages = len(bbu_beam.stage)
print(f"Number of stages and elements: {n_stages} {lat.n_ele_track}")
# Use the initial lattice
lat = lat0
# --------------------------------------------------------------------------
# Tracking Execution
# --------------------------------------------------------------------------
print("Starting bbu_track_all...")
(hom_voltage_gain, growth_rate, lost, irep) = pb.bbu_track_all(
lat=lat,
bbu_beam=bbu_beam,
bbu_param=bbu_param,
beam_init=beam_init,
)
print("bbu_track_all complete !!!")
print(f"HOM VOLT GAIN: {hom_voltage_gain}")
print(f"growth_rate: {growth_rate}")
output_file = pathlib.Path("for_py.txt")
with output_file.open("w") as f:
f.write(f"lostbool = {lost}\n")
# Format: Scientific notation with specifically padded output
f.write(f"v_gain = {hom_voltage_gain:.8E}\n")
f.write(f"bunch_dt = {beam_init.dt_bunch:.6E}\n")
# Check against a 'garbage' constant, in Python usually math.nan check or checking initialization
# Fortran used real_garbage$. Assuming -1 or specific check.
# Here we just print validity.
valid_growth = growth_rate != -1.0 # Simplify garbage check
f.write(f"growth_rate_set = {valid_growth}\n")
f.write(f"growth_rate = {growth_rate:.6E}\n")
return {
"lat": lat,
"bbu_beam": bbu_beam,
"hom_voltage_gain": hom_voltage_gain,
"growth_rate": growth_rate,
"lost": lost,
"irep": irep,
}
if __name__ == "__main__":
res = main()
lat = res["lat"]
Geometry¶
wall_generator.py¶
Extract wall geometry from wall3d lattice sections. Computes wall radius and
surface normals at angular sample points and writes output files.
"""
Program wall_generator
Parses lattice file and outputs wall.out file with points that can be linearly
interpolated to make wall geometry.
"""
from __future__ import annotations
import argparse
import math
import pathlib
import sys
from typing import IO
import pybmad as pb
from pybmad import BranchStruct, CoordStruct, RealAlloc1D
def write_header(fp: IO[str]) -> None:
"""
Write the header to the output file.
Parameters
----------
fp : IO[str]
The open file object to write to.
"""
# Mimicking Fortran (9a18) format
headers = ["x", "normal_x", "y", "normal_y", "z", "normal_z", "ix_ele", "angle_index", "s"]
units = ["m", "1", "m", "1", "m", "1", "1", "1", "m"]
# Format string for 18-character width columns
fmt = "{:>18}" * 9
fp.write("#" + fmt.format(*headers) + "\n")
fp.write("#" + fmt.format(*units) + "\n")
def write_line(f_out: IO[str], orb_out: CoordStruct, ix_ele: int, i_angle: int, s: float):
vec_str = "".join([f"{val:18.10E}" for val in orb_out.vec])
f_out.write(f"{vec_str}{ix_ele:>18}{i_angle:>18}{s:18.10E}\n")
def print_logo() -> None:
print(r"")
print(r" __ ___ ___ __ ___ __ __ ")
print(r"| | /\ | | / _` |__ |\ | |__ |__) /\ | / \ |__) ")
print(r"|/\| /~~\ |___ |___ ___ \__> |___ | \| |___ | \ /~~\ | \__/ | \ ")
print(r"")
def get_wall(branch: BranchStruct, *, n_angles: int = 2):
wall3d = branch.wall3d[0]
for section in wall3d.section:
ele = branch.ele[section.ix_ele]
# Calculate relative s position within the element
# section.s is absolute s, ele.s is total s at exit of element
s_rel = section.s - (ele.s - ele.value[pb.EleAttribute.L])
for i_angle in range(n_angles):
theta = i_angle * 2 * math.pi / n_angles
if ele.key == pb.SBEND:
theta = theta - ele.value[pb.EleAttribute.REF_TILT_TOT]
r_wall, _dr_dtheta, _ix_vertex = pb.calc_wall_radius(section.v, math.cos(theta), math.sin(theta))
dummy_orb = pb.CoordStruct(
ix_ele=ele.ix_ele,
# WARNING: NORMAL IS NOT COMPUTED!!!
# x, px, y, py, z, pz
vec=[
r_wall * math.cos(theta),
0.0,
r_wall * math.sin(theta),
0.0,
s_rel,
0.0,
],
)
orb_out = pb.particle_in_global_frame(
dummy_orb,
branch,
in_time_coordinates=True,
in_body_frame=False,
)
yield section, i_angle, ele, orb_out
def get_wall_contour(
branch: BranchStruct,
n_angles: int = 2,
ds: float = 0.0,
r0=1e-3,
):
point = pb.CoordStruct()
point.s = 0.0
# Initial dummy vector setup from Fortran: vec(6) = 1.0 (index 5)
# x, px, y, py, z, pz
X, _PX, Y, _PY, Z, _PZ = range(6)
point_vec = RealAlloc1D([0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
last_ele_s = branch.ele[branch.n_ele_track].s
point_s = 0.0
while point_s <= last_ele_s:
ix_ele = pb.element_at_s(branch, point_s, True).ix_ele
ele = branch.ele[ix_ele]
point_vec[Z] = point_s - (ele.s - ele.value[pb.L])
for i_angle in range(n_angles):
theta = i_angle * 2 * math.pi / n_angles
if ele.key == pb.SBEND:
theta = theta - ele.value[pb.EleAttribute.REF_TILT_TOT]
point_vec[X] = r0 * math.cos(theta)
point_vec[Y] = r0 * math.sin(theta)
# Calculate d_radius
w_res = pb.wall3d_d_radius(point_vec, ele, 1)
d_radius = w_res.d_radius
norm_x, norm_y, norm_z = w_res.perp
r_wall = r0 - d_radius
dummy_orb = pb.CoordStruct(
ix_ele=ele.ix_ele,
# Indices in global frame: [x, normal_x, y, normal_y, z, normal_z]
vec=[
# x
r_wall * math.cos(theta),
norm_x,
# y
r_wall * math.sin(theta),
norm_y,
# z
point_vec[Z], # The longitudinal pos
norm_z,
],
)
orb_out = pb.particle_in_global_frame(
dummy_orb, branch, in_time_coordinates=True, in_body_frame=False
)
yield point_s, i_angle, ele, orb_out
point_s += ds
def main() -> None:
parser = argparse.ArgumentParser(
description="Wall Generator: Create wall geometry from lattice file.",
prog="wall_generator",
usage="%(prog)s <lattice> [n_angles] [ds]",
)
parser.add_argument("lattice", help="The input Bmad lattice file")
parser.add_argument(
"n_angles", type=int, nargs="?", default=2, help="Number of angles in a section (default: 2)"
)
parser.add_argument("ds", type=float, nargs="?", default=0.0, help="Step size for contour (default: 0)")
parser.add_argument("--terse", action="store_true")
if len(sys.argv) == 1:
parser.print_help()
sys.exit(1)
args = parser.parse_args()
lat_name: str = args.lattice
n_angles: int = args.n_angles
ds: float | None = args.ds
verbose: bool = not args.terse
if verbose:
print_logo()
if n_angles < 1:
print(f"Not enough angles: {n_angles}", file=sys.stderr)
sys.exit(1)
if verbose:
print(f"Creating wall for lattice file: {lat_name}")
print(f"Using number of angles: {n_angles}")
if ds != 0:
print(f"using ds: {ds}")
lat = pb.bmad_parser(lat_name).lat
with pathlib.Path("wall.out").open("w") as f_out:
branch = lat.branch[0]
if not branch.wall3d:
print("No branch.wall3d, exiting...", file=sys.stderr)
sys.exit(1)
write_header(f_out)
for section, i_angle, ele, orb_out in get_wall(branch, n_angles=n_angles):
write_line(f_out, orb_out=orb_out, ix_ele=ele.ix_ele, i_angle=i_angle, s=section.s)
if verbose:
print("Wrote: wall.out")
if not ds:
return
with pathlib.Path("wall_contour.out").open("w") as f_out:
write_header(f_out)
for point_s, i_angle, ele, orb_out in get_wall_contour(branch, n_angles=n_angles, ds=ds):
write_line(f_out, orb_out=orb_out, ix_ele=ele.ix_ele, i_angle=i_angle, s=point_s)
if verbose:
print("Wrote: wall_contour.out")
if __name__ == "__main__":
main()
Tao¶
pybmad is also able to interface with Tao's data structures directly. Loading a lattice with PyTao, we can poke into all of the Bmad and Tao structures directly from Python.
Interaction with pytao.Tao instances is supported natively, but
pytao.SubprocessTao requires more effort to work with - as Tao's data is no
longer in the same process. See the tao-subproc.py example below for how to
work around this.
tao.py¶
Initialize PyTao and access universe/lattice data structures via get_super_universe.
from __future__ import annotations
from pytao import Tao
import pybmad
tao = Tao(init_file="$ACC_ROOT_DIR/bmad-doc/tao_examples/optics_matching/tao.init", noplot=True)
print(tao)
print("\n".join(tao.cmd("show uni")))
s = pybmad.get_super_universe()
print("Lattice:", s.u[0].model.lat.use_name)
print("Model element #0 name", s.u[0].model.lat.ele[0].name)
print("Some orbit value", s.u[0].design.tao_branch[0].orbit[0].vec)
tao-plot.py¶
Extract graph and curve data from a Tao session and plot with matplotlib.
#!/usr/bin/env python3
from __future__ import annotations
import matplotlib.pyplot as plt
from pytao import Tao
import pybmad
tao = Tao(init_file="$ACC_ROOT_DIR/bmad-doc/tao_examples/optics_matching/tao.init", plot="mpl")
s = pybmad.get_super_universe()
regions = [r for r in s.plot_page.region if r.plot.description and r.plot.graph.is_valid()]
plt.ion()
for region in regions:
for graph in region.plot.graph:
curves = [curve for curve in graph.curve if curve.valid]
if not curves:
continue
plt.figure()
for curve in curves:
plt.plot(curve.x_line, curve.y_line, label=curve.legend_text)
plt.title(graph.title)
plt.legend()
plt.ioff()
tao-subproc.py¶
Run Tao in a subprocess, interact with Bmad/Tao structures using pybmad in that subprocess, and return data to the parent process.
This example is in 2 parts tao-subproc.py and tao_subproc_helper.py. This
is due to a limitation in PyTao - it requires that custom functions be in an
importable module (which __main__ is not).
#!/usr/bin/env python3
from __future__ import annotations
from pytao import SubprocessTao
from tao_subproc_helper import get_info_example
def test():
with SubprocessTao(
init_file="$ACC_ROOT_DIR/bmad-doc/tao_examples/optics_matching/tao.init",
noplot=True,
) as tao:
print("Tao is:", tao)
print("\n".join(tao.cmd("show ele 1")[:10]))
data = tao.subprocess_call(get_info_example)
print()
print("Retrieved info from subprocess with pybmad:")
print(f"data = {data}")
if __name__ == "__main__":
test()
tao_subproc_helper.py
Helper module for tao-subproc.py -- provides a callable that accesses
universe and lattice data from within the subprocess.
from __future__ import annotations
import numpy as np
import pybmad
def get_info_example(_tao):
"""
This is meant to be used by tao-subproc.py.
SubprocessTao may only call importable functions (i.e., those outside of
__main__).
"""
s = pybmad.get_super_universe()
univ = s.u[0]
lattice = univ.model.lat.use_name
ele1_name = univ.model.lat.ele[1].name
orbit = univ.design.tao_branch[0].orbit[0].vec
# Note that we are also somewhat limited on supported return types here.
# Native types: float, int, str, bytes, booll
# Native containers: list, tuple, set, dict
# And np.ndarray, which is recommended for passing back large arrays.
return {
"lattice": lattice,
"ele1_name": ele1_name,
"orbit": np.asarray(orbit),
}