Update dataloader.py

dataloader.py

@@ -1,46 +1,39 @@
 import torch
 import torch.nn.functional as F
 from torch.utils.data import Dataset
-from sklearn.decomposition import TruncatedSVD
+import pandas as pd
 import numpy as np
-from random import shuffle
 import os
-import pandas as pd
 
-class VRP_Dataset(Dataset):
 
+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
+        # Load CSV data
         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)
+        if len(raw_data) < dataset_size * num_nodes:
+            raise ValueError("Not enough rows in CSV to build required dataset")
 
-        # Assign node positions
+        # Randomly sample and reshape
+        coords = torch.tensor(raw_data[['longitude', 'latitude']].values[:dataset_size * num_nodes], dtype=torch.float32)
         node_positions = coords.view(dataset_size, num_nodes, 2)
         self.node_positions = node_positions
 
-        # Generate fleet data
+        # 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
+            'car_start_node': car_start_node,
         }
 
-        # Generate graph data
+        # 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)
@@ -63,13 +56,13 @@ class VRP_Dataset(Dataset):
     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
+        distance = torch.sqrt(((x - y) ** 2).sum(dim=3))
         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
+        graph = {key: self.graph_data[key][idx].unsqueeze(0).to(self.device) for key in self.graph_data}
+        fleet = {key: self.fleet_data[key][idx].unsqueeze(0).to(self.device) for key in self.fleet_data}
+        return graph, fleet
 
     def __len__(self):
         return self.dataset_size