Add classes for discontinuous axis

docs/usage.md

@@ -15,6 +15,7 @@ examples/hamiltonians
 examples/timestructure
 examples/animations
 examples/colors
+examples/utilities
 ```
 
 :::{tip}

tests/test_examples.py

@@ -174,3 +174,8 @@ def test_timestructure():
 def test_twiss():
     t = NotebookTester(f"{dir}/twiss.ipynb")
     t.execute(check_outputs=True)
+
+
+def test_utilities():
+    t = NotebookTester(f"{dir}/utilities.ipynb")
+    t.execute(check_outputs=True)

xplt/__init__.py

@@ -5,7 +5,7 @@
 __contact__ = "[email protected]"
 
 
-__version__ = "0.11.1"
+__version__ = "0.11.2"
 
 
 # expose the following in global namespace

xplt/base.py

@@ -15,6 +15,10 @@
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
+import matplotlib.scale
+import matplotlib.transforms
+import matplotlib.patches
+import matplotlib.path
 import numpy as np
 import pint
 
@@ -156,6 +160,135 @@ def tick_values(self, vmin, vmax):
         return ticks
 
 
+class DiscontinuousLinearScale(mpl.scale.LinearScale):
+    name = "discontinuous_linear"
+
+    def __init__(self, axis, *, breaks=(), space=0, hide=False):
+        super().__init__(axis)
+        self.breaks = breaks
+        self.space = space
+        self.hide = hide
+
+    def set_default_locators_and_formatters(self, axis):
+        axis.set(
+            major_locator=DiscontinuousLocator(self.breaks, self.space),
+            major_formatter=DiscontinuousFormatter(self.breaks),
+        )
+        self.add_discontinuity_markers(axis)
+
+    def get_transform(self):
+        return DiscontinuousTransform(self.breaks, self.space, hide=self.hide)
+
+    def add_discontinuity_markers(self, axis):
+        ax = axis.axes
+        xy = lambda x, y: (x, y) if axis == ax.xaxis else (y, x)
+        trans = mpl.transforms.blended_transform_factory(*xy(ax.transData, ax.transAxes))
+        kwargs = dict(color="k", linewidth=0.8, transform=trans, clip_on=False, zorder=9)
+        for br in self.breaks:
+            if self.hide:
+                ax.add_artist(
+                    mpl.patches.Rectangle(xy(br[0], 0), *xy(br[1] - br[0], 1), fc="w", **kwargs)
+                )
+            else:
+                for i in range(2):
+                    x, w, h = np.mean(br), np.diff(br).item() / 4, 0.02
+                    points = list(zip(*xy((x - w, x, x + w, x), (i - h, i + h, i + h, i - h))))
+                    ax.add_artist(
+                        mpl.patches.PathPatch(mpl.path.Path(points), ec="none", fc="w", **kwargs)
+                    )
+                    ax.add_artist(
+                        mpl.patches.PathPatch(
+                            mpl.path.Path(
+                                points,
+                                [
+                                    mpl.path.Path.MOVETO,
+                                    mpl.path.Path.LINETO,
+                                    mpl.path.Path.MOVETO,
+                                    mpl.path.Path.LINETO,
+                                ],
+                            ),
+                            **kwargs,
+                        )
+                    )
+
+
+class DiscontinuousLocator(mpl.ticker.AutoLocator):
+    def __init__(self, breaks, space, **kwargs):
+        super().__init__(**kwargs)
+        self.breaks = breaks
+        self.space = space
+
+    def set_params(self, **kwargs):
+        if "nbins" in kwargs:
+            self._nbins_real = kwargs.pop("nbins")
+        super().set_params(**kwargs)
+
+    def _raw_ticks(self, vmin, vmax):
+        visible = total = vmax - vmin
+        for br in self.breaks:
+            left, right = np.clip(br, vmin, vmax)
+            visible -= right - left - self.space
+
+        nbins = self._nbins_real
+        if nbins == "auto":
+            if self.axis is not None:
+                nbins = np.clip(self.axis.get_tick_space(), max(1, self._min_n_ticks - 1), 9)
+            else:
+                nbins = 9
+
+        self._nbins = nbins * total / visible
+        ticks = super()._raw_ticks(vmin, vmax)
+
+        # Remove ticks inside breaks
+        # for left, right in self.breaks:
+        #    ticks = ticks[(ticks <= left) | (ticks >= right)]
+
+        return ticks
+
+
+class DiscontinuousFormatter(mpl.ticker.ScalarFormatter):
+    def __init__(self, breaks, **kwargs):
+        super().__init__(**kwargs)
+        self.breaks = breaks
+
+    def __call__(self, x, pos=None):
+        """Return the format for tick value *x* at position *pos*."""
+        for left, right in self.breaks:
+            if left < x < right:
+                return None
+        return super().__call__(x, pos)
+
+
+class DiscontinuousTransform(mpl.transforms.Transform):
+    input_dims = output_dims = 1
+
+    def __init__(self, breaks, space, *, hide=True):
+        super().__init__()
+        self.breaks = breaks
+        self.spaces = space * (np.ones(len(breaks)) if np.isscalar(space) else 1)
+        self.hide = hide
+
+    def transform_non_affine(self, values):
+        a = np.asanyarray(values, dtype="float").copy()
+        breaks = np.asanyarray(self.breaks, dtype="float").copy()
+        spaces = np.asanyarray(self.spaces, dtype="float")
+        m = np.zeros_like(a)
+        for (l, r), s in zip(breaks, spaces):
+            gap, over = (a > l) & (a < r), a >= r
+            m[gap] = 1
+            a[gap] = l + s * (a[gap] - l) / (r - l)
+            a[over] -= r - l - s
+            breaks[breaks >= r] -= r - l - s
+        if self.hide:
+            a = np.ma.masked_where(m, a)
+        return a
+
+    def inverted(self):
+        spaces = np.diff(self.breaks, axis=1)[:, 0]
+        breaks = self.transform_non_affine(self.breaks)
+        return DiscontinuousTransform(breaks, spaces, hide=False)
+
+
 class XPlot:
     def __init__(
         self,

xplt/hooks.py

@@ -11,7 +11,7 @@
 
 import types
 import warnings
-from xplt.colors import (
+from .colors import (
     cmap_petroff,
     cmap_petroff_gradient,
     cmap_petroff_bipolar,
@@ -55,6 +55,11 @@ def register_matplotlib_options():
         cmap_r.name = cmap.name + "_r"
         mpl.colormaps.register(cmap=cmap_r)
 
+    # register custom scales
+    from .base import DiscontinuousLinearScale
+
+    mpl.scale.register_scale(DiscontinuousLinearScale)
+
 
 def register_pint_options():
     """Register default options for pint"""