Interactive Bokeh plot of Elegant Twiss output

A. Petrenko (CERN, 2017)

This notebook explains how to plot Elegant-produced Twiss-output using Bokeh library. The notebook can be executed at the MS Azure Cloud service (file Twiss_plot.ipynb).

Similar approach can be also used to plot the MAD-X output.

First download some sample twiss_output data:

In [1]:
# Downloading example Elegant twiss_output data:
!wget -Nq https://apetrenko.blob.core.windows.net/sdds-jupyter/twiss.txt

This file is the result of this SDDS command

sdds2stream twiss.twi -col=s,ElementName,ElementType,betax,betay,etax > twiss.txt

applied to the result of twiss_output. Also here is the corresponding .lte and .ele files for details.

In [1]:
#!head twiss.txt

Such file is easy to load into the pandas DataFrame for later use:

In [2]:
import pandas as pd
In [4]:
df = pd.read_table('twiss.txt', names=['s','ElementName','ElementType','betax','betay','etax'], delim_whitespace=True)
In [5]:
df.head()
Out[5]:
s ElementName ElementType betax betay etax
0 0.0 _BEG_ MARK 4.915917 1.275074 9.063422e-07
1 0.0 Q CHARGE 4.915917 1.275074 9.063422e-07
2 0.0 W2 DRIF 4.915917 1.275074 9.063422e-07
3 0.0 MA_80 MAXAMP 4.915917 1.275074 9.063422e-07
4 0.1 Q3F3 QUAD 4.531497 1.522587 8.679051e-07
In [6]:
from bokeh.plotting import figure, show
from bokeh.io import output_notebook
output_notebook()
Loading BokehJS ...

Simple plot:

In [7]:
p = figure(tools="save,box_zoom,pan,reset", toolbar_location="right",
            logo="grey", plot_width=800, plot_height=300, active_drag="box_zoom")

p.line(df.s, df.betax, color='red', line_width=2, line_alpha=0.5, legend='Beta x')
p.line(df.s, df.betay, color='blue', line_width=2, line_alpha=0.5, legend='Beta y')
p.legend.background_fill_alpha = 0.5
p.y_range.start=0; #p.y_range.end = df[['betax','betay']].values.max() + 0.5
p.xaxis.axis_label='s (m)'
p.yaxis.axis_label='Beta x,y (m)'

show(p)

More complicated plot with lattice layout and hover annotations:

In [8]:
df['L'] = df['s'] - df['s'].shift(+1) # lengths of all elements
df.set_value(0, 'L', 0) # replace NaN with 0 in the first row
df['s0'] = df['s'] - df['L']/2 # centers of all elements
df.head()
Out[8]:
s ElementName ElementType betax betay etax L s0
0 0.0 _BEG_ MARK 4.915917 1.275074 9.063422e-07 0.0 0.00
1 0.0 Q CHARGE 4.915917 1.275074 9.063422e-07 0.0 0.00
2 0.0 W2 DRIF 4.915917 1.275074 9.063422e-07 0.0 0.00
3 0.0 MA_80 MAXAMP 4.915917 1.275074 9.063422e-07 0.0 0.00
4 0.1 Q3F3 QUAD 4.531497 1.522587 8.679051e-07 0.1 0.05
In [9]:
from bokeh.models import (HoverTool, ColumnDataSource,
                          CustomJS)
from bokeh.layouts import gridplot

def pic_h(row):
    if row.ElementType.endswith('BEN'): return 4
    if row.ElementType.startswith('SEXT'): return 5
    if row.ElementType.endswith('QUAD'): return 6
    if row.ElementName.startswith('BPM'): return 5
    return 3

def pic_L(row):
    if row.L < 0.01:
        return 0.01
    else:
        return row.L

def pic_color(row):
    if row.ElementType.endswith('BEN'): return 'green'
    if row.ElementType.endswith('SEXT'): return 'yellow'
    if row.ElementType.endswith('QUAD'):
        if row.ElementName[-2] == 'F': return 'red'
        return 'blue'
    if row.ElementName.startswith('BPM'): return 'magenta'
    return 'white'

df["pic_h"] = df.apply(pic_h, axis = 1)
df["pic_L"] = df.apply(pic_L, axis = 1)
df["pic_color"] = df.apply(pic_color, axis = 1)

all_elements_src = ColumnDataSource(data=dict(
    NAME=df.ElementName.values,
    s = df.s0.values,
    L = df.pic_L.values,
    h = df.pic_h.values,
    color = df.pic_color.values,
))

p = figure(tools="save,xbox_zoom,xpan,reset", toolbar_location="above",
            logo="grey", plot_width=800, plot_height=110, active_drag="xbox_zoom",
            y_axis_type=None)

#p.x_range.start = df.s.values.min() - 1
#p.x_range.end   = df.s.values.max() + 1

r1 = p.rect('s', 0, width='L', height='h', fill_color='color',
            fill_alpha=0.5, line_width=1.0, line_alpha=0.2,
            line_color='black', source=all_elements_src)

tips = [
    ('Name', '@NAME'),      
]

hover = HoverTool(tooltips=tips)
p.add_tools(hover)

#hover = p.select_one(HoverTool)
#hover.tooltips = tips

#hover.renderers=[c1,c2,r1,r2]
hover.renderers=[r1]

p.y_range.start = -7
p.y_range.end   = +7
p.xaxis.axis_label='s (m)'

p_layout = p
show(p_layout)
In [10]:
p = figure(tools="save,xbox_zoom,xpan,reset", x_range=p_layout.x_range, toolbar_location="right",
            logo="grey", plot_width=800, plot_height=300, active_drag="xbox_zoom")

p.line(df.s, df.betax, color='red', line_width=2, line_alpha=0.5, legend='Beta x')
p.line(df.s, df.betay, color='blue', line_width=2, line_alpha=0.5, legend='Beta y')
p.legend.background_fill_alpha = 0.5
p.y_range.start=0; #p.y_range.end = df[['betax','betay']].values.max() + 0.5

p.xaxis.major_label_text_alpha = 0
p.xaxis.major_label_text_font_size = '1pt'

p.yaxis.axis_label='Beta x,y (m)'

lay=gridplot([
        [p],
        [p_layout]
            ], toolbar_location='right')
show(lay)

Generated p_layout can be similarly attached to any other plot, like dispersion function for example:

In [11]:
p = figure(tools="save,xbox_zoom,xpan,reset", x_range=p_layout.x_range, toolbar_location="right",
            logo="grey", plot_width=800, plot_height=300, active_drag="xbox_zoom")

p.line(df.s, df.etax, color='green', line_width=2, line_alpha=0.5)
p.y_range.start=-0.05

p.xaxis.major_label_text_alpha = 0
p.xaxis.major_label_text_font_size = '1pt'

p.yaxis.axis_label='Dx (m)'

lay=gridplot([
        [p],
        [p_layout]
            ], toolbar_location='right')
show(lay)
In [ ]: