optimisation_models.py

import math
import torch
import gurobipy as gp
from gurobipy import GRB

from pyepo.model.grb.grbmodel import optGrbModel

from ortools.constraint_solver import routing_enums_pb2
from ortools.constraint_solver import pywrapcp
from pyepo.model.opt import optModel
from pyepo import EPO

import numpy as np


def compute_euclidean_distance_matrix(locations):
    """Creates callback to return distance between points."""
    distances = {}
    for from_counter, from_node in enumerate(locations):
        distances[from_counter] = {}
        for to_counter, to_node in enumerate(locations):
            if from_counter == to_counter:
                distances[from_counter][to_counter] = 0
            else:
                # Euclidean distance
                distances[from_counter][to_counter] = int(
                    math.hypot((from_node[0] - to_node[0]), (from_node[1] - to_node[1]))
                )
                # Manhattan distance
                distances[from_counter][to_counter] = int(
                    abs(from_node[0] - to_node[0]) + abs(from_node[1] - to_node[1])
                )
    return distances

# Create a console solution printer.
def print_solution(manager, routing, assignment):
    """Prints assignment on console."""
    output = {}
    output["Objective"] = assignment.ObjectiveValue()
#     print(f'Objective: {assignment.ObjectiveValue()}')
   # Display dropped nodes.
#     dropped_nodes = 'Dropped nodes:'
#     for index in range(routing.Size()):
#         if routing.IsStart(index) or routing.IsEnd(index):
#             continue
#         if assignment.Value(routing.NextVar(index)) == index:
#             node = manager.IndexToNode(index)
#             dropped_nodes += f' {node}({VISIT_VALUES[node]})'
#     print(dropped_nodes)
    # Display routes
    index = routing.Start(0)
#     plan_output = 'Route for vehicle 0:\n'
    route_distance = 0
    value_collected = 0
    nodes_visited = []
    while not routing.IsEnd(index):
        node = manager.IndexToNode(index)
        value_collected += routing.GetDisjunctionPenalty(node)
        nodes_visited.append(node)
        previous_index = index
        index = assignment.Value(routing.NextVar(index))
        route_distance += routing.GetArcCostForVehicle(previous_index, index, 0)

#     plan_output += f' {manager.IndexToNode(index)}\n'
#     plan_output += f'Distance of the route: {route_distance}m\n'
#     plan_output += f'Value collected: {value_collected}/{sum(VISIT_VALUES)}\n'
    output['Distance of the route'] = route_distance
    output['Nodes visited'] = nodes_visited
    output['Value collected'] = value_collected
    # output['Proportion collected'] = (value_collected,sum(VISIT_VALUES))

    
    return output


def construct_prize_collecting_tsp(k, h):
    """Entry point of the program.
    https://github.com/google/or-tools/blob/stable/examples/python/prize_collecting_tsp.py
    
    k (int): size of grid
    h (int): limit on distance travelled
    VALUE_MATRIX (kxk array of (int)): reward for visiting node
    """
    nodes = list(range(k*k))
    node_index_dct = {node: (node // k,node % k) for node in nodes}
    DISTANCE_MATRIX = compute_euclidean_distance_matrix(list(node_index_dct.values()))
    MAX_DISTANCE = int(h)
    
    num_nodes = len(DISTANCE_MATRIX)
    num_vehicles = 1
    depot = 0
    all_nodes = range(num_nodes)

    # Create the routing index manager.
    manager = pywrapcp.RoutingIndexManager(
            num_nodes,
            num_vehicles,
            depot)

    # Create routing model.
    routing = pywrapcp.RoutingModel(manager)

    # Create and register a transit callback.
    def distance_callback(from_index, to_index):
        """Returns the distance between the two nodes."""
        # Convert from routing variable Index to distance matrix NodeIndex.
        from_node = manager.IndexToNode(from_index)
        to_node = manager.IndexToNode(to_index)
        return DISTANCE_MATRIX[from_node][to_node]

    transit_callback_index = routing.RegisterTransitCallback(distance_callback)

    # Define cost of each arc.
    routing.SetArcCostEvaluatorOfAllVehicles(transit_callback_index)

    # Limit Vehicle distance.
    dimension_name = 'Distance'
    routing.AddDimension(
        transit_callback_index,
        0,  # no slack
        MAX_DISTANCE,  # vehicle maximum travel distance
        True,  # start cumul to zero
        dimension_name)
    #distance_dimension = routing.GetDimensionOrDie(dimension_name)
    #distance_dimension.SetGlobalSpanCostCoefficient(100)


        
    return manager, routing

def set_objective_prize_collecting_tsp(manager, routing, VISIT_VALUES, num_nodes):
    # Allow to drop nodes.
    for node in range(0, int(num_nodes)):
        routing.AddDisjunction(
                [manager.NodeToIndex(node)],
                int(VISIT_VALUES[node]))
        
def solve_prize_collecting_tsp(manager, routing, seconds_per_solve = 10):

    # Setting first solution heuristic.
    search_parameters = pywrapcp.DefaultRoutingSearchParameters()
    
    search_parameters.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PARALLEL_CHEAPEST_INSERTION)
    search_parameters.local_search_metaheuristic = (
        routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH)
    
    search_parameters.time_limit.FromSeconds(seconds_per_solve)
    
    search_parameters.savings_parallel_routes = True

#     search_parameters.log_search = True

    # Solve the problem.
    assignment = routing.SolveWithParameters(search_parameters)

    # Print solution on console.
    solution_dct = print_solution(manager, routing, assignment)
    return solution_dct

class PrizeCollectingTSP(optModel):

    def __init__(self, k, h, seconds_per_solve = 10):
        """
        k (int): size of grid
        h (int): limit on distance travelled
        """
        self.k = k
        self.h = h
        self.sps = seconds_per_solve
        super().__init__()

    def _getModel(self):
        """
        A method to build model

        Returns:
            tuple: optimization model and variables
        """
        # build ortools model
        self.modelSense = EPO.MAXIMIZE
        
        return construct_prize_collecting_tsp, np.zeros(self.k*self.k)
    
    def num_cost(self):
        return int(self.k**2)


    def setObj(self, c):
        """
        c (tensor): vector of benefits associated with stacked rows of grid
        """
        num_nodes = self.num_cost()

        manager, routing = construct_prize_collecting_tsp(k = self.k, h = self.h)

        self.manager = manager
        self.routing = routing

        if not torch.is_tensor(c):
            c = torch.tensor(c) 
        c = 1000 * c# make it pop, scaling does not affect optimal solution but makes it work in integer context of ortools
        c = torch.nn.ReLU()(torch.round(c)) + self.h # So the routing does not interfere with the prize collection
        set_objective_prize_collecting_tsp(manager = self.manager, 
                                           routing = self.routing, 
                                           VISIT_VALUES = c, 
                                           num_nodes = num_nodes)



    def solve(self):
        """
        A method to solve model

        Returns:
            tuple: optimal solution (list) and objective value (float)
        """
        solution = solve_prize_collecting_tsp(manager = self.manager, 
                                              routing = self.routing,
                                              seconds_per_solve= self.sps
                                              )
        
        nodes_visited = np.array([1 if i in solution['Nodes visited'] else 0 for i in range(self.num_cost())])
        return nodes_visited, solution['Value collected']

class shortestPathModel(optGrbModel):
    """
    This class is optimization model for shortest path problem on 2D grid with 8 neighbors

    Attributes:
        _model (GurobiPy model): Gurobi model
        grid (tuple of int): Size of grid network
        nodes (list): list of vertex
        edges (list): List of arcs
        nodes_map (ndarray): 2D array for node index
    """

    def __init__(self, grid):
        """
        Args:
            grid (tuple of int): size of grid network
        """
        self.grid = grid
        self.nodes, self.edges, self.nodes_map = self._getEdges()
        super().__init__()

    def _getEdges(self):
        """
        A method to get list of edges for grid network

        Returns:
            list: arcs
        """
        # init list
        nodes, edges = [], []
        # init map from coord to ind
        nodes_map = {}
        for i in range(self.grid[0]):
            for j in range(self.grid[1]):
                u = self._calNode(i, j)
                nodes_map[u] = (i,j)
                nodes.append(u)
                # edge to 8 neighbors
                # up
                if i != 0:
                    v = self._calNode(i-1, j)
                    edges.append((u,v))
                    # up-right
                    if j != self.grid[1] - 1:
                        v = self._calNode(i-1, j+1)
                        edges.append((u,v))
                # right
                if j != self.grid[1] - 1:
                    v = self._calNode(i, j+1)
                    edges.append((u,v))
                    # down-right
                    if i != self.grid[0] - 1:
                        v = self._calNode(i+1, j+1)
                        edges.append((u,v))
                # down
                if i != self.grid[0] - 1:
                    v = self._calNode(i+1, j)
                    edges.append((u,v))
                    # down-left
                    if j != 0:
                        v = self._calNode(i+1, j-1)
                        edges.append((u,v))
                # left
                if j != 0:
                    v = self._calNode(i, j-1)
                    edges.append((u,v))
                    # top-left
                    if i != 0:
                        v = self._calNode(i-1, j-1)
                        edges.append((u,v))
        return nodes, edges, nodes_map

    def _calNode(self, x, y):
        """
        A method to calculate index of node
        """
        v = x * self.grid[1] + y
        return v

    def _getModel(self):
        """
        A method to build Gurobi model

        Returns:
            tuple: optimization model and variables
        """
        # ceate a model
        m = gp.Model("shortest path")
        # varibles
        x = m.addVars(self.edges, ub=1, name="x")
        # sense
        m.modelSense = GRB.MINIMIZE
        # constraints
        for i in range(self.grid[0]):
            for j in range(self.grid[1]):
                v = self._calNode(i, j)
                expr = 0
                for e in self.edges:
                    # flow in
                    if v == e[1]:
                        expr += x[e]
                    # flow out
                    elif v == e[0]:
                        expr -= x[e]
                # source
                if i == 0 and j == 0:
                    m.addConstr(expr == -1)
                # sink
                elif i == self.grid[0] - 1 and j == self.grid[0] - 1:
                    m.addConstr(expr == 1)
                # transition
                else:
                    m.addConstr(expr == 0)
        return m, x

    def setObj(self, c):
        """
        A method to set objective function

        Args:
            c (np.ndarray): cost of objective function
        """
        # vector to matrix
        c = c.reshape(self.grid)
        # sum up vector cost
        obj = float(c[0,0]) + gp.quicksum(c[self.nodes_map[j]] * self.x[i,j] for i, j in self.x)
        self._model.setObjective(obj)

    def solve(self):
        """
        A method to solve model

        Returns:
            tuple: optimal solution (list) and objective value (float)
        """
        # update gurobi model
        self._model.update()
        # solve
        self._model.optimize()
        # kxk solution map
        sol = np.zeros(self.grid)
        for i, j in self.edges:
            # active edge
            if abs(1 - self.x[i,j].x) < 1e-3:
                # node on active edge
                sol[self.nodes_map[i]] = 1
                sol[self.nodes_map[j]] = 1
        # matrix to vector
        sol = sol.reshape(-1)
        return sol, self._model.objVal