com2014-template/main.py

99 lines
2.8 KiB
Python

import os
import argparse
import json
import matplotlib
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from utils import visualize_tsp
from model.anneal_model import SimAnneal
from model.my_model import MyModel
def write_to_file(solution, fitness):
with open('output/' + os.path.splitext(os.path.basename(args.tsp_file))[0] + '.txt', "w") as outfile:
outfile.write(", ".join(str(item) for item in solution))
outfile.write("\n")
outfile.write(str(fitness))
def plot_learning(fitness_list):
"""
Plot the fitness through iterations.
"""
plt.plot([i for i in range(len(fitness_list))], fitness_list)
plt.ylabel("Fitness")
plt.xlabel("Iteration")
plt.show()
def load_data(file):
coords = []
with open(file, "r") as infile:
line = infile.readline()
# Skip instance header
while "NODE_COORD_SECTION" not in line:
line = infile.readline()
for line in infile.readlines():
line = line.replace("\n", "")
if line and 'EOF' not in line:
line = [float(x) for x in line.lstrip().split(" ", 1)[1].split(" ") if x]
coords.append(line)
return coords
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='TSP Solver')
parser.add_argument(
'tsp_file',
type=str,
help='Path to tsp file.'
)
parser.add_argument(
'--plot',
action="store_true",
help='Plot results'
)
args = parser.parse_args()
coords = load_data(args.tsp_file)
# model = MyModel(coords)
# best_solution, fitness_list = model.fit(max_it=100000)
sa = SimAnneal(coords)
best_solution, fitness_list = sa.fit(max_it=100000)
if best_solution:
if args.plot:
visualize_tsp.plotTSP([best_solution], coords)
if fitness_list:
plot_learning(fitness_list)
print("Writing the best solution to file" )
write_to_file(best_solution, sa.fitness(best_solution))
data = {}
data['nodes'] = []
for i in range(len(coords)):
data['nodes'].append({
'title': str(i),
'id': i,
'x': int(coords[i][0]),
'y': int(coords[i][1])
})
data['edges'] = []
for i in range(len(best_solution)):
if i == len(best_solution)-1:
data['edges'].append({
'source': best_solution[i],
'target': best_solution[0]
})
else:
data['edges'].append({
'source': best_solution[i],
'target': best_solution[i+1]
})
with open('output/' + os.path.splitext(os.path.basename(args.tsp_file))[0] + '.json', 'w') as outfile:
json.dump(data, outfile)