Plot bokeh particles

Floor plan plot with ParticleGroup plot overlays¶

In [1]:

import dataclasses
import logging

import bokeh.plotting
import matplotlib.pyplot as plt
import numpy as np
import pytao
from pmd_beamphysics import ParticleGroup
from pytao.plotting.bokeh import BokehAppCreator, BokehFloorPlanGraph, FloorPlanElement

import dataclasses import logging import bokeh.plotting import matplotlib.pyplot as plt import numpy as np import pytao from pmd_beamphysics import ParticleGroup from pytao.plotting.bokeh import BokehAppCreator, BokehFloorPlanGraph, FloorPlanElement

Configuration¶

In [2]:

# The number of particles to generate.
N_PARTICLE = 10000
# Change this glob string ("*") to tell bmad when to save particles:
BEAM_SAVED_AT = "*"
# The matplotlib figure size for ParticleGroup:
PLOT_FIGSIZE = (4.0, 4.0)
# The dots per inch (DPI) for ParticleGroup plots:
PLOT_DPI = 100
# The size of the overlay image, in floor coordinates:
OVERLAY_DISPLAY_WIDTH = 1.0
# Calculated aspect ratio and floor coordinate height of the image:
OVERLAY_ASPECT = float(PLOT_FIGSIZE[1]) / PLOT_FIGSIZE[0]
OVERLAY_DISPLAY_HEIGHT = OVERLAY_ASPECT * OVERLAY_DISPLAY_WIDTH
# For each of matching key here, an offset will be added to the floor position:
NAME_TO_OFFSET = {
    # This matches all elements - it's our default offset:
    "": (-OVERLAY_DISPLAY_WIDTH / 2.0, OVERLAY_DISPLAY_HEIGHT / 2.0),
    # For this example, P5 overlaps - so move it up:
    "P5": (0, OVERLAY_DISPLAY_HEIGHT),
    # The remainder are examples of other offsets you could use for multipass elements
    ".END": (0, OVERLAY_DISPLAY_HEIGHT),
    r"\2": (0, OVERLAY_DISPLAY_HEIGHT),
    r"\3": (0, OVERLAY_DISPLAY_HEIGHT),
}
# Skip these elements (by name)
SKIP_ELEMENTS = []
# Use a transparent background for the plots:
TRANSPARENT_BACKGROUND = True
# The width and height of the bokeh plot:
BOKEH_WIDTH = 1000
BOKEH_HEIGHT = 500

# The number of particles to generate. N_PARTICLE = 10000 # Change this glob string ("*") to tell bmad when to save particles: BEAM_SAVED_AT = "*" # The matplotlib figure size for ParticleGroup: PLOT_FIGSIZE = (4.0, 4.0) # The dots per inch (DPI) for ParticleGroup plots: PLOT_DPI = 100 # The size of the overlay image, in floor coordinates: OVERLAY_DISPLAY_WIDTH = 1.0 # Calculated aspect ratio and floor coordinate height of the image: OVERLAY_ASPECT = float(PLOT_FIGSIZE[1]) / PLOT_FIGSIZE[0] OVERLAY_DISPLAY_HEIGHT = OVERLAY_ASPECT * OVERLAY_DISPLAY_WIDTH # For each of matching key here, an offset will be added to the floor position: NAME_TO_OFFSET = { # This matches all elements - it's our default offset: "": (-OVERLAY_DISPLAY_WIDTH / 2.0, OVERLAY_DISPLAY_HEIGHT / 2.0), # For this example, P5 overlaps - so move it up: "P5": (0, OVERLAY_DISPLAY_HEIGHT), # The remainder are examples of other offsets you could use for multipass elements ".END": (0, OVERLAY_DISPLAY_HEIGHT), r"\2": (0, OVERLAY_DISPLAY_HEIGHT), r"\3": (0, OVERLAY_DISPLAY_HEIGHT), } # Skip these elements (by name) SKIP_ELEMENTS = [] # Use a transparent background for the plots: TRANSPARENT_BACKGROUND = True # The width and height of the bokeh plot: BOKEH_WIDTH = 1000 BOKEH_HEIGHT = 500

Tao setup¶

In [3]:

tao = pytao.Tao(
    init_file="$ACC_ROOT_DIR/bmad-doc/tao_examples/optics_matching/tao.init",
    plot="bokeh",
)

for cmd in [
    f"set beam_init beam_saved_at = {BEAM_SAVED_AT}",
    f"set beam_init n_particle = {N_PARTICLE}",
    "set global track_type = beam",
    "set beam_init distribution_type(1) = ran_gauss",
    "set beam_init distribution_type(2) = ran_gauss",
    "set beam_init distribution_type(3) = ran_gauss",
]:
    print(f"Tao> {cmd}")
    print("\n".join(tao.cmd(cmd)))

tao = pytao.Tao( init_file="$ACC_ROOT_DIR/bmad-doc/tao_examples/optics_matching/tao.init", plot="bokeh", ) for cmd in [ f"set beam_init beam_saved_at = {BEAM_SAVED_AT}", f"set beam_init n_particle = {N_PARTICLE}", "set global track_type = beam", "set beam_init distribution_type(1) = ran_gauss", "set beam_init distribution_type(2) = ran_gauss", "set beam_init distribution_type(3) = ran_gauss", ]: print(f"Tao> {cmd}") print("\n".join(tao.cmd(cmd)))

Tao> set beam_init beam_saved_at = *

Tao> set beam_init n_particle = 10000

Tao> set global track_type = beam

Tao> set beam_init distribution_type(1) = ran_gauss
[WARNING] tao_beam_track:
    Beam parameters not computed at: BEGINNING  (0)
    [This will happen, for example, with round beams. Ignore if the beam parameters at problem locations are not needed.]
    The singular sigma matrix is:
        2.5583371E-07 -4.1364916E-08  1.4395908E-07 -5.1162180E-08  0.0000000E+00  2.4771946E-06
       -4.1364916E-08  1.0233348E-08  1.6194173E-09  8.3526758E-09  0.0000000E+00 -7.8951558E-07
        1.4395908E-07  1.6194173E-09  2.5583371E-07 -2.8224437E-08  0.0000000E+00 -1.3376786E-06
       -5.1162180E-08  8.3526758E-09 -2.8224437E-08  1.0233348E-08  0.0000000E+00 -5.0421971E-07
        0.0000000E+00  0.0000000E+00  0.0000000E+00  0.0000000E+00  0.0000000E+00  0.0000000E+00
        2.4771946E-06 -7.8951558E-07 -1.3376786E-06 -5.0421971E-07  0.0000000E+00  6.6666667E-05
    Will not print any more singular sigma matrices for this track...
Tao> set beam_init distribution_type(2) = ran_gauss
[WARNING] tao_beam_track:
    Beam parameters not computed at: BEGINNING  (0)
    [This will happen, for example, with round beams. Ignore if the beam parameters at problem locations are not needed.]
    The singular sigma matrix is:
        2.5583371E-07  4.0753901E-08 -2.1032010E-07 -4.4602661E-08  0.0000000E+00  8.7494201E-07
        4.0753901E-08  1.0233348E-08 -5.1118087E-08 -4.0731011E-09  0.0000000E+00 -3.4870725E-07
       -2.1032010E-07 -5.1118087E-08  2.5583371E-07  2.2392654E-08  0.0000000E+00  1.5786534E-06
       -4.4602661E-08 -4.0731011E-09  2.2392654E-08  1.0233348E-08  0.0000000E+00 -5.4809175E-07
        0.0000000E+00  0.0000000E+00  0.0000000E+00  0.0000000E+00  0.0000000E+00  0.0000000E+00
        8.7494201E-07 -3.4870725E-07  1.5786534E-06 -5.4809175E-07  0.0000000E+00  6.6666667E-05
    Will not print any more singular sigma matrices for this track...
Tao> set beam_init distribution_type(3) = ran_gauss

Helper dataclasses¶

In [4]:

@dataclasses.dataclass
class Position:
    x: float
    y: float
    w: float
    h: float
    floor_x: float
    floor_y: float


@dataclasses.dataclass
class Element:
    index: int
    name: str
    head: dict
    bunch_params: dict
    particles: ParticleGroup
    particles_filename: str
    position: Position
    floor_ele: FloorPlanElement | None = None
    image: np.ndarray | None = None

@dataclasses.dataclass class Position: x: float y: float w: float h: float floor_x: float floor_y: float @dataclasses.dataclass class Element: index: int name: str head: dict bunch_params: dict particles: ParticleGroup particles_filename: str position: Position floor_ele: FloorPlanElement | None = None image: np.ndarray | None = None

Helper functions¶

In [5]:

def position_overlay_for_element(floor_ele: FloorPlanElement, ele_name: str) -> Position:
    """
    Calculate the position of an overlay image on the floor plan.

    Parameters
    ----------
    floor_ele : FloorPlanElement
    ele_name : str
        The name of the element to be overlayed.

    Returns
    -------
    Position
    """
    dw = OVERLAY_DISPLAY_WIDTH
    dh = OVERLAY_DISPLAY_HEIGHT

    floor_x = floor_ele.info["end1_r1"]
    floor_y = floor_ele.info["end1_r2"]
    image_x = floor_x
    image_y = floor_y

    for match, (offset_x, offset_y) in NAME_TO_OFFSET.items():
        if match in ele_name:
            image_x += offset_x
            image_y += offset_y

    return Position(x=image_x, y=image_y, w=dw, h=dh, floor_x=floor_x, floor_y=floor_y)


def fig_to_image(fig, dpi: int) -> np.ndarray:
    """
    Convert a Matplotlib figure to an RGBA numpy array.

    Parameters
    ----------
    fig : matplotlib.figure.Figure
        The figure to convert.
    dpi : int
        The resolution in dots per inch.

    Returns
    -------
    np.ndarray
        The RGBA image array of the figure.
    """
    fig.set_tight_layout(True)
    fig.set_dpi(dpi)
    fig.canvas.draw()
    return np.array(fig.canvas.renderer.buffer_rgba())


def image_array_to_bokeh(img: np.ndarray) -> np.ndarray:
    """
    Convert a numpy image array to a format compatible with Bokeh.

    Mirrors horizontally.
    """
    return img.view(dtype=np.uint32).reshape(img.shape[0], img.shape[1])[::-1, :]

def position_overlay_for_element(floor_ele: FloorPlanElement, ele_name: str) -> Position: """ Calculate the position of an overlay image on the floor plan. Parameters ---------- floor_ele : FloorPlanElement ele_name : str The name of the element to be overlayed. Returns ------- Position """ dw = OVERLAY_DISPLAY_WIDTH dh = OVERLAY_DISPLAY_HEIGHT floor_x = floor_ele.info["end1_r1"] floor_y = floor_ele.info["end1_r2"] image_x = floor_x image_y = floor_y for match, (offset_x, offset_y) in NAME_TO_OFFSET.items(): if match in ele_name: image_x += offset_x image_y += offset_y return Position(x=image_x, y=image_y, w=dw, h=dh, floor_x=floor_x, floor_y=floor_y) def fig_to_image(fig, dpi: int) -> np.ndarray: """ Convert a Matplotlib figure to an RGBA numpy array. Parameters ---------- fig : matplotlib.figure.Figure The figure to convert. dpi : int The resolution in dots per inch. Returns ------- np.ndarray The RGBA image array of the figure. """ fig.set_tight_layout(True) fig.set_dpi(dpi) fig.canvas.draw() return np.array(fig.canvas.renderer.buffer_rgba()) def image_array_to_bokeh(img: np.ndarray) -> np.ndarray: """ Convert a numpy image array to a format compatible with Bokeh. Mirrors horizontally. """ return img.view(dtype=np.uint32).reshape(img.shape[0], img.shape[1])[::-1, :]

Create the floor plan¶

We need to interact directly with the Bokeh graph manager built into PyTao, rather than using the tao.plot functionality here.

prepare_graphs_by_name asks Tao to place a floor plan for us, and then PyTao inspects the graph to generate the floor_graph object. That object contains information about all elements in the floor plan, what shape they are, where they are located, and so on.

In [6]:

(floor_graph,) = tao.bokeh.prepare_graphs_by_name("floor")
bokeh_app = BokehAppCreator(
    graphs=(floor_graph,),
    include_variables=False,
    manager=tao.bokeh,
    width=BOKEH_WIDTH,
    height=BOKEH_HEIGHT,
)

index_to_floor_plan_element = {
    floor_ele.index: floor_ele for floor_ele in floor_graph.elements
}

(floor_graph,) = tao.bokeh.prepare_graphs_by_name("floor") bokeh_app = BokehAppCreator( graphs=(floor_graph,), include_variables=False, manager=tao.bokeh, width=BOKEH_WIDTH, height=BOKEH_HEIGHT, ) index_to_floor_plan_element = { floor_ele.index: floor_ele for floor_ele in floor_graph.elements }

Loading BokehJS ...

Gather all of the elements¶

For all elements which:

• Have particles saved by bmad ('beam_saved' bunch parameter)
• Are not in the skip list
• Are included in the Tao floor plan

We save OpenPMD standard particle data to a corresponding HDF5 file. We create an Element() instance for each of these with all relevant information.

In [7]:

elements = []
for idx in tao.lat_list("*", "ele.ix_ele", flags="-array_out -track_only"):
    head = tao.ele_head(idx)
    ele_name = head["name"]

    if ele_name in SKIP_ELEMENTS:
        print(f"Skipping {ele_name} ({idx}) as configured")
        continue

    bunch_params = tao.bunch_params(idx)
    if not bunch_params["beam_saved"]:
        print(f"Skipping {ele_name} ({idx}) has no particles")
        continue

    floor_ele = index_to_floor_plan_element.get(idx, None)
    if floor_ele is None:
        print(f"Skipping {ele_name} ({idx}) as it's not in the floor plan")
        continue

    fn = f"{ele_name}.h5".replace("#", "_").replace("\\", "_pass")

    # Write the particles at this element to 'fn':
    tao.cmd(f"write beam -at {idx} {fn}")

    print(f"{ele_name} ({idx}) particles in {fn}")

    elements.append(
        Element(
            index=idx,
            name=ele_name,
            head=head,
            bunch_params=bunch_params,
            particles=ParticleGroup(h5=fn),
            particles_filename=fn,
            position=position_overlay_for_element(floor_ele, ele_name),
            floor_ele=floor_ele,
        )
    )

elements = [] for idx in tao.lat_list("*", "ele.ix_ele", flags="-array_out -track_only"): head = tao.ele_head(idx) ele_name = head["name"] if ele_name in SKIP_ELEMENTS: print(f"Skipping {ele_name} ({idx}) as configured") continue bunch_params = tao.bunch_params(idx) if not bunch_params["beam_saved"]: print(f"Skipping {ele_name} ({idx}) has no particles") continue floor_ele = index_to_floor_plan_element.get(idx, None) if floor_ele is None: print(f"Skipping {ele_name} ({idx}) as it's not in the floor plan") continue fn = f"{ele_name}.h5".replace("#", "_").replace("\\", "_pass") # Write the particles at this element to 'fn': tao.cmd(f"write beam -at {idx} {fn}") print(f"{ele_name} ({idx}) particles in {fn}") elements.append( Element( index=idx, name=ele_name, head=head, bunch_params=bunch_params, particles=ParticleGroup(h5=fn), particles_filename=fn, position=position_overlay_for_element(floor_ele, ele_name), floor_ele=floor_ele, ) )

BEGINNING (0) particles in BEGINNING.h5
MAR.BEG (1) particles in MAR.BEG.h5
P1 (2) particles in P1.h5
B1 (3) particles in B1.h5
Skipping P2#1 (4) as it's not in the floor plan
Skipping P2\Q1 (5) as it's not in the floor plan
Skipping P2#2 (6) as it's not in the floor plan
Skipping P2\Q2 (7) as it's not in the floor plan
Skipping P2#3 (8) as it's not in the floor plan
B2 (9) particles in B2.h5
Skipping P3#1 (10) as it's not in the floor plan
Skipping P3\Q3 (11) as it's not in the floor plan
Skipping P3#2 (12) as it's not in the floor plan
Skipping P3\Q4 (13) as it's not in the floor plan
Skipping P3#3 (14) as it's not in the floor plan
B3 (15) particles in B3.h5
Skipping P4#1 (16) as it's not in the floor plan
Skipping P4\Q5 (17) as it's not in the floor plan
Skipping P4#2 (18) as it's not in the floor plan
Skipping P4\Q6 (19) as it's not in the floor plan
Skipping P4#3 (20) as it's not in the floor plan
B4 (21) particles in B4.h5
P5 (22) particles in P5.h5
MAR.END (23) particles in MAR.END.h5
END (24) particles in END.h5

Create the bokeh floor plan¶

For each element, we overlay the generated image into its position. Optionally, the white background of the particle plot is removed.

In [8]:

white = 0xff_ff_ff_ff  # 32-bit RGBA all bits set -> white
transparent = 0  # transparent color -> alpha channel all 0

# Create the base bokeh floor plan graph:
bokeh_fig = bokeh_app.create_state().figures[0]

for ele in elements:
    fig = ele.particles.plot("t", "energy", return_figure=True, figsize=PLOT_FIGSIZE)
    fig.axes[0].set_title(ele.name)
    ele.image = fig_to_image(fig, dpi=PLOT_DPI)
    plt.close()

    img = image_array_to_bokeh(ele.image)
    if TRANSPARENT_BACKGROUND:
        img[img == white] = transparent

    pos = ele.position

    bokeh_fig.image_rgba(image=[img], x=pos.x, y=pos.y, dw=pos.w, dh=pos.h)
    bokeh_fig.line(
        [pos.floor_x, pos.x + pos.w / 2.0], [pos.floor_y, pos.y], line_dash="dotted"
    )

# bokeh.plotting.output_file("gen-overlay.html")
# bokeh.plotting.save(bokeh_fig)
bokeh.plotting.output_notebook()
bokeh.plotting.show(bokeh_fig)

white = 0xff_ff_ff_ff # 32-bit RGBA all bits set -> white transparent = 0 # transparent color -> alpha channel all 0 # Create the base bokeh floor plan graph: bokeh_fig = bokeh_app.create_state().figures[0] for ele in elements: fig = ele.particles.plot("t", "energy", return_figure=True, figsize=PLOT_FIGSIZE) fig.axes[0].set_title(ele.name) ele.image = fig_to_image(fig, dpi=PLOT_DPI) plt.close() img = image_array_to_bokeh(ele.image) if TRANSPARENT_BACKGROUND: img[img == white] = transparent pos = ele.position bokeh_fig.image_rgba(image=[img], x=pos.x, y=pos.y, dw=pos.w, dh=pos.h) bokeh_fig.line( [pos.floor_x, pos.x + pos.w / 2.0], [pos.floor_y, pos.y], line_dash="dotted" ) # bokeh.plotting.output_file("gen-overlay.html") # bokeh.plotting.save(bokeh_fig) bokeh.plotting.output_notebook() bokeh.plotting.show(bokeh_fig)

Loading BokehJS ...