<jupyter_start><jupyter_text>Movie example using write_beam Here we insert write_beam elements into an existing lattice, run, save the beams to an h5 file, and plot using openPMD-beamphysics tools<jupyter_code>from impact import Impact
from distgen import Generator
import numpy as np
import matplotlib.pyplot as plt
import os
IMPACT_IN  = 'templates/apex_gun/ImpactT.in'
DISTGEN_IN = 'templates/apex_gun/distgen.yaml'
os.path.exists(IMPACT_IN)
G = Generator(DISTGEN_IN)
G['n_particle'] = 10000
G.run()
P0 = G.particles
P0.plot('x', 'y')
# 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 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)
I.timeout = 1000
I.run()
len(I.particles)<jupyter_output><empty_output><jupyter_text>Plot<jupyter_code>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
# 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 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')

plot = figure(#x_range=[xmin,xmax], y_range=[ymin,ymax], 
              #    x_axis_label = xlabel,  y_axis_label = ylabel,
               plot_width=500, plot_height=500)
plot.image(image='image', x='xmin', y='ymin', dw='dw', dh='dh', source=ds,palette=pal)
show(plot)<jupyter_output><empty_output><jupyter_text>Interactive<jupyter_code>from bokeh.models.widgets import Slider
from bokeh import palettes, colors

# 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,
           plot_width=500, plot_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)
# If there are multiple 
import os
os.environ['BOKEH_ALLOW_WS_ORIGIN'] = 'localhost:8888'
%%time
I.archive()<jupyter_output>Archiving to file impact_4713be609d6c5ed64c6dd62272833320.h5
CPU times: user 1.33 s, sys: 134 ms, total: 1.46 s
Wall time: 1.53 s
{'fraction_non_alphanumeric': 0.1280104144065958, 'fraction_numerical': 0.03124321978737253, 'mean_word_length': 3.767321613236815, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '30115', 'license': 'Apache-2.0', 'n_tokens_mistral': 1864, 'n_tokens_neox': 1751, 'n_words': 379, 'path': 'examples/movie_example.ipynb', 'repo_name': 'jacquelinegarrahan/lume-impact', 'size': '4609'}