Merge pull request #159 from lukelbd/dataframe-standardize_2d

Fix standardize_2d issue with pandas dataframes

Author: lukelbd

proplot/axes/base.py

@@ -1297,20 +1297,37 @@ def get_tightbbox(self, renderer, *args, **kwargs):
         self._tightbbox = bbox
         return bbox
 
-    def heatmap(self, *args, **kwargs):
+    def heatmap(self, *args, aspect=None, **kwargs):
         """
         Pass all arguments to `~matplotlib.axes.Axes.pcolormesh` then apply
-        settings that are suitable for heatmaps: no gridlines, no minor ticks,
-        and major ticks at the center of each grid box.
+        settings that are suitable for heatmaps: square grid boxes by default,
+        major ticks at the center of each grid box, no minor ticks, and no gridlines.
+
+        Parameters
+        ----------
+        aspect : {'equal', 'auto'} or float, optional
+            Controls the aspect ratio of the axes. The aspect is of particular
+            relevance for heatmaps since it may distort the heatmap, i.e. a grid box
+            will not be square. This parameter is a shortcut for explicitly calling
+            `~matplotlib.axes.set_aspect`.
+
+            The default is :rc:`image.heatmap`. The options are:
+
+            - ``'equal'``: Ensures an aspect ratio of 1. Grid boxes will be square.
+            - ``'auto'``: The axes is kept fixed and the aspect is adjusted so
+              that the data fit in the axes. In general, this will result in non-square
+              grid boxes.
         """
         obj = self.pcolormesh(*args, **kwargs)
+        aspect = _not_none(aspect, rc['image.aspect'])
         xlocator, ylocator = None, None
         if hasattr(obj, '_coordinates'):
             coords = obj._coordinates
             coords = (coords[1:, ...] + coords[:-1, ...]) / 2
             coords = (coords[:, 1:, :] + coords[:, :-1, :]) / 2
             xlocator, ylocator = coords[0, :, 0], coords[:, 0, 1]
         self.format(
+            aspect=aspect,
             xgrid=False, ygrid=False, xtickminor=False, ytickminor=False,
             xlocator=xlocator, ylocator=ylocator,
         )

proplot/wrappers.py

@@ -467,25 +467,29 @@ def standardize_2d(self, func, *args, order='C', globe=False, **kwargs):
 
     * If *x* and *y* or *latitude* and *longitude* coordinates were not
       provided, and a `~pandas.DataFrame` or `~xarray.DataArray` is passed, we
-      try to infer them from the metadata. Otherwise,
-      ``np.arange(0, data.shape[0])`` and ``np.arange(0, data.shape[1])``
-      are used.
+      try to infer them from the metadata. Otherwise, ``np.arange(0, data.shape[0])``
+      and ``np.arange(0, data.shape[1])`` are used.
     * For ``pcolor`` and ``pcolormesh``, coordinate *edges* are calculated
       if *centers* were provided. For all other methods, coordinate *centers*
       are calculated if *edges* were provided.
 
-    For `~proplot.axes.CartopyAxes` and `~proplot.axes.BasemapAxes`, the
-    `globe` keyword arg is added, suitable for plotting datasets with global
-    coverage. Passing ``globe=True`` does the following:
-
-    1. "Interpolates" input data to the North and South poles.
-    2. Makes meridional coverage "circular", i.e. the last longitude coordinate
-       equals the first longitude coordinate plus 360\N{DEGREE SIGN}.
-
-    For `~proplot.axes.BasemapAxes`, 1d longitude vectors are also cycled to
-    fit within the map edges. For example, if the projection central longitude
-    is 90\N{DEGREE SIGN}, the data is shifted so that it spans
-    -90\N{DEGREE SIGN} to 270\N{DEGREE SIGN}.
+    Parameters
+    ----------
+    order : {'C', 'F'}, optional
+        If ``'C'``, arrays should be shaped as ``(y, x)``. If ``'F'``, arrays
+        should be shaped as ``(x, y)``. Default is ``'C'``.
+    globe : bool, optional
+        Whether to ensure global coverage for `~proplot.axes.GeoAxes` plots.
+        Default is ``False``. When set to ``True`` this does the following:
+
+        #. Interpolates input data to the North and South poles by setting the data
+           values at the poles to the mean from latitudes nearest each pole.
+        #. Makes meridional coverage "circular", i.e. the last longitude coordinate
+           equals the first longitude coordinate plus 360\N{DEGREE SIGN}.
+        #. For `~proplot.axes.BasemapAxes`, 1D longitude vectors are also cycled to
+           fit within the map edges. For example, if the projection central longitude
+           is 90\N{DEGREE SIGN}, the data is shifted so that it spans -90\N{DEGREE SIGN}
+           to 270\N{DEGREE SIGN}.
 
     See also
     --------
@@ -532,8 +536,12 @@ def standardize_2d(self, func, *args, order='C', globe=False, **kwargs):
             x = Z.coords[Z.dims[idx]]
             y = Z.coords[Z.dims[idy]]
         else:  # DataFrame; never Series or Index because these are 1d
-            x = Z.index
-            y = Z.columns
+            if order == 'C':
+                x = Z.columns
+                y = Z.index
+            else:
+                x = Z.index
+                y = Z.columns
 
     # Check coordinates
     x, y = _to_array(x), _to_array(y)