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]:
Copied!
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]:
Copied!
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]:
Copied!
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]:
Copied!
%%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]:
Copied!
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]:
Copied!
%%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]:
Copied!
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]:
Copied!
%%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]:
Copied!
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]:
Copied!
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]:
Copied!
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