def to_sparse_coo_matrix(
data: Union[GraphData, GraphBatch],
fill_value: Optional[Union[int, float, torch.Tensor]] = None,
dtype=torch.float,
):
"""Return the sparse coo representation of the edge attribute data of the.
:class:`caldera.data.GraphData` instance. If provided with the optional
fill value, will return a sparse matrix filled with the value instead of
the edge attribute data.
.. code-block::
to_sparse_coo_matrix(data)
to_sparse_coo_matrix(data, fill_value=1)
:param data:
:param fill_value:
:param dtype:
:return:
"""
if fill_value is not None:
return scatter_coo(
data.edges,
fill_value,
expand=True,
dtype=dtype,
size=(data.num_nodes, data.num_nodes, ...),
)
else:
return scatter_coo(
data.edges,
data.e,
expand=False,
dtype=dtype,
size=(data.num_nodes, data.num_nodes, ...),
)
def validate_scatter_coo(self, indices, values, size=None):
if size is None:
expected_size = tuple((indices.max(dim=1).values + 1).tolist())
elif size[-1] is ...:
expected_size = size[:-1] + values.shape[1:]
else:
expected_size = size
expected = torch.zeros(expected_size)
expected[indices.unbind()] = values
result = scatter_coo(indices, values, size)
if expected_size is not None:
assert result.size() == expected_size
assert torch.allclose(result.to_dense(), expected)
def validate_scatter_coo(self, indices, values, size=None):
if torch.is_tensor(values):
value_shape = values.shape
else:
value_shape = torch.Size([])
if size is None:
expected_size = tuple((indices.max(dim=1).values + 1).tolist())
elif size[-1] is ...:
expected_size = size[:-1] + value_shape
else:
expected_size = size
expected = torch.zeros(expected_size)
expected[indices.to(torch.long).unbind()] = values
result = scatter_coo(indices, values, size=size, expand=True)
if expected_size is not None:
assert result.size() == expected_size
assert torch.all(result.to_dense() == expected)
def test_scatter_coo_infer_second_dim(self):
indices = torch.LongTensor([0, 1, 2, 3])
values = torch.tensor([10, 20, 30, 40])
matrix = scatter_coo(indices, values)
assert matrix.size() == torch.Size([1, 4])
def test_scatter_expand_fail(self):
indices = torch.LongTensor([[0, 1, 2, 3], [4, 3, 2, 1], [3, 4, 5, 10]])
values = torch.tensor([[1, 2, 3, 1], [3, 4, 5, 1], [3, 4, 5, 5]],
dtype=torch.float)
scatter_coo(indices, values, expand=True, size=(50, 50, ...))