add: class VrptwAcoFigure

1 files changed, 47 insertions, 0 deletions

diff --git a/vprtw_aco_figure.py b/vprtw_aco_figure.py
new file mode 100644
index 0000000..3d1a269
--- /dev/null
+++ b/vprtw_aco_figure.py
@@ -0,0 +1,47 @@
+import matplotlib.pyplot as plt
+from vrptw_base import VrptwGraph
+import time
+
+
+class VrptwAcoFigure:
+    def __init__(self, graph: VrptwGraph):
+        self.graph = graph
+        self.figure = plt.figure(figsize=(10, 10))
+        self.figure_ax = self.figure.add_subplot(1, 1, 1)
+
+        self._dot_color = 'k'
+        self._line_color_list = ['r', 'y', 'g', 'c', 'b', 'm']
+
+    def init_figure(self, path):
+        self.figure_ax.plot(list(self.graph.nodes[index].x for index in path),
+                            list(self.graph.nodes[index].y for index in path), '%s.' % self._dot_color)
+        self._draw_line(path)
+        self.figure.show()
+        time.sleep(0.2)
+
+    def update_figure(self, path):
+        for line in self.figure_ax.lines:
+            if line._color != self._dot_color:
+                self.figure_ax.lines.remove(line)
+
+        self._draw_line(path)
+
+        self.figure.canvas.draw()
+        time.sleep(0.2)
+
+    def _draw_line(self, path):
+        color_ind = 0
+        x_list = []
+        y_list = []
+        i = 0
+        while i < len(path):
+            x_list.append(self.graph.nodes[path[i]].x)
+            y_list.append(self.graph.nodes[path[i]].y)
+
+            if path[i] == 0 and len(x_list) > 1:
+                self.figure_ax.plot(x_list, y_list, '%s-' % self._line_color_list[color_ind])
+                color_ind = (color_ind + 1) % len(self._line_color_list)
+                x_list.clear()
+                y_list.clear()
+                i -= 1
+            i += 1