logitscale: rewrite of LogitLocator

jb-leger · jb-leger · commit 5847a840676a · 2019-07-23T11:21:05.000+02:00
diff --git a/lib/matplotlib/ticker.py b/lib/matplotlib/ticker.py
@@ -2502,64 +2502,116 @@ def view_limits(self, vmin, vmax):
         return result
 
 
-class LogitLocator(Locator):
+class LogitLocator(MaxNLocator):
     """
     Determine the tick locations for logit axes
     """
 
-    def __init__(self, minor=False):
-        """Place ticks on the logit locations."""
-        self.minor = minor
+    def __init__(self, minor=False, *, nbins="auto"):
+        """
+        Place ticks on the logit locations
+
+        Parameters
+        ----------
+        nbins : int or 'auto', optional
+            Number of ticks. Only used if minor is False.
+        minor : bool, default: False
+            Indicate if this locator is for minor ticks or not.
+        """
+
+        self._minor = minor
+        MaxNLocator.__init__(self, nbins=nbins, steps=[1, 2, 5, 10])
 
-    def set_params(self, minor=None):
+    def set_params(self, minor=None, **kwargs):
         """Set parameters within this locator."""
         if minor is not None:
-            self.minor = minor
+            self._minor = minor
+        MaxNLocator.set_params(self, **kwargs)
 
-    def __call__(self):
-        """Return the locations of the ticks."""
-        vmin, vmax = self.axis.get_view_interval()
-        return self.tick_values(vmin, vmax)
+    @property
+    def minor(self):
+        return self._minor
+
+    @minor.setter
+    def minor(self, value):
+        self.set_params(minor=value)
 
     def tick_values(self, vmin, vmax):
         # dummy axis has no axes attribute
-        if hasattr(self.axis, 'axes') and self.axis.axes.name == 'polar':
-            raise NotImplementedError('Polar axis cannot be logit scaled yet')
+        if hasattr(self.axis, "axes") and self.axis.axes.name == "polar":
+            raise NotImplementedError("Polar axis cannot be logit scaled yet")
 
-        vmin, vmax = self.nonsingular(vmin, vmax)
-        vmin = np.log10(vmin / (1 - vmin))
-        vmax = np.log10(vmax / (1 - vmax))
-
-        decade_min = np.floor(vmin)
-        decade_max = np.ceil(vmax)
-
-        # major ticks
-        if not self.minor:
-            ticklocs = []
-            if decade_min <= -1:
-                expo = np.arange(decade_min, min(0, decade_max + 1))
-                ticklocs.extend(10**expo)
-            if decade_min <= 0 <= decade_max:
-                ticklocs.append(0.5)
-            if decade_max >= 1:
-                expo = -np.arange(max(1, decade_min), decade_max + 1)
-                ticklocs.extend(1 - 10**expo)
-
-        # minor ticks
+        if self._nbins == "auto":
+            if self.axis is not None:
+                nbins = self.axis.get_tick_space()
+                if nbins < 2:
+                    nbins = 2
+            else:
+                nbins = 9
         else:
-            ticklocs = []
-            if decade_min <= -2:
-                expo = np.arange(decade_min, min(-1, decade_max))
-                newticks = np.outer(np.arange(2, 10), 10**expo).ravel()
-                ticklocs.extend(newticks)
-            if decade_min <= 0 <= decade_max:
-                ticklocs.extend([0.2, 0.3, 0.4, 0.6, 0.7, 0.8])
-            if decade_max >= 2:
-                expo = -np.arange(max(2, decade_min), decade_max + 1)
-                newticks = 1 - np.outer(np.arange(2, 10), 10**expo).ravel()
-                ticklocs.extend(newticks)
+            nbins = self._nbins
 
-        return self.raise_if_exceeds(np.array(ticklocs))
+        # We define ideal ticks with their index:
+        # linscale: ... 1e-3 1e-2 1e-1 1/2 1-1e-1 1-1e-2 1-1e-3 ...
+        # b-scale : ... -3   -2   -1   0   1      2      3      ...
+        def ideal_ticks(x):
+            return 10 ** x if x < 0 else 1 - (10 ** (-x)) if x > 0 else 1 / 2
+
+        vmin, vmax = self.nonsingular(vmin, vmax)
+        binf = int(
+            np.floor(np.log10(vmin))
+            if vmin < 0.5
+            else 0
+            if vmin < 0.9
+            else -np.ceil(np.log10(1 - vmin))
+        )
+        bsup = int(
+            np.ceil(np.log10(vmax))
+            if vmax <= 0.5
+            else 1
+            if vmax <= 0.9
+            else -np.floor(np.log10(1 - vmax))
+        )
+        numideal = bsup - binf - 1
+        if numideal >= 2:
+            # have 2 or more wanted ideal ticks, so use them as major ticks
+            if numideal > nbins:
+                # to many ideal ticks, subsampling ideals for major ticks, and
+                # take others for minor ticks
+                subsampling_factor = math.ceil(numideal / nbins)
+                if self._minor:
+                    ticklocs = [
+                        ideal_ticks(b)
+                        for b in range(binf, bsup + 1)
+                        if (b % subsampling_factor) != 0
+                    ]
+                else:
+                    ticklocs = [
+                        ideal_ticks(b)
+                        for b in range(binf, bsup + 1)
+                        if (b % subsampling_factor) == 0
+                    ]
+                return self.raise_if_exceeds(np.array(ticklocs))
+            if self._minor:
+                ticklocs = []
+                for b in range(binf, bsup):
+                    if b < -1:
+                        ticklocs.extend(np.arange(2, 10) * 10 ** b)
+                    elif b == -1:
+                        ticklocs.extend(np.arange(2, 5) / 10)
+                    elif b == 0:
+                        ticklocs.extend(np.arange(6, 9) / 10)
+                    else:
+                        ticklocs.extend(
+                            1 - np.arange(2, 10)[::-1] * 10 ** (-b - 1)
+                        )
+                return self.raise_if_exceeds(np.array(ticklocs))
+            ticklocs = [ideal_ticks(b) for b in range(binf, bsup + 1)]
+            return self.raise_if_exceeds(np.array(ticklocs))
+        # the scale is zoomed so same ticks as linear scale can be used
+        if self._minor:
+            return []
+        return MaxNLocator.tick_values(self, vmin, vmax)
 
     def nonsingular(self, vmin, vmax):
         initial_range = (1e-7, 1 - 1e-7)
