diff --git a/mbtrack2/impedance/wakefield.py b/mbtrack2/impedance/wakefield.py index 5de4a8fffb9ecf70a7feb9f14a1a1347ba8d1e5a..147e4451778454170ff49de934ec890d60a65695 100644 --- a/mbtrack2/impedance/wakefield.py +++ b/mbtrack2/impedance/wakefield.py @@ -764,7 +764,7 @@ class WakeField: """ Return an array of the impedance component names for the element. """ - valid = ["Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad"] + valid = ["Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad", "Zxcst", "Zycst"] return np.array([comp for comp in dir(self) if comp in valid]) @property @@ -772,7 +772,7 @@ class WakeField: """ Return an array of the wake function component names for the element. """ - valid = ["Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad"] + valid = ["Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad", "Wxcst", "Wycst"] return np.array([comp for comp in dir(self) if comp in valid]) @property @@ -780,8 +780,8 @@ class WakeField: """ Return an array of the component names for the element. """ - valid = ["Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad", - "Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad"] + valid = ["Wlong", "Wxdip", "Wydip", "Wxquad", "Wyquad", "Wxcst", "Wycst", + "Zlong", "Zxdip", "Zydip", "Zxquad", "Zyquad", "Zxcst", "Zycst"] return np.array([comp for comp in dir(self) if comp in valid]) def drop(self, component): diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py index 14ac15167630cda3e042d5b623fe3e9b80a3532a..e1cfc04916451afcd4c7ad88ef1e927171cc1f76 100644 --- a/mbtrack2/tracking/wakepotential.py +++ b/mbtrack2/tracking/wakepotential.py @@ -273,7 +273,7 @@ class WakePotential(Element): profile0 = np.interp(tau0, self.tau, self.rho, 0, 0) - if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad": + if wake_type == "Wlong" or wake_type == "Wxquad" or wake_type == "Wyquad" or wake_type == "Wxcst" or wake_type == "Wycst": Wp = signal.convolve(profile0, W0*-1, mode='same')*dtau0 elif wake_type == "Wxdip": dipole0 = self.dipole_moment(bunch, "x", tau0) @@ -318,6 +318,10 @@ class WakePotential(Element): elif wake_type == "Wyquad": bunch["yp"] += (bunch["y"] * Wp_interp * bunch.charge / self.ring.E0) + elif wake_type == "Wxcst": + bunch["xp"] += Wp_interp * bunch.charge / self.ring.E0 + elif wake_type == "Wycst": + bunch["yp"] += Wp_interp * bunch.charge / self.ring.E0 def plot_last_wake(self, wake_type, plot_rho=True, plot_dipole=False, plot_wake_function=True): @@ -346,7 +350,9 @@ class WakePotential(Element): "Wxdip" : r"$W_{p,x}^{D} (V/pC)$", "Wydip" : r"$W_{p,y}^{D} (V/pC)$", "Wxquad" : r"$W_{p,x}^{Q} (V/pC/m)$", - "Wyquad" : r"$W_{p,y}^{Q} (V/pC/m)$"} + "Wyquad" : r"$W_{p,y}^{Q} (V/pC/m)$", + "Wxcst" : r"$W_{p,x}^{M}$ (V/pC)", + "Wycst" : r"$W_{p,y}^{M}$ (V/pC)",} Wp = getattr(self, wake_type) tau0 = getattr(self, "tau0_" + wake_type)