import numpy as np
from matplotlib import ticker as mticker
from matplotlib.transforms import Bbox, Transform
def _find_line_box_crossings(xys, bbox):
"""
Find the points where a polyline crosses a bbox, and the crossing angles.
Parameters
----------
xys : (N, 2) array
The polyline coordinates.
bbox : `.Bbox`
The bounding box.
Returns
-------
list of ((float, float), float)
Four separate lists of crossings, for the left, right, bottom, and top
sides of the bbox, respectively. For each list, the entries are the
``((x, y), ccw_angle_in_degrees)`` of the crossing, where an angle of 0
means that the polyline is moving to the right at the crossing point.
The entries are computed by linearly interpolating at each crossing
between the nearest points on either side of the bbox edges.
"""
crossings = []
dxys = xys[1:] - xys[:-1]
for sl in [slice(None), slice(None, None, -1)]:
us, vs = xys.T[sl] # "this" coord, "other" coord
dus, dvs = dxys.T[sl]
umin, vmin = bbox.min[sl]
umax, vmax = bbox.max[sl]
for u0, inside in [(umin, us > umin), (umax, us < umax)]:
crossings.append([])
idxs, = (inside[:-1] ^ inside[1:]).nonzero()
for idx in idxs:
v = vs[idx] + (u0 - us[idx]) * dvs[idx] / dus[idx]
if not vmin <= v <= vmax:
continue
crossing = (u0, v)[sl]
theta = np.degrees(np.arctan2(*dxys[idx][::-1]))
crossings[-1].append((crossing, theta))
return crossings
class ExtremeFinderSimple:
"""
A helper class to figure out the range of grid lines that need to be drawn.
"""
def __init__(self, nx, ny):
"""
Parameters
----------
nx, ny : int
The number of samples in each direction.
"""
self.nx = nx
self.ny = ny
def __call__(self, transform_xy, x1, y1, x2, y2):
"""
Compute an approximation of the bounding box obtained by applying
*transform_xy* to the box delimited by ``(x1, y1, x2, y2)``.
The intended use is to have ``(x1, y1, x2, y2)`` in axes coordinates,
and have *transform_xy* be the transform from axes coordinates to data
coordinates; this method then returns the range of data coordinates
that span the actual axes.
The computation is done by sampling ``nx * ny`` equispaced points in
the ``(x1, y1, x2, y2)`` box and finding the resulting points with
extremal coordinates; then adding some padding to take into account the
finite sampling.
As each sampling step covers a relative range of *1/nx* or *1/ny*,
the padding is computed by expanding the span covered by the extremal
coordinates by these fractions.
"""
x, y = np.meshgrid(
np.linspace(x1, x2, self.nx), np.linspace(y1, y2, self.ny))
xt, yt = transform_xy(np.ravel(x), np.ravel(y))
return self._add_pad(xt.min(), xt.max(), yt.min(), yt.max())
def _add_pad(self, x_min, x_max, y_min, y_max):
"""Perform the padding mentioned in `__call__`."""
dx = (x_max - x_min) / self.nx
dy = (y_max - y_min) / self.ny
return x_min - dx, x_max + dx, y_min - dy, y_max + dy
class _User2DTransform(Transform):
"""A transform defined by two user-set functions."""
input_dims = output_dims = 2
def __init__(self, forward, backward):
"""
Parameters
----------
forward, backward : callable
The forward and backward transforms, taking ``x`` and ``y`` as
separate arguments and returning ``(tr_x, tr_y)``.
"""
# The normal Matplotlib convention would be to take and return an
# (N, 2) array but axisartist uses the transposed version.
super().__init__()
self._forward = forward
self._backward = backward
def transform_non_affine(self, values):
# docstring inherited
return np.transpose(self._forward(*np.transpose(values)))
def inverted(self):
# docstring inherited
return type(self)(self._backward, self._forward)
class GridFinder:
"""
Internal helper for `~.grid_helper_curvelinear.GridHelperCurveLinear`, with
the same constructor parameters; should not be directly instantiated.
"""
def __init__(self,
transform,
extreme_finder=None,
grid_locator1=None,
grid_locator2=None,
tick_formatter1=None,
tick_formatter2=None):
if extreme_finder is None:
extreme_finder = ExtremeFinderSimple(20, 20)
if grid_locator1 is None:
grid_locator1 = MaxNLocator()
if grid_locator2 is None:
grid_locator2 = MaxNLocator()
if tick_formatter1 is None:
tick_formatter1 = FormatterPrettyPrint()
if tick_formatter2 is None:
tick_formatter2 = FormatterPrettyPrint()
self.extreme_finder = extreme_finder
self.grid_locator1 = grid_locator1
self.grid_locator2 = grid_locator2
self.tick_formatter1 = tick_formatter1
self.tick_formatter2 = tick_formatter2
self.set_transform(transform)
def get_grid_info(self, x1, y1, x2, y2):
"""
lon_values, lat_values : list of grid values. if integer is given,
rough number of grids in each direction.
"""
extremes = self.extreme_finder(self.inv_transform_xy, x1, y1, x2, y2)
# min & max rage of lat (or lon) for each grid line will be drawn.
# i.e., gridline of lon=0 will be drawn from lat_min to lat_max.
lon_min, lon_max, lat_min, lat_max = extremes
lon_levs, lon_n, lon_factor = self.grid_locator1(lon_min, lon_max)
lon_levs = np.asarray(lon_levs)
lat_levs, lat_n, lat_factor = self.grid_locator2(lat_min, lat_max)
lat_levs = np.asarray(lat_levs)
lon_values = lon_levs[:lon_n] / lon_factor
lat_values = lat_levs[:lat_n] / lat_factor
lon_lines, lat_lines = self._get_raw_grid_lines(lon_values,
lat_values,
lon_min, lon_max,
lat_min, lat_max)
ddx = (x2-x1)*1.e-10
ddy = (y2-y1)*1.e-10
bb = Bbox.from_extents(x1-ddx, y1-ddy, x2+ddx, y2+ddy)
grid_info = {
"extremes": extremes,
"lon_lines": lon_lines,
"lat_lines": lat_lines,
"lon": self._clip_grid_lines_and_find_ticks(
lon_lines, lon_values, lon_levs, bb),
"lat": self._clip_grid_lines_and_find_ticks(
lat_lines, lat_values, lat_levs, bb),
}
tck_labels = grid_info["lon"]["tick_labels"] = {}
for direction in ["left", "bottom", "right", "top"]:
levs = grid_info["lon"]["tick_levels"][direction]
tck_labels[direction] = self.tick_formatter1(
direction, lon_factor, levs)
tck_labels = grid_info["lat"]["tick_labels"] = {}
for direction in ["left", "bottom", "right", "top"]:
levs = grid_info["lat"]["tick_levels"][direction]
tck_labels[direction] = self.tick_formatter2(
direction, lat_factor, levs)
return grid_info
def _get_raw_grid_lines(self,
lon_values, lat_values,
lon_min, lon_max, lat_min, lat_max):
lons_i = np.linspace(lon_min, lon_max, 100) # for interpolation
lats_i = np.linspace(lat_min, lat_max, 100)
lon_lines = [self.transform_xy(np.full_like(lats_i, lon), lats_i)
for lon in lon_values]
lat_lines = [self.transform_xy(lons_i, np.full_like(lons_i, lat))
for lat in lat_values]
return lon_lines, lat_lines
def _clip_grid_lines_and_find_ticks(self, lines, values, levs, bb):
gi = {
"values": [],
"levels": [],
"tick_levels": dict(left=[], bottom=[], right=[], top=[]),
"tick_locs": dict(left=[], bottom=[], right=[], top=[]),
"lines": [],
}
tck_levels = gi["tick_levels"]
tck_locs = gi["tick_locs"]
for (lx, ly), v, lev in zip(lines, values, levs):
tcks = _find_line_box_crossings(np.column_stack([lx, ly]), bb)
gi["levels"].append(v)
gi["lines"].append([(lx, ly)])
for tck, direction in zip(tcks,
["left", "right", "bottom", "top"]):
for t in tck:
tck_levels[direction].append(lev)
tck_locs[direction].append(t)
return gi
def set_transform(self, aux_trans):
if isinstance(aux_trans, Transform):
self._aux_transform = aux_trans
elif len(aux_trans) == 2 and all(map(callable, aux_trans)):
self._aux_transform = _User2DTransform(*aux_trans)
else:
raise TypeError("'aux_trans' must be either a Transform "
"instance or a pair of callables")
def get_transform(self):
return self._aux_transform
update_transform = set_transform # backcompat alias.
def transform_xy(self, x, y):
return self._aux_transform.transform(np.column_stack([x, y])).T
def inv_transform_xy(self, x, y):
return self._aux_transform.inverted().transform(
np.column_stack([x, y])).T
def update(self, **kwargs):
for k, v in kwargs.items():
if k in ["extreme_finder",
"grid_locator1",
"grid_locator2",
"tick_formatter1",
"tick_formatter2"]:
setattr(self, k, v)
else:
raise ValueError(f"Unknown update property {k!r}")
class MaxNLocator(mticker.MaxNLocator):
def __init__(self, nbins=10, steps=None,
trim=True,
integer=False,
symmetric=False,
prune=None):
# trim argument has no effect. It has been left for API compatibility
super().__init__(nbins, steps=steps, integer=integer,
symmetric=symmetric, prune=prune)
self.create_dummy_axis()
def __call__(self, v1, v2):
locs = super().tick_values(v1, v2)
return np.array(locs), len(locs), 1 # 1: factor (see angle_helper)
class FixedLocator:
def __init__(self, locs):
self._locs = locs
def __call__(self, v1, v2):
v1, v2 = sorted([v1, v2])
locs = np.array([l for l in self._locs if v1 <= l <= v2])
return locs, len(locs), 1 # 1: factor (see angle_helper)
# Tick Formatter
class FormatterPrettyPrint:
def __init__(self, useMathText=True):
self._fmt = mticker.ScalarFormatter(
useMathText=useMathText, useOffset=False)
self._fmt.create_dummy_axis()
def __call__(self, direction, factor, values):
return self._fmt.format_ticks(values)
class DictFormatter:
def __init__(self, format_dict, formatter=None):
"""
format_dict : dictionary for format strings to be used.
formatter : fall-back formatter
"""
super().__init__()
self._format_dict = format_dict
self._fallback_formatter = formatter
def __call__(self, direction, factor, values):
"""
factor is ignored if value is found in the dictionary
"""
if self._fallback_formatter:
fallback_strings = self._fallback_formatter(
direction, factor, values)
else:
fallback_strings = [""] * len(values)
return [self._format_dict.get(k, v)
for k, v in zip(values, fallback_strings)]