299 lines
11 KiB
Python
299 lines
11 KiB
Python
|
"""
|
||
|
|
An experimental support for curvilinear grid.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# TODO :
|
||
|
|
# see if tick_iterator method can be simplified by reusing the parent method.
|
||
|
|
|
||
|
|
import functools
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
import matplotlib as mpl
|
||
|
|
from matplotlib import _api, cbook
|
||
|
|
import matplotlib.patches as mpatches
|
||
|
|
from matplotlib.path import Path
|
||
|
|
|
||
|
|
from mpl_toolkits.axes_grid1.parasite_axes import host_axes_class_factory
|
||
|
|
|
||
|
|
from . import axislines, grid_helper_curvelinear
|
||
|
|
from .axis_artist import AxisArtist
|
||
|
|
from .grid_finder import ExtremeFinderSimple
|
||
|
|
|
||
|
|
|
||
|
|
class FloatingAxisArtistHelper(
|
||
|
|
grid_helper_curvelinear.FloatingAxisArtistHelper):
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
class FixedAxisArtistHelper(grid_helper_curvelinear.FloatingAxisArtistHelper):
|
||
|
|
|
||
|
|
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
|
||
|
|
"""
|
||
|
|
lon1, lon2, lat1, lat2 = grid_helper.grid_finder.extreme_finder(*[None] * 5)
|
||
|
|
value, nth_coord = _api.check_getitem(
|
||
|
|
dict(left=(lon1, 0), right=(lon2, 0), bottom=(lat1, 1), top=(lat2, 1)),
|
||
|
|
side=side)
|
||
|
|
super().__init__(grid_helper, nth_coord, value, axis_direction=side)
|
||
|
|
if nth_coord_ticks is None:
|
||
|
|
nth_coord_ticks = nth_coord
|
||
|
|
self.nth_coord_ticks = nth_coord_ticks
|
||
|
|
|
||
|
|
self.value = value
|
||
|
|
self.grid_helper = grid_helper
|
||
|
|
self._side = side
|
||
|
|
|
||
|
|
def update_lim(self, axes):
|
||
|
|
self.grid_helper.update_lim(axes)
|
||
|
|
self._grid_info = self.grid_helper._grid_info
|
||
|
|
|
||
|
|
def get_tick_iterators(self, axes):
|
||
|
|
"""tick_loc, tick_angle, tick_label, (optionally) tick_label"""
|
||
|
|
|
||
|
|
grid_finder = self.grid_helper.grid_finder
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
extremes = self.grid_helper.grid_finder.extreme_finder(*[None] * 5)
|
||
|
|
xmin, xmax = sorted(extremes[:2])
|
||
|
|
ymin, ymax = sorted(extremes[2:])
|
||
|
|
|
||
|
|
def trf_xy(x, y):
|
||
|
|
trf = grid_finder.get_transform() + axes.transData
|
||
|
|
return trf.transform(np.column_stack(np.broadcast_arrays(x, y))).T
|
||
|
|
|
||
|
|
if self.nth_coord == 0:
|
||
|
|
mask = (ymin <= yy0) & (yy0 <= ymax)
|
||
|
|
(xx1, yy1), (dxx1, dyy1), (dxx2, dyy2) = \
|
||
|
|
grid_helper_curvelinear._value_and_jacobian(
|
||
|
|
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(
|
||
|
|
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))
|
||
|
|
|
||
|
|
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(self, axes):
|
||
|
|
self.update_lim(axes)
|
||
|
|
k, v = dict(left=("lon_lines0", 0),
|
||
|
|
right=("lon_lines0", 1),
|
||
|
|
bottom=("lat_lines0", 0),
|
||
|
|
top=("lat_lines0", 1))[self._side]
|
||
|
|
xx, yy = self._grid_info[k][v]
|
||
|
|
return Path(np.column_stack([xx, yy]))
|
||
|
|
|
||
|
|
|
||
|
|
class ExtremeFinderFixed(ExtremeFinderSimple):
|
||
|
|
# docstring inherited
|
||
|
|
|
||
|
|
def __init__(self, extremes):
|
||
|
|
"""
|
||
|
|
This subclass always returns the same bounding box.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
extremes : (float, float, float, float)
|
||
|
|
The bounding box that this helper always returns.
|
||
|
|
"""
|
||
|
|
self._extremes = extremes
|
||
|
|
|
||
|
|
def __call__(self, transform_xy, x1, y1, x2, y2):
|
||
|
|
# docstring inherited
|
||
|
|
return self._extremes
|
||
|
|
|
||
|
|
|
||
|
|
class GridHelperCurveLinear(grid_helper_curvelinear.GridHelperCurveLinear):
|
||
|
|
|
||
|
|
def __init__(self, aux_trans, extremes,
|
||
|
|
grid_locator1=None,
|
||
|
|
grid_locator2=None,
|
||
|
|
tick_formatter1=None,
|
||
|
|
tick_formatter2=None):
|
||
|
|
# docstring inherited
|
||
|
|
super().__init__(aux_trans,
|
||
|
|
extreme_finder=ExtremeFinderFixed(extremes),
|
||
|
|
grid_locator1=grid_locator1,
|
||
|
|
grid_locator2=grid_locator2,
|
||
|
|
tick_formatter1=tick_formatter1,
|
||
|
|
tick_formatter2=tick_formatter2)
|
||
|
|
|
||
|
|
@_api.deprecated("3.8")
|
||
|
|
def get_data_boundary(self, side):
|
||
|
|
"""
|
||
|
|
Return v=0, nth=1.
|
||
|
|
"""
|
||
|
|
lon1, lon2, lat1, lat2 = self.grid_finder.extreme_finder(*[None] * 5)
|
||
|
|
return dict(left=(lon1, 0),
|
||
|
|
right=(lon2, 0),
|
||
|
|
bottom=(lat1, 1),
|
||
|
|
top=(lat2, 1))[side]
|
||
|
|
|
||
|
|
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
|
||
|
|
# This is not the same as the FixedAxisArtistHelper class used by
|
||
|
|
# grid_helper_curvelinear.GridHelperCurveLinear.new_fixed_axis!
|
||
|
|
helper = FixedAxisArtistHelper(
|
||
|
|
self, loc, nth_coord_ticks=nth_coord)
|
||
|
|
axisline = AxisArtist(axes, helper, axis_direction=axis_direction)
|
||
|
|
# Perhaps should be moved to the base class?
|
||
|
|
axisline.line.set_clip_on(True)
|
||
|
|
axisline.line.set_clip_box(axisline.axes.bbox)
|
||
|
|
return axisline
|
||
|
|
|
||
|
|
# new_floating_axis will inherit the grid_helper's extremes.
|
||
|
|
|
||
|
|
# def new_floating_axis(self, nth_coord,
|
||
|
|
# value,
|
||
|
|
# axes=None,
|
||
|
|
# axis_direction="bottom"
|
||
|
|
# ):
|
||
|
|
|
||
|
|
# axis = super(GridHelperCurveLinear,
|
||
|
|
# self).new_floating_axis(nth_coord,
|
||
|
|
# value, axes=axes,
|
||
|
|
# axis_direction=axis_direction)
|
||
|
|
|
||
|
|
# # set extreme values of the axis helper
|
||
|
|
# if nth_coord == 1:
|
||
|
|
# axis.get_helper().set_extremes(*self._extremes[:2])
|
||
|
|
# elif nth_coord == 0:
|
||
|
|
# axis.get_helper().set_extremes(*self._extremes[2:])
|
||
|
|
|
||
|
|
# return axis
|
||
|
|
|
||
|
|
def _update_grid(self, x1, y1, x2, y2):
|
||
|
|
if self._grid_info is None:
|
||
|
|
self._grid_info = dict()
|
||
|
|
|
||
|
|
grid_info = self._grid_info
|
||
|
|
|
||
|
|
grid_finder = self.grid_finder
|
||
|
|
extremes = grid_finder.extreme_finder(grid_finder.inv_transform_xy,
|
||
|
|
x1, y1, x2, y2)
|
||
|
|
|
||
|
|
lon_min, lon_max = sorted(extremes[:2])
|
||
|
|
lat_min, lat_max = sorted(extremes[2:])
|
||
|
|
grid_info["extremes"] = lon_min, lon_max, lat_min, lat_max # extremes
|
||
|
|
|
||
|
|
lon_levs, lon_n, lon_factor = \
|
||
|
|
grid_finder.grid_locator1(lon_min, lon_max)
|
||
|
|
lon_levs = np.asarray(lon_levs)
|
||
|
|
lat_levs, lat_n, lat_factor = \
|
||
|
|
grid_finder.grid_locator2(lat_min, lat_max)
|
||
|
|
lat_levs = np.asarray(lat_levs)
|
||
|
|
|
||
|
|
grid_info["lon_info"] = lon_levs, lon_n, lon_factor
|
||
|
|
grid_info["lat_info"] = lat_levs, lat_n, lat_factor
|
||
|
|
|
||
|
|
grid_info["lon_labels"] = grid_finder.tick_formatter1(
|
||
|
|
"bottom", lon_factor, lon_levs)
|
||
|
|
grid_info["lat_labels"] = grid_finder.tick_formatter2(
|
||
|
|
"bottom", lat_factor, lat_levs)
|
||
|
|
|
||
|
|
lon_values = lon_levs[:lon_n] / lon_factor
|
||
|
|
lat_values = lat_levs[:lat_n] / lat_factor
|
||
|
|
|
||
|
|
lon_lines, lat_lines = grid_finder._get_raw_grid_lines(
|
||
|
|
lon_values[(lon_min < lon_values) & (lon_values < lon_max)],
|
||
|
|
lat_values[(lat_min < lat_values) & (lat_values < lat_max)],
|
||
|
|
lon_min, lon_max, lat_min, lat_max)
|
||
|
|
|
||
|
|
grid_info["lon_lines"] = lon_lines
|
||
|
|
grid_info["lat_lines"] = lat_lines
|
||
|
|
|
||
|
|
lon_lines, lat_lines = grid_finder._get_raw_grid_lines(
|
||
|
|
# lon_min, lon_max, lat_min, lat_max)
|
||
|
|
extremes[:2], extremes[2:], *extremes)
|
||
|
|
|
||
|
|
grid_info["lon_lines0"] = lon_lines
|
||
|
|
grid_info["lat_lines0"] = lat_lines
|
||
|
|
|
||
|
|
def get_gridlines(self, which="major", axis="both"):
|
||
|
|
grid_lines = []
|
||
|
|
if axis in ["both", "x"]:
|
||
|
|
grid_lines.extend(self._grid_info["lon_lines"])
|
||
|
|
if axis in ["both", "y"]:
|
||
|
|
grid_lines.extend(self._grid_info["lat_lines"])
|
||
|
|
return grid_lines
|
||
|
|
|
||
|
|
|
||
|
|
class FloatingAxesBase:
|
||
|
|
|
||
|
|
def __init__(self, *args, grid_helper, **kwargs):
|
||
|
|
_api.check_isinstance(GridHelperCurveLinear, grid_helper=grid_helper)
|
||
|
|
super().__init__(*args, grid_helper=grid_helper, **kwargs)
|
||
|
|
self.set_aspect(1.)
|
||
|
|
|
||
|
|
def _gen_axes_patch(self):
|
||
|
|
# docstring inherited
|
||
|
|
x0, x1, y0, y1 = self.get_grid_helper().grid_finder.extreme_finder(*[None] * 5)
|
||
|
|
patch = mpatches.Polygon([(x0, y0), (x1, y0), (x1, y1), (x0, y1)])
|
||
|
|
patch.get_path()._interpolation_steps = 100
|
||
|
|
return patch
|
||
|
|
|
||
|
|
def clear(self):
|
||
|
|
super().clear()
|
||
|
|
self.patch.set_transform(
|
||
|
|
self.get_grid_helper().grid_finder.get_transform()
|
||
|
|
+ self.transData)
|
||
|
|
# The original patch is not in the draw tree; it is only used for
|
||
|
|
# clipping purposes.
|
||
|
|
orig_patch = super()._gen_axes_patch()
|
||
|
|
orig_patch.set_figure(self.figure)
|
||
|
|
orig_patch.set_transform(self.transAxes)
|
||
|
|
self.patch.set_clip_path(orig_patch)
|
||
|
|
self.gridlines.set_clip_path(orig_patch)
|
||
|
|
self.adjust_axes_lim()
|
||
|
|
|
||
|
|
def adjust_axes_lim(self):
|
||
|
|
bbox = self.patch.get_path().get_extents(
|
||
|
|
# First transform to pixel coords, then to parent data coords.
|
||
|
|
self.patch.get_transform() - self.transData)
|
||
|
|
bbox = bbox.expanded(1.02, 1.02)
|
||
|
|
self.set_xlim(bbox.xmin, bbox.xmax)
|
||
|
|
self.set_ylim(bbox.ymin, bbox.ymax)
|
||
|
|
|
||
|
|
|
||
|
|
floatingaxes_class_factory = cbook._make_class_factory(
|
||
|
|
FloatingAxesBase, "Floating{}")
|
||
|
|
FloatingAxes = floatingaxes_class_factory(
|
||
|
|
host_axes_class_factory(axislines.Axes))
|
||
|
|
FloatingSubplot = FloatingAxes
|