Let Proj accept objects, misc minor changes

lukelbd · Sep 18, 2019 · Sep 18, 2019 · Sep 18, 2019 · Sep 29, 2019 · Sep 29, 2019
commit 707d377483e7d82f6791366c422876c528270df3
diff --git a/proplot/projs.py b/proplot/projs.py
@@ -27,15 +27,18 @@
     'WinkelTripel',
     ]
 try:
-    from cartopy.crs import (_WarpedRectangularProjection,
+    from mpl_toolkits.basemap import Basemap
+except:
+    Basemap = object
+try:
+    from cartopy.crs import (CRS, _WarpedRectangularProjection,
         LambertAzimuthalEqualArea, AzimuthalEquidistant, Gnomonic)
-    _cartopy_installed = True
 except ModuleNotFoundError:
+    CRS = object
     _WarpedRectangularProjection = object
     LambertAzimuthalEqualArea = object
     AzimuthalEquidistant = object
     Gnomonic = object
-    _cartopy_installed = False
 
 def Circle(N=100):
     """Returns a circle `~matplotlib.path.Path` used as the outline
@@ -58,10 +61,13 @@ def Proj(name, basemap=False, **kwargs):
 
     Parameters
     ----------
-    name : str
-        The projection name. Like basemap, we use the PROJ.4 shorthands. The
-        following table lists the valid projection names, their full names
-        (with links to the relevant `PROJ.4 documentation
+    name : str, `~mpl_toolkits.basemap.Basemap`, or `cartopy.crs.Projection`
+        The projection name or projection class instance. If the latter, it
+        is simply returned. If the former, it must correspond to one of the
+        PROJ.4 projection name shorthands, like in basemap.
+
+        The following table lists the valid projection name shorthands, their
+        full names (with links to the relevant `PROJ.4 documentation
         <https://proj4.org/operations/projections/index.html>`__),
         and whether they are available in the cartopy and basemap packages.
         "``(added)``" indicates a cartopy projection that ProPlot has
@@ -122,8 +128,8 @@ def Proj(name, basemap=False, **kwargs):
         ====================================  ===========================================================================================  =========  =======
 
     basemap : bool, optional
-        Whether to use the basemap or cartopy package. Default is
-        ``False``.
+        Whether to use the basemap package as opposed to the cartopy package.
+        Default is ``False``.
     **kwargs
         Passed to the `~mpl_toolkits.basemap.Basemap` or
         cartopy `~cartopy.crs.Projection` class.
@@ -141,8 +147,14 @@ def Proj(name, basemap=False, **kwargs):
     --------
     `~proplot.axes.CartopyProjectionAxes`, `~proplot.axes.BasemapProjectionAxes`
     """
+    # Class instances
+    if ((CRS is not object and isinstance(name, CRS))
+        or (Basemap is not object and isinstance(name, Basemap))):
+        proj = name
+    elif not isinstance(proj, str):
+        raise ValueError(f'Unexpected Proj() argument {proj!r}. Must be name, mpl_toolkits.basemap.Basemap instance, or cartopy.crs.CRS instance.')
     # Basemap
-    if basemap:
+    elif basemap:
         import mpl_toolkits.basemap as mbasemap # verify package is available
         name = BASEMAP_TRANSLATE.get(name, name)
         kwproj = basemap_rc.get(name, {})
@@ -313,7 +325,7 @@ def __init__(self, central_longitude=0.0, globe=None):
 
 cartopy_projs = {}
 """Mapping of "projection names" to cartopy `~cartopy.crs.Projection` classes."""
-if _cartopy_installed:
+if CRS is not object:
     # Custom ones, these are always present
     import cartopy.crs as ccrs # verify package is available
     cartopy_projs = { # interpret string, create cartopy projection

diff --git a/proplot/styletools.py b/proplot/styletools.py
@@ -1429,13 +1429,15 @@ def __setitem__(self, key, item):
         a matplotlib `~matplotlib.colors.ListedColormap` or
         `~matplotlib.colors.LinearSegmentedColormap`, it is converted to the
         ProPlot `ListedColormap` or `LinearSegmentedColormap` subclass."""
-        if type(item) is mcolors.LinearSegmentedColormap:
+        if isinstance(item, (ListedColormap, LinearSegmentedColormap)):
+            pass
+        elif isinstance(item, mcolors.LinearSegmentedColormap):
             item = LinearSegmentedColormap(
                 item.name, item._segmentdata, item.N, item._gamma)
-        elif type(item) is mcolors.ListedColormap:
+        elif isinstance(item, mcolors.ListedColormap):
             item = ListedColormap(
                 item.colors, item.name, item.N)
-        elif not isinstance(item, (ListedColormap, LinearSegmentedColormap)):
+        else:
             raise ValueError(f'Invalid colormap {item}. Must be instance of matplotlib.colors.ListedColormap or matplotlib.colors.LinearSegmentedColormap.')
         key = self._sanitize_key(key, mirror=False)
         return super().__setitem__(key, item)
@@ -1664,10 +1666,10 @@ def Colormap(*args, name=None, listmode='perceptual',
         # Convert matplotlib colormaps to subclasses
         if isinstance(cmap, (ListedColormap, LinearSegmentedColormap)):
             pass
-        elif type(cmap) is mcolors.LinearSegmentedColormap:
+        elif isinstance(cmap, mcolors.LinearSegmentedColormap):
             cmap = LinearSegmentedColormap(
                 cmap.name, cmap._segmentdata, cmap.N, cmap._gamma)
-        elif type(cmap) is mcolors.ListedColormap:
+        elif isinstance(cmap, mcolors.ListedColormap):
             cmap = ListedColormap(
                 cmap.colors, cmap.name, cmap.N)
         # Dictionary of hue/sat/luminance values or 2-tuples representing linear transition

diff --git a/proplot/subplots.py b/proplot/subplots.py
@@ -323,9 +323,11 @@ def __init__(self, nrows=1, ncols=1, **kwargs):
             scales a 2-column gridspec so that the second column is twice as
             wide as the first column.
         left, right, top, bottom : float or str, optional
-            Denotes the margin *widths* in physical units. These are *not*
-            the margin coordinates. Units are interpreted by
-            `~proplot.utils.units`.
+            Denotes the margin *widths* in physical units. Units are
+            interpreted by `~proplot.utils.units`. These are *not* the
+            margin coordinates -- for example, ``left=0.1`` and ``right=0.9``
+            corresponds to a left-hand margin of 0.1 inches and a right-hand
+            margin of 0.9 inches.
         **kwargs
             Passed to `~matplotlib.gridspec.GridSpec`.
         """
@@ -385,34 +387,37 @@ def _normalize(key, size):
 
     def _spaces_as_ratios(self,
         hspace=None, wspace=None, # spacing between axes
+        hratios=None, wratios=None,
         height_ratios=None, width_ratios=None,
         **kwargs):
         """For keyword arg usage, see `GridSpec`."""
         # Parse flexible input
         nrows, ncols = self.get_active_geometry()
-        hratios = np.atleast_1d(_notNone(height_ratios, 1))
-        wratios = np.atleast_1d(_notNone(width_ratios,  1))
+        hratios = np.atleast_1d(_notNone(hratios, height_ratios, 1,
+            names=('hratios', 'height_ratios')))
+        wratios = np.atleast_1d(_notNone(wratios, width_ratios,  1,
+            names=('wratios', 'width_ratios')))
         hspace = np.atleast_1d(_notNone(hspace, np.mean(hratios)*0.10)) # this is relative to axes
         wspace = np.atleast_1d(_notNone(wspace, np.mean(wratios)*0.10))
-        if len(wspace) == 1:
-            wspace = np.repeat(wspace, (ncols-1,)) # note: may be length 0
-        if len(hspace) == 1:
-            hspace = np.repeat(hspace, (nrows-1,))
         if len(wratios) == 1:
             wratios = np.repeat(wratios, (ncols,))
         if len(hratios) == 1:
             hratios = np.repeat(hratios, (nrows,))
+        if len(wspace) == 1:
+            wspace = np.repeat(wspace, (ncols-1,)) # note: may be length 0
+        if len(hspace) == 1:
+            hspace = np.repeat(hspace, (nrows-1,))
 
         # Verify input ratios and spacings
         # Translate height/width spacings, implement as extra columns/rows
         if len(hratios) != nrows:
-            raise ValueError(f'Got {nrows} rows, but {len(hratios)} hratios.')
+            raise ValueError(f'Got {nrows} rows, but {len(hratios)} height_ratios.')
         if len(wratios) != ncols:
-            raise ValueError(f'Got {ncols} columns, but {len(wratios)} wratios.')
+            raise ValueError(f'Got {ncols} columns, but {len(wratios)} width_ratios.')
         if len(wspace) != ncols-1:
-            raise ValueError(f'Require {ncols-1} width spacings for {ncols} columns, got {len(wspace)}.')
+            raise ValueError(f'Require length {ncols-1} wspace vector for {ncols} columns, but got {len(wspace)}.')
         if len(hspace) != nrows-1:
-            raise ValueError(f'Require {nrows-1} height spacings for {nrows} rows, got {len(hspace)}.')
+            raise ValueError(f'Require length {nrows-1} hspace vector for {nrows} rows, but got {len(hspace)}.')
 
         # Assign spacing as ratios
         nrows, ncols = self.get_geometry()
@@ -426,16 +431,11 @@ def _spaces_as_ratios(self,
             hratios_final[1::2] = [*hspace]
         return wratios_final, hratios_final, kwargs # bring extra kwargs back
 
-    def add_figure(self, fig):
+    def add_figure(self, figure):
         """Adds `~matplotlib.figure.Figure` to the list of figures that are
         using this gridspec. This is done automatically when calling
         `~Figure.add_subplot` with a subplotspec generated by this gridspec."""
-        self._figures.add(fig)
-
-    def remove_figure(self, fig):
-        """Removes `~matplotlib.figure.Figure` from the list of figures that
-        are using this gridspec."""
-        self._figures.discard(fig)
+        self._figures.add(figure)
 
     def get_margins(self):
         """Returns left, bottom, right, top values. Not sure why this method
@@ -463,6 +463,16 @@ def get_active_geometry(self):
         columns that aren't skipped by `~GridSpec.__getitem__`."""
         return self._nrows_active, self._ncols_active
 
+    def remove_figure(self, figure):
+        """Removes `~matplotlib.figure.Figure` from the list of figures that
+        are using this gridspec."""
+        self._figures.discard(figure)
+
+    def tight_layout(self):
+        """Method is disabled because ProPlot has its own simplified
+        tight layout algorithm."""
+        raise RuntimeError(f'Native matplotlib tight layout is disabled.')
+
     def update(self, figure=None, **kwargs):
         """
         Updates the width and height ratios, gridspec margins, and spacing
@@ -491,12 +501,12 @@ def update(self, figure=None, **kwargs):
             raise ValueError(f'Unknown keyword arg(s): {kwargs}.')
 
         # Apply to figure(s) and all axes
-        for fig in self._figures:
-            fig.subplotpars.update(self.left, self.bottom, self.right, self.top)
-            for ax in fig.axes:
+        for figure in self._figures:
+            figure.subplotpars.update(self.left, self.bottom, self.right, self.top)
+            for ax in figure.axes:
                 ax.update_params()
                 ax.set_position(ax.figbox)
-            fig.stale = True
+            figure.stale = True
 
 #-----------------------------------------------------------------------------#
 # Helper funcs
@@ -545,6 +555,7 @@ def __init__(self,
         axwidth=None, axheight=None, aspect=1,
         tight=None, pad=None, axpad=None, panelpad=None,
         left=None, right=None, bottom=None, top=None,
+        wspace=None, hspace=None,
         share=None, sharex=None, sharey=None,
         span=None, spanx=None, spany=None,
         align=None, alignx=None, aligny=None,
@@ -714,10 +725,15 @@ def __init__(self,
             spec = ', '.join(spec)
             names = ('axwidth', 'axheight')
             values = (axwidth, axheight)
-        # Raise warning
+        # Save dimensions
         for name,value in zip(names,values):
             if value is not None:
                 warnings.warn(f'You specified both {spec} and {name}={value!r}. Ignoring {name!r}.')
+        width, height = units(width), units(height)
+        axwidth, axheight = units(axwidth), units(axheight)
+        self._width, self._height = width, height
+        self._ref_width, self._ref_height = axwidth, axheight
+
         # Input border spacing
         left, right = units(left), units(right)
         bottom, top = units(bottom), units(top)
@@ -729,11 +745,7 @@ def __init__(self,
         right  = _notNone(right, units(rc['subplots.innerspace']))
         top    = _notNone(top, units(rc['subplots.titlespace']))
         bottom = _notNone(bottom, units(rc['subplots.xlabspace']))
-        width, height = units(width), units(height)
-        axwidth, axheight = units(axwidth), units(axheight)
         self._ref_aspect = aspect
-        self._width, self._height = width, height
-        self._ref_width, self._ref_height = axwidth, axheight
         self._left, self._right = left, right
         self._bottom, self._top = bottom, top
         self._wspace, self._hspace = wspace, hspace
@@ -938,25 +950,20 @@ def _adjust_tight_layout(self, renderer):
         instance."""
         # Initial stuff
         axs = self._iter_axes()
+        pad = self._pad
         obox = self.bbox_inches # original bbox
         bbox = self.get_tightbbox(renderer)
         gridspec = self._gridspec
         if not axs or not gridspec:
             return
 
         # Tight box *around* figure
-        # Get bounds from old bounding box
-        pad = self._pad
-        left = bbox.xmin
-        bottom = bbox.ymin
-        right = obox.xmax - bbox.xmax
-        top = obox.ymax - bbox.ymax
         # Apply new bounds, permitting user overrides
         # TODO: Account for bounding box NaNs?
-        self._left = _notNone(self._left_orig, self._left - offset + pad)
-        self._right = _notNone(self._right_orig, self._right - offset + pad)
-        self._bottom = _notNone(self._bottom_orig, self._bottom - offset + pad)
-        self._top = _notNone(self._top_orig, self._top - offset + pad)
+        self._left = _notNone(self._left_orig, self._left - bbox.xmin + pad)
+        self._right = _notNone(self._right_orig, self._right - bbox.ymin + pad)
+        self._bottom = _notNone(self._bottom_orig, self._bottom - (obox.xmax - bbox.xmax) + pad)
+        self._top = _notNone(self._top_orig, self._top - (obox.ymax - bbox.ymax) + pad)
 
         # Get arrays storing gridspec spacing args
         axpad = self._axpad
@@ -1577,9 +1584,10 @@ def add_subplot(self, *args,
               or divide the "main" gridspec geometry.
             * A tuple indicating (nrows, ncols, index). Geometry must be
               equivalent to or divide the "main" gridspec geometry.
-        proj, projection : str, optional
-            The registered matplotlib projection name, or a basemap or cartopy
-            map projection name. Passed to `~proplot.projs.Proj`. See
+        proj, projection : str, `~mpl_toolkits.basemap.Basemap`, or `~cartopy.crs.CRS`, optional
+            A registered matplotlib projection name, a basemap or cartopy
+            map projection name, a `~mpl_toolkits.basemap.Basemap` instance, or
+            a `~cartpoy.crs.CRS` instance. Passed to `~proplot.projs.Proj`. See
             `~proplot.projs.Proj` for a table of map projection names.
         proj_kw, projection_kw : dict-like, optional
             Dictionary of keyword args for the projection class. Passed to
@@ -2142,7 +2150,7 @@ def subplots(array=None, ncols=1, nrows=1,
     # TODO: Fix xref and yref so update geometry gets them automatically
     # from the reference axes! Then user can also change the reference axes
     # which is pretty neat!
-    fig = plt.figure(FigureClass=Figure, figsize=figsize, ref=ref, **kwargs)
+    fig = plt.figure(FigureClass=Figure, ref=ref, **kwargs)
     fig._wspace_orig = wspace_orig
     fig._hspace_orig = hspace_orig
     gridspec = fig._update_geometry(nrows=nrows, ncols=ncols,