Ticket #14141: trac_14141-plot-edges-fs.patch

sage/combinat/knutson_tao_puzzles.py

@@ -327 +327 @@
         else:
             return Graphics()
 
-    def _plot_piece(self, coords, border_color=(0.5,0.5,0.5), border_thickness=1):
+    def _plot_piece(self, coords, border_color=(0.5,0.5,0.5), border_thickness=1, style='fill'):
         """
         TESTS::
 
@@ -335 +335 @@
             sage: delta = DeltaPiece('2','K','2')
             sage: delta._plot_piece([(1,1),(1,2),(2,2)]) # not tested
         """
-        P = polygon(coords, color=self.color())
-        P += polygon(coords, fill=False, color=border_color, thickness=border_thickness)
-        return P
+        if style == 'fill':
+            P = polygon(coords, color=self.color())
+            P += polygon(coords, fill=False, color=border_color, thickness=border_thickness)
+            return P
+        elif style == 'edges':
+            if isinstance(self, DeltaPiece):
+                edges = ('north_west', 'south', 'north_east')
+            elif isinstance(self, NablaPiece):
+                edges = ('south_west', 'north', 'south_east')
+            else:
+                edges = self.edges()
+            P = Graphics()
+            for (i, edge) in enumerate(edges):
+                P += line([coords[i], coords[(i+1)%3]], color=self.edge_color(edge), thickness=border_thickness)
+            return P
+        else:
+            return NotImplemented
+
+    def edge_color(self, edge):
+        """
+        Color of the specified edge of ``self`` (to be used when plotting the
+        piece).
+
+        EXAMPLES::
+
+            sage: from sage.combinat.knutson_tao_puzzles import DeltaPiece
+            sage: delta = DeltaPiece('1','0','10')
+            sage: delta.edge_color('south')
+            'blue'
+            sage: delta.edge_color('north_west')
+            'red'
+            sage: delta.edge_color('north_east')
+            (1, 0.65, 0)
+        """
+        edge_label = self.edge_label(edge)
+        colors = {'1': 'blue', '0': 'red'}
+        if edge_label in colors:
+            color = colors[edge_label]
+        elif 'K' in edge_label:
+            color = 'orange'
+        elif edge_label.startswith('T'):
+            color = 'yellow'
+        else:
+            color = 'white'
+        return color
 
     def edge_label(self, edge):
         """
@@ -1371 +1413 @@
             for k in range(d + 1):
                 yield self[k + 1, self._n - d + k]
 
-    def plot(self, labels=True, engine="fill triangles"):
+    def plot(self, labels=True, style="fill"):
         """
         Plots completed puzzle.
 
@@ -1380 +1422 @@
             sage: ps = PuzzleSolver("H")
             sage: puzzle = ps('0101','1001')[0]
             sage: puzzle.plot()  #not tested
-        """
-        if engine == "fill triangles":
-            return self._plot_fill_triangles(labels=labels)
-        elif engine == "color edges":
-            return self._plot_color_edges(labels=labels)
-
-    def _plot_fill_triangles(self, labels=True):
-        """
-        Plots completed puzzle by filling triangles.
-
-        EXAMPLES::
-
-            sage: ps = PuzzleSolver("H")
-            sage: puzzle = ps('0101','1001')[0]
-            sage: puzzle._plot_fill_triangles()  #not tested
+            sage: puzzle.plot(style='fill')  #not tested
+            sage: puzzle.plot(style='edges')  #not tested
         """
         P = Graphics()
         coords = [(k, -d) for d in range(self._n) for k in range(-d, d + 1, 2)]
         for ((k, d), piece) in zip(coords, self):
             if isinstance(piece, RhombusPiece):
                 for (i, triangle) in enumerate(piece):
-                    P += triangle._plot_piece([(k, d - 2*i), (k - 1, d - 1), (k + 1, d - 1)])
+                    P += triangle._plot_piece([(k, d - 2*i), (k - 1, d - 1), (k + 1, d - 1)], style=style)
                 if labels:
                     P += piece._plot_label(piece['north_west'], (k - 0.5, d - 0.5), rotation=60)
                     P += piece._plot_label(piece['north_east'], (k + 0.5, d - 0.5), rotation=-60)
                     P += piece._plot_label(piece.north_piece()['south'], (k, d - 1))
             else:
-                P += piece._plot_piece([(k, d), (k - 1, d - 1), (k + 1, d - 1)])
+                P += piece._plot_piece([(k, d), (k - 1, d - 1), (k + 1, d - 1)], style=style)
                 if labels:
                     P += piece._plot_label(piece['north_west'], (k - 0.5, d - 0.5), rotation=60)
                     P += piece._plot_label(piece['north_east'], (k + 0.5, d - 0.5), rotation=-60)
@@ -1416 +1445 @@
         P.axes(False)
         return P
 
-    # FIXME: I don't think this method survived the refactoring
-    def _plot_color_edges(self, labels=True):
-        P = Graphics()
-        text_color = (0, 0, 0)
-        colors = {10: (1, 0.65, 0), 1: 'blue', 0: 'red'}
-        line_thickness = 3
-        coords = [(k, -d) for d in range(self._n) for k in range(-d, d + 1, 2)]
-        for ((k, d), piece) in zip(coords, self):
-            if isinstance(piece, RhombusPiece):
-                triangle = piece.north_piece()
-                P += line([(k, d), (k - 1, d - 1)], color=colors[triangle['north_west']], thickness=line_thickness)
-                P += line([(k - 1, d - 1), (k + 1, d - 1)], color=colors[triangle['south']], thickness=line_thickness)
-                P += line([(k + 1, d - 1), (k, d)], color=colors[triangle['north_east']], thickness=line_thickness)
-
-                triangle = piece.south_piece()
-                P += line([(k, d - 2), (k - 1, d - 1)], color=colors[triangle['south_west']], thickness=line_thickness)
-                P += line([(k - 1, d - 1), (k + 1, d - 1)], color=colors[triangle['north']], thickness=line_thickness)
-                P += line([(k + 1, d - 1), (k, d - 2)], color=colors[triangle['south_east']], thickness=line_thickness)
-
-                if labels:
-                    P += text("%s" % piece['north_west'], (k - 0.5, d - 0.5), color=text_color, rotation=60)
-                    P += text("%s" % piece['north_east'], (k + 0.5, d - 0.5), color=text_color, rotation=-60)
-                    P += text("%s" % piece.north_piece()['south'], (k, d - 1), color=text_color)
-            else:
-                P += line([(k, d), (k - 1, d - 1)], color=colors[piece['north_west']], thickness=line_thickness)
-                P += line([(k - 1, d - 1), (k + 1, d - 1)], color=colors[piece['south']], thickness=line_thickness)
-                P += line([(k + 1, d - 1), (k, d)], color=colors[piece['north_east']], thickness=line_thickness)
-                if labels:
-                    P += text("%s" % piece['north_west'], (k - 0.5, d - 0.5), color=text_color, rotation=60)
-                    P += text("%s" % piece['north_east'], (k + 0.5, d - 0.5), color=text_color, rotation=-60)
-                    P += text("%s" % piece['south'], (k, d - 1), color=text_color)
-
-        P.set_aspect_ratio(1.73)
-        P.axes(False)
-        return P
-
     def _latex_(self):
         r"""
         Return latex version of ``self``.