Hi everyone,
I am trying to check that copy_u does what I want, but it doesn’t seem to be the case. I’d like to sum the features of the 5 nearest neighbours to a given node. However, although the edges seem to be correct, when I sum up the features I don’t get the same answer as I do it within a for loop (as can be seen in my dirty_graph function). Should I be using a different function rather than copy_u?
import numpy as np
import dgl
import dgl.function as fn
from scipy.spatial import cKDTree
def generate_fake_data(n_particles):
pos = np.random.random((n_particles,3)).astype(np.float32)
mass = np.random.normal(loc=13, scale = 3, size=n_particles).astype(np.float32)
return pos , mass
def build_test_graph_distance(positions,
features,
distance = 1.,
n_neighbors = 5
):
tree = cKDTree(positions,)
distances, edges = tree.query(positions,
k=n_neighbors+1,
)
edges = [(idx, dest) for idx, destination in enumerate(edges) for dest in destination \
if ((idx!=dest)&(dest<len(positions)))]
print(edges)
src, dst = zip(*edges)
G = dgl.DGLGraph()
G.add_nodes(len(positions))
G.add_edges(src, dst)
G.ndata["masses"] = features
G.update_all(fn.copy_u('masses', 'm'), fn.sum('m', 'sum_mass'))
labels = G.ndata['sum_mass']
return G, labels
def dirty_graph(positions,
features,
n_neighbors=5
):
labels = []
for i in range(len(positions)):
difference = positions - positions[i]
distance = np.linalg.norm(difference, axis=-1)
idx = np.argsort(distance)[1:n_neighbors+1]
labels.append(np.sum(features[idx]))
return labels
if __name__=='__main__':
pos, mass = generate_fake_data(10)
graph, labels= build_test_graph_distance(pos, mass, )
dirty_labels = dirty_graph(pos, mass)
np.testing.assert_almost_equal(labels.numpy(), dirty_labels, decimal=1)