More fun with plots

[linpy.git] / pypol / polyhedra.py

diff --git a/pypol/polyhedra.py b/pypol/polyhedra.py
index 5d1bfa1..a5048d1 100644 (file)
--- a/pypol/polyhedra.py
+++ b/pypol/polyhedra.py
@@ -290,79 +290,66 @@ class Polyhedron(Domain):
             faces.append(face)
         return faces
 
-    def plot(self):
-        """
-        Display 3D plot of set. 
-        """
+    def _plot_2d(self, plot=None, **kwargs):
+        from matplotlib import pylab
         import matplotlib.pyplot as plt
-        import matplotlib.patches as patches
-
-        if len(self.symbols)> 3:
-            raise TypeError
-
-        elif len(self.symbols) == 2:
-            import pylab
-            points = []  
-            for verts in self.vertices():
-                    pairs=()
-                    for coordinate, point in verts.coordinates():
-                        pairs = pairs + (float(point),)
-                    points.append(pairs)
-            cent=(sum([p[0] for p in points])/len(points),sum([p[1] for p in points])/len(points))
-            points.sort(key=lambda p: math.atan2(p[1]-cent[1],p[0]-cent[0]))
-            pylab.scatter([p[0] for p in points],[p[1] for p in points])
-            pylab.gca().add_patch(patches.Polygon(points,closed=True,fill=True))
-            pylab.grid()
-            pylab.show()
-
-        elif len(self.symbols)==3:
-            from mpl_toolkits.mplot3d import Axes3D
-            from mpl_toolkits.mplot3d.art3d import Poly3DCollection
-            faces = self.faces()
+        from matplotlib.axes import Axes
+        from matplotlib.patches import Polygon
+        vertices = self._sort_polygon_2d(self.vertices())
+        xys = [tuple(vertex.values()) for vertex in vertices]
+        if plot is None:
             fig = plt.figure()
-            ax = Axes3D(fig)
-            for face in faces:
-                points = []
-                vertices = Polyhedron._sort_polygon_3d(face)
-                for verts in vertices:
-                    pairs=()
-                    for coordinate, point in verts.coordinates():
-                        pairs = pairs + (float(point),)
-                    points.append(pairs)
-                collection = Poly3DCollection([points], alpha=0.7)
-                face_color = [0.5, 0.5, 1] # alternative: matplotlib.colors.rgb2hex([0.5, 0.5, 1])
-                collection.set_facecolor(face_color)
-                ax.add_collection3d(collection)
-            ax.set_xlabel('X')   
-            ax.set_xlim(0, 5)
-            ax.set_ylabel('Y')
-            ax.set_ylim(0, 5)
-            ax.set_zlabel('Z')
-            ax.set_zlim(0, 5)
-            plt.grid()      
-            plt.show()
-        return points
-    
-    @classmethod
-    def limit(cls, faces, variable, lim):
-        sym = []
-        if variable is 'x':
-            n = 0
-        elif variable is 'y':
-            n = 1
-        elif variable is 'z':
-            n = 2
-        for face in faces:
-            for vert in face:
-                coordinates = vert.coordinates()
-                for point in enumerate(coordinates):
-                        coordinates.get(n)
-                        sym.append(points)
-        if lim == 0:
-            value = min(sym)
+            plot = fig.add_subplot(1, 1, 1)
+            xs, ys = zip(*xys)
+            plot.set_xlim(float(min(xs)), float(max(xs)))
+            plot.set_ylim(float(min(ys)), float(max(ys)))
+        plot.add_patch(Polygon(xys, closed=True, **kwargs))
+        return plot
+
+    def _plot_3d(self, plot=None, **kwargs):
+        import matplotlib.pyplot as plt
+        from mpl_toolkits.mplot3d import Axes3D
+        from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+        if plot is None:
+            fig = plt.figure()
+            axes = Axes3D(fig)
+            xmin, xmax = float('inf'), float('-inf')
+            ymin, ymax = float('inf'), float('-inf')
+            zmin, zmax = float('inf'), float('-inf')
         else:
-            value = max(sym)
-        return value
+            axes = plot
+        poly_xyzs = []
+        for vertices in self.faces():
+            if len(vertices) == 0:
+                continue
+            vertices = Polyhedron._sort_polygon_3d(vertices)
+            vertices.append(vertices[0])
+            face_xyzs = [tuple(vertex.values()) for vertex in vertices]
+            if plot is None:
+                xs, ys, zs = zip(*face_xyzs)
+                xmin, xmax = min(xmin, float(min(xs))), max(xmax, float(max(xs)))
+                ymin, ymax = min(ymin, float(min(ys))), max(ymax, float(max(ys)))
+                zmin, zmax = min(zmin, float(min(zs))), max(zmax, float(max(zs)))
+            poly_xyzs.append(face_xyzs)
+        collection = Poly3DCollection(poly_xyzs, **kwargs)
+        axes.add_collection3d(collection)
+        if plot is None:
+            axes.set_xlim(xmin, xmax)
+            axes.set_ylim(ymin, ymax)
+            axes.set_zlim(zmin, zmax)
+        return axes
+
+    def plot(self, plot=None, **kwargs):
+        """
+        Display 3D plot of set.
+        """
+        if self.dimension == 2:
+            return self._plot_2d(plot=plot, **kwargs)
+        elif self.dimension == 3:
+            return self._plot_3d(plot=plot, **kwargs)
+        else:
+            raise ValueError('polyhedron must be 2 or 3-dimensional')
+
 
 def _polymorphic(func):
     @functools.wraps(func)