Spaces:
Running
Running
File size: 2,086 Bytes
719d0db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import torch.nn as nn
from models.solvers.ortools.ortools import ORTools
from models.solvers.lkh.lkh import LKH
from models.solvers.concorde.concorde import ConcordeTSP
class GeneralSolver(nn.Module):
def __init__(self, problem, solver_type, large_value=1e+6, scaling=True):
super().__init__()
self.problem = problem
self.large_value = large_value
self.scaling = scaling
self.solver_type = solver_type
supported_problem = {
"ortools": ["tsp", "tsptw", "pctsp", "pctsptw", "cvrp", "cvrptw"],
"lkh": ["tsp", "tsptw", "cvrp", "cvrptw"],
"concorde": ["tsp"]
}
# validate solver_type & problem
assert solver_type in supported_problem.keys(), f"Invalid solver type: {solver_type}. Please select from {supported_problem.keys()}"
assert problem in supported_problem[solver_type], f"{solver_type} does not support {problem}."
self.solver = self.get_solver(problem, solver_type)
def change_solver(self, problem, solver_type):
if self.solver_type != solver_type or self.problem != problem:
self.problem = problem
self.solver_type = solver_type
self.solver = self.get_solver(problem, solver_type)
def get_solver(self, problem, solver_type):
if solver_type == "ortools":
return ORTools(problem, self.large_value, self.scaling)
elif solver_type == "lkh":
return LKH(problem, self.large_value, self.scaling)
elif solver_type == "concorde":
assert problem == "tsp", "Concorde solver supports only TSP."
return ConcordeTSP(self.large_value, self.scaling)
else:
assert False, f"Invalid solver type: {solver_type}"
def solve(self, node_feats, fixed_paths=None, dist_matrix=None, instance_name=None):
if isinstance(self.solver, ORTools):
return self.solver.solve(node_feats, fixed_paths, dist_matrix, instance_name)
else:
return self.solver.solve(node_feats, fixed_paths, instance_name) |