diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py index 47581c48bcd7a1cac3f4cb5347fbbb17b9ccb2a1..63e3e7d460116962f7fa016bd9e2b607c528ead4 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" or wake_type == "Wxcst" or wake_type == "Wycst": + if wake_type == "Wlong" 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) @@ -281,6 +281,8 @@ class WakePotential(Element): elif wake_type == "Wydip": dipole0 = self.dipole_moment(bunch, "y", tau0) Wp = signal.convolve(profile0 * dipole0, W0, mode='same') * dtau0 + elif wake_type == "Wxquad" or wake_type == "Wyquad": + Wp = signal.convolve(profile0, W0, mode='same') * dtau0 else: raise ValueError("This type of wake is not taken into account.") @@ -387,7 +389,7 @@ class WakePotential(Element): if plot_wake_function is True: W0 = getattr(self, "W0_" + wake_type) - W0_rescaled = W0 / max(W0) * max(np.abs(Wp)) + W0_rescaled = W0 / max(np.abs(W0)) * max(np.abs(Wp)) ax.plot(tau0 * 1e12, W0_rescaled * 1e-12, label=r"$W_{function}$ (a.u.)")