File size: 3,327 Bytes
a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce a3d45a8 f2809ce |
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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import torch
class Graph(object):
def __init__(self, graph_data, device='cpu'):
self.device = device
# Load and move data to the correct device
self.start_time = graph_data['start_times'].to(device)
self.end_time = graph_data['end_times'].to(device)
self.depot = graph_data['depot'].to(device)
self.node_positions = graph_data['node_vector'].to(device)
self.distance_matrix = graph_data['distance_matrix'].to(device)
self.time_matrix = graph_data['time_matrix'].to(device)
self.num_nodes = self.distance_matrix.shape[1]
self.batch_size = self.distance_matrix.shape[0]
self.correct_depot_features()
self.time_window_compatibility = self.compute_time_window_compatibility()
self.max_dist = self.distance_matrix.max()
self.max_drive_time = self.time_matrix.max()
def correct_depot_features(self):
# Remove time windows for depot nodes
self.start_time = self.start_time.clone() * (1 - self.depot)
self.end_time = self.end_time.clone() * (1 - self.depot)
def construct_vector(self):
# Concatenate node features for model input
L = [self.node_positions, self.start_time, self.end_time, self.depot]
self.vector = torch.cat(L, dim=2)
return self.vector
def get_drive_times(self, from_nodes, to_nodes):
"""
Get drive times between from_nodes and to_nodes.
"""
num_elements = from_nodes.shape[1]
assert num_elements == to_nodes.shape[1]
# Extract relevant entries from the time matrix
ind_1 = from_nodes.unsqueeze(2).expand(-1, -1, self.num_nodes)
dist = torch.gather(self.time_matrix, dim=1, index=ind_1)
ind_2 = to_nodes.unsqueeze(2)
drive_times = torch.gather(dist, dim=2, index=ind_2)
return drive_times
def get_distances(self, from_nodes, to_nodes):
"""
Get Euclidean distances between from_nodes and to_nodes.
"""
num_elements = from_nodes.shape[1]
assert num_elements == to_nodes.shape[1]
# Extract relevant entries from the distance matrix
ind_1 = from_nodes.unsqueeze(2).expand(-1, -1, self.num_nodes)
dist = torch.gather(self.distance_matrix, dim=1, index=ind_1)
ind_2 = to_nodes.unsqueeze(2)
distances = torch.gather(dist, dim=2, index=ind_2)
return distances
def compute_time_window_compatibility(self):
"""
Determine if traveling from node i to node j respects time window constraints:
i.e., start_time[i] + drive_time[i][j] <= end_time[j]
"""
x = self.start_time.reshape(self.batch_size, self.num_nodes, 1).repeat(1, 1, self.num_nodes)
y = self.end_time.reshape(self.batch_size, 1, self.num_nodes).repeat(1, self.num_nodes, 1)
time_mask = (x + self.time_matrix <= y).float()
return time_mask
def to(self, device):
"""
Move all tensors to the specified device.
"""
self.device = device
for attr in ['start_time', 'end_time', 'depot', 'node_positions', 'distance_matrix', 'time_matrix']:
setattr(self, attr, getattr(self, attr).to(device))
self.time_window_compatibility = self.time_window_compatibility.to(device)
|