def test_radius_graph(): assert RadiusGraph(1).__repr__() == 'RadiusGraph(r=1)' pos = [[0, 0], [1, 0], [2, 0], [0, 1], [-2, 0], [0, -2]] pos = torch.tensor(pos, dtype=torch.float) data = Data(pos=pos) edge_index = RadiusGraph(1)(data).edge_index assert edge_index.tolist() == [[0, 0, 1, 1, 2, 3], [1, 3, 0, 2, 1, 0]]
def test_radius_graph(): assert RadiusGraph(r=1).__repr__() == 'RadiusGraph(r=1)' pos = torch.Tensor([[0, 0], [1, 0], [2, 0], [0, 1], [-2, 0], [0, -2]]) data = Data(pos=pos) data = RadiusGraph(r=1)(data) assert len(data) == 2 assert data.pos.tolist() == pos.tolist() assert data.edge_index.tolist() == [[0, 0, 1, 1, 2, 3], [1, 3, 0, 2, 1, 0]]
def __init__( self, root: str = "data/", n_segments: int = 100, max_num_neighbors: int = 8, r: float = 10, loop: bool = True, transform: Optional[Callable] = None, pre_transform: Optional[Callable] = None, pre_filter: Optional[Callable] = None, **kwargs, ): self.data_dir = root self.n_segments = n_segments self.max_num_neighbors = max_num_neighbors self.r = r self.loop = loop self.slic_kwargs = kwargs self.base_transform = T.Compose([ T.ToTensor(), ToSLIC(n_segments=n_segments, **kwargs), RadiusGraph(r=r, max_num_neighbors=max_num_neighbors, loop=loop), ]) super().__init__(os.path.join(root, "FashionMNIST"), transform, pre_transform, pre_filter) self.data, self.slices = torch.load(self.processed_paths[0])
def process_rollout(self, positions): data_list = [] num_steps = len(positions) # Calculate velocities velocities = np.concatenate(([np.zeros(positions[0].shape)], positions[1:] - positions[0:-1]),axis=0) # Calculate accelerations accelerations = np.concatenate(([np.zeros(velocities[0].shape)], velocities[1:] - velocities[0:-1]),axis=0) # Material properties (using one-hot encoding for now) m = np.zeros((len(positions[0]), 2)) m[0:64] = [0,1] # First 64 particles are solid m[64:] = [1,0] # TODO: Global forces # Drop the first 5 and the last step since we don't have accurate velocities/accelerations for t in range(6,num_steps-1): # Distance to the 5 walls d = np.stack([ positions[t][:,1], # bottom positions[t][:,0], # left positions[t][:,2], # back 1.2 - positions[t][:,0], # right 0.4 - positions[t][:,2] # front ], axis=1) d = np.clip(d, 0, R) x = np.concatenate((positions[t], m, np.concatenate(velocities[t-5:t], axis=1), d), axis=1) y = torch.tensor(accelerations[t]).float() data = Data(x=torch.tensor(x).float(), y=y, pos=torch.as_tensor(positions[t])) # print(f'Step {t}:', data) # Apply pre-transform to get edges calculate_edges = self.pre_transform or RadiusGraph(R) data = calculate_edges(data) data_list.append(data) return data_list
def print_dataset(train_dataset, test_dataset): num_nodes = num_edges = 0 for data in train_dataset: data = RadiusGraph(0.2)(data) num_nodes += data.num_nodes num_edges += data.num_edges for data in test_dataset: data = RadiusGraph(0.2)(data) num_nodes += data.num_nodes num_edges += data.num_edges num_graphs = len(train_dataset) + len(test_dataset) print('Graphs', num_graphs) print('Nodes', num_nodes / num_graphs) print('Edges', (num_edges // 2) / num_graphs) print('Label rate', len(train_dataset) / num_graphs) print()
def make_graph(self, positions, properties, velocities): d = torch.stack( [ positions[:, 1], # bottom positions[:, 0], # left positions[:, 2], # back 1.2 - positions[:, 0], # right 0.4 - positions[:, 2] # front ], dim=1) d = torch.clamp(d, min=0, max=self.R) x = torch.cat([positions, properties, velocities, d], 1) data = Data(x=x, pos=positions) find_edges = RadiusGraph(self.R) data = find_edges(data) return data
def __call__(self, data): r = self.calculate_radius(data.area, len(data.x)) return RadiusGraph(r)(data)