Improve handling of degenerate jacobians in non-rectilinear grids.

anntzer · anntzer · commit 3c6efe2e1415 · 2024-12-18T22:43:02.000+01:00
grid_helper_curvelinear and floating_axes have code to specifically
handle the case where the transform from rectlinear to non-rectilinear
axes has null derivatives in one of the directions, inferring the angle
of the jacobian from the derivative in the other direction.  (This angle
defines the rotation applied to axis labels, ticks, and tick labels.)

This approach, however, is insufficient if the derivatives in both
directions are zero.  A classical example is e.g. the ``exp(-1/x**2)``
transform, for which all derivatives are zero.  To handle this case more
robustly (and also to better encapsulate the angle calculation, which is
currently repeated at a few places), instead, one can increase the step
size of the numerical differentiation until the gradient becomes
nonzero.  This amounts to moving along the corresponding gridline until
one actually leaves the position of the tick, and thus is indeed a
justifiable approach to compute the tick rotation.

Full examples:

    import numpy as np
    from matplotlib import pyplot as plt
    from matplotlib.projections.polar import PolarTransform
    from matplotlib.transforms import Affine2D
    from mpl_toolkits.axisartist import (
        angle_helper, GridHelperCurveLinear, HostAxes)
    import mpl_toolkits.axisartist.floating_axes as floating_axes

    # def tr(x, y): return x - y, x + y
    # def inv_tr(u, v): return (u + v) / 2, (v - u) / 2

    @np.errstate(divide="ignore")  # at x=0, exp(-1/x**2)=0; div-by-zero can be ignored.
    def tr(x, y):
        return np.exp(-x**-2) - np.exp(-y**-2), np.exp(-x**-2) + np.exp(-y**-2)
    def inv_tr(u, v):
        return (-np.log((u+v)/2))**(1/2), (-np.log((v-u)/2))**(1/2)

    plt.subplot(
        121, axes_class=floating_axes.FloatingAxes,
        grid_helper=floating_axes.GridHelperCurveLinear(
            (tr, inv_tr), extremes=(0, 10, 0, 10)))

    ax = plt.subplot(
        122, axes_class=HostAxes,
        grid_helper=GridHelperCurveLinear(
            Affine2D().scale(np.pi / 180, 1)
            + PolarTransform()
            + Affine2D().scale(2, 1),
            extreme_finder=angle_helper.ExtremeFinderCycle(
                20, 20,
                lon_cycle=360, lat_cycle=None,
                lon_minmax=None, lat_minmax=(0, np.inf),
            ),
            grid_locator1=angle_helper.LocatorDMS(12),
            tick_formatter1=angle_helper.FormatterDMS(),
        ),
        aspect=1, xlim=(-5, 12), ylim=(-5, 10))
    ax.axis["lat"] = axis = ax.new_floating_axis(0, 40)
    axis.label.set_text(r"$\theta = 40^{\circ}$")
    axis.label.set_visible(True)
    ax.grid(True)

    plt.show()
diff --git a/lib/mpl_toolkits/axisartist/floating_axes.py b/lib/mpl_toolkits/axisartist/floating_axes.py
@@ -71,25 +71,19 @@ def trf_xy(x, y):
 
         if self.nth_coord == 0:
             mask = (ymin <= yy0) & (yy0 <= ymax)
-            (xx1, yy1), (dxx1, dyy1), (dxx2, dyy2) = \
-                grid_helper_curvelinear._value_and_jacobian(
+            (xx1, yy1), angle_normal, angle_tangent = \
+                grid_helper_curvelinear._value_and_jac_angle(
                     trf_xy, self.value, yy0[mask], (xmin, xmax), (ymin, ymax))
             labels = self._grid_info["lat_labels"]
 
         elif self.nth_coord == 1:
             mask = (xmin <= xx0) & (xx0 <= xmax)
-            (xx1, yy1), (dxx2, dyy2), (dxx1, dyy1) = \
-                grid_helper_curvelinear._value_and_jacobian(
+            (xx1, yy1), angle_tangent, angle_normal = \
+                grid_helper_curvelinear._value_and_jac_angle(
                     trf_xy, xx0[mask], self.value, (xmin, xmax), (ymin, ymax))
             labels = self._grid_info["lon_labels"]
 
         labels = [l for l, m in zip(labels, mask) if m]
-
-        angle_normal = np.arctan2(dyy1, dxx1)
-        angle_tangent = np.arctan2(dyy2, dxx2)
-        mm = (dyy1 == 0) & (dxx1 == 0)  # points with degenerate normal
-        angle_normal[mm] = angle_tangent[mm] + np.pi / 2
-
         tick_to_axes = self.get_tick_transform(axes) - axes.transAxes
         in_01 = functools.partial(
             mpl.transforms._interval_contains_close, (0, 1))
diff --git a/lib/mpl_toolkits/axisartist/grid_helper_curvelinear.py b/lib/mpl_toolkits/axisartist/grid_helper_curvelinear.py
@@ -16,29 +16,74 @@
 from .grid_finder import GridFinder
 
 
-def _value_and_jacobian(func, xs, ys, xlims, ylims):
+def _value_and_jac_angle(func, xs, ys, xlim, ylim):
     """
-    Compute *func* and its derivatives along x and y at positions *xs*, *ys*,
-    while ensuring that finite difference calculations don't try to evaluate
-    values outside of *xlims*, *ylims*.
+    Parameters
+    ----------
+    func : callable
+        A function that transforms the coordinates of a point (x, y) to a new coordinate
+        system (u, v), and which can also take x and y as arrays of shape *shape* and
+        returns (u, v) as a ``(2, shape)`` array.
+    xs, ys : array-likes
+        Points where *func* and its derivatives will be evaluated.
+    xlim, ylim : pairs of floats
+        (min, max) beyond which *func* should not be evaluated.
+
+    Returns
+    -------
+    val
+        Value of *func* at each point of ``(xs, ys)``.
+    thetas_dx
+        Angles (in radians) defined by the (u, v) components of the numerically
+        differentiated df/dx vector, at each point of ``(xs, ys)``.  If needed, the
+        differentiation step size is increased until at least one component of df/dx
+        is nonzero, under the constraint of not going out of the *xlims*, *ylims*
+        bounds.  If the gridline at a point is actually null (and the angle is thus not
+        well defined), the derivatives are evaluated after taking a small step along y;
+        this ensures e.g. that the tick at r=0 on a radial axis of a polar plot is
+        parallel with the ticks at r!=0.
+    thetas_dy
+        Like *thetas_dx*, but for df/dy.
     """
-    eps = np.finfo(float).eps ** (1/2)  # see e.g. scipy.optimize.approx_fprime
+
+    shape = np.broadcast_shapes(np.shape(xs), np.shape(ys))
     val = func(xs, ys)
-    # Take the finite difference step in the direction where the bound is the
-    # furthest; the step size is min of epsilon and distance to that bound.
-    xlo, xhi = sorted(xlims)
-    dxlo = xs - xlo
-    dxhi = xhi - xs
-    xeps = (np.take([-1, 1], dxhi >= dxlo)
-            * np.minimum(eps, np.maximum(dxlo, dxhi)))
-    val_dx = func(xs + xeps, ys)
-    ylo, yhi = sorted(ylims)
-    dylo = ys - ylo
-    dyhi = yhi - ys
-    yeps = (np.take([-1, 1], dyhi >= dylo)
-            * np.minimum(eps, np.maximum(dylo, dyhi)))
-    val_dy = func(xs, ys + yeps)
-    return (val, (val_dx - val) / xeps, (val_dy - val) / yeps)
+
+    eps0 = np.finfo(float).eps ** (1/2)  # cf. scipy.optimize.approx_fprime
+    # Take finite difference steps towards the furthest bound; the step size will be the
+    # min of epsilon and the distance to that bound.
+    xlo, xhi = sorted(xlim)
+    xdlo = xs - xlo
+    xdhi = xhi - xs
+    xeps_max = np.maximum(xdlo, xdhi)
+    xeps0 = np.where(xdhi >= xdlo, 1, -1) * np.minimum(eps0, xeps_max)
+    ylo, yhi = sorted(ylim)
+    ydlo = ys - ylo
+    ydhi = yhi - ys
+    yeps_max = np.maximum(ydlo, ydhi)
+    yeps0 = np.where(ydhi >= ydlo, 1, -1) * np.minimum(eps0, yeps_max)
+
+    def calc_thetas(dfunc, ps, eps0, eps_max, eps_q):
+        thetas_dp = np.full(shape, np.nan)
+        missing = np.full(shape, True)
+        eps_p = eps0
+        for it, eps_q in enumerate([0, eps_q]):
+            while missing.any() and (abs(eps_p) < eps_max).any():
+                if it == 0 and (eps_p > 1).any():
+                    break  # Degenerate derivative, move a bit along the other coord.
+                eps_p = np.minimum(eps_p, eps_max)
+                df_x, df_y = (dfunc(eps_p, eps_q) - dfunc(0, eps_q)) / eps_p
+                good = missing & ((df_x != 0) | (df_y != 0))
+                thetas_dp[good] = np.arctan2(df_y, df_x)[good]
+                missing &= ~good
+                eps_p *= 2
+        return thetas_dp
+
+    thetas_dx = calc_thetas(lambda eps_p, eps_q: func(xs + eps_p, ys + eps_q),
+                            xs, xeps0, xeps_max, yeps0)
+    thetas_dy = calc_thetas(lambda eps_p, eps_q: func(xs + eps_q, ys + eps_p),
+                            ys, yeps0, yeps_max, xeps0)
+    return (val, thetas_dx, thetas_dy)
 
 
 class FixedAxisArtistHelper(_FixedAxisArtistHelperBase):
@@ -167,12 +212,11 @@ def trf_xy(x, y):
         elif self.nth_coord == 1:
             xx0 = (xmin + xmax) / 2
             yy0 = self.value
-        xy1, dxy1_dx, dxy1_dy = _value_and_jacobian(
+        xy1, angle_dx, angle_dy = _value_and_jac_angle(
             trf_xy, xx0, yy0, (xmin, xmax), (ymin, ymax))
         p = axes.transAxes.inverted().transform(xy1)
         if 0 <= p[0] <= 1 and 0 <= p[1] <= 1:
-            d = [dxy1_dy, dxy1_dx][self.nth_coord]
-            return xy1, np.rad2deg(np.arctan2(*d[::-1]))
+            return xy1, np.rad2deg([angle_dy, angle_dx][self.nth_coord])
         else:
             return None, None
 
@@ -197,23 +241,17 @@ def trf_xy(x, y):
         # find angles
         if self.nth_coord == 0:
             mask = (e0 <= yy0) & (yy0 <= e1)
-            (xx1, yy1), (dxx1, dyy1), (dxx2, dyy2) = _value_and_jacobian(
+            (xx1, yy1), angle_normal, angle_tangent = _value_and_jac_angle(
                 trf_xy, self.value, yy0[mask], (-np.inf, np.inf), (e0, e1))
             labels = self._grid_info["lat_labels"]
 
         elif self.nth_coord == 1:
             mask = (e0 <= xx0) & (xx0 <= e1)
-            (xx1, yy1), (dxx2, dyy2), (dxx1, dyy1) = _value_and_jacobian(
+            (xx1, yy1), angle_tangent, angle_normal = _value_and_jac_angle(
                 trf_xy, xx0[mask], self.value, (-np.inf, np.inf), (e0, e1))
             labels = self._grid_info["lon_labels"]
 
         labels = [l for l, m in zip(labels, mask) if m]
-
-        angle_normal = np.arctan2(dyy1, dxx1)
-        angle_tangent = np.arctan2(dyy2, dxx2)
-        mm = (dyy1 == 0) & (dxx1 == 0)  # points with degenerate normal
-        angle_normal[mm] = angle_tangent[mm] + np.pi / 2
-
         tick_to_axes = self.get_tick_transform(axes) - axes.transAxes
         in_01 = functools.partial(
             mpl.transforms._interval_contains_close, (0, 1))
diff --git a/lib/mpl_toolkits/axisartist/tests/test_floating_axes.py b/lib/mpl_toolkits/axisartist/tests/test_floating_axes.py
@@ -1,5 +1,7 @@
 import numpy as np
 
+import pytest
+
 import matplotlib.pyplot as plt
 import matplotlib.projections as mprojections
 import matplotlib.transforms as mtransforms
@@ -113,3 +115,29 @@ def test_axis_direction():
     ax.axis['y'] = ax.new_floating_axis(nth_coord=1, value=0,
                                         axis_direction='left')
     assert ax.axis['y']._axis_direction == 'left'
+
+
+def test_transform_with_zero_derivatives():
+    # The transform is really a 45° rotation
+    #     tr(x, y) = x-y, x+y; inv_tr(u, v) = (u+v)/2, (u-v)/2
+    # with an additional x->exp(-x**-2) on each coordinate.
+    # Therefore all ticks should be at +/-45°, even the one at zero where the
+    # transform derivatives are zero.
+
+    # at x=0, exp(-x**-2)=0; div-by-zero can be ignored.
+    @np.errstate(divide="ignore")
+    def tr(x, y):
+        return np.exp(-x**-2) - np.exp(-y**-2), np.exp(-x**-2) + np.exp(-y**-2)
+
+    def inv_tr(u, v):
+        return (-np.log((u+v)/2))**(1/2), (-np.log((v-u)/2))**(1/2)
+
+    fig = plt.figure()
+    ax = fig.add_subplot(
+        axes_class=FloatingAxes, grid_helper=GridHelperCurveLinear(
+            (tr, inv_tr), extremes=(0, 10, 0, 10)))
+    fig.canvas.draw()
+
+    for k in ax.axis:
+        for l, a in ax.axis[k].major_ticks.locs_angles:
+            assert a % 90 == pytest.approx(45, abs=1e-3)
