Changing line color in Line2d - Python GIF-CodePudding

I created an algorithm to solve the TSP problem. I want to present the result in the form of a GIF that shows the distance traveled. I use line2d from the Matplotlib library for this purpose. Every time the car reaches the starting point, I want to change the color of the line. I am currently using the line2d.set_color () functions for this. Unfortunately, it causes a color change of the entire road. My goal is to change the color at any given moment while keeping the previous colors. Can you help me with this problem?

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import random
from gurobipy import Model, GRB, quicksum
from matplotlib import collections  as mc

rnd = np.random
rnd.seed()

n = 20  # numbre of clients
xc = rnd.rand(n   1) * 200
yc = rnd.rand(n   1) * 100
xc[0] = 100
yc[0] = 50

N = [i for i in range(1, n   1)]
V = [0]   N
A = [(i, j) for i in V for j in V if i != j]
c = {(i, j): np.hypot(xc[i] - xc[j], yc[i] - yc[j]) for i, j in A}
Q = 30
# q = {i: random.choice([-5, -4, -3, -2, -1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) for i in N}
# q = {i: random.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) for i in N}
q = {i: rnd.randint(1, 10) for i in N}

# a = [['x', xc], ['y', yc], ['Capacity', q]]
# for i in range(len(a)) :
#     for j in range(len(a[i])):
#         print(a[i][j], end=" ")
#     print()

mdl = Model('CVRP')

x = mdl.addVars(A, vtype=GRB.BINARY)
u = mdl.addVars(N, vtype=GRB.CONTINUOUS)

mdl.modelSense = GRB.MINIMIZE
mdl.setObjective(quicksum(x[i, j] * c[i, j] for i, j in A))

mdl.addConstrs(quicksum(x[i, j] for j in V if j != i) == 1 for i in N)
mdl.addConstrs(quicksum(x[i, j] for i in V if i != j) == 1 for j in N)
mdl.addConstrs((x[i, j] == 1) >> (u[i]   q[j] == u[j])
               for i, j in A if i != 0 and j != 0)
mdl.addConstrs(u[i] >= q[i] for i in N)
mdl.addConstrs(u[i] <= Q for i in N)

mdl.Params.MIPGap = 0.1
mdl.Params.TimeLimit = 30  # seconds
mdl.optimize()

active_arcs = [a for a in A if x[a].x > 0.99]
final_arcs = []
roads = 0
while len(active_arcs) != 0:
    temp_arcs = [active_arcs[0]]
    for i, j in temp_arcs:
        temp_j = j

    i2 = 0
    while temp_j != 0:
        for i, j in active_arcs:
            if temp_j == i and temp_j != 0:
                temp_arcs.append(active_arcs[i2])
                temp_j = j
            i2  = 1
        i2 = 0

    i = 0
    i2 = 0
    while i < len(temp_arcs):
        while i2 < len(active_arcs):
            if temp_arcs[i] == active_arcs[i2]:
                active_arcs.pop(i2)
            i2  = 1
        i  = 1
        i2 = 0

    final_i = 0
    while final_i < len(temp_arcs):
        final_arcs.append(temp_arcs[final_i])
        final_i  = 1
    temp_arcs.clear()
    roads  = 1

plt.plot(xc[0], yc[0], c='r', marker='s')
for i in N:
    if q[i] > 0:
        plt.scatter(xc[i], yc[i], c='b')
    else:
        plt.scatter(xc[i], yc[i], c='y')

plt.show()
print('Liczba wykonanych tras: ', roads)

# for i, j in final_arcs:
#     plt.plot([xc[i], xc[j]], [yc[i], yc[j]], color, zorder=0)
#     if j == 0:
#         if len(colors_base) != 0:
#             c = random.randint(0, len(colors_base) - 1)
#             color = colors_base[c]
#             colors_base.pop(c)
#
# plt.show()


fig, ax = plt.subplots()
ax.set_xlim(0, 200)
ax.set_ylim(0, 100)
x_data, y_data = [], []
ln1, ln2 = [], []
ln, = plt.plot([], [], zorder=0)
colors_base = ['r', 'cyan', 'black', 'orange', 'm', 'darkorchid', 'dimgray', 'indigo', 'lawngreen', 'darkred']
ln.set_color('g')

plt.plot(xc[0], yc[0], 'r', marker='s')
for i in N:
    if q[i] > 0:
        plt.scatter(xc[i], yc[i], c='b')
    else:
        plt.scatter(xc[i], yc[i], c='y')

tmp = 0
for i, j in final_arcs:
    x_data_tmp = [xc[i], xc[j]]
    y_data_tmp = [yc[i], yc[j]]
    x_data.append(x_data_tmp)
    y_data.append(y_data_tmp)


def update(p):
    ln1.append(x_data[p])
    ln2.append(y_data[p])
    ln.set_data(ln1, ln2)
    for i, j in final_arcs:
        if yc[i] == 50 and y_data[p] == [yc[i], yc[j]] and p != 0:
            if len(colors_base) != 0:
                color = colors_base[0]
                ln.set_color(color)
                colors_base.pop(0)
    return ln,


ani = FuncAnimation(fig, update, frames=len(final_arcs), interval=1000)

ani.save('animation.gif', fps=3)
# ani.save('animation.mp4', fps=3)

plt.show()

This is how the received gif looks like:

enter image description here

CodePudding user response：

One possible solution:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import random
from gurobipy import Model, GRB, quicksum
from matplotlib import collections  as mc

rnd = np.random
rnd.seed()

n = 20  # numbre of clients
xc = rnd.rand(n   1) * 200
yc = rnd.rand(n   1) * 100
xc[0] = 100
yc[0] = 50

N = [i for i in range(1, n   1)]
V = [0]   N
A = [(i, j) for i in V for j in V if i != j]
c = {(i, j): np.hypot(xc[i] - xc[j], yc[i] - yc[j]) for i, j in A}
Q = 30
# q = {i: random.choice([-5, -4, -3, -2, -1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) for i in N}
# q = {i: random.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) for i in N}
q = {i: rnd.randint(1, 10) for i in N}


mdl = Model('CVRP')

x = mdl.addVars(A, vtype=GRB.BINARY)
u = mdl.addVars(N, vtype=GRB.CONTINUOUS)

mdl.modelSense = GRB.MINIMIZE
mdl.setObjective(quicksum(x[i, j] * c[i, j] for i, j in A))

mdl.addConstrs(quicksum(x[i, j] for j in V if j != i) == 1 for i in N)
mdl.addConstrs(quicksum(x[i, j] for i in V if i != j) == 1 for j in N)
mdl.addConstrs((x[i, j] == 1) >> (u[i]   q[j] == u[j])
               for i, j in A if i != 0 and j != 0)
mdl.addConstrs(u[i] >= q[i] for i in N)
mdl.addConstrs(u[i] <= Q for i in N)

mdl.Params.MIPGap = 0.1
mdl.Params.TimeLimit = 30  # seconds
mdl.optimize()

active_arcs = [a for a in A if x[a].x > 0.99]
final_arcs = []
roads = 0
while len(active_arcs) != 0:
    temp_arcs = [active_arcs[0]]
    for i, j in temp_arcs:
        temp_j = j

    i2 = 0
    while temp_j != 0:
        for i, j in active_arcs:
            if temp_j == i and temp_j != 0:
                temp_arcs.append(active_arcs[i2])
                temp_j = j
            i2  = 1
        i2 = 0

    i = 0
    i2 = 0
    while i < len(temp_arcs):
        while i2 < len(active_arcs):
            if temp_arcs[i] == active_arcs[i2]:
                active_arcs.pop(i2)
            i2  = 1
        i  = 1
        i2 = 0

    final_i = 0
    while final_i < len(temp_arcs):
        final_arcs.append(temp_arcs[final_i])
        final_i  = 1
    temp_arcs.clear()
    roads  = 1



fig, ax = plt.subplots()
ax.set_xlim(0, 200)
ax.set_ylim(0, 100)
x_data, y_data = [], []
ln1, ln2 = [], []
ln, = ax.plot([], [], zorder=0)
colors_base = ['r', 'cyan', 'black', 'orange', 'm', 'darkorchid', 'dimgray', 'indigo', 'lawngreen', 'darkred']
ln.set_color('g')

plt.plot(xc[0], yc[0], 'r', marker='s')
for i in N:
    if q[i] > 0:
        plt.scatter(xc[i], yc[i], c='b')
    else:
        plt.scatter(xc[i], yc[i], c='y')

# precompute each line coordinates
lines_data = []
xd, yd = [], []
for i, j in final_arcs:
    xd.append(xc[i])
    yd.append(yc[i])
    if j == 0:
        xd.append(xc[0])
        yd.append(yc[0])
        lines_data.append([xd, yd])
        xd, yd = [], []

# add empty lines with the specified colors
lines = []
for i in range(len(lines_data)):
    lines.append(ax.plot([], [], color=colors_base[i])[0]) 


def update(p):
    global current_line, up_to_point
    
    # reset after each animation cycle is complete
    if current_line > len(lines_data) - 1:
        current_line = 0
        for l in lines:
            l.set_data([], [])
    
    up_to_point  = 1
    xd, yd = lines_data[current_line]
    lines[current_line].set_data(xd[:up_to_point], yd[:up_to_point])
    if len(xd[:up_to_point]) == len(xd):
        current_line  = 1
        up_to_point = 1

# index of the line to animate
current_line = 0
# draw the current line up to the point of this index
up_to_point = 1

ani = FuncAnimation(fig, update, frames=len(final_arcs), interval=1000)

ani.save('animation.gif', fps=3)
# ani.save('animation.mp4', fps=3)

plt.show()