File size: 3,615 Bytes
c389cb0 |
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 83 84 85 86 87 88 89 90 91 92 93 94 95 |
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from sklearn.decomposition import TruncatedSVD
import numpy as np
from random import shuffle
import os
import pandas as pd
class VRP_Dataset(Dataset):
def __init__(self, dataset_size, num_nodes, num_depots, dataset_path, device='cpu', *args, **kwargs):
super().__init__()
self.device = device
self.dataset_size = dataset_size
self.num_nodes = num_nodes
self.num_depots = num_depots
# Load external CSV data from Hugging Face
raw_data = pd.read_csv(dataset_path)
if len(raw_data) < dataset_size:
raise ValueError("Dataset size requested exceeds available data")
sampled_data = raw_data.sample(n=dataset_size, random_state=42).reset_index(drop=True)
# Extract coordinates (assuming columns named 'longitude', 'latitude')
coords = torch.tensor(sampled_data[['longitude', 'latitude']].values, dtype=torch.float32)
# Assign node positions
node_positions = coords.view(dataset_size, num_nodes, 2)
self.node_positions = node_positions
# Generate fleet data
num_cars = num_nodes
launch_time = torch.zeros(dataset_size, num_cars, 1)
car_start_node = torch.randint(low=0, high=num_depots, size=(dataset_size, num_cars, 1))
self.fleet_data = {
'start_time': launch_time,
'car_start_node': car_start_node
}
# Generate graph data
a = torch.arange(num_nodes).reshape(1, 1, -1).repeat(dataset_size, num_cars, 1)
b = car_start_node.repeat(1, 1, num_nodes)
depot = ((a == b).sum(dim=1) > 0).float().unsqueeze(2)
start_times = (torch.rand(dataset_size, num_nodes, 1) * 2 + 3) * (1 - depot)
end_times = start_times + (0.1 + 0.5 * torch.rand(dataset_size, num_nodes, 1)) * (1 - depot)
distance_matrix = self.compute_distance_matrix(node_positions)
time_matrix = distance_matrix.clone()
self.graph_data = {
'start_times': start_times,
'end_times': end_times,
'depot': depot,
'node_vector': node_positions,
'distance_matrix': distance_matrix,
'time_matrix': time_matrix
}
def compute_distance_matrix(self, node_positions):
x = node_positions.unsqueeze(1).repeat(1, self.num_nodes, 1, 1)
y = node_positions.unsqueeze(2).repeat(1, 1, self.num_nodes, 1)
distance = (((x - y) ** 2).sum(dim=3)) ** 0.5
return distance
def __getitem__(self, idx):
A = {key: self.graph_data[key][idx].unsqueeze(0).to(self.device) for key in self.graph_data}
B = {key: self.fleet_data[key][idx].unsqueeze(0).to(self.device) for key in self.fleet_data}
return A, B
def __len__(self):
return self.dataset_size
def collate(self, batch):
graph_data = {key: torch.cat([item[0][key] for item in batch], dim=0) for key in self.graph_data}
fleet_data = {key: torch.cat([item[1][key] for item in batch], dim=0) for key in self.fleet_data}
return graph_data, fleet_data
def get_batch(self, idx, batch_size=10):
return self.collate([self.__getitem__(i) for i in range(idx, idx + batch_size)])
def get_data(self):
return self.graph_data, self.fleet_data
def model_input_length(self):
return 3 + self.graph_data['node_vector'].shape[2]
def save_data(self, fp):
data = (self.graph_data, self.fleet_data)
with open(fp, 'wb') as f:
torch.save(data, f)
|