Genesis4 FODO Scan¶

LUME-Genesis makes it easy to explore a parameter space within a notebook.

Here we will make a lattice file within the notebook, and scan the quadrupole k to find the best final power for a simple benchmark example.

In [1]:

import matplotlib.pyplot as plt
import numpy as np

import genesis.version4 as g4
from genesis.version4 import Genesis4

%config InlineBackend.figure_format = 'retina'

import matplotlib.pyplot as plt import numpy as np import genesis.version4 as g4 from genesis.version4 import Genesis4 %config InlineBackend.figure_format = 'retina'

Lattice¶

Here we will use Genesis4 SASE Benchmark as the basis for a lattice, but create it here in code.

Note that the FODO cell length is 9.5 m, and there are 13 cells, for a total length of 123.5 m.

In [2]:

def make_lat(k1=2):
    return g4.Lattice(
        {
            "D1": g4.Drift(L=0.445),
            "D2": g4.Drift(L=0.24),
            "QF": g4.Quadrupole(L=0.08, k1=k1),
            "QD": g4.Quadrupole(L=0.08, k1=-k1),
            "UND": g4.Undulator(aw=0.84853, lambdau=0.015, nwig=266),
            "FODO": g4.Line(
                elements=["UND", "D1", "QF", "D2", "UND", "D1", "QD", "D2"]
            ),
            "ARAMIS": g4.Line(elements=[g4.DuplicatedLineItem(label="FODO", count=13)]),
        },
    )


make_lat()

def make_lat(k1=2): return g4.Lattice( { "D1": g4.Drift(L=0.445), "D2": g4.Drift(L=0.24), "QF": g4.Quadrupole(L=0.08, k1=k1), "QD": g4.Quadrupole(L=0.08, k1=-k1), "UND": g4.Undulator(aw=0.84853, lambdau=0.015, nwig=266), "FODO": g4.Line( elements=["UND", "D1", "QF", "D2", "UND", "D1", "QD", "D2"] ), "ARAMIS": g4.Line(elements=[g4.DuplicatedLineItem(label="FODO", count=13)]), }, ) make_lat()

Out[2]:

Lattice(
  elements={
    'D1': Drift(L=0.445),
    'D2': Drift(L=0.24),
    'QF': Quadrupole(L=0.08, k1=2.0),
    'QD': Quadrupole(L=0.08, k1=-2.0),
    'UND': Undulator(aw=0.84853, lambdau=0.015, nwig=266),
    'FODO': Line(elements=['UND', 'D1', 'QF', 'D2', 'UND', 'D1', 'QD', 'D2']),
    'ARAMIS': Line(elements=[DuplicatedLineItem(label='FODO', count=13)]),
  },
)

In [3]:

main = g4.MainInput(
    [
        g4.Setup(
            rootname="Benchmark",
            beamline="ARAMIS",
            gamma0=11357.82,
            delz=0.045,
            shotnoise=False,
            beam_global_stat=True,
            field_global_stat=True,
        ),
        g4.LatticeNamelist(zmatch=9.5),
        g4.Field(power=5000.0, waist_size=3e-05, dgrid=0.0002, ngrid=255),
        g4.Beam(delgam=1.0, current=3000.0, ex=4e-07, ey=4e-07),
        g4.Track(zstop=123.5),
    ],
)

G = Genesis4(main, make_lat())

main = g4.MainInput( [ g4.Setup( rootname="Benchmark", beamline="ARAMIS", gamma0=11357.82, delz=0.045, shotnoise=False, beam_global_stat=True, field_global_stat=True, ), g4.LatticeNamelist(zmatch=9.5), g4.Field(power=5000.0, waist_size=3e-05, dgrid=0.0002, ngrid=255), g4.Beam(delgam=1.0, current=3000.0, ex=4e-07, ey=4e-07), g4.Track(zstop=123.5), ], ) G = Genesis4(main, make_lat())

In [4]:

%%time
G.nproc = 0  # Auto-select
output = G.run()

%%time G.nproc = 0 # Auto-select output = G.run()

CPU times: user 23 ms, sys: 8.26 ms, total: 31.3 ms
Wall time: 27.4 s

In [5]:

G.plot(
    "field_peak_power", yscale="log", y2=["beam_xsize", "beam_ysize"], ylim2=(0, 50e-6)
)

G.plot( "field_peak_power", yscale="log", y2=["beam_xsize", "beam_ysize"], ylim2=(0, 50e-6) )

Run1 function¶

Make a simple function to run a complete simulation and return a Genesis4 object.

In [6]:

%%time


def run1(k):
    main.track.zstop = 20
    G = Genesis4(main, make_lat(k))
    G.nproc = 0
    G.run()
    return G


G2 = run1(4)

%%time def run1(k): main.track.zstop = 20 G = Genesis4(main, make_lat(k)) G.nproc = 0 G.run() return G G2 = run1(4)

CPU times: user 18.5 ms, sys: 9.72 ms, total: 28.2 ms
Wall time: 4.09 s

In [7]:

G2.plot(
    "field_peak_power", yscale="log", y2=["beam_xsize", "beam_ysize"], ylim2=(0, 50e-6)
)

G2.plot( "field_peak_power", yscale="log", y2=["beam_xsize", "beam_ysize"], ylim2=(0, 50e-6) )

Scan k1¶

Now scan use this function to scan various quadrupole k1. Note that Genesis4 is doing matching internally with zmatch in the input.

In [8]:

%%time
klist = np.linspace(1, 3, 10)
Glist = [run1(k) for k in klist]

%%time klist = np.linspace(1, 3, 10) Glist = [run1(k) for k in klist]

CPU times: user 178 ms, sys: 70.5 ms, total: 249 ms
Wall time: 45.9 s

In [9]:

fig, ax = plt.subplots()
for k, g in zip(klist, Glist):
    x = g.stat("zplot")
    y = g.stat("field_peak_power")
    ax.plot(x, y / 1e6, label=f"{k:0.1f}")
ax.set_yscale("log")
ax.set_xlabel(r"$z$ (m)")
ax.set_ylabel("power (MW)")
plt.legend(title=r"$k$ (1/m$^2$)")

fig, ax = plt.subplots() for k, g in zip(klist, Glist): x = g.stat("zplot") y = g.stat("field_peak_power") ax.plot(x, y / 1e6, label=f"{k:0.1f}") ax.set_yscale("log") ax.set_xlabel(r"$z$ (m)") ax.set_ylabel("power (MW)") plt.legend(title=r"$k$ (1/m$^2$)")

Out[9]:

<matplotlib.legend.Legend at 0x14ab161e0>

In [10]:

fig, ax = plt.subplots()
y = np.array([g.stat("field_peak_power")[-1] for g in Glist])
ixbest = y.argmax()
Gbest = Glist[ixbest]
kbest = klist[ixbest]
ybest = y[ixbest]
ax.plot(klist, y / 1e6)
ax.scatter(kbest, ybest / 1e6, marker="*", label=rf"$k$= {kbest:0.1f} 1/m$^2$")
ax.set_ylabel("end power (MW)")
ax.set_xlabel(r"$k$ (1/m$^2$)")
plt.legend()

fig, ax = plt.subplots() y = np.array([g.stat("field_peak_power")[-1] for g in Glist]) ixbest = y.argmax() Gbest = Glist[ixbest] kbest = klist[ixbest] ybest = y[ixbest] ax.plot(klist, y / 1e6) ax.scatter(kbest, ybest / 1e6, marker="*", label=rf"$k$= {kbest:0.1f} 1/m$^2$") ax.set_ylabel("end power (MW)") ax.set_xlabel(r"$k$ (1/m$^2$)") plt.legend()

Out[10]:

<matplotlib.legend.Legend at 0x14af86ff0>

In [11]:

fig, ax = plt.subplots()
for k, g in zip(klist, Glist):
    x = g.stat("zplot")
    y = g.stat("beam_xsize")
    if k == kbest:
        color = "black"
    else:
        color = None
    ax.plot(x, y * 1e6, label=f"{k:0.1f}", color=color)

ax.set_xlabel(r"$z$ (m)")
ax.set_ylabel("Beam xsize (µm)")
# ax.set_ylim(0, None)
plt.legend(title=r"$k$ (1/m$^2$)")

fig, ax = plt.subplots() for k, g in zip(klist, Glist): x = g.stat("zplot") y = g.stat("beam_xsize") if k == kbest: color = "black" else: color = None ax.plot(x, y * 1e6, label=f"{k:0.1f}", color=color) ax.set_xlabel(r"$z$ (m)") ax.set_ylabel("Beam xsize (µm)") # ax.set_ylim(0, None) plt.legend(title=r"$k$ (1/m$^2$)")

Out[11]:

<matplotlib.legend.Legend at 0x14b0ee1e0>

Cleanup