Matplotlib in wxPython - Fixing Equal Axis problem for 3D and 2D plots

The matplotlib plot I embedded in wxpython frame, did not have equal axis, making a circle appear as an ellipse.

I wrote a function Plot_Equal_Axis() which fixed the problem of getting an equal axis for 3D plotting (here, I got the help) and pyplot function, axis('equal') works fine for 2D plotting. The entire source code of the entire work:

1:  import wx  
2:  import matplotlib.pyplot as plt  
3:  from mpl_toolkits.mplot3d import Axes3D  
4:  from matplotlib.backends.backend_wxagg import FigureCanvasWxAgg as FigureCanvas  
5:  from matplotlib.backends.backend_wxagg import NavigationToolbar2WxAgg  
6:  import numpy as np  
7:    
8:  # INTPUTS  
9:  dimension = 'three'  
10:  task = 'tracking'  
11:  if dimension == 'three'and task == 'tracking':  
12:    n_plot_inputs = 6  
13:  elif dimension == 'two' and task == 'tracking':  
14:    n_plot_inputs = 4  
15:  elif dimension == 'three' and task == 'goal':  
16:    n_plot_inputs = 3  
17:  elif dimension == 'two' and task == 'goal':  
18:    n_plot_inputs = 2  
19:    
20:  x_UAV = [2,3,4,5,6]  
21:  y_UAV = [3,4,2,6,4]  
22:  z_UAV = [7,3,2,2,6]  
23:    
24:  x_target = [6,4,3,4,3]  
25:  y_target = [6,2,6,2,6]  
26:  z_target = [6,2,2,6,2]  
27:    
28:    
29:  class basic_computation(object):  
30:    def __init__(self):  
31:      return None  
32:    
33:    def plot_axis_equal_input(self):  
34:      if n_plot_inputs == 3:  
35:        return [x_UAV, y_UAV, z_UAV]  
36:      elif n_plot_inputs == 6:  
37:        (x,y,z) = (x_UAV, y_UAV, z_UAV)  
38:        x.extend(x_target)  
39:        y.extend(y_target)  
40:        z.extend(z_target)  
41:        return [x, y, z]  
42:    
43:  # Make the axis of equal size  
44:  class Plot_Equal_Axis(object):  
45:    def __init__(self, axes):  
46:      p = basic_computation()  
47:      point = p.plot_axis_equal_input()  
48:      self.X = np.array(point[0])  
49:      self.Y = np.array(point[1])  
50:      self.Z = np.array(point[2])  
51:      max_range = self.max_range()  
52:      mean_x = self.X.mean()  
53:      mean_y = self.Y.mean()  
54:      mean_z = self.Z.mean()  
55:      axes.set_xlim(mean_x - max_range, mean_x + max_range)  
56:      axes.set_ylim(mean_y - max_range, mean_y + max_range)  
57:      axes.set_zlim(mean_z - max_range, mean_z + max_range)  
58:    
59:    def max_range(self):  
60:      return np.array([self.X.max()-self.X.min(), self.Y.max()-self.Y.min(), self.Z.max()-self.Z.min()]).max() / 2.0  
61:    
62:  class Plotting(object):  
63:    def __init__(self, axes, type, *args):  
64:      x = args[0]  
65:      y = args[1]  
66:      if dimension == 'two':  
67:        axes.plot(x,y,type)  
68:      else:  
69:        z = args[2]  
70:        axes.plot(np.array(x), np.array(y),np.array(z),type)  
71:    
72:    
73:  class PlotFrame(wx.Frame):  
74:    def __init__(self):  
75:      wx.Frame.__init__(self, None, -1, 'Simulation', (756,100), (600,532))  
76:      self.figure = plt.figure()  
77:      plt.ion()  
78:      if dimension == 'two':  
79:        self.axes = plt.axes()  
80:      else:  
81:        self.axes = plt.axes(projection = '3d')  
82:        self.axes.set_zlabel('z axis')  
83:      plot_UAV = Plotting(self.axes, '-.g', x_UAV, y_UAV, z_UAV)  
84:      if task == 'tracking':  
85:        plot_target = Plotting(self.axes, '-or', x_target, y_target, z_target)  
86:      self.axes.margins(0.05)  
87:      self.axes.set_xlabel('x axis')  
88:      self.axes.set_ylabel('y axis')  
89:      self.canvas = FigureCanvas(self, -1, self.figure)  
90:      self.sizer = wx.BoxSizer(wx.VERTICAL)  
91:      self.sizer.Add(self.canvas, 1, wx.LEFT | wx.TOP | wx.GROW)  
92:      self.SetSizer(self.sizer)  
93:      self.add_toolbar()  
94:      if dimension == 'three':  
95:        plot_axis_equal = Plot_Equal_Axis(self.axes)  
96:      else:  
97:        plt.axis('equal')  
98:    
99:    def add_toolbar(self):  
100:      self.toolbar = NavigationToolbar2WxAgg(self.canvas)  
101:      self.toolbar.Realize()  
102:      if wx.Platform == '__WXMAC__':  
103:        self.SetToolBar(self.toolbar)  
104:      else:  
105:        tw, th = self.toolbar.GetSizeTuple()  
106:        fw, fh = self.canvas.GetSizeTuple()  
107:        self.toolbar.SetSize(wx.Size(fw, th))  
108:        self.sizer.Add(self.toolbar, 0, wx.LEFT | wx.EXPAND)  
109:      self.toolbar.update()  
110:    
111:    
112:    
113:  class App(wx.App):  
114:    def OnInit(self):  
115:      self.plotframe = PlotFrame()  
116:      self.plotframe.Center()  
117:      self.plotframe.Show()  
118:      return True  
119:    
120:    
121:  if __name__ == '__main__':  
122:    app = App(False)  
123:    app.MainLoop()