diff --git a/traveling-salesman-portrait/traveling-salesman-portrait.py b/traveling-salesman-portrait/traveling-salesman-portrait.py
index 14f86cc..77a0fe4 100644
--- a/traveling-salesman-portrait/traveling-salesman-portrait.py
+++ b/traveling-salesman-portrait/traveling-salesman-portrait.py
@@ -39,11 +39,15 @@
 
 optimized_path_points = [chosen_black_indices[x] for x in optimized_path]
 
-plt.figure(figsize=(8, 10), dpi=100)
-plt.plot([x[1] for x in optimized_path_points], [x[0] for x in optimized_path_points], color='black', lw=1)
-plt.xlim(0, 600)
-plt.ylim(0, 800)
-plt.gca().invert_yaxis()
-plt.xticks([])
-plt.yticks([])
-plt.savefig('traveling-salesman-portrait.png', bbox_inches='tight')
+def draw_plot(optimized_path_points):
+    """Draw the plot using matplotlib
+    params: optimized_path_points"""
+    plt.figure(figsize=(8, 10), dpi=100)
+    plt.plot([x[1] for x in optimized_path_points], [x[0] for x in optimized_path_points], color='black', lw=1)
+    plt.xlim(0, 600)
+    plt.ylim(0, 800)
+    plt.gca().invert_yaxis()
+    plt.xticks([])
+    plt.yticks([])
+    plt.savefig('traveling-salesman-portrait.png', bbox_inches='tight')
+