Skip to content

Commit

Permalink
Merge pull request #120 from fsoubelet/ibs_gpu_fix
Browse files Browse the repository at this point in the history
IBS CUDA GPUs
  • Loading branch information
giadarol authored Jun 13, 2024
2 parents 3206ff1 + eda71bd commit 92b9235
Show file tree
Hide file tree
Showing 9 changed files with 186 additions and 113 deletions.
12 changes: 10 additions & 2 deletions examples/005_ibs/000_growth_rates_from_parameters_no_vdisp.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,6 @@
# ########################################### #
import xtrack as xt

from xfields.ibs import get_intrabeam_scattering_growth_rates

##########################
# Load xt.Line from file #
##########################
Expand Down Expand Up @@ -63,6 +61,11 @@
print(f"Nagaitsev: {nag_growth_rates}")
print(f"Bjorken-Mtingwa: {bm_growth_rates}")

# Computed from normalized emittances:
# ------------------------------------
# Nagaitsev: IBSGrowthRates(Tx=0.00104, Ty=0.011, Tz=0.00803)
# Bjorken-Mtingwa: IBSGrowthRates(Tx=0.00111, Ty=0.0109, Tz=0.008)

#####################
# Define parameters #
#####################
Expand Down Expand Up @@ -108,3 +111,8 @@
print("------------------------------------------------------")
print(f"Nagaitsev: {nag_growth_rates2}")
print(f"Bjorken-Mtingwa: {bm_growth_rates2}")

# Computed from geometric emittances (rough equivalent):
# ------------------------------------------------------
# Nagaitsev: IBSGrowthRates(Tx=0.00103, Ty=0.0111, Tz=0.00803)
# Bjorken-Mtingwa: IBSGrowthRates(Tx=0.00111, Ty=0.0109, Tz=0.008)
50 changes: 4 additions & 46 deletions examples/005_ibs/001_growth_rates_from_parameters_with_vdisp.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
# ########################################### #
import json

import numpy as np
import xtrack as xt

##########################
Expand Down Expand Up @@ -60,7 +59,6 @@
bunched=True,
)


##########################################################
# Compare: we expect Nagaitsev to be wrong in horizontal #
##########################################################
Expand All @@ -71,47 +69,7 @@
print(f"Nagaitsev: {nag_growth_rates}")
print(f"Bjorken-Mtingwa: {bm_growth_rates}")

#####################
# Define parameters #
#####################

gemitt_x: float = 2.598e-10
gemitt_y: float = 2.598e-10

###################################
# Get growth rates with Nagaitsev #
###################################

nag_growth_rates2 = tw.get_ibs_growth_rates(
formalism="nagaitsev",
total_beam_intensity=bunch_intensity,
gemitt_x=gemitt_x,
gemitt_y=gemitt_y,
sigma_delta=sigma_delta,
bunch_length=bunch_length,
bunched=True,
)

#########################################
# Get growth rates with Bjorken-Mtingwa #
#########################################

bm_growth_rates2 = tw.get_ibs_growth_rates(
formalism="bjorken-mtingwa", # also accepts "b&m"
total_beam_intensity=bunch_intensity,
gemitt_x=gemitt_x,
gemitt_y=gemitt_y,
sigma_delta=sigma_delta,
bunch_length=bunch_length,
bunched=True,
)

##########################################################
# Compare: we expect Nagaitsev to be wrong in horizontal #
##########################################################

print()
print("Computed from geometric emittances (rough equivalent):")
print("------------------------------------------------------")
print(f"Nagaitsev: {nag_growth_rates2}")
print(f"Bjorken-Mtingwa: {bm_growth_rates2}")
# Computed from normalized emittances:
# ------------------------------------
# Nagaitsev: IBSGrowthRates(Tx=6.24e-05, Ty=-2.27e-09, Tz=0.00031)
# Bjorken-Mtingwa: IBSGrowthRates(Tx=6.21e-05, Ty=1.1e-06, Tz=0.00031)
10 changes: 10 additions & 0 deletions examples/005_ibs/002_growth_rates_from_particles.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,11 @@
print(f"Nagaitsev: {nag_growth_rates}")
print(f"Bjorken-Mtingwa: {bm_growth_rates}")

# Computed from particles object:
# -------------------------------
# Nagaitsev: IBSGrowthRates(Tx=1.54e-06, Ty=-1.46e-07, Tz=1.65e-06)
# Bjorken-Mtingwa: IBSGrowthRates(Tx=1.54e-06, Ty=-1.48e-07, Tz=1.65e-06)

###################################
# Get growth rates with Nagaitsev #
###################################
Expand Down Expand Up @@ -106,3 +111,8 @@
print("-------------------------------------------------------")
print(f"Nagaitsev: {nag_growth_rates}")
print(f"Bjorken-Mtingwa: {bm_growth_rates}")

# Computed from normalized emittances (rough equivalent):
# -------------------------------------------------------
# Nagaitsev: IBSGrowthRates(Tx=1.54e-06, Ty=-1.46e-07, Tz=1.65e-06)
# Bjorken-Mtingwa: IBSGrowthRates(Tx=1.54e-06, Ty=-1.48e-07, Tz=1.65e-06)
40 changes: 20 additions & 20 deletions examples/005_ibs/003_tracking_with_kicks.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,26 @@
fname_line_particles = "../../../xtrack/test_data/sps_w_spacecharge/"\
"line_no_spacecharge_and_particle.json"
line: xt.Line = xt.Line.from_json(fname_line_particles)
line.build_tracker()
line.build_tracker(_context=context)

#######################################
# Create and Install IBS Kick Element #
#######################################

# For the analytical kick formalism: kicks are computed based
# on the analytical growth rates (so it needs a formalism)
# ibs_kick = xf.IBSAnalyticalKick(formalism="nagaitsev", num_slices=50)

# For the kinetic formalism: kicks are computed based on the
# friction and diffusion terms of the kinetic theory of gases
ibs_kick = xf.IBSKineticKick(num_slices=50)

# By default the element is off until configuration. Let's install
# the kick at the end of the line and configure it. This internally
# provides the necessary information to the element
line.configure_intrabeam_scattering(
element=ibs_kick, name="ibskick", index=-1, update_every=50
)

############################################
# Define parameters and Generate Particles #
Expand All @@ -40,25 +59,6 @@
_context=context,
)

#######################################
# Create and Install IBS Kick Element #
#######################################

# For the analytical kick formalism: kicks are computed based
# on the analytical growth rates (so it needs a formalism)
# ibs_kick = xf.IBSAnalyticalKick(formalism="nagaitsev", num_slices=50)

# For the kinetic formalism: kicks are computed based on the
# friction and diffusion terms of the kinetic theory of gases
ibs_kick = xf.IBSKineticKick(num_slices=50)

# By default the element is off until configuration. Let's install
# the kick at the end of the line and configure it. This internally
# provides the necessary information to the element
line.configure_intrabeam_scattering(
element=ibs_kick, name="ibskick", index=-1, update_every=50
)

##############################################
# Track now applies an IBS kick at each turn #
##############################################
Expand Down
94 changes: 94 additions & 0 deletions examples/005_ibs/004_tracking_with_kicks_tbt_emittances.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
# copyright ################################# #
# This file is part of the Xfields Package. #
# Copyright (c) CERN, 2021. #
# ########################################### #
import xfields as xf
import xobjects as xo
import xpart as xp
import xtrack as xt
from xfields.ibs._formulary import _gemitt_x, _gemitt_y, _sigma_delta, _bunch_length
import numpy as np
import matplotlib.pyplot as plt

# context = xo.ContextCupy()
context = xo.ContextCpu(omp_num_threads="auto")

##########################
# Load xt.Line from file #
##########################

fname_line_particles = "../../../xtrack/test_data/clic_dr/line.json"
line: xt.Line = xt.Line.from_json(fname_line_particles)
line.build_tracker(_context=context)
cavities = [element for element in line.elements if isinstance(element, xt.Cavity)]
for cavity in cavities:
cavity.lag = 180
tw = line.twiss(method="4d")

#######################################
# Create and Install IBS Kick Element #
#######################################

# ibs_kick = xf.IBSKineticKick(num_slices=50)
ibs_kick = xf.IBSAnalyticalKick(formalism="nagaitsev", num_slices=50)
line.configure_intrabeam_scattering(
element=ibs_kick, name="ibskick", index=-1, update_every=50
)

################################
# Generate Particles and Track #
################################

nturns: int = 2000
epsx, epsy, sigd, bl = [], [], [], []

particles = xp.generate_matched_gaussian_bunch(
num_particles=2500,
total_intensity_particles=int(4.5e9),
nemitt_x=5.66e-7,
nemitt_y=3.7e-9,
sigma_z=1.58e-3,
line=line,
_context=context,
)

for turn in range(nturns):
line.track(particles, num_turns=1)
epsx.append(_gemitt_x(particles, tw.betx[0], tw.dx[0]))
epsy.append(_gemitt_y(particles, tw.bety[0], tw.dy[0]))
sigd.append(_sigma_delta(particles))
bl.append(_bunch_length(particles))

#############################
# Plot turn-by-turn results #
#############################

turns = np.arange(nturns) + 1 # start from 1
epsx = np.array(epsx)
epsy = np.array(epsy)
sigd = np.array(sigd)
bl = np.array(bl)

fig, axx = plt.subplots()
axx.plot(turns, 1e10 * epsx, c="C0", label=r"$\varepsilon_x$")
axx.set_xlabel("Turn")
axx.set_ylabel(r"$\varepsilon_x$ [$10^{-10}$m]")
axy = axx.twinx()
axy.plot(turns, 1e13 * epsy, c="C1", label=r"$\varepsilon_y$")
axy.yaxis.set_label_position("right")
axy.set_ylabel(r"$\varepsilon_y$ [$10^{-13}$m]")
fig.legend(loc="upper center", bbox_to_anchor=(0.5, 0.95), ncols=2)
fig.tight_layout()
fig.show()

fig2, axd = plt.subplots()
axd.plot(turns, 1e3 * sigd, c="C2", label=r"$\sigma_{\delta}$")
axd.set_xlabel("Turn")
axd.set_ylabel(r"$\sigma_{\delta}$ [$10^{-3}$]")
axb = axd.twinx()
axb.plot(turns, 1e3 * bl, c="C3", label=r"$\sigma_z$")
axb.yaxis.set_label_position("right")
axb.set_ylabel(r"$\sigma_z$ [$10^{-3}$m]")
fig2.legend(loc="upper center", bbox_to_anchor=(0.5, 0.95), ncols=2)
fig2.tight_layout()
fig2.show()
Loading

0 comments on commit 92b9235

Please sign in to comment.