Movie example using write_beam¶
TODO: Fix bokeh
In [ ]:
Copied!
from impact import Impact
from distgen import Generator
import numpy as np
import matplotlib.pyplot as plt
import os
from impact import Impact
from distgen import Generator
import numpy as np
import matplotlib.pyplot as plt
import os
In [ ]:
Copied!
# Nicer plotting
import matplotlib
%config InlineBackend.figure_format = 'retina'
matplotlib.rcParams['figure.figsize'] = (8,4)
# Nicer plotting
import matplotlib
%config InlineBackend.figure_format = 'retina'
matplotlib.rcParams['figure.figsize'] = (8,4)
Here we insert write_beam elements into an existing lattice, run, save the beams to an h5 file, and plot using openPMD-beamphysics tools
In [ ]:
Copied!
IMPACT_IN = 'templates/apex_gun/ImpactT.in'
DISTGEN_IN = 'templates/apex_gun/distgen.yaml'
os.path.exists(IMPACT_IN)
IMPACT_IN = 'templates/apex_gun/ImpactT.in'
DISTGEN_IN = 'templates/apex_gun/distgen.yaml'
os.path.exists(IMPACT_IN)
In [ ]:
Copied!
G = Generator(DISTGEN_IN)
G['n_particle'] = 10000
G.run()
P0 = G.particles
P0.plot('x', 'y')
G = Generator(DISTGEN_IN)
G['n_particle'] = 10000
G.run()
P0 = G.particles
P0.plot('x', 'y')
In [ ]:
Copied!
# Make Impact object
I = Impact(IMPACT_IN, initial_particles = P0, verbose=True)
# Change some things
I.header['Nx'] = 32
I.header['Ny'] = 32
I.header['Nz'] = 32
I.header['Dt'] = 1e-13
I.total_charge = P0['charge']
# Change stop location
I.stop = 0.1
# Make Impact object
I = Impact(IMPACT_IN, initial_particles = P0, verbose=True)
# Change some things
I.header['Nx'] = 32
I.header['Ny'] = 32
I.header['Nz'] = 32
I.header['Dt'] = 1e-13
I.total_charge = P0['charge']
# Change stop location
I.stop = 0.1
In [ ]:
Copied!
# Make new write_beam elements and add them to the lattice.
from impact.lattice import new_write_beam
# Make a list of s
for s in np.linspace(0.001, 0.1, 98):
ele = new_write_beam(s=s, ref_eles=I.lattice) # ref_eles will ensure that there are no naming conflicts
I.add_ele(ele)
# Make new write_beam elements and add them to the lattice.
from impact.lattice import new_write_beam
# Make a list of s
for s in np.linspace(0.001, 0.1, 98):
ele = new_write_beam(s=s, ref_eles=I.lattice) # ref_eles will ensure that there are no naming conflicts
I.add_ele(ele)
In [ ]:
Copied!
I.timeout = 1000
I.run()
I.timeout = 1000
I.run()
In [ ]:
Copied!
len(I.particles)
len(I.particles)
Plot¶
In [ ]:
Copied!
from bokeh.plotting import show, figure, output_notebook
from bokeh.layouts import column, row
from bokeh.models import ColumnDataSource
from bokeh import palettes, colors
pal = palettes.Viridis[256]
white=colors.named.white
pal = list(pal)
pal[0] = white # replace 0 with white
pal = tuple(pal)
output_notebook(verbose=False, hide_banner=True)
import os
from bokeh.plotting import show, figure, output_notebook
from bokeh.layouts import column, row
from bokeh.models import ColumnDataSource
from bokeh import palettes, colors
pal = palettes.Viridis[256]
white=colors.named.white
pal = list(pal)
pal[0] = white # replace 0 with white
pal = tuple(pal)
output_notebook(verbose=False, hide_banner=True)
import os
In [ ]:
Copied!
# Prepare histogram function
PL = I.particles
ilist = []
for k in PL:
if k.startswith('write_beam_'):
ilist.append(int(k.strip('write_beam_')))
def bin_particles(i, key1='x', key2='y', bins=40):
P = I.particles[f'write_beam_{i}']
return np.histogram2d(P[key1], P[key2], weights=P.weight, bins=bins)
bin_particles(100)
# Prepare histogram function
PL = I.particles
ilist = []
for k in PL:
if k.startswith('write_beam_'):
ilist.append(int(k.strip('write_beam_')))
def bin_particles(i, key1='x', key2='y', bins=40):
P = I.particles[f'write_beam_{i}']
return np.histogram2d(P[key1], P[key2], weights=P.weight, bins=bins)
bin_particles(100)
In [ ]:
Copied!
# Prepare a datasource for Bokeh
def bin_bunch_datasource_h5(i, key1, key2, bins=20, nice=True, liveOnly=True, liveStatus=1):
H, xedges, yedges = bin_particles(i, key1, key2, bins=bins)
xmin = min(xedges)
xmax = max(xedges)
ymin = min(yedges)
ymax = max(yedges)
#if nice:
# f1 = nice_phase_space_factor[component1]
# f2 = nice_phase_space_factor[component2]
# xlabel = nice_phase_space_label[component1]
# ylabel = nice_phase_space_label[component2]
# xmin *= f1
# xmax *= f1
# ymin *= f2
# ymax *= f2
#else:
# xlabel = component1
# ylabel = component2
# Form datasource
dat = {'image':[H.transpose()], 'xmin':[xmin], 'ymin':[ymin], 'dw':[xmax-xmin], 'dh':[ymax-ymin]}
dat['xmax'] = [xmax]
dat['ymax'] = [ymax]
ds = ColumnDataSource(data=dat)
return ds
ds = bin_bunch_datasource_h5(100, 'x', 'y')
# Prepare a datasource for Bokeh
def bin_bunch_datasource_h5(i, key1, key2, bins=20, nice=True, liveOnly=True, liveStatus=1):
H, xedges, yedges = bin_particles(i, key1, key2, bins=bins)
xmin = min(xedges)
xmax = max(xedges)
ymin = min(yedges)
ymax = max(yedges)
#if nice:
# f1 = nice_phase_space_factor[component1]
# f2 = nice_phase_space_factor[component2]
# xlabel = nice_phase_space_label[component1]
# ylabel = nice_phase_space_label[component2]
# xmin *= f1
# xmax *= f1
# ymin *= f2
# ymax *= f2
#else:
# xlabel = component1
# ylabel = component2
# Form datasource
dat = {'image':[H.transpose()], 'xmin':[xmin], 'ymin':[ymin], 'dw':[xmax-xmin], 'dh':[ymax-ymin]}
dat['xmax'] = [xmax]
dat['ymax'] = [ymax]
ds = ColumnDataSource(data=dat)
return ds
ds = bin_bunch_datasource_h5(100, 'x', 'y')
In [ ]:
Copied!
plot = figure(#x_range=[xmin,xmax], y_range=[ymin,ymax],
# x_axis_label = xlabel, y_axis_label = ylabel,
width=500, height=500)
plot.image(image='image', x='xmin', y='ymin', dw='dw', dh='dh', source=ds,palette=pal)
show(plot)
plot = figure(#x_range=[xmin,xmax], y_range=[ymin,ymax],
# x_axis_label = xlabel, y_axis_label = ylabel,
width=500, height=500)
plot.image(image='image', x='xmin', y='ymin', dw='dw', dh='dh', source=ds,palette=pal)
show(plot)
Interactive¶
In [ ]:
Copied!
from bokeh.models.widgets import Slider
from bokeh import palettes, colors
from bokeh.models.widgets import Slider
from bokeh import palettes, colors
In [ ]:
Copied!
# interactive
def myapp2(doc):
bunches = ilist
doc.bunchi = bunches[0]
doc.component1 = 'z'
doc.component2 = 'x'
doc.xlabel = doc.component1
doc.ylabel = doc.component2
doc.bins = 100
#doc.range = FULLRANGE
ds = bin_bunch_datasource_h5(doc.bunchi, doc.component1, doc.component2,bins=doc.bins)
def refresh():
ds.data = dict(bin_bunch_datasource_h5(doc.bunchi, doc.component1, doc.component2,bins=doc.bins).data )
# Default plot
plot = figure(title='',
x_axis_label = doc.xlabel, y_axis_label = doc.ylabel,
width=500, height=500)
plot.image(image='image', x='xmin', y='ymin', dw='dw', dh='dh', source=ds, palette=pal)
def slider_handler(attr, old, new):
doc.bunchi = bunches[new]
refresh()
slider = Slider(start=0, end=len(bunches)-1, value=0, step=1, title='x')
slider.on_change('value', slider_handler)
# Add plot to end
doc.add_root(column(slider, plot))
show(myapp2)# , notebook_url=remote_jupyter_proxy_url)
# interactive
def myapp2(doc):
bunches = ilist
doc.bunchi = bunches[0]
doc.component1 = 'z'
doc.component2 = 'x'
doc.xlabel = doc.component1
doc.ylabel = doc.component2
doc.bins = 100
#doc.range = FULLRANGE
ds = bin_bunch_datasource_h5(doc.bunchi, doc.component1, doc.component2,bins=doc.bins)
def refresh():
ds.data = dict(bin_bunch_datasource_h5(doc.bunchi, doc.component1, doc.component2,bins=doc.bins).data )
# Default plot
plot = figure(title='',
x_axis_label = doc.xlabel, y_axis_label = doc.ylabel,
width=500, height=500)
plot.image(image='image', x='xmin', y='ymin', dw='dw', dh='dh', source=ds, palette=pal)
def slider_handler(attr, old, new):
doc.bunchi = bunches[new]
refresh()
slider = Slider(start=0, end=len(bunches)-1, value=0, step=1, title='x')
slider.on_change('value', slider_handler)
# Add plot to end
doc.add_root(column(slider, plot))
show(myapp2)# , notebook_url=remote_jupyter_proxy_url)
In [ ]:
Copied!
# If there are multiple
import os
os.environ['BOKEH_ALLOW_WS_ORIGIN'] = 'localhost:8888'
# If there are multiple
import os
os.environ['BOKEH_ALLOW_WS_ORIGIN'] = 'localhost:8888'