Merge pull request #4310 from insertroar/master

tacaswell · tacaswell · commit 42a852db5bcc · 2015-06-15T22:02:03.000-04:00
ENH: 'square' argument to `Axes.axis`
diff --git a/doc/users/whats_new/square_plot_feature.rst b/doc/users/whats_new/square_plot_feature.rst
@@ -0,0 +1,11 @@
+Square Plot 
+-------------------------
+
+Square plot:
+
+Implemented square plots feature as a new parameter in the axis function. When argument 'square' is specified, equal scaling is set, and the limits are set such that xmax-xmin == ymax-ymin.  
+
+Example:
+
+``ax.axis('square')``
+
diff --git a/lib/matplotlib/axes/_base.py b/lib/matplotlib/axes/_base.py
@@ -1301,26 +1301,67 @@ def apply_aspect(self, position=None):
                 self.set_xbound((x0, x1))
 
     def axis(self, *v, **kwargs):
-        """
-        Convenience method for manipulating the x and y view limits
-        and the aspect ratio of the plot. For details, see
-        :func:`~matplotlib.pyplot.axis`.
+        """Set axis properties.
+
+        Valid signatures::
+
+          xmin, xmax, ymin, ymax = axis()
+          xmin, xmax, ymin, ymax = axis(list_arg)
+          xmin, xmax, ymin, ymax = axis(string_arg)
+          xmin, xmax, ymin, ymax = axis(**kwargs)
+
+        Parameters
+        ----------
+        v : list of float or {'on', 'off', 'equal', 'tight', 'scaled',\
+            'normal', 'auto', 'image', 'square'}
+            Optional positional argument
+
+            Axis data limits set from a list; or a command relating to axes:
+
+                ========== ================================================
+                Value      Description
+                ========== ================================================
+                'on'       Toggle axis lines and labels on
+                'off'      Toggle axis lines and labels off
+                'equal'    Equal scaling by changing limits
+                'scaled'   Equal scaling by changing box dimensions
+                'tight'    Limits set such that all data is shown
+                'auto'     Automatic scaling, fill rectangle with data
+                'normal'   Same as 'auto'; deprecated
+                'image'    'scaled' with axis limits equal to data limits
+                'square'   Square plot; similar to 'scaled', but initially\
+                           forcing xmax-xmin = ymax-ymin
+                ========== ================================================
+
+        emit : bool, optional
+            Passed to set_{x,y}lim functions, if observers
+            are notified of axis limit change
+
+        xmin, ymin, xmax, ymax : float, optional
+            The axis limits to be set
+
+        Returns
+        -------
+        xmin, xmax, ymin, ymax : float
+            The axis limits
 
-        *kwargs* are passed on to :meth:`set_xlim` and
-        :meth:`set_ylim`
         """
+
         if len(v) == 0 and len(kwargs) == 0:
             xmin, xmax = self.get_xlim()
             ymin, ymax = self.get_ylim()
             return xmin, xmax, ymin, ymax
 
+        emit = kwargs.get('emit', True)
+
         if len(v) == 1 and is_string_like(v[0]):
             s = v[0].lower()
             if s == 'on':
                 self.set_axis_on()
             elif s == 'off':
                 self.set_axis_off()
-            elif s in ('equal', 'tight', 'scaled', 'normal', 'auto', 'image'):
+            elif s in ('equal', 'tight', 'scaled', 'normal',
+                       'auto', 'image', 'square'):
                 self.set_autoscale_on(True)
                 self.set_aspect('auto')
                 self.autoscale_view(tight=False)
@@ -1337,15 +1378,23 @@ def axis(self, *v, **kwargs):
                     self.autoscale_view(tight=True)
                     self.set_autoscale_on(False)
                     self.set_aspect('equal', adjustable='box', anchor='C')
-
+                elif s == 'square':
+                    self.set_aspect('equal', adjustable='box', anchor='C')
+                    self.set_autoscale_on(False)
+                    xlim = self.get_xlim()
+                    ylim = self.get_ylim()
+                    edge_size = max(np.diff(xlim), np.diff(ylim))
+                    self.set_xlim([xlim[0], xlim[0] + edge_size],
+                                  emit=emit, auto=False)
+                    self.set_ylim([ylim[0], ylim[0] + edge_size],
+                                  emit=emit, auto=False)
             else:
                 raise ValueError('Unrecognized string %s to axis; '
                                  'try on or off' % s)
             xmin, xmax = self.get_xlim()
             ymin, ymax = self.get_ylim()
             return xmin, xmax, ymin, ymax
 
-        emit = kwargs.get('emit', True)
         try:
             v[0]
         except IndexError:
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -1480,6 +1480,12 @@ def axis(*v, **kwargs):
     as kwargs selectively to alter just those limits without changing
     the others.
 
+      >>> axis('square')
+
+    changes the limit ranges (*xmax*-*xmin*) and (*ymax*-*ymin*) of
+    the *x* and *y* axes to be the same, and have the same scaling,
+    resulting in a square plot.
+
     The xmin, xmax, ymin, ymax tuple is returned
 
     .. seealso::
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3679,6 +3679,16 @@ def test_bar_negative_width():
         assert_equal(b._width, 1)
         assert_equal(b._height, indx + 1)
 
+@cleanup
+def test_square_plot():
+    x = np.arange(4)
+    y = np.array([1., 3., 5., 7.])
+    fig, ax = plt.subplots()
+    ax.plot(x, y, 'mo')
+    ax.axis('square')
+    xlim, ylim = ax.get_xlim(), ax.get_ylim()
+    assert_true(np.diff(xlim) == np.diff(ylim))
+    assert_true(ax.get_aspect() == 'equal')
 
 @cleanup
 def test_no_None():
