ticker.EngFormatter: base upon ScalarFormatter

doronbehar · doronbehar · commit 6a2460244568 · 2024-07-04T18:21:43.000+03:00
Allows us to use many order of magnitude and offset related routines from ScalarFormatter, and removes a bit usetex related duplicated code. Solves #28463.
diff --git a/doc/users/next_whats_new/EngFormatter-offset.rst b/doc/users/next_whats_new/EngFormatter-offset.rst
@@ -0,0 +1,12 @@
+ticker.EngFormatter now computes offset by default
+--------------------------------------------------
+
+:class:`matplotlib.ticker.EngFormatter` has gained the ability to show an offset text near the
+axis. With shared logic with :class:`matplotlib.ticker.ScalarFormatter`, it is capable of
+deciding whether the data qualifies having an offset and show it with an
+appropriate SI quantity prefix, and with the supplied ``unit``.
+
+To enable this new behavior, simply pass ``useOffset=True`` when you instantiate
+:class:`matplotlib.ticker.EngFormatter`. Example is available here_.
+
+.. _here: ../../gallery/ticks/engformatter_offset.html
diff --git a/galleries/examples/ticks/engformatter_offset.py b/galleries/examples/ticks/engformatter_offset.py
@@ -0,0 +1,31 @@
+"""
+===================================================
+SI prefixed offsets and natural order of magnitudes
+===================================================
+
+:class:`matplotlib.ticker.EngFormatter` is capable of computing a natural
+offset for your axis data, and presenting it with a standard SI prefix
+automatically calculated.
+
+Below is an examples of such a plot:
+
+"""
+
+import matplotlib.pyplot as plt
+import matplotlib.ticker as mticker
+import numpy as np
+
+np.random.seed(19680801)
+
+UNIT = "Hz"
+
+fig, ax = plt.subplots()
+ax.yaxis.set_major_formatter(mticker.EngFormatter(
+    useOffset=True,
+    unit=UNIT
+))
+size = 100
+measurement = np.full(size, 1)*1e9
+noise = np.random.uniform(low=-2e3, high=2e3, size=(size))
+ax.plot(measurement + noise)
+plt.show()
diff --git a/lib/matplotlib/tests/test_ticker.py b/lib/matplotlib/tests/test_ticker.py
@@ -1594,6 +1594,78 @@ def test_engformatter_usetex_useMathText():
         assert x_tick_label_text == ['$0$', '$500$', '$1$ k']
 
 
+@pytest.mark.parametrize(
+    'oom_center, oom_noise, oom_center_desired, oom_noise_desired', [
+        (11, 1, 9, 0),
+        (13, 7, 12, 6),
+        (1, -2, 0, -3),
+        (3, -2, 3, -3),
+        (5, -3, 3, -3),
+        (2, -3, 0, -3),
+        # The following sets of parameters demonstrates that when oom_center-1
+        # and oom_noise-2 equal a standard 3*N oom, we get that
+        # oom_noise_desired < oom_noise
+        (10, 2, 9, 3),
+        (1, -7, 0, -6),
+        (2, -4, 0, -3),
+        (1, -4, 0, -3),
+        # Tests where oom_center <= oom_noise
+        (4,  4, 0, 3),
+        (1,  4, 0, 3),
+        (1,  3, 0, 3),
+        (1,  2, 0, 0),
+        (1,  1, 0, 0),
+    ]
+)
+def test_engformatter_offset_oom(
+    oom_center,
+    oom_noise,
+    oom_center_desired,
+    oom_noise_desired
+):
+    UNIT = "eV"
+    # Doesn't really matter here, but should be of order of magnitude ~= 1
+    r = range(-5, 7)
+    fig, ax = plt.subplots()
+    # Use some random ugly number
+    data_offset = 2.7149*10**oom_center
+    ydata = data_offset + np.array(r, dtype=float)*10**oom_noise
+    ax.plot(ydata)
+    formatter = mticker.EngFormatter(useOffset=True, unit=UNIT)
+    # So that offset strings will always have the same size
+    formatter.ENG_PREFIXES[0] = "_"
+    ax.yaxis.set_major_formatter(formatter)
+    fig.canvas.draw()
+    offsetGot = formatter.get_offset()
+    ticksGot = [labl.get_text() for labl in ax.get_yticklabels()]
+    # Predicting whether offset should be 0 or not is essentially testing
+    # ScalarFormatter._compute_offset . This function is pretty complex and it
+    # would be nice to test it, but this is out of scope for this test which
+    # only makes sure that offset text and the ticks gets the correct unit
+    # prefixes and the ticks.
+    if formatter.offset:
+        prefix_noise_got = offsetGot[2]
+        prefix_noise_desired = formatter.ENG_PREFIXES[oom_noise_desired]
+        prefix_center_got = offsetGot[-1-len(UNIT)]
+        prefix_center_desired = formatter.ENG_PREFIXES[oom_center_desired]
+        assert prefix_noise_desired == prefix_noise_got
+        assert prefix_center_desired == prefix_center_got
+        # Make sure the ticks didn't get the UNIT
+        for tick in ticksGot:
+            assert UNIT not in tick
+    else:
+        assert oom_center_desired == 0
+        assert offsetGot == ""
+        # Make sure the ticks contain now the prefixes
+        for tick in ticksGot:
+            # 0 is zero on all orders of magnitudes, no
+            if tick[0] == "0":
+                prefixIdx = 0
+            else:
+                prefixIdx = oom_noise_desired
+            assert tick.endswith(formatter.ENG_PREFIXES[prefixIdx] + UNIT)
+
+
 class TestPercentFormatter:
     percent_data = [
         # Check explicitly set decimals over different intervals and values
diff --git a/lib/matplotlib/ticker.py b/lib/matplotlib/ticker.py
@@ -1340,7 +1340,7 @@ def format_data_short(self, value):
         return f"1-{1 - value:e}"
 
 
-class EngFormatter(Formatter):
+class EngFormatter(ScalarFormatter):
     """
     Format axis values using engineering prefixes to represent powers
     of 1000, plus a specified unit, e.g., 10 MHz instead of 1e7.
@@ -1372,7 +1372,7 @@ class EngFormatter(Formatter):
     }
 
     def __init__(self, unit="", places=None, sep=" ", *, usetex=None,
-                 useMathText=None):
+                 useMathText=None, useOffset=False):
         r"""
         Parameters
         ----------
@@ -1406,76 +1406,123 @@ def __init__(self, unit="", places=None, sep=" ", *, usetex=None,
         useMathText : bool, default: :rc:`axes.formatter.use_mathtext`
             To enable/disable the use mathtext for rendering the numbers in
             the formatter.
+        useOffset : bool or float, default: False
+            Whether to use offset notation with :math:`10^{3*N}` based prefixes.
+            This features allows showing an offset with standard SI order of
+            magnitude prefix near the axis. Offset is computed similarly to
+            how `ScalarFormatter` computes it internally, but here you are
+            guaranteed to get an offset which will make the tick labels exceed
+            3 digits. See also `.set_useOffset`.
+
+        .. versionadded:: 3.10
         """
         self.unit = unit
         self.places = places
         self.sep = sep
-        self.set_usetex(usetex)
-        self.set_useMathText(useMathText)
-
-    def get_usetex(self):
-        return self._usetex
-
-    def set_usetex(self, val):
-        if val is None:
-            self._usetex = mpl.rcParams['text.usetex']
-        else:
-            self._usetex = val
-
-    usetex = property(fget=get_usetex, fset=set_usetex)
-
-    def get_useMathText(self):
-        return self._useMathText
+        super().__init__(
+            useOffset=useOffset,
+            useMathText=useMathText,
+            useLocale=False,
+            usetex=usetex,
+        )
 
-    def set_useMathText(self, val):
-        if val is None:
-            self._useMathText = mpl.rcParams['axes.formatter.use_mathtext']
+    def __call__(self, x, pos=None):
+        """
+        Return the format for tick value *x* at position *pos*. If there is no
+        currently offset in the data, it returns the best engineering formatting
+        that fits the given argument, independently.
+        """
+        if len(self.locs) == 0 or self.offset == 0:
+            return self.fix_minus(self.format_data(x))
         else:
-            self._useMathText = val
+            xp = (x - self.offset) / (10. ** self.orderOfMagnitude)
+            if abs(xp) < 1e-8:
+                xp = 0
+            return self._format_maybe_minus_and_locale(self.format, xp)
 
-    useMathText = property(fget=get_useMathText, fset=set_useMathText)
+    def set_locs(self, locs):
+        # docstring inherited
+        self.locs = locs
+        if len(self.locs) > 0:
+            vmin, vmax = sorted(self.axis.get_view_interval())
+            if self._useOffset:
+                self._compute_offset()
+                if self.offset != 0:
+                    # We don't want to use the offset computed by
+                    # self._compute_offset because it rounds the offset unaware
+                    # of our engineering prefixes preference, and this can
+                    # cause ticks with 4+ digits to appear. These ticks are
+                    # slightly less readable, so if offset is justified
+                    # (decided by self._compute_offset) we set it to better
+                    # value:
+                    self.offset = round((vmin + vmax)/2, 3)
+            # Use log1000 to use engineers' oom standards
+            self.orderOfMagnitude = math.floor(math.log(vmax - vmin, 1000))*3
+            self._set_format()
 
-    def __call__(self, x, pos=None):
-        s = f"{self.format_eng(x)}{self.unit}"
-        # Remove the trailing separator when there is neither prefix nor unit
-        if self.sep and s.endswith(self.sep):
-            s = s[:-len(self.sep)]
-        return self.fix_minus(s)
+    # Simplify a bit ScalarFormatter.get_offset: We always want to use
+    # self.format_data. Also we want to return a non-empty string only if there
+    # is an offset, no matter what is self.orderOfMagnitude. if there is an
+    # offset OTH, self.orderOfMagnitude is consulted. This behavior is verified
+    # in `test_ticker.py`.
+    def get_offset(self):
+        # docstring inherited
+        if len(self.locs) == 0:
+            return ''
+        if self.offset:
+            offsetStr = ''
+            if self.offset:
+                offsetStr = self.format_data(self.offset)
+                if self.offset > 0:
+                    offsetStr = '+' + offsetStr
+            sciNotStr = self.format_data(10 ** self.orderOfMagnitude)
+            if self._useMathText or self._usetex:
+                if sciNotStr != '':
+                    sciNotStr = r'\times%s' % sciNotStr
+                s = fr'${sciNotStr}{offsetStr}$'
+            else:
+                s = ''.join((sciNotStr, offsetStr))
+            return self.fix_minus(s)
+        return ''
 
     def format_eng(self, num):
+        """Alias to EngFormatter.format_data"""
+        return self.format_data(num)
+
+    def format_data(self, value):
         """
         Format a number in engineering notation, appending a letter
         representing the power of 1000 of the original number.
         Some examples:
 
-        >>> format_eng(0)        # for self.places = 0
+        >>> format_data(0)        # for self.places = 0
         '0'
 
-        >>> format_eng(1000000)  # for self.places = 1
+        >>> format_data(1000000)  # for self.places = 1
         '1.0 M'
 
-        >>> format_eng(-1e-6)  # for self.places = 2
+        >>> format_data(-1e-6)  # for self.places = 2
         '-1.00 \N{MICRO SIGN}'
         """
         sign = 1
         fmt = "g" if self.places is None else f".{self.places:d}f"
 
-        if num < 0:
+        if value < 0:
             sign = -1
-            num = -num
+            value = -value
 
-        if num != 0:
-            pow10 = int(math.floor(math.log10(num) / 3) * 3)
+        if value != 0:
+            pow10 = int(math.floor(math.log10(value) / 3) * 3)
         else:
             pow10 = 0
-            # Force num to zero, to avoid inconsistencies like
+            # Force value to zero, to avoid inconsistencies like
             # format_eng(-0) = "0" and format_eng(0.0) = "0"
             # but format_eng(-0.0) = "-0.0"
-            num = 0.0
+            value = 0.0
 
         pow10 = np.clip(pow10, min(self.ENG_PREFIXES), max(self.ENG_PREFIXES))
 
-        mant = sign * num / (10.0 ** pow10)
+        mant = sign * value / (10.0 ** pow10)
         # Taking care of the cases like 999.9..., which may be rounded to 1000
         # instead of 1 k.  Beware of the corner case of values that are beyond
         # the range of SI prefixes (i.e. > 'Y').
@@ -1484,13 +1531,15 @@ def format_eng(self, num):
             mant /= 1000
             pow10 += 3
 
-        prefix = self.ENG_PREFIXES[int(pow10)]
+        unitPrefix = self.ENG_PREFIXES[int(pow10)]
+        if self.unit or unitPrefix:
+            suffix = f"{self.sep}{unitPrefix}{self.unit}"
+        else:
+            suffix = ""
         if self._usetex or self._useMathText:
-            formatted = f"${mant:{fmt}}${self.sep}{prefix}"
+            return rf"${mant:{fmt}}${suffix}"
         else:
-            formatted = f"{mant:{fmt}}{self.sep}{prefix}"
-
-        return formatted
+            return rf"{mant:{fmt}}{suffix}"
 
 
 class PercentFormatter(Formatter):
diff --git a/lib/matplotlib/ticker.pyi b/lib/matplotlib/ticker.pyi
@@ -130,7 +130,7 @@ class LogitFormatter(Formatter):
     def set_minor_number(self, minor_number: int) -> None: ...
     def format_data_short(self, value: float) -> str: ...
 
-class EngFormatter(Formatter):
+class EngFormatter(ScalarFormatter):
     ENG_PREFIXES: dict[int, str]
     unit: str
     places: int | None
@@ -142,20 +142,9 @@ class EngFormatter(Formatter):
         sep: str = ...,
         *,
         usetex: bool | None = ...,
-        useMathText: bool | None = ...
+        useMathText: bool | None = ...,
+        useOffset: bool | float | None = ...,
     ) -> None: ...
-    def get_usetex(self) -> bool: ...
-    def set_usetex(self, val: bool | None) -> None: ...
-    @property
-    def usetex(self) -> bool: ...
-    @usetex.setter
-    def usetex(self, val: bool | None) -> None: ...
-    def get_useMathText(self) -> bool: ...
-    def set_useMathText(self, val: bool | None) -> None: ...
-    @property
-    def useMathText(self) -> bool: ...
-    @useMathText.setter
-    def useMathText(self, val: bool | None) -> None: ...
     def format_eng(self, num: float) -> str: ...
 
 class PercentFormatter(Formatter):
