Handling edgecases of n_bin=1 and n_bin=2

mbtrack2/tracking/particles.py

@@ -481,16 +481,22 @@ class Bunch:
         bin_max = max(bin_max * 0.99, bin_max * 1.01)
 
         bins = np.linspace(bin_min, bin_max, n_bin)
+        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)
         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)
 
-        # Compute bin indices directly (O(1) per element)
-        sorted_index = ((self[dimension] - bin_min) / bin_width).astype(int)
-
-        # Clip indices to ensure they fall within the valid range
+        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)
 
-        # Compute profile using bincount
         profile = np.bincount(sorted_index, minlength=n_bin - 1)
         return bins, sorted_index, profile, center