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)