diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py index f8afc9aa64ee942fdf3b8a9d3335b2e82bc1f050..ceb5c4cd7a614238ff693758778a1288fb2892fb 100644 --- a/mbtrack2/tracking/particles.py +++ b/mbtrack2/tracking/particles.py @@ -480,24 +480,24 @@ class Bunch: bin_max = self[dimension].max() bin_max = max(bin_max * 0.99, bin_max * 1.01) - bins = np.linspace(bin_min, bin_max, n_bin) + bins = np.linspace(bin_min, bin_max, n_bin + 1) if n_bin == 1: return np.array([bin_min, bin_max]), np.zeros_like( self[dimension], dtype=int), np.array([len(self[dimension]) ]), np.array([(bin_min+bin_max) / 2]) - bin_width = (bin_max-bin_min) / (n_bin-1) + bin_width = (bin_max-bin_min) / n_bin center = (bins[1:] + bins[:-1]) / 2 sorted_index = np.floor( (self[dimension] - bin_min) / bin_width).astype(np.int64) - sorted_index = np.clip(sorted_index, 0, n_bin - 2) + sorted_index = np.clip(sorted_index, 0, n_bin - 1) if self.track_alive and return_full_length: sorted_index = np.floor((self.particles[dimension] - bin_min) / bin_width).astype(np.int64) - sorted_index = np.clip(sorted_index, 0, n_bin - 2) + sorted_index = np.clip(sorted_index, 0, n_bin - 1) - profile = np.bincount(sorted_index, minlength=n_bin - 1) + profile = np.bincount(sorted_index, minlength=n_bin) return bins, sorted_index, profile, center def plot_profile(self, dimension="tau", n_bin=75): diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py index 2533689c0aa17785822f72261bd3288ec646e4c3..fbacc0ec1e5b5d2b6b8ee4e026c1775b6bdaab49 100644 --- a/mbtrack2/tracking/wakepotential.py +++ b/mbtrack2/tracking/wakepotential.py @@ -91,10 +91,20 @@ class WakePotential(Element): np.seterr(invalid='ignore') def interp_regular_numpy(self, x_new, x_min, dx, y): - indices = ((x_new-x_min) / dx).astype(int) - indices = np.clip(indices, 0, len(y) - 2) - weights = (x_new - (x_min + indices*dx)) / dx - return y[indices] + weights * (y[indices + 1] - y[indices]) + x_max = x_min + (len(y) - 1) * dx # Compute the maximum valid x value + + result = np.zeros_like(x_new, dtype=y.dtype) + + valid_mask = (x_new >= x_min) & (x_new <= x_max) + + indices = ((x_new[valid_mask] - x_min) / dx).astype(int) + indices = np.clip(indices, 0, len(y) - 2) # Ensure within valid range + weights = (x_new[valid_mask] - (x_min + indices*dx)) / dx + + result[valid_mask] = y[indices] + weights * (y[indices + 1] - + y[indices]) + + return result def charge_density(self, bunch): """ @@ -159,7 +169,7 @@ class WakePotential(Element): """ dipole = np.bincount(self.sorted_index, weights=bunch[plane], - minlength=self.n_bin)[:-1] + minlength=self.n_bin + 1)[:-1] dipole = np.divide(dipole, self.profile, where=self.profile != 0) dipole = np.nan_to_num(dipole)