"""
An experimental support for curvilinear grid.
"""
import functools
from itertools import chain
import numpy as np
import matplotlib as mpl
from matplotlib.path import Path
from matplotlib.transforms import Affine2D, IdentityTransform
from .axislines import (
_FixedAxisArtistHelperBase, _FloatingAxisArtistHelperBase, GridHelperBase)
from .axis_artist import AxisArtist
from .grid_finder import GridFinder
def _value_and_jacobian(func, xs, ys, xlims, ylims):
"""
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*.
"""
eps = np.finfo(float).eps ** (1/2) # see e.g. scipy.optimize.approx_fprime
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)
class FixedAxisArtistHelper(_FixedAxisArtistHelperBase):
"""
Helper class for a fixed axis.
"""
def __init__(self, grid_helper, side, nth_coord_ticks=None):
"""
nth_coord = along which coordinate value varies.
nth_coord = 0 -> x axis, nth_coord = 1 -> y axis
"""
super().__init__(loc=side)
self.grid_helper = grid_helper
if nth_coord_ticks is None:
nth_coord_ticks = self.nth_coord
self.nth_coord_ticks = nth_coord_ticks
self.side = side
def update_lim(self, axes):
self.grid_helper.update_lim(axes)
def get_tick_transform(self, axes):
return axes.transData
def get_tick_iterators(self, axes):
"""tick_loc, tick_angle, tick_label"""
v1, v2 = axes.get_ylim() if self.nth_coord == 0 else axes.get_xlim()
if v1 > v2: # Inverted limits.
side = {"left": "right", "right": "left",
"top": "bottom", "bottom": "top"}[self.side]
else:
side = self.side
g = self.grid_helper
ti1 = g.get_tick_iterator(self.nth_coord_ticks, side)
ti2 = g.get_tick_iterator(1-self.nth_coord_ticks, side, minor=True)
return chain(ti1, ti2), iter([])
class FloatingAxisArtistHelper(_FloatingAxisArtistHelperBase):
def __init__(self, grid_helper, nth_coord, value, axis_direction=None):
"""
nth_coord = along which coordinate value varies.
nth_coord = 0 -> x axis, nth_coord = 1 -> y axis
"""
super().__init__(nth_coord, value)
self.value = value
self.grid_helper = grid_helper
self._extremes = -np.inf, np.inf
self._line_num_points = 100 # number of points to create a line
def set_extremes(self, e1, e2):
if e1 is None:
e1 = -np.inf
if e2 is None:
e2 = np.inf
self._extremes = e1, e2
def update_lim(self, axes):
self.grid_helper.update_lim(axes)
x1, x2 = axes.get_xlim()
y1, y2 = axes.get_ylim()
grid_finder = self.grid_helper.grid_finder
extremes = grid_finder.extreme_finder(grid_finder.inv_transform_xy,
x1, y1, x2, y2)
lon_min, lon_max, lat_min, lat_max = extremes
e_min, e_max = self._extremes # ranges of other coordinates
if self.nth_coord == 0:
lat_min = max(e_min, lat_min)
lat_max = min(e_max, lat_max)
elif self.nth_coord == 1:
lon_min = max(e_min, lon_min)
lon_max = min(e_max, lon_max)
lon_levs, lon_n, lon_factor = \
grid_finder.grid_locator1(lon_min, lon_max)
lat_levs, lat_n, lat_factor = \
grid_finder.grid_locator2(lat_min, lat_max)
if self.nth_coord == 0:
xx0 = np.full(self._line_num_points, self.value)
yy0 = np.linspace(lat_min, lat_max, self._line_num_points)
xx, yy = grid_finder.transform_xy(xx0, yy0)
elif self.nth_coord == 1:
xx0 = np.linspace(lon_min, lon_max, self._line_num_points)
yy0 = np.full(self._line_num_points, self.value)
xx, yy = grid_finder.transform_xy(xx0, yy0)
self._grid_info = {
"extremes": (lon_min, lon_max, lat_min, lat_max),
"lon_info": (lon_levs, lon_n, np.asarray(lon_factor)),
"lat_info": (lat_levs, lat_n, np.asarray(lat_factor)),
"lon_labels": grid_finder.tick_formatter1(
"bottom", lon_factor, lon_levs),
"lat_labels": grid_finder.tick_formatter2(
"bottom", lat_factor, lat_levs),
"line_xy": (xx, yy),
}
def get_axislabel_transform(self, axes):
return Affine2D() # axes.transData
def get_axislabel_pos_angle(self, axes):
def trf_xy(x, y):
trf = self.grid_helper.grid_finder.get_transform() + axes.transData
return trf.transform([x, y]).T
xmin, xmax, ymin, ymax = self._grid_info["extremes"]
if self.nth_coord == 0:
xx0 = self.value
yy0 = (ymin + ymax) / 2
elif self.nth_coord == 1:
xx0 = (xmin + xmax) / 2
yy0 = self.value
xy1, dxy1_dx, dxy1_dy = _value_and_jacobian(
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]))
else:
return None, None
def get_tick_transform(self, axes):
return IdentityTransform() # axes.transData
def get_tick_iterators(self, axes):
"""tick_loc, tick_angle, tick_label, (optionally) tick_label"""
lat_levs, lat_n, lat_factor = self._grid_info["lat_info"]
yy0 = lat_levs / lat_factor
lon_levs, lon_n, lon_factor = self._grid_info["lon_info"]
xx0 = lon_levs / lon_factor
e0, e1 = self._extremes
def trf_xy(x, y):
trf = self.grid_helper.grid_finder.get_transform() + axes.transData
return trf.transform(np.column_stack(np.broadcast_arrays(x, y))).T
# find angles
if self.nth_coord == 0:
mask = (e0 <= yy0) & (yy0 <= e1)
(xx1, yy1), (dxx1, dyy1), (dxx2, dyy2) = _value_and_jacobian(
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(
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))
def f1():
for x, y, normal, tangent, lab \
in zip(xx1, yy1, angle_normal, angle_tangent, labels):
c2 = tick_to_axes.transform((x, y))
if in_01(c2[0]) and in_01(c2[1]):
yield [x, y], *np.rad2deg([normal, tangent]), lab
return f1(), iter([])
def get_line_transform(self, axes):
return axes.transData
def get_line(self, axes):
self.update_lim(axes)
x, y = self._grid_info["line_xy"]
return Path(np.column_stack([x, y]))
class GridHelperCurveLinear(GridHelperBase):
def __init__(self, aux_trans,
extreme_finder=None,
grid_locator1=None,
grid_locator2=None,
tick_formatter1=None,
tick_formatter2=None):
"""
Parameters
----------
aux_trans : `.Transform` or tuple[Callable, Callable]
The transform from curved coordinates to rectilinear coordinate:
either a `.Transform` instance (which provides also its inverse),
or a pair of callables ``(trans, inv_trans)`` that define the
transform and its inverse. The callables should have signature::
x_rect, y_rect = trans(x_curved, y_curved)
x_curved, y_curved = inv_trans(x_rect, y_rect)
extreme_finder
grid_locator1, grid_locator2
Grid locators for each axis.
tick_formatter1, tick_formatter2
Tick formatters for each axis.
"""
super().__init__()
self._grid_info = None
self.grid_finder = GridFinder(aux_trans,
extreme_finder,
grid_locator1,
grid_locator2,
tick_formatter1,
tick_formatter2)
def update_grid_finder(self, aux_trans=None, **kwargs):
if aux_trans is not None:
self.grid_finder.update_transform(aux_trans)
self.grid_finder.update(**kwargs)
self._old_limits = None # Force revalidation.
def new_fixed_axis(self, loc,
nth_coord=None,
axis_direction=None,
offset=None,
axes=None):
if axes is None:
axes = self.axes
if axis_direction is None:
axis_direction = loc
helper = FixedAxisArtistHelper(self, loc, nth_coord_ticks=nth_coord)
axisline = AxisArtist(axes, helper, axis_direction=axis_direction)
# Why is clip not set on axisline, unlike in new_floating_axis or in
# the floating_axig.GridHelperCurveLinear subclass?
return axisline
def new_floating_axis(self, nth_coord,
value,
axes=None,
axis_direction="bottom"
):
if axes is None:
axes = self.axes
helper = FloatingAxisArtistHelper(
self, nth_coord, value, axis_direction)
axisline = AxisArtist(axes, helper)
axisline.line.set_clip_on(True)
axisline.line.set_clip_box(axisline.axes.bbox)
# axisline.major_ticklabels.set_visible(True)
# axisline.minor_ticklabels.set_visible(False)
return axisline
def _update_grid(self, x1, y1, x2, y2):
self._grid_info = self.grid_finder.get_grid_info(x1, y1, x2, y2)
def get_gridlines(self, which="major", axis="both"):
grid_lines = []
if axis in ["both", "x"]:
for gl in self._grid_info["lon"]["lines"]:
grid_lines.extend(gl)
if axis in ["both", "y"]:
for gl in self._grid_info["lat"]["lines"]:
grid_lines.extend(gl)
return grid_lines
def get_tick_iterator(self, nth_coord, axis_side, minor=False):
# axisnr = dict(left=0, bottom=1, right=2, top=3)[axis_side]
angle_tangent = dict(left=90, right=90, bottom=0, top=0)[axis_side]
# angle = [0, 90, 180, 270][axisnr]
lon_or_lat = ["lon", "lat"][nth_coord]
if not minor: # major ticks
for (xy, a), l in zip(
self._grid_info[lon_or_lat]["tick_locs"][axis_side],
self._grid_info[lon_or_lat]["tick_labels"][axis_side]):
angle_normal = a
yield xy, angle_normal, angle_tangent, l
else:
for (xy, a), l in zip(
self._grid_info[lon_or_lat]["tick_locs"][axis_side],
self._grid_info[lon_or_lat]["tick_labels"][axis_side]):
angle_normal = a
yield xy, angle_normal, angle_tangent, ""
# for xy, a, l in self._grid_info[lon_or_lat]["ticks"][axis_side]:
# yield xy, a, ""