337 lines
12 KiB
Python
337 lines
12 KiB
Python
|
"""
|
||
|
|
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, ""
|