Skip to content

Commit 5bf295a

Browse files
lukelbdlukelbd
committed
Fix basemap and cartopy wrappers
1 parent 819b525 commit 5bf295a

File tree

3 files changed

+256
-251
lines changed

3 files changed

+256
-251
lines changed

‎proplot/axes/base.py

Lines changed: 125 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -350,35 +350,43 @@ def __init__(self, *args, number=None, main=False, **kwargs):
350350
# Ensure isDefault_minloc enabled at start, needed for dual axes
351351
self.xaxis.isDefault_minloc = self.yaxis.isDefault_minloc = True
352352

353-
# Properties
354353
self._auto_format = None # manipulated by wrapper functions
354+
355355
self._abc_loc = None
356356
self._abc_text = None
357357
self._abc_border_kwargs = {} # abs border properties
358+
358359
self._title_loc = None # location of main title
359360
self._title_pad = rc['axes.titlepad'] # format() can overwrite
360361
self._title_pad_active = None
361362
self._title_border_kwargs = {} # title border properties
363+
362364
self._above_top_panels = True # TODO: add rc prop?
363365
self._bottom_panels = []
364366
self._top_panels = []
365367
self._left_panels = []
366368
self._right_panels = []
369+
367370
self._tightbbox = None # bounding boxes are saved
371+
368372
self._panel_hidden = False # True when "filled" with cbar/legend
369373
self._panel_parent = None
370374
self._panel_share = False
371375
self._panel_sharex_group = False
372376
self._panel_sharey_group = False
373377
self._panel_side = None
378+
374379
self._inset_parent = None
375380
self._inset_zoom = False
376381
self._inset_zoom_data = None
382+
377383
self._alty_child = None
378384
self._altx_child = None
379385
self._alty_parent = None
380386
self._altx_parent = None
387+
381388
self.number = number # for abc numbering
389+
382390
if main:
383391
self.figure._axes_main.append(self)
384392

@@ -681,6 +689,15 @@ def inset_locator(ax, renderer):
681689
return bb
682690
return inset_locator
683691

692+
def _plot_redirect(self, name, *args, **kwargs):
693+
"""
694+
Redirect to the associated basemap method if possible.
695+
"""
696+
if getattr(self, 'name', '') == 'basemap':
697+
return getattr(self.projection, name)(*args, ax=self, **kwargs)
698+
else:
699+
return getattr(maxes.Axes, name)(self, *args, **kwargs)
700+
684701
def _range_gridspec(self, x):
685702
"""
686703
Return the column or row gridspec range for the axes.
@@ -1338,7 +1355,7 @@ def barbs(self, *args, **kwargs):
13381355
"""
13391356
args = plot._parse_2d(*args)
13401357
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
1341-
mappable = super().barbs(*args, **kwargs)
1358+
mappable = self._plot_redirect('barbs', *args, **kwargs)
13421359
plot._auto_colorbar(mappable, **kwargs_colorbar)
13431360
return mappable
13441361

@@ -1701,7 +1718,7 @@ def contour(self, *args, **kwargs):
17011718
"""
17021719
args = plot._parse_2d(*args)
17031720
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
1704-
mappable = super().contour(*args, **kwargs)
1721+
mappable = self._plot_redirect('contour', *args, **kwargs)
17051722
plot._auto_colorbar(mappable, **kwargs_colorbar)
17061723
return mappable
17071724

@@ -1720,10 +1737,15 @@ def contourf(self, *args, **kwargs):
17201737
"""
17211738
args = plot._parse_2d(*args)
17221739
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
1723-
mappable = super().contourf(*args, **kwargs)
1740+
mappable = self._plot_redirect('contourf', *args, **kwargs)
17241741
plot._auto_colorbar(mappable, **kwargs_colorbar)
17251742
return mappable
17261743

1744+
def draw(self, renderer=None, *args, **kwargs):
1745+
# Perform extra post-processing steps
1746+
self._reassign_title()
1747+
super().draw(renderer, *args, **kwargs)
1748+
17271749
# fill_between = _fill_between_wrapper(_standardize_1d(_cycle_changer(
17281750
def _fill_between_apply(
17291751
self, xy, *args,
@@ -1834,101 +1856,6 @@ def fill_betweenx(self, *args, **kwargs):
18341856
"""
18351857
return self._fill_between_apply('yx', *args, **kwargs)
18361858

1837-
@plot._concatenate_docstrings
1838-
@docstring.add_snippets
1839-
def legend(self, *args, loc=None, width=None, space=None, **kwargs):
1840-
"""
1841-
Add an *inset* legend or *outer* legend along the edge of the axes.
1842-
1843-
Parameters
1844-
----------
1845-
%(plot.legend_args)s
1846-
loc : int or str, optional
1847-
The legend location. The following location keys are valid:
1848-
1849-
.. _legend_table:
1850-
1851-
================== =======================================
1852-
Location Valid keys
1853-
================== =======================================
1854-
outer left ``'left'``, ``'l'``
1855-
outer right ``'right'``, ``'r'``
1856-
outer bottom ``'bottom'``, ``'b'``
1857-
outer top ``'top'``, ``'t'``
1858-
"best" inset ``'best'``, ``'inset'``, ``'i'``, ``0``
1859-
upper right inset ``'upper right'``, ``'ur'``, ``1``
1860-
upper left inset ``'upper left'``, ``'ul'``, ``2``
1861-
lower left inset ``'lower left'``, ``'ll'``, ``3``
1862-
lower right inset ``'lower right'``, ``'lr'``, ``4``
1863-
center left inset ``'center left'``, ``'cl'``, ``5``
1864-
center right inset ``'center right'``, ``'cr'``, ``6``
1865-
lower center inset ``'lower center'``, ``'lc'``, ``7``
1866-
upper center inset ``'upper center'``, ``'uc'``, ``8``
1867-
center inset ``'center'``, ``'c'``, ``9``
1868-
"filled" ``'fill'``
1869-
================== =======================================
1870-
1871-
width : float or str, optional
1872-
For outer legends only. The space allocated for the legend box.
1873-
This does nothing if :rcraw:`tight` is ``True``. Units are
1874-
interpreted by `~proplot.utils.units`.
1875-
space : float or str, optional
1876-
For outer legends only. The space between the axes and the legend
1877-
box. Units are interpreted by `~proplot.utils.units`.
1878-
When :rcraw:`tight` is ``True``, this is adjusted automatically.
1879-
Otherwise, the default is :rc:`subplots.panelpad`.
1880-
%(plot.legend_kwargs)s
1881-
1882-
Other parameters
1883-
----------------
1884-
**kwargs
1885-
Passed to `~matplotlib.axes.Axes.legend`.
1886-
"""
1887-
if loc != 'fill':
1888-
loc = self._loc_translate(loc, 'legend')
1889-
if isinstance(loc, np.ndarray):
1890-
loc = loc.tolist()
1891-
1892-
# Generate panel
1893-
if loc in ('left', 'right', 'top', 'bottom'):
1894-
ax = self.panel_axes(loc, width=width, space=space, filled=True)
1895-
return ax.legend(*args, loc='fill', **kwargs)
1896-
1897-
# Fill
1898-
if loc == 'fill':
1899-
# Hide content
1900-
self._hide_panel()
1901-
1902-
# Try to make handles and stuff flush against the axes edge
1903-
kwargs.setdefault('borderaxespad', 0)
1904-
frameon = _not_none(
1905-
kwargs.get('frame', None), kwargs.get('frameon', None),
1906-
rc['legend.frameon']
1907-
)
1908-
if not frameon:
1909-
kwargs.setdefault('borderpad', 0)
1910-
1911-
# Apply legend location
1912-
side = self._panel_side
1913-
if side == 'bottom':
1914-
loc = 'upper center'
1915-
elif side == 'right':
1916-
loc = 'center left'
1917-
elif side == 'left':
1918-
loc = 'center right'
1919-
elif side == 'top':
1920-
loc = 'lower center'
1921-
else:
1922-
raise ValueError(f'Invalid panel side {side!r}.')
1923-
1924-
# Draw legend
1925-
return plot._add_legend(self, *args, loc=loc, **kwargs)
1926-
1927-
def draw(self, renderer=None, *args, **kwargs):
1928-
# Perform extra post-processing steps
1929-
self._reassign_title()
1930-
super().draw(renderer, *args, **kwargs)
1931-
19321859
def get_size_inches(self):
19331860
# Return the width and height of the axes in inches.
19341861
width, height = self.figure.get_size_inches()
@@ -2000,7 +1927,7 @@ def hexbin(self, *args, **kwargs):
20001927
"""
20011928
args = plot._parse_1d(*args)
20021929
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
2003-
mappable = super().hexbin(*args, **kwargs)
1930+
mappable = self._plot_redirect('hexbin', *args, **kwargs)
20041931
plot._auto_colorbar(mappable, **kwargs_colorbar)
20051932
return mappable
20061933

@@ -2140,7 +2067,7 @@ def imshow(self, *args, **kwargs):
21402067
%(plot.cmap_args)s
21412068
"""
21422069
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
2143-
mappable = super().imshow(*args, **kwargs)
2070+
mappable = self._plot_redirect('imshow', *args, **kwargs)
21442071
plot._auto_colorbar(mappable, **kwargs_colorbar)
21452072
return mappable
21462073

@@ -2292,6 +2219,96 @@ def indicate_inset_zoom(
22922219
self._inset_zoom_data = (rectpatch, connects)
22932220
return rectpatch, connects
22942221

2222+
@plot._concatenate_docstrings
2223+
@docstring.add_snippets
2224+
def legend(self, *args, loc=None, width=None, space=None, **kwargs):
2225+
"""
2226+
Add an *inset* legend or *outer* legend along the edge of the axes.
2227+
2228+
Parameters
2229+
----------
2230+
%(plot.legend_args)s
2231+
loc : int or str, optional
2232+
The legend location. The following location keys are valid:
2233+
2234+
.. _legend_table:
2235+
2236+
================== =======================================
2237+
Location Valid keys
2238+
================== =======================================
2239+
outer left ``'left'``, ``'l'``
2240+
outer right ``'right'``, ``'r'``
2241+
outer bottom ``'bottom'``, ``'b'``
2242+
outer top ``'top'``, ``'t'``
2243+
"best" inset ``'best'``, ``'inset'``, ``'i'``, ``0``
2244+
upper right inset ``'upper right'``, ``'ur'``, ``1``
2245+
upper left inset ``'upper left'``, ``'ul'``, ``2``
2246+
lower left inset ``'lower left'``, ``'ll'``, ``3``
2247+
lower right inset ``'lower right'``, ``'lr'``, ``4``
2248+
center left inset ``'center left'``, ``'cl'``, ``5``
2249+
center right inset ``'center right'``, ``'cr'``, ``6``
2250+
lower center inset ``'lower center'``, ``'lc'``, ``7``
2251+
upper center inset ``'upper center'``, ``'uc'``, ``8``
2252+
center inset ``'center'``, ``'c'``, ``9``
2253+
"filled" ``'fill'``
2254+
================== =======================================
2255+
2256+
width : float or str, optional
2257+
For outer legends only. The space allocated for the legend box.
2258+
This does nothing if :rcraw:`tight` is ``True``. Units are
2259+
interpreted by `~proplot.utils.units`.
2260+
space : float or str, optional
2261+
For outer legends only. The space between the axes and the legend
2262+
box. Units are interpreted by `~proplot.utils.units`.
2263+
When :rcraw:`tight` is ``True``, this is adjusted automatically.
2264+
Otherwise, the default is :rc:`subplots.panelpad`.
2265+
%(plot.legend_kwargs)s
2266+
2267+
Other parameters
2268+
----------------
2269+
**kwargs
2270+
Passed to `~matplotlib.axes.Axes.legend`.
2271+
"""
2272+
if loc != 'fill':
2273+
loc = self._loc_translate(loc, 'legend')
2274+
if isinstance(loc, np.ndarray):
2275+
loc = loc.tolist()
2276+
2277+
# Generate panel
2278+
if loc in ('left', 'right', 'top', 'bottom'):
2279+
ax = self.panel_axes(loc, width=width, space=space, filled=True)
2280+
return ax.legend(*args, loc='fill', **kwargs)
2281+
2282+
# Fill
2283+
if loc == 'fill':
2284+
# Hide content
2285+
self._hide_panel()
2286+
2287+
# Try to make handles and stuff flush against the axes edge
2288+
kwargs.setdefault('borderaxespad', 0)
2289+
frameon = _not_none(
2290+
kwargs.get('frame', None), kwargs.get('frameon', None),
2291+
rc['legend.frameon']
2292+
)
2293+
if not frameon:
2294+
kwargs.setdefault('borderpad', 0)
2295+
2296+
# Apply legend location
2297+
side = self._panel_side
2298+
if side == 'bottom':
2299+
loc = 'upper center'
2300+
elif side == 'right':
2301+
loc = 'center left'
2302+
elif side == 'left':
2303+
loc = 'center right'
2304+
elif side == 'top':
2305+
loc = 'lower center'
2306+
else:
2307+
raise ValueError(f'Invalid panel side {side!r}.')
2308+
2309+
# Draw legend
2310+
return plot._add_legend(self, *args, loc=loc, **kwargs)
2311+
22952312
@plot._concatenate_docstrings
22962313
@docstring.add_snippets
22972314
def matshow(self, *args, **kwargs):
@@ -2466,7 +2483,7 @@ def pcolor(self, *args, **kwargs):
24662483
"""
24672484
args = plot._parse_2d(*args)
24682485
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
2469-
mappable = super().pcolor(*args, **kwargs)
2486+
mappable = self._plot_redirect('pcolor', *args, **kwargs)
24702487
plot._auto_colorbar(mappable, **kwargs_colorbar)
24712488
return mappable
24722489

@@ -2504,7 +2521,7 @@ def pcolormesh(self, *args, **kwargs):
25042521
"""
25052522
args = plot._parse_2d(*args)
25062523
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
2507-
mappable = super().pcolormesh(*args, **kwargs)
2524+
mappable = self._plot_redirect('pcolormesh', *args, **kwargs)
25082525
plot._auto_colorbar(mappable, **kwargs_colorbar)
25092526
return mappable
25102527

@@ -2565,7 +2582,7 @@ def plot(self, *args, cmap=None, values=None, **kwargs):
25652582
kwargs, kwargs_legend_colorbar = plot._parse_cycle(**kwargs)
25662583

25672584
# Draw lines
2568-
objs = super().plot(x, y, values=values, **kwargs)
2585+
objs = self._plot_redirect('plot', x, y, values=values, **kwargs)
25692586

25702587
# Add sticky edges? No because there is no way to check whether "dependent
25712588
# variable" is x or y axis like with area/areax and bar/barh. Better to always
@@ -2598,7 +2615,7 @@ def quiver(self, *args, **kwargs):
25982615
"""
25992616
args = plot._parse_2d(*args)
26002617
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
2601-
mappable = super().quiver(*args, **kwargs)
2618+
mappable = self._plot_redirect('quiver', *args, **kwargs)
26022619
plot._auto_colorbar(mappable, **kwargs_colorbar)
26032620
return mappable
26042621

@@ -2725,8 +2742,8 @@ def scatter(
27252742
)
27262743

27272744
# Draw scatterplot
2728-
objs = super().scatter(
2729-
*args, c=c, s=s, cmap=cmap, norm=norm,
2745+
objs = self._plot_redirect(
2746+
'scatter', *args, c=c, s=s, cmap=cmap, norm=norm,
27302747
linewidths=lw, edgecolors=ec, **kwargs
27312748
)
27322749
if ticks is not None:
@@ -2813,7 +2830,7 @@ def streamplot(self, *args, **kwargs):
28132830
"""
28142831
args = plot._parse_2d(*args)
28152832
kwargs, kwargs_colorbar = plot._parse_cmap(**kwargs)
2816-
mappable = super().streamplot(*args, **kwargs)
2833+
mappable = self._plot_redirect('streamplot', *args, **kwargs)
28172834
plot._auto_colorbar(mappable, **kwargs_colorbar)
28182835
return mappable
28192836

0 commit comments

Comments
 (0)