From 97432be9b2cbb4dae93559a4230ee6808ef97688 Mon Sep 17 00:00:00 2001
From: Alexis Gamelin <alexis.gamelin@synchrotron-soleil.fr>
Date: Tue, 23 Jan 2024 12:08:59 +0100
Subject: [PATCH] Modification from Vadim faster branch
Change the way Bunch particles coordinates are stored and accessed.
Speed-up TransverseMap.
Refactor WakePotential.
---
mbtrack2/tracking/element.py | 28 +++++++++++++++-------------
mbtrack2/tracking/particles.py | 15 +++++++++------
mbtrack2/tracking/wakepotential.py | 28 ++++++++++++----------------
3 files changed, 36 insertions(+), 35 deletions(-)
diff --git a/mbtrack2/tracking/element.py b/mbtrack2/tracking/element.py
index 7ccfa4d..7fab92b 100644
--- a/mbtrack2/tracking/element.py
+++ b/mbtrack2/tracking/element.py
@@ -201,27 +201,29 @@ class TransverseMap(Element):
self.adts_poly[1](Jx) + self.adts_poly[3](Jy))
# 6x6 matrix corresponding to (x, xp, delta, y, yp, delta)
- matrix = np.zeros((6, 6, len(bunch)))
+ matrix = np.zeros((6, 6, len(bunch)), dtype=np.float64)
# Horizontal
- matrix[0, 0, :] = np.cos(
- phase_advance_x) + self.alpha[0] * np.sin(phase_advance_x)
- matrix[0, 1, :] = self.beta[0] * np.sin(phase_advance_x)
+ c_x = np.cos(phase_advance_x)
+ s_x = np.sin(phase_advance_x)
+
+ matrix[0, 0, :] = c_x + self.alpha[0] * s_x
+ matrix[0, 1, :] = self.beta[0] * s_x
matrix[0, 2, :] = self.dispersion[0]
- matrix[1, 0, :] = -1 * self.gamma[0] * np.sin(phase_advance_x)
- matrix[1, 1, :] = np.cos(
- phase_advance_x) - self.alpha[0] * np.sin(phase_advance_x)
+ matrix[1, 0, :] = -1 * self.gamma[0] * s_x
+ matrix[1, 1, :] = c_x - self.alpha[0] * s_x
matrix[1, 2, :] = self.dispersion[1]
matrix[2, 2, :] = 1
# Vertical
- matrix[3, 3, :] = np.cos(
- phase_advance_y) + self.alpha[1] * np.sin(phase_advance_y)
- matrix[3, 4, :] = self.beta[1] * np.sin(phase_advance_y)
+ c_y = np.cos(phase_advance_y)
+ s_y = np.sin(phase_advance_y)
+
+ matrix[3, 3, :] = c_y + self.alpha[1] * s_y
+ matrix[3, 4, :] = self.beta[1] * s_y
matrix[3, 5, :] = self.dispersion[2]
- matrix[4, 3, :] = -1 * self.gamma[1] * np.sin(phase_advance_y)
- matrix[4, 4, :] = np.cos(
- phase_advance_y) - self.alpha[1] * np.sin(phase_advance_y)
+ matrix[4, 3, :] = -1 * self.gamma[1] * s_y
+ matrix[4, 4, :] = c_y - self.alpha[1] * s_y
matrix[4, 5, :] = self.dispersion[3]
matrix[5, 5, :] = 1
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index 1c3f20b..04f26b8 100644
--- a/mbtrack2/tracking/particles.py
+++ b/mbtrack2/tracking/particles.py
@@ -128,11 +128,14 @@ class Bunch:
current = 0
self._mp_number = int(mp_number)
- self.dtype = np.dtype([('x', float), ('xp', float), ('y', float),
- ('yp', float), ('tau', float),
- ('delta', float)])
-
- self.particles = np.zeros(self.mp_number, self.dtype)
+ self.particles = {
+ "x": np.empty(self.mp_number, dtype=np.float64),
+ "xp": np.empty(self.mp_number, dtype=np.float64),
+ "y": np.empty(self.mp_number, dtype=np.float64),
+ "yp": np.empty(self.mp_number, dtype=np.float64),
+ "tau": np.empty(self.mp_number, dtype=np.float64),
+ "delta": np.empty(self.mp_number, dtype=np.float64),
+ }
self.track_alive = track_alive
self.alive = np.ones((self.mp_number, ), dtype=bool)
self.current = current
@@ -144,7 +147,7 @@ class Bunch:
def __len__(self):
"""Return the number of alive particles"""
- return len(self[:])
+ return self.alive.sum()
def __getitem__(self, label):
"""Return the columns label for alive particles"""
diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py
index 5881a2d..eb6a105 100644
--- a/mbtrack2/tracking/wakepotential.py
+++ b/mbtrack2/tracking/wakepotential.py
@@ -110,18 +110,12 @@ class WakePotential(Element):
self.dtau = self.tau[1] - self.tau[0]
# Add N values before and after rho and tau
- if self.n_bin % 2 == 0:
- N = int(self.n_bin / 2)
- self.tau = np.arange(self.tau[0] - self.dtau * N,
- self.tau[-1] + self.dtau * N, self.dtau)
- self.rho = np.append(self.rho, np.zeros(N))
- self.rho = np.insert(self.rho, 0, np.zeros(N))
- else:
- N = int(np.floor(self.n_bin / 2))
- self.tau = np.arange(self.tau[0] - self.dtau * N,
- self.tau[-1] + self.dtau * (N+1), self.dtau)
- self.rho = np.append(self.rho, np.zeros(N))
- self.rho = np.insert(self.rho, 0, np.zeros(N + 1))
+ N = int(self.n_bin / 2)
+ self.tau = np.arange(self.tau[0] - self.dtau * N,
+ self.tau[-1] + self.dtau * N, self.dtau)
+ self.rho = np.pad(self.rho, (N, N + self.n_bin % 2),
+ mode="constant",
+ constant_values=(0, 0))
if len(self.tau) != len(self.rho):
self.tau = np.append(self.tau, self.tau[-1] + self.dtau)
@@ -148,7 +142,7 @@ class WakePotential(Element):
Dipole moment of the bunch.
"""
- dipole = np.zeros((self.n_bin - 1, ))
+ dipole = np.empty((self.n_bin - 1, ))
for i in range(self.n_bin - 1):
dipole[i] = bunch[plane][self.sorted_index == i].sum()
dipole = dipole / self.profile
@@ -219,14 +213,14 @@ class WakePotential(Element):
tau0 = np.arange(tau0[0] - dtau0*n_assym, tau0[-1] + dtau0,
dtau0)
n_to_add = len(tau0) - len(W0)
- W0 = np.insert(W0, 0, np.zeros(n_to_add))
+ W0 = np.concatenate((np.zeros(n_to_add), W0))
# add at tail
elif np.abs(tau0[0]) > np.abs(tau0[-1]):
tau0 = np.arange(tau0[0], tau0[-1] + dtau0 * (n_assym+1),
dtau0)
n_to_add = len(tau0) - len(W0)
- W0 = np.insert(W0, 0, np.zeros(n_to_add))
+ W0 = np.concatenate((np.zeros(n_to_add), W0))
# Check is the wf is shorter than rho then add zeros
if (tau0[0] > tau[0]) or (tau0[-1] < tau[-1]):
@@ -237,7 +231,9 @@ class WakePotential(Element):
dtau0)
W0 = np.insert(W0, 0, np.zeros(n))
n_to_add = len(tau0) - len(W0)
- W0 = np.insert(W0, len(W0), np.zeros(n_to_add))
+ W0 = np.pad(W0, (n, n_to_add),
+ mode="constant",
+ constant_values=(0, 0))
if save_data:
setattr(self, "tau0_" + wake_type, tau0)
--
GitLab