ENH: add arbitrary-scale functionality

jklymak · jklymak · commit 1e2d23e6d338 · 2018-11-17T18:54:08.000-08:00
API: add ability to set formatter and locators from scale init
diff --git a/.flake8 b/.flake8
@@ -197,6 +197,7 @@ per-file-ignores =
     examples/pyplots/whats_new_99_spines.py: E231, E402
     examples/recipes/placing_text_boxes.py: E501
     examples/scales/power_norm.py: E402
+    examples/scales/scales.py: E402
     examples/shapes_and_collections/artist_reference.py: E402
     examples/shapes_and_collections/collections.py: E402
     examples/shapes_and_collections/compound_path.py: E402
diff --git a/examples/scales/custom_scale.py b/examples/scales/custom_scale.py
@@ -5,6 +5,12 @@
 
 Create a custom scale, by implementing the scaling use for latitude data in a
 Mercator Projection.
+
+Unless you are making special use of the `~.Transform` class, you probably
+don't need to use this verbose method, and instead can use
+`~.matplotlib.scale.ArbitraryScale` and the ``'arbitrary'`` option of
+`~.matplotlib.axes.Axes.set_xscale` and `~.matplotlib.axes.Axes.set_yscale`.
+See the last example in :doc:`/gallery/scales/scales`.
 """
 
 import numpy as np
diff --git a/examples/scales/scales.py b/examples/scales/scales.py
@@ -4,10 +4,13 @@
 ======
 
 Illustrate the scale transformations applied to axes, e.g. log, symlog, logit.
+
+The last two examples are examples of using the ``'arbitrary'`` scale by
+supplying forward and inverse functions for the scale transformation.
 """
 import numpy as np
 import matplotlib.pyplot as plt
-from matplotlib.ticker import NullFormatter
+from matplotlib.ticker import NullFormatter, FixedLocator
 
 # Fixing random state for reproducibility
 np.random.seed(19680801)
@@ -19,8 +22,8 @@
 x = np.arange(len(y))
 
 # plot with various axes scales
-fig, axs = plt.subplots(2, 2, sharex=True)
-fig.subplots_adjust(left=0.08, right=0.98, wspace=0.3)
+fig, axs = plt.subplots(3, 2, figsize=(6, 8),
+                        constrained_layout=True)
 
 # linear
 ax = axs[0, 0]
@@ -54,4 +57,73 @@
 ax.yaxis.set_minor_formatter(NullFormatter())
 
 
+# Arbitrary x**(1/2)
+def forward(x):
+    return x**(1/2)
+
+
+def inverse(x):
+    return x**2
+
+
+ax = axs[2, 0]
+ax.plot(x, y)
+ax.set_yscale('arbitrary', functions=(forward, inverse))
+ax.set_title('arbitrary: $x^{1/2}$')
+ax.grid(True)
+ax.yaxis.set_major_locator(FixedLocator(np.arange(0, 1, 0.2)**2))
+ax.yaxis.set_major_locator(FixedLocator(np.arange(0, 1, 0.2)))
+
+
+# Arbitrary Mercator transform
+def forward(a):
+    a = np.deg2rad(a)
+    # mask values too close to +/- 90, which are ill-defined.
+    masked = np.ma.masked_where((a <= -np.pi/2 + 0.02) |
+                                (a >= np.pi/2 - 0.02), a)
+    if masked.mask.any():
+        return np.rad2deg(
+                np.ma.log(np.abs(np.ma.tan(masked) + 1.0 / np.ma.cos(masked))))
+    else:
+        return np.rad2deg(np.log(np.abs(np.tan(a) + 1.0 / np.cos(a))))
+
+
+def inverse(a):
+    a = np.deg2rad(a)
+    return np.rad2deg(np.arctan(np.sinh(a)))
+
+ax = axs[2, 1]
+
+t = np.arange(-170.0, 170.0, 0.1)
+s = t / 2.
+
+ax.plot(t, s, '-', lw=2)
+
+ax.set_yscale('arbitrary', functions=(forward, inverse))
+ax.set_title('arbitrary: Mercator')
+ax.grid(True)
+ax.set_xlim([-180, 180])
+ax.yaxis.set_minor_formatter(NullFormatter())
+ax.yaxis.set_major_locator(FixedLocator(np.arange(-90, 90, 30)))
+
 plt.show()
+
+#############################################################################
+#
+# ------------
+#
+# References
+# """"""""""
+#
+# The use of the following functions, methods, classes and modules is shown
+# in this example:
+
+import matplotlib
+matplotlib.axes.Axes.set_yscale
+matplotlib.axes.Axes.set_xscale
+matplotlib.axis.Axis.set_major_locator
+matplotlib.scale.LogitScale
+matplotlib.scale.LogScale
+matplotlib.scale.LinearScale
+matplotlib.scale.SymmetricalLogScale
+matplotlib.scale.ArbitraryScale
diff --git a/lib/matplotlib/scale.py b/lib/matplotlib/scale.py
@@ -82,6 +82,84 @@ def get_transform(self):
         return IdentityTransform()
 
 
+class ArbitraryTransform(Transform):
+     """
+     A simple transform that takes and arbitrary function for the
+     forward and inverse transform.
+     """
+
+     input_dims = 1
+     output_dims = 1
+     is_separable = True
+     has_inverse = True
+
+     def __init__(self, forward, inverse):
+         """
+         Parameters
+         ----------
+
+         forward: The forward function for the transform
+
+         inverse: The inverse of the forward function.
+         """
+         super().__init__()
+         if callable(forward) and callable(inverse):
+             self._forward = forward
+             self._inverse = inverse
+         else:
+             raise ValueError('arguments to ArbitraryTransform must '
+                              'be functions')
+
+     def transform_non_affine(self, values):
+         return self._forward(values)
+
+     def inverted(self):
+         return ArbitraryTransform(self._inverse, self._forward)
+
+
+class ArbitraryScale(ScaleBase):
+     """
+     Provide an arbitrary scale for the axis.
+     """
+
+     name = 'arbitrary'
+
+     def __init__(self, axis, functions):
+         """
+         Parameters
+         ----------
+
+         axis: the axis for the scale
+
+         functions: (forward, inverse)
+            two-tuple of the forward and inverse functions for the scale.
+         """
+         forward, inverse = functions
+         transform = ArbitraryTransform(forward, inverse)
+         self._transform = transform
+
+     def get_transform(self):
+         """
+         The transform for linear scaling is just the
+         :class:`~matplotlib.transforms.IdentityTransform`.
+         """
+         return self._transform
+
+     def set_default_locators_and_formatters(self, axis):
+         """
+         Set the locators and formatters to reasonable defaults for
+         linear scaling.
+         """
+         axis.set_major_locator(AutoLocator())
+         axis.set_major_formatter(ScalarFormatter())
+         axis.set_minor_formatter(NullFormatter())
+         # update the minor locator for x and y axis based on rcParams
+         if rcParams['xtick.minor.visible']:
+             axis.set_minor_locator(AutoMinorLocator())
+         else:
+             axis.set_minor_locator(NullLocator())
+
+
 class LogTransformBase(Transform):
     input_dims = 1
     output_dims = 1
@@ -545,6 +623,7 @@ def limit_range_for_scale(self, vmin, vmax, minpos):
     'log':    LogScale,
     'symlog': SymmetricalLogScale,
     'logit':  LogitScale,
+    'arbitrary': ArbitraryScale,
     }
 
 
diff --git a/lib/matplotlib/tests/baseline_images/test_scale/arbitrary_scales.png b/lib/matplotlib/tests/baseline_images/test_scale/arbitrary_scales.png
diff --git a/lib/matplotlib/tests/test_scale.py b/lib/matplotlib/tests/test_scale.py
@@ -1,6 +1,8 @@
 from matplotlib.testing.decorators import image_comparison
 import matplotlib.pyplot as plt
 from matplotlib.scale import Log10Transform, InvertedLog10Transform
+import matplotlib.ticker as mticker
+
 import numpy as np
 import io
 import platform
@@ -148,3 +150,26 @@ def test_invalid_log_lims():
     with pytest.warns(UserWarning):
         ax.set_ylim(top=-1)
     assert ax.get_ylim() == original_ylim
+
+
+def inverse(x):
+    return x**2
+
+
+def forward(x):
+    good = x >= 0
+    y = np.full_like(x, np.NaN)
+    y[good] = x[good]**(1/2)
+    return y
+
+
+@image_comparison(baseline_images=['arbitrary_scales'], remove_text=True,
+                  extensions=['png'], style='mpl20')
+def test_arbitrary_scale():
+    fig, ax = plt.subplots()
+
+    x = np.arange(1, 1000)
+
+    ax.plot(x, x)
+    ax.set_xscale('arbitrary', functions=(forward, inverse))
+    ax.set_xlim(1, 1000)
