Movie example using write_beam¶

TODO: Fix bokeh

In [ ]:

from impact import Impact
from distgen import Generator

# Make new write_beam elements and add them to the lattice.
from impact.lattice import new_write_beam


import numpy as np
import os


from bokeh.plotting import show, figure, output_notebook
from bokeh.layouts import column
from bokeh.models import ColumnDataSource
from bokeh import palettes, colors
from bokeh.models.widgets import Slider

output_notebook(verbose=False, hide_banner=True)

from impact import Impact from distgen import Generator # Make new write_beam elements and add them to the lattice. from impact.lattice import new_write_beam import numpy as np import os from bokeh.plotting import show, figure, output_notebook from bokeh.layouts import column from bokeh.models import ColumnDataSource from bokeh import palettes, colors from bokeh.models.widgets import Slider output_notebook(verbose=False, hide_banner=True)

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 [ ]:

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 [ ]:

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 [ ]:

# 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 [ ]:

# 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 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 [ ]:

I.timeout = 1000
I.run()

I.timeout = 1000 I.run()

In [ ]:

len(I.particles)

len(I.particles)

Plot¶

In [ ]:

pal = palettes.Viridis[256]
white = colors.named.white
pal = list(pal)
pal[0] = white  # replace 0 with white
pal = tuple(pal)

pal = palettes.Viridis[256] white = colors.named.white pal = list(pal) pal[0] = white # replace 0 with white pal = tuple(pal)

In [ ]:

# 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 [ ]:

# 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 [ ]:

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 [ ]:

# 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 [ ]:

# If there are multiple


os.environ["BOKEH_ALLOW_WS_ORIGIN"] = "localhost:8888"

# If there are multiple os.environ["BOKEH_ALLOW_WS_ORIGIN"] = "localhost:8888"