CSR Zeuthen Benchmark - Comparing to Bmad¶

This tries to replicate the Zeuthern CSR benchmark according to: https://journals.aps.org/prab/abstract/10.1103/PhysRevSTAB.16.060703

In [ ]:

import matplotlib.pyplot as plt
import numpy as np
from pmd_beamphysics import ParticleGroup
from pytao import Tao
from scipy.constants import c

import impact.z as IZ
from impact.z.interfaces.bmad import plot_impactz_and_tao_stats
from impact.tests.z.conftest import bmad_files

import matplotlib.pyplot as plt import numpy as np from pmd_beamphysics import ParticleGroup from pytao import Tao from scipy.constants import c import impact.z as IZ from impact.z.interfaces.bmad import plot_impactz_and_tao_stats from impact.tests.z.conftest import bmad_files

Parameters¶

In [ ]:

TRACK_START = "BEGINNING"
TRACK_END = "END"

NX = 32
NY = 32
NZ = 128
N_PARTICLE = 100_000
DS_STEP = 0.01

CSR_ON = True
DRIFT_CSR_ON = True

BMAD_SC_ON = False

CHIRP = -36  # 1/m

TRACK_START = "BEGINNING" TRACK_END = "END" NX = 32 NY = 32 NZ = 128 N_PARTICLE = 100_000 DS_STEP = 0.01 CSR_ON = True DRIFT_CSR_ON = True BMAD_SC_ON = False CHIRP = -36 # 1/m

Tao¶

In [ ]:

# Gaussian
def set_gaussian(
    tao,
    n_particle=N_PARTICLE,
    a_norm_emit=1.0e-6,
    b_norm_emit=1.0e-6,
    bunch_charge=1e-9,
    sig_pz0=2e-6,
    sig_z=200e-6,
    center_pz=0,
    chirp=0,  # 1/m
):
    sig_pz = np.hypot(sig_pz0, chirp * sig_z)

    cmds = [
        f"set beam_init n_particle = {n_particle}",
        "set beam_init random_engine = quasi",
        "set beam_init saved_at = MARKER::*, BEGINNING, END",
        f"set beam_init a_norm_emit = {a_norm_emit}",
        f"set beam_init b_norm_emit = {b_norm_emit}",
        f"set beam_init bunch_charge = {bunch_charge}",
        f"set beam_init sig_pz = {sig_pz}",
        f"set beam_init sig_z = {sig_z}",
        f"set beam_init dpz_dz = {chirp}",
        f"set beam_init center(6) = {center_pz}",
    ]
    tao.cmds(cmds)
    tao.cmd("set global lattice_calc_on = T")


def get_particles(tao, ele_id):
    return ParticleGroup(data=tao.bunch_data(ele_id))

# Gaussian def set_gaussian( tao, n_particle=N_PARTICLE, a_norm_emit=1.0e-6, b_norm_emit=1.0e-6, bunch_charge=1e-9, sig_pz0=2e-6, sig_z=200e-6, center_pz=0, chirp=0, # 1/m ): sig_pz = np.hypot(sig_pz0, chirp * sig_z) cmds = [ f"set beam_init n_particle = {n_particle}", "set beam_init random_engine = quasi", "set beam_init saved_at = MARKER::*, BEGINNING, END", f"set beam_init a_norm_emit = {a_norm_emit}", f"set beam_init b_norm_emit = {b_norm_emit}", f"set beam_init bunch_charge = {bunch_charge}", f"set beam_init sig_pz = {sig_pz}", f"set beam_init sig_z = {sig_z}", f"set beam_init dpz_dz = {chirp}", f"set beam_init center(6) = {center_pz}", ] tao.cmds(cmds) tao.cmd("set global lattice_calc_on = T") def get_particles(tao, ele_id): return ParticleGroup(data=tao.bunch_data(ele_id))

In [ ]:

tao = Tao(lattice_file=bmad_files / "csr_zeuthen.bmad", plot="mpl")

# tao.cmd('set ele beginning e_tot = 500e6') # TEST

set_gaussian(tao, n_particle=N_PARTICLE, chirp=CHIRP)

tao = Tao(lattice_file=bmad_files / "csr_zeuthen.bmad", plot="mpl") # tao.cmd('set ele beginning e_tot = 500e6') # TEST set_gaussian(tao, n_particle=N_PARTICLE, chirp=CHIRP)

In [ ]:

tao.cmd("set bmad_com csr_and_space_charge_on = T")
tao.cmd("set space_charge particle_bin_span = 1")
tao.cmd(f"set space_charge n_bin = {NZ}")
tao.cmd(f"set space_charge ds_track_step = {DS_STEP}")
tao.cmd("set bmad_com radiation_damping_on  = T")  # off by default
tao.cmd("set bmad_com radiation_fluctuations_on = T")  # off by default

if BMAD_SC_ON:
    tao.cmd("set ele * space_charge_method = fft_3d")
    tao.cmd(f"set space_charge space_charge_mesh_size = {NX} {NY} {NZ}")
else:
    tao.cmd("set ele * space_charge_method = off")
tao.cmd(f"set ele * DS_STEP = {DS_STEP}")

tao.cmd("set bmad_com csr_and_space_charge_on = T") tao.cmd("set space_charge particle_bin_span = 1") tao.cmd(f"set space_charge n_bin = {NZ}") tao.cmd(f"set space_charge ds_track_step = {DS_STEP}") tao.cmd("set bmad_com radiation_damping_on = T") # off by default tao.cmd("set bmad_com radiation_fluctuations_on = T") # off by default if BMAD_SC_ON: tao.cmd("set ele * space_charge_method = fft_3d") tao.cmd(f"set space_charge space_charge_mesh_size = {NX} {NY} {NZ}") else: tao.cmd("set ele * space_charge_method = off") tao.cmd(f"set ele * DS_STEP = {DS_STEP}")

In [ ]:

if not CSR_ON:
    tao.cmd("set ele * csr_method = off")
elif DRIFT_CSR_ON:
    tao.cmd("set ele * csr_method = 1_dim")
else:
    tao.cmd("set ele * csr_method = off")
    tao.cmd("set ele sbend::* csr_method = 1_dim")

if not CSR_ON: tao.cmd("set ele * csr_method = off") elif DRIFT_CSR_ON: tao.cmd("set ele * csr_method = 1_dim") else: tao.cmd("set ele * csr_method = off") tao.cmd("set ele sbend::* csr_method = 1_dim")

In [ ]:

%%time
tao.track_beam()

P0 = get_particles(tao, "beginning")
P1 = get_particles(tao, "end")

P1.plot("delta_t", "energy"), P1["sigma_t"] * c * 1e6

%%time tao.track_beam() P0 = get_particles(tao, "beginning") P1 = get_particles(tao, "end") P1.plot("delta_t", "energy"), P1["sigma_t"] * c * 1e6

In [ ]:

sigma_z0 = P0["sigma_t"] * c
sigma_z1 = P1["sigma_t"] * c

sigma_z1 * 1e6

sigma_z0 = P0["sigma_t"] * c sigma_z1 = P1["sigma_t"] * c sigma_z1 * 1e6

In [ ]:

# Compression factor
sigma_z0 / sigma_z1

# Compression factor sigma_z0 / sigma_z1

Impact-Z¶

In [ ]:

input = IZ.ImpactZInput.from_tao(tao, verbose=False, write_beam_eles="")
I = IZ.ImpactZ(input)
I.initial_particles = P0

input = IZ.ImpactZInput.from_tao(tao, verbose=False, write_beam_eles="") I = IZ.ImpactZ(input) I.initial_particles = P0

In [ ]:

I.input.integrator_type = 1  # default

I.input.integrator_type = 1 # default

In [ ]:

%%time
I.nproc = 0
I.run(verbose=False);

%%time I.nproc = 0 I.run(verbose=False);

In [ ]:

P2 = I.output.particles["final_particles"]
P2.plot("delta_t", "energy")

P2 = I.output.particles["final_particles"] P2.plot("delta_t", "energy")

Compare¶

In [ ]:

plot_impactz_and_tao_stats(I, tao)

plot_impactz_and_tao_stats(I, tao)

In [ ]:

def compare(
    xkey="x",
    ykey="y",
    skip=1,
):
    fig, ax = plt.subplots()

    plist = (
        (P1, "Bmad", "X", "blue"),
        (P2, "Impact-Z", ".", "orange"),
    )

    for p, label, marker, color in plist:
        ax.scatter(
            p[xkey][::skip],
            p[ykey][::skip],
            label=label,
            marker=marker,
            alpha=0.5,
            color=color,
        )

    for p, label, marker, color in plist:
        ax.scatter(
            p["mean_" + xkey],
            p["mean_" + ykey],
            marker="+",
            color=color,
            facecolor="black",
        )

    plt.legend()

    ax.set_xlabel(xkey)
    ax.set_ylabel(ykey)


compare()

def compare( xkey="x", ykey="y", skip=1, ): fig, ax = plt.subplots() plist = ( (P1, "Bmad", "X", "blue"), (P2, "Impact-Z", ".", "orange"), ) for p, label, marker, color in plist: ax.scatter( p[xkey][::skip], p[ykey][::skip], label=label, marker=marker, alpha=0.5, color=color, ) for p, label, marker, color in plist: ax.scatter( p["mean_" + xkey], p["mean_" + ykey], marker="+", color=color, facecolor="black", ) plt.legend() ax.set_xlabel(xkey) ax.set_ylabel(ykey) compare()

In [ ]:

P1["mean_x"], P2["mean_x"]

P1["mean_x"], P2["mean_x"]

In [ ]:

compare("delta_t", "energy")

compare("delta_t", "energy")

In [ ]:

compare("delta_t", "xp")

compare("delta_t", "xp")

In [ ]:

energy0 = P0["mean_energy"]
P1["mean_energy"] - energy0, P2["mean_energy"] - energy0

energy0 = P0["mean_energy"] P1["mean_energy"] - energy0, P2["mean_energy"] - energy0

Impact-Z Lorentz integrator¶

ImpactZInput.integrator_type = 2

Note that this is not expected to work as well because the dipole fringe fields are not able to be specified.

In [ ]:

I2 = I.copy()

I2 = I.copy()

In [ ]:

I2.input.integrator_type = 2
I2.run()
P3 = I2.output.particles["final_particles"]

I2.input.integrator_type = 2 I2.run() P3 = I2.output.particles["final_particles"]

In [ ]:

stats1 = I.output.stats
stats2 = I2.output.stats

eref = I.output.stats.energy_ref[0]  # should be constant

fig, ax = plt.subplots(1, figsize=(12, 4))

ax.plot(stats1.z, (stats1.mean_energy - eref) / 1e6, label="Impact-T, map")
ax.plot(stats2.z, (stats2.mean_energy - eref) / 1e6, "--", label="Impact-T, nonlinear")

ax.set_xlabel(r"$s$ (m)")
ax.set_ylabel(r"$dE$ (MeV)")

plt.legend()

stats1 = I.output.stats stats2 = I2.output.stats eref = I.output.stats.energy_ref[0] # should be constant fig, ax = plt.subplots(1, figsize=(12, 4)) ax.plot(stats1.z, (stats1.mean_energy - eref) / 1e6, label="Impact-T, map") ax.plot(stats2.z, (stats2.mean_energy - eref) / 1e6, "--", label="Impact-T, nonlinear") ax.set_xlabel(r"$s$ (m)") ax.set_ylabel(r"$dE$ (MeV)") plt.legend()

In [ ]:

def compare(
    xkey="x",
    ykey="y",
    skip=10,
):
    fig, ax = plt.subplots()

    for p, label, marker in (
        #  (P1, "Bmad", "X"),
        (P2, "Impact-Z map", "x"),
        (P3, "Impact-Z nonlinear", "."),
    ):
        ax.scatter(
            p[xkey][::skip], p[ykey][::skip], label=label, marker=marker, alpha=0.5
        )

    plt.legend()

    ax.set_xlabel(xkey)
    ax.set_ylabel(ykey)


compare("delta_t", "xp")

def compare( xkey="x", ykey="y", skip=10, ): fig, ax = plt.subplots() for p, label, marker in ( # (P1, "Bmad", "X"), (P2, "Impact-Z map", "x"), (P3, "Impact-Z nonlinear", "."), ): ax.scatter( p[xkey][::skip], p[ykey][::skip], label=label, marker=marker, alpha=0.5 ) plt.legend() ax.set_xlabel(xkey) ax.set_ylabel(ykey) compare("delta_t", "xp")