Merge pull request #68 from stefraynaud/master

lukelbd · web-flow · commit 3224bbddabf9 · 2019-11-27T01:26:42.000-07:00
Add support for 2D coordinates
diff --git a/proplot/utils.py b/proplot/utils.py
@@ -14,7 +14,7 @@
     from icecream import ic
 except ImportError:  # graceful fallback if IceCream isn't installed
     ic = lambda *a: None if not a else (a[0] if len(a) == 1 else a) # noqa
-__all__ = ['arange', 'edges', 'units']
+__all__ = ['arange', 'edges', 'edges2d', 'units']
 
 # Change this to turn on benchmarking
 BENCHMARK = False
@@ -153,22 +153,50 @@ def edges(array, axis=-1):
     array = np.array(array)
     array = np.swapaxes(array, axis, -1)
     # Next operate
-    flip = False
-    idxs = [[0] for _ in range(array.ndim-1)] # must be list because we use it twice
-    if array[np.ix_(*idxs, [1])] < array[np.ix_(*idxs, [0])]:
-        flip = True
-        array = np.flip(array, axis=-1)
     array = np.concatenate((
         array[...,:1]  - (array[...,1]-array[...,0])/2,
         (array[...,1:] + array[...,:-1])/2,
         array[...,-1:] + (array[...,-1]-array[...,-2])/2,
         ), axis=-1)
-    if flip:
-        array = np.flip(array, axis=-1)
     # Permute back and return
     array = np.swapaxes(array, axis, -1)
     return array
 
+def edges2d(z):
+    """
+    Like :func:`edges` but for 2D arrays.
+    The size of both axes are increased of one.
+
+    Parameters
+    ----------
+    array : array-like
+        Two-dimensional array.
+
+    Returns
+    -------
+    `~numpy.ndarray`
+        Array of "edge" coordinates.
+    """
+    z = np.asarray(z)
+    ny, nx = z.shape
+    zzb = np.zeros((ny+1, nx+1))
+
+    # Inner
+    zzb[1:-1, 1:-1] = 0.25 * (z[1:, 1:] + z[:-1, 1:] +
+                              z[1:, :-1] + z[:-1, :-1])
+    # Lower and upper
+    zzb[0] += edges(1.5*z[0]-0.5*z[1])
+    zzb[-1] += edges(1.5*z[-1]-0.5*z[-2])
+
+    # Left and right
+    zzb[:, 0] += edges(1.5*z[:, 0]-0.5*z[:, 1])
+    zzb[:, -1] += edges(1.5*z[:, -1]-0.5*z[:, -2])
+
+    # Corners
+    zzb[[0, 0, -1, -1], [0, -1, -1, 0]] *= 0.5
+
+    return zzb
+
 def units(value, numeric='in'):
     """
     Flexible units -- this function is used internally all over ProPlot, so
diff --git a/proplot/wrappers.py b/proplot/wrappers.py
@@ -473,6 +473,13 @@ def standardize_2d(self, func, *args, order='C', globe=False, **kwargs):
                 if all(z.ndim == 1 and z.size > 1 and z.dtype != 'object' for z in (x,y)):
                     x = utils.edges(x)
                     y = utils.edges(y)
+                else:
+                    if (x.ndim == 2 and x.shape[0] > 1 and x.shape[1] > 1 and
+                            x.dtype != 'object'):
+                        x = utils.edges2d(x)
+                    if (y.ndim == 2 and y.shape[0] > 1 and y.shape[1] > 1 and
+                            y.dtype != 'object'):
+                        y = utils.edges2d(y)
             elif Z.shape[1] != xlen-1 or Z.shape[0] != ylen-1:
                 raise ValueError(f'Input shapes x {x.shape} and y {y.shape} must match Z centers {Z.shape} or Z borders {tuple(i+1 for i in Z.shape)}.')
         # Optionally re-order
@@ -491,10 +498,19 @@ def standardize_2d(self, func, *args, order='C', globe=False, **kwargs):
         for Z in Zs:
             if Z.ndim != 2:
                 raise ValueError(f'Input arrays must be 2D, instead got shape {Z.shape}.')
-            elif Z.shape[1] == xlen-1 and Z.shape[0] == ylen-1 and x.ndim == 1 and y.ndim == 1:
+            elif Z.shape[1] == xlen-1 and Z.shape[0] == ylen-1:
                 if all(z.ndim == 1 and z.size > 1 and z.dtype != 'object' for z in (x,y)):
                     x = (x[1:] + x[:-1])/2
                     y = (y[1:] + y[:-1])/2
+                else:
+                    if (x.ndim == 2 and x.shape[0] > 1 and x.shape[1] > 1 and
+                            x.dtype != 'object'):
+                        x = 0.25 * (x[:-1, :-1] + x[:-1, 1:] +
+                                    x[1:, :-1] + x[1:, 1:])
+                    if (y.ndim == 2 and y.shape[0] > 1 and y.shape[1] > 1 and
+                            y.dtype != 'object'):
+                        y = 0.25 * (y[:-1, :-1] + y[:-1, 1:] +
+                                    y[1:, :-1] + y[1:, 1:])
             elif Z.shape[1] != xlen or Z.shape[0] != ylen:
                 raise ValueError(f'Input shapes x {x.shape} and y {y.shape} must match Z centers {Z.shape} or Z borders {tuple(i+1 for i in Z.shape)}.')
         # Optionally re-order
