import matplotlib.pyplot as plt
import numpy as np
import xpart as xp
import xtrack as xt
import xwakes as xw

%config InlineBackend.figure_format = "retina"

/home/runner/work/Journees_Accelerateurs/Journees_Accelerateurs/.venv/lib/python3.13/site-packages/xwakes/wit/landau_damping.py:22: SyntaxWarning: invalid escape sequence '\s'
  :param b_direct: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in
/home/runner/work/Journees_Accelerateurs/Journees_Accelerateurs/.venv/lib/python3.13/site-packages/xwakes/wit/landau_damping.py:76: SyntaxWarning: invalid escape sequence '\s'
  :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in
/home/runner/work/Journees_Accelerateurs/Journees_Accelerateurs/.venv/lib/python3.13/site-packages/xwakes/wit/landau_damping.py:140: SyntaxWarning: invalid escape sequence '\s'
  :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in

env = xt.Environment.from_json("pimms.json")
env.vars.load_json("pimms_strengths.json")
ring = env.ring

# Modèle de dipôle adapté aux petits anneaux
ring.configure_bend_model(edge="full", core="adaptive", num_multipole_kicks=5)
twiss4d = ring.twiss(method="4d")

# Installation d'une limite d'ouverture horizontale pour obtenir des pertes
ring.append_element(name="aperture", element=xt.LimitRect(min_x=-0.05, max_x=0.05));

# Alimentation de la cavité pour 1 bucket RF
env.vars["vrf"] = 10e3  # 10kV
env.vars["frf"] = 1 / twiss4d.T_rev0  # correspond à h=1

# Twiss fonctionne maintenant en six dimensions
twiss0 = ring.twiss()
twiss0.qs  # nombre d'onde synchrotron

np.float64(0.0013332617169061756)

# Génération de particules sur l'axe longitudinal
p = ring.build_particles(delta=np.linspace(-8e-3, 8e-3, 81))
ring.track(p, num_turns=1500, turn_by_turn_monitor=True, with_progress=75)
mon = ring.record_last_track

# On peut visualiser notre bucket RF
plt.figure()
plt.plot(mon.zeta.T, mon.delta.T * 1e3, c="C0", lw=1)
plt.xlim(-40, 40)
plt.xlabel(r"$\zeta$ [m]")
plt.ylabel(r"$\delta$ [10$^{-3}$]")
plt.title("Espace des phases longitudinal")
plt.show()

# Champ de sillage représentant l'impédance de la cavité RF
wakefield = xw.WakeResonator(kind="dipolar_x", r=100e6, f_r=1.3e6, q=1.0)  # Shunt impedance
wakefield.configure_for_tracking(zeta_range=(-20, 20), num_slices=20)
ring.append("wf", wakefield)

%%capture
ring.build_tracker()
bunch = xp.generate_matched_gaussian_bunch(
    line=ring,
    num_particles=700,
    total_intensity_particles=1e11,
    nemitt_x=2e-6,
    nemitt_y=2e-6,
    sigma_z=10.0,
)
bunch.circumference = twiss0.circumference  # requis par xwakes

# Application d'un décallage horizontal initial de 1 mm
bunch.x += 1e-3
bunch0 = bunch.copy()

# Definition de quantités à enregistrer pendant le tracking
def compute_mean_x(line: xt.Line, particles: xt.Particles) -> float:
    """Return the bunch's average x position."""
    nplike = particles._context.nplike_lib
    particles.hide_lost_particles()
    x_average = nplike.mean(particles.x)
    particles.unhide_lost_particles()
    return x_average


def compute_intensity(line: xt.Line, particles: xt.Particles) -> float:
    """Return the bunch's intensity, in charges."""
    nplike = particles._context.nplike_lib
    particles.hide_lost_particles()
    intensity = particles.q0 * nplike.sum(particles.weight)
    particles.unhide_lost_particles()
    return intensity


# Création d'un logger pour enregistrer ces quantités à chaque tour
track_log = xt.Log(mean_x=compute_mean_x, intensity=compute_intensity)

# Track!
ring.track(bunch, log=track_log, num_turns=850, with_progress=5)

Compiling ContextCpu kernels...

Done compiling ContextCpu kernels.
Compiling ContextCpu kernels...

Done compiling ContextCpu kernels.

# Visualisation des quantités enregistrées
mean_x = np.array(ring.log_last_track["mean_x"])
intensity = np.array(ring.log_last_track["intensity"])

plt.figure(figsize=(7, 5))
ax1 = plt.subplot(2, 1, 1)
ax1.plot(mean_x, label=rf"$Q_x^' = {twiss0.dqx:.2f}, Q_y^' = {twiss0.dqy:.2f}$")
ax1.set_ylabel(r"$x_{\mathrm{centroid}}$ [m]")

ax2 = plt.subplot(2, 1, 2, sharex=ax1)
ax2.plot(intensity * 1e-11, label=rf"$Q_x^' = {twiss0.dqx:.2f}, Q_y^' = {twiss0.dqy:.2f}$")
ax2.set_ylim(bottom=0)
ax2.set_ylabel("Intensité [$10^{11}$ppb]")
ax2.set_xlabel("Tour")

plt.gcf().align_ylabels((ax1, ax2))
ax1.legend(ncols=2, loc="lower center", bbox_to_anchor=(0.5, 1))
plt.subplots_adjust(left=0.2, hspace=0.3, top=0.9)
plt.show()

%%capture
opt = ring.match(
    solve=False,
    method="4d",
    vary=xt.VaryList(["ksf", "ksd"], step=1e-3),
    targets=xt.TargetSet(dqx=-4.0, dqy=-4.0, tol=1e-3, tag="chrom"),
)
opt.solve()
twiss1 = ring.twiss()

ring.track(bunch0, log=track_log, num_turns=850, with_progress=5)

# Visualisation des quantités enregistrées
mean_x_corr = np.array(ring.log_last_track["mean_x"])
intensity_corr = np.array(ring.log_last_track["intensity"])

plt.figure(figsize=(7, 5))
ax1 = plt.subplot(2, 1, 1)
ax1.plot(mean_x, label=rf"$Q_x^' = {twiss0.dqx:.2f}, Q_y^' = {twiss0.dqy:.2f}$")
ax1.plot(mean_x_corr, label=rf"$Q_x^' = {twiss1.dqx:.2f}, Q_y^' = {twiss1.dqy:.2f}$")
ax1.set_ylabel(r"$x_{\mathrm{centroid}}$ [m]")

ax2 = plt.subplot(2, 1, 2, sharex=ax1)
ax2.plot(intensity * 1e-11, label=rf"$Q_x^' = {twiss0.dqx:.2f}, Q_y^' = {twiss0.dqy:.2f}$")
ax2.plot(intensity_corr * 1e-11, label=rf"$Q_x^' = {twiss1.dqx:.2f}, Q_y^' = {twiss1.dqy:.2f}$")
ax2.set_ylim(bottom=0)
ax2.set_ylabel("Intensité [$10^{11}$ppb]")
ax2.set_xlabel("Tour")

plt.subplots_adjust(left=0.2, hspace=0.3, top=0.9)
ax1.legend(ncols=2, loc="lower center", bbox_to_anchor=(0.5, 1))
plt.show()

Tracking de particules et mitigation d'instabilité dûe à l'impédance de la cavité RF¶

Chargement de la maille¶

Activation de la cavité RF¶

Espace des phases longitudinal¶

Installation d'un champ de sillage transverse¶

Génération d'un paquet adapté au bucket RF et à l’optique¶

Tracking des particules pour 1000 tours¶

Ajustement de la chromaticité via les sextupôles¶

Tracking à nouveau avec la nouvelle chromaticité¶