Merge pull request #24290 from anntzer/bft

Cleanup Barbs._find_tails.

lib/matplotlib/quiver.py

@@ -890,9 +890,11 @@ class Barbs(mcollections.PolyCollection):
     From there :meth:`_make_barbs` is used to find the vertices of the
     polygon to represent the barb based on this information.
     """
+
     # This may be an abuse of polygons here to render what is essentially maybe
     # 1 triangle and a series of lines.  It works fine as far as I can tell
     # however.
+
     @_docstring.interpd
     def __init__(self, ax, *args,
                  pivot='tip', length=7, barbcolor=None, flagcolor=None,
@@ -952,35 +954,35 @@ def __init__(self, ax, *args,
 
     def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):
         """
-        Find how many of each of the tail pieces is necessary.  Flag
-        specifies the increment for a flag, barb for a full barb, and half for
-        half a barb. Mag should be the magnitude of a vector (i.e., >= 0).
-
-        This returns a tuple of:
+        Find how many of each of the tail pieces is necessary.
 
-            (*number of flags*, *number of barbs*, *half_flag*, *empty_flag*)
+        Parameters
+        ----------
+        mag : array
+            Vector magnitudes; must be non-negative (and an actual ndarray).
+        rounding : bool, default: True
+            Whether to round or to truncate to the nearest half-barb.
+        half, full, flag : float, defaults: 5, 10, 50
+            Increments for a half-barb, a barb, and a flag.
 
-        The bool *half_flag* indicates whether half of a barb is needed,
-        since there should only ever be one half on a given
-        barb. *empty_flag* flag is an array of flags to easily tell if
-        a barb is empty (too low to plot any barbs/flags.
+        Returns
+        -------
+        n_flags, n_barbs : int array
+            For each entry in *mag*, the number of flags and barbs.
+        half_flag : bool array
+            For each entry in *mag*, whether a half-barb is needed.
+        empty_flag : bool array
+            For each entry in *mag*, whether nothing is drawn.
         """
-
         # If rounding, round to the nearest multiple of half, the smallest
         # increment
         if rounding:
-            mag = half * (mag / half + 0.5).astype(int)
-
-        num_flags = np.floor(mag / flag).astype(int)
-        mag = mag % flag
-
-        num_barb = np.floor(mag / full).astype(int)
-        mag = mag % full
-
+            mag = half * np.around(mag / half)
+        n_flags, mag = divmod(mag, flag)
+        n_barb, mag = divmod(mag, full)
         half_flag = mag >= half
-        empty_flag = ~(half_flag | (num_flags > 0) | (num_barb > 0))
-
-        return num_flags, num_barb, half_flag, empty_flag
+        empty_flag = ~(half_flag | (n_flags > 0) | (n_barb > 0))
+        return n_flags.astype(int), n_barb.astype(int), half_flag, empty_flag
 
     def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
                     pivot, sizes, fill_empty, flip):
@@ -1152,9 +1154,8 @@ def set_UVC(self, U, V, C=None):
             _check_consistent_shapes(x, y, u, v, flip)
 
         magnitude = np.hypot(u, v)
-        flags, barbs, halves, empty = self._find_tails(magnitude,
-                                                       self.rounding,
-                                                       **self.barb_increments)
+        flags, barbs, halves, empty = self._find_tails(
+            magnitude, self.rounding, **self.barb_increments)
 
         # Get the vertices for each of the barbs