Python torch_geometric.data.Data() Examples

def test_in_memory_dataset():
    class TestDataset(InMemoryDataset):
        def __init__(self, data_list):
            super(TestDataset, self).__init__('/tmp/TestDataset')
            self.data, self.slices = self.collate(data_list)

    x = torch.Tensor([[1], [1], [1]])
    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
    face = torch.tensor([[0], [1], [2]])
    i = 1
    s = '1'

    data1 = Data(x=x, edge_index=edge_index, face=face, test_int=i, test_str=s)
    data1.num_nodes = 10

    data2 = Data(x=x, edge_index=edge_index, face=face, test_int=i, test_str=s)
    data2.num_nodes = 5

    dataset = TestDataset([data1, data2])
    assert len(dataset) == 2
    assert dataset[0].num_nodes == 10
    assert len(dataset[0]) == 5
    assert dataset[1].num_nodes == 5
    assert len(dataset[1]) == 5 

def test_local_cartesian():
    assert LocalCartesian().__repr__() == 'LocalCartesian()'

    pos = torch.Tensor([[-1, 0], [0, 0], [2, 0]])
    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
    edge_attr = torch.Tensor([1, 1, 1, 1])

    data = Data(edge_index=edge_index, pos=pos)
    data = LocalCartesian()(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[1, 0.5], [0.25, 0.5], [1, 0.5],
                                       [0, 0.5]]

    data = Data(edge_index=edge_index, pos=pos, edge_attr=edge_attr)
    data = LocalCartesian()(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[1, 1, 0.5], [1, 0.25, 0.5],
                                       [1, 1, 0.5], [1, 0, 0.5]] 

def from_data_list(data_list, follow_batch=[]):
        laplacians = None
        v_plus = None

        if 'laplacians' in data_list[0]:
            laplacians = [d.laplacians[:] for d in data_list]
            v_plus = [d.v_plus[:] for d in data_list]

        copy_data = []
        for d in data_list:
            copy_data.append(Data(x=d.x,
                                  y=d.y,
                                  edge_index=d.edge_index,
                                  edge_attr=d.edge_attr,
                                  v_outs=d.v_outs,
                                  g_outs=d.g_outs,
                                  e_outs=d.e_outs,
                                  o_outs=d.o_outs)
                             )

        batch = data.Batch.from_data_list(copy_data, follow_batch=follow_batch)
        batch['laplacians'] = laplacians
        batch['v_plus'] = v_plus

        return batch 

def test_distance():
    assert Distance().__repr__() == 'Distance(norm=True, max_value=None)'

    pos = torch.Tensor([[-1, 0], [0, 0], [2, 0]])
    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
    edge_attr = torch.Tensor([1, 1, 1, 1])

    data = Data(edge_index=edge_index, pos=pos)
    data = Distance(norm=False)(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[1], [1], [2], [2]]

    data = Data(edge_index=edge_index, pos=pos, edge_attr=edge_attr)
    data = Distance(norm=True)(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[1, 0.5], [1, 0.5], [1, 1], [1, 1]] 

def test_sample_points():
    assert SamplePoints(1024).__repr__() == 'SamplePoints(1024)'

    pos = torch.Tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0]])
    face = torch.tensor([[0, 1], [1, 2], [2, 3]])

    data = Data(pos=pos)
    data.face = face
    data = SamplePoints(8)(data)
    assert len(data) == 1
    assert pos[:, 0].min().item() >= 0 and pos[:, 0].max().item() <= 1
    assert pos[:, 1].min().item() >= 0 and pos[:, 1].max().item() <= 1
    assert pos[:, 2].abs().sum().item() == 0

    data = Data(pos=pos)
    data.face = face
    data = SamplePoints(8, include_normals=True)(data)
    assert len(data) == 2
    assert data.norm[:, :2].abs().sum().item() == 0
    assert data.norm[:, 2].abs().sum().item() == 8 

def collate(self, batch):
        elem = batch[0]
        if isinstance(elem, Data):
            return Batch.from_data_list(batch, self.follow_batch)
        elif isinstance(elem, torch.Tensor):
            return default_collate(batch)
        elif isinstance(elem, float):
            return torch.tensor(batch, dtype=torch.float)
        elif isinstance(elem, int_classes):
            return torch.tensor(batch)
        elif isinstance(elem, string_classes):
            return batch
        elif isinstance(elem, container_abcs.Mapping):
            return {key: self.collate([d[key] for d in batch]) for key in elem}
        elif isinstance(elem, tuple) and hasattr(elem, '_fields'):
            return type(elem)(*(self.collate(s) for s in zip(*batch)))
        elif isinstance(elem, container_abcs.Sequence):
            return [self.collate(s) for s in zip(*batch)]

        raise TypeError('DataLoader found invalid type: {}'.format(type(elem))) 

def _load_split_star_data():
    res = []
    for filepath in ["./data/aclImdb/train",
                     "./data/aclImdb/test"]:
        bert_encode_res = np.load(os.path.join(filepath, "split_bert_encode_res.npy"))  # 25000,768
        y = np.load(os.path.join(filepath, "split_y.npy"))  # 25000
        edge = np.load(os.path.join(filepath, "split_edge.npy"))  # 2,num_edge*2
        datas = []
        for x, y, e in zip(bert_encode_res, y, edge):
            x = [_.tolist() for _ in x]
            x = torch.tensor(x, dtype=torch.float32)
            y = torch.tensor(y, dtype=torch.long)
            star = x.mean(dim=0).view(1, -1)
            x = torch.cat([x, star], dim=0)
            y = torch.cat([y, torch.tensor([-1], dtype=torch.long)])
            a = torch.full((1, len(x) - 1), len(x) - 1).long()
            b = torch.tensor(range(len(x) - 1)).view(1, -1).long()
            e = torch.cat([torch.cat([a, b], dim=0), torch.cat([b, a], dim=0)], dim=1)
            datas.append(Data(x=x, edge_index=e, y=y))
        res.append(datas)
    return res[0], res[1] 

def process(self):
        mat = loadmat(self.raw_paths[0])['Problem'][0][0][2].tocsr().tocoo()

        row = torch.from_numpy(mat.row).to(torch.long)
        col = torch.from_numpy(mat.col).to(torch.long)
        edge_index = torch.stack([row, col], dim=0)

        edge_attr = torch.from_numpy(mat.data).to(torch.float)
        if torch.all(edge_attr == 1.):
            edge_attr = None

        size = torch.Size(mat.shape)
        if mat.shape[0] == mat.shape[1]:
            size = None

        num_nodes = mat.shape[0]

        data = Data(edge_index=edge_index, edge_attr=edge_attr, size=size,
                    num_nodes=num_nodes)

        if self.pre_transform is not None:
            data = self.pre_transform(data)

        torch.save(self.collate([data]), self.processed_paths[0]) 

def _load_split_2k_data():
    res = []
    for filepath in ["./data/aclImdb/train",
                     "./data/aclImdb/test"]:
        bert_encode_res = np.load(os.path.join(filepath, "split_2k_bert_large_encode_res.npy"),
                                  allow_pickle=True)  # 25000,768
        y = np.load(os.path.join(filepath, "split_2k_y.npy"), allow_pickle=True)  # 25000
        edge = np.load(os.path.join(filepath, "split_2k_edge.npy"), allow_pickle=True)  # 2,num_edge*2
        datas = []
        for x, y, e in zip(bert_encode_res, y, edge):
            x = [_.tolist() for _ in x]
            x = torch.tensor(x, dtype=torch.float32)
            y = torch.tensor(y, dtype=torch.long)
            if len(e) == 0:
                e = torch.empty((0, 2), dtype=torch.long).t()
            else:
                e = torch.tensor(e, dtype=torch.long).t()
            datas.append(Data(x=x, edge_index=e, y=y))
        res.append(datas)
    return res[0], res[1] 

def _load_split_data():
    res = []
    for filepath in ["./data/aclImdb/train",
                     "./data/aclImdb/test"]:
        bert_encode_res = np.load(os.path.join(filepath, "split_bert_large_encode_res.npy"),
                                  allow_pickle=True)  # 25000,768
        y = np.load(os.path.join(filepath, "split_y.npy"), allow_pickle=True)  # 25000
        edge = np.load(os.path.join(filepath, "split_edge.npy"), allow_pickle=True)  # 2,num_edge*2
        datas = []
        for x, y, e in zip(bert_encode_res, y, edge):
            x = [_.tolist() for _ in x]
            x = torch.tensor(x, dtype=torch.float32)
            y = torch.tensor(y, dtype=torch.long)
            if len(e) == 0:
                e = torch.empty((0, 2), dtype=torch.long).t()
            else:
                e = torch.tensor(e, dtype=torch.long).t()
            datas.append(Data(x=x, edge_index=e, y=y))
        res.append(datas)
    return res[0], res[1] 

def test_generate_normals():
    assert GenerateMeshNormals().__repr__() == 'GenerateMeshNormals()'

    pos = torch.Tensor([
        [0, 0, 0],
        [-2, 1, 0],
        [-1, 1, 0],
        [0, 1, 0],
        [1, 1, 0],
        [2, 1, 0],
    ])
    face = torch.tensor([
        [0, 0, 0, 0],
        [1, 2, 3, 4],
        [2, 3, 4, 5],
    ])

    data = GenerateMeshNormals()(Data(pos=pos, face=face))
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.face.tolist() == face.tolist()
    assert data.norm.tolist() == [[0, 0, -1]] * 6 

def test_constant():
    assert Constant().__repr__() == 'Constant(value=1)'

    x = torch.Tensor([[-1, 0], [0, 0], [2, 0]])
    edge_index = torch.tensor([[0, 1], [1, 2]])

    data = Data(edge_index=edge_index, num_nodes=3)
    data = Constant()(data)
    assert len(data) == 2
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.x.tolist() == [[1], [1], [1]]

    data = Data(edge_index=edge_index, x=x)
    data = Constant()(data)
    assert len(data) == 2
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.x.tolist() == [[-1, 0, 1], [0, 0, 1], [2, 0, 1]] 

def test_polar():
    assert Polar().__repr__() == 'Polar(norm=True, max_value=None)'

    pos = torch.Tensor([[0, 0], [1, 0]])
    edge_index = torch.tensor([[0, 1], [1, 0]])
    edge_attr = torch.Tensor([1, 1])

    data = Data(edge_index=edge_index, pos=pos)
    data = Polar(norm=False)(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert torch.allclose(
        data.edge_attr, torch.Tensor([[1, 0], [1, PI]]), atol=1e-04)

    data = Data(edge_index=edge_index, pos=pos, edge_attr=edge_attr)
    data = Polar(norm=True)(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert torch.allclose(
        data.edge_attr, torch.Tensor([[1, 1, 0], [1, 1, 0.5]]), atol=1e-04) 

def read_obj(in_file):
    vertices = []
    faces = []

    for k, v in yield_file(in_file):
        if k == 'v':
            vertices.append(v)
        elif k == 'f':
            for i in v:
                faces.append(i)

    if not len(faces) or not len(vertices):
        return None

    pos = torch.tensor(vertices, dtype=torch.float)
    face = torch.tensor(faces, dtype=torch.long).t().contiguous()

    data = Data(pos=pos, face=face)

    return data 

def test_static_graph():
    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
    x1, x2 = torch.randn(3, 8), torch.randn(3, 8)

    data1 = Data(edge_index=edge_index, x=x1)
    data2 = Data(edge_index=edge_index, x=x2)
    batch = Batch.from_data_list([data1, data2])

    x = torch.stack([x1, x2], dim=0)
    for conv in [MyConv(), GCNConv(8, 16), ChebConv(8, 16, K=2)]:
        out1 = conv(batch.x, batch.edge_index)
        assert out1.size(0) == 6
        conv.node_dim = 1
        out2 = conv(x, edge_index)
        assert out2.size()[:2] == (2, 3)
        assert torch.allclose(out1, out2.view(-1, out2.size(-1))) 

def test_compose():
    transform = T.Compose([T.Center(), T.AddSelfLoops()])
    assert transform.__repr__() == ('Compose([\n'
                                    '    Center(),\n'
                                    '    AddSelfLoops(),\n'
                                    '])')

    pos = torch.Tensor([[0, 0], [2, 0], [4, 0]])
    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])

    data = Data(edge_index=edge_index, pos=pos)
    data = transform(data)
    assert len(data) == 2
    assert data.pos.tolist() == [[-2, 0], [0, 0], [2, 0]]
    assert data.edge_index.tolist() == [[0, 0, 1, 1, 1, 2, 2],
                                        [0, 1, 0, 1, 2, 1, 2]] 

def parse_sdf(src):
    src = src.split('\n')[3:]
    num_atoms, num_bonds = [int(item) for item in src[0].split()[:2]]

    atom_block = src[1:num_atoms + 1]
    pos = parse_txt_array(atom_block, end=3)
    x = torch.tensor([elems[item.split()[3]] for item in atom_block])
    x = F.one_hot(x, num_classes=len(elems))

    bond_block = src[1 + num_atoms:1 + num_atoms + num_bonds]
    row, col = parse_txt_array(bond_block, end=2, dtype=torch.long).t() - 1
    row, col = torch.cat([row, col], dim=0), torch.cat([col, row], dim=0)
    edge_index = torch.stack([row, col], dim=0)
    edge_attr = parse_txt_array(bond_block, start=2, end=3) - 1
    edge_attr = torch.cat([edge_attr, edge_attr], dim=0)
    edge_index, edge_attr = coalesce(edge_index, edge_attr, num_atoms,
                                     num_atoms)

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, pos=pos)
    return data 

def test_cartesian():
    assert Cartesian().__repr__() == 'Cartesian(norm=True, max_value=None)'

    pos = torch.Tensor([[-1, 0], [0, 0], [2, 0]])
    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
    edge_attr = torch.Tensor([1, 1, 1, 1])

    data = Data(edge_index=edge_index, pos=pos)
    data = Cartesian(norm=False)(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[1, 0], [-1, 0], [2, 0], [-2, 0]]

    data = Data(edge_index=edge_index, pos=pos, edge_attr=edge_attr)
    data = Cartesian(norm=True)(data)
    assert len(data) == 3
    assert data.pos.tolist() == pos.tolist()
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[1, 0.75, 0.5], [1, 0.25, 0.5],
                                       [1, 1, 0.5], [1, 0, 0.5]] 

def test_to_networkx_undirected():
    x = torch.Tensor([[1, 2], [3, 4]])
    pos = torch.Tensor([[0, 0], [1, 1]])
    edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
    edge_attr = torch.Tensor([1, 2, 3])
    data = Data(x=x, pos=pos, edge_index=edge_index, weight=edge_attr)

    for remove_self_loops in [True, False]:
        G = to_networkx(data, node_attrs=['x', 'pos'], edge_attrs=['weight'],
                        remove_self_loops=remove_self_loops,
                        to_undirected=True)

        assert G.nodes[0]['x'] == [1, 2]
        assert G.nodes[1]['x'] == [3, 4]
        assert G.nodes[0]['pos'] == [0, 0]
        assert G.nodes[1]['pos'] == [1, 1]

        if remove_self_loops:
            assert nx.to_numpy_matrix(G).tolist() == [[0, 2], [2, 0]]
        else:
            assert nx.to_numpy_matrix(G).tolist() == [[3, 2], [2, 0]] 

def test_random_translate():
    assert RandomTranslate(0.1).__repr__() == 'RandomTranslate(0.1)'

    pos = torch.Tensor([[0, 0], [0, 0], [0, 0], [0, 0]])

    data = Data(pos=pos)
    data = RandomTranslate(0)(data)
    assert len(data) == 1
    assert data.pos.tolist() == pos.tolist()

    data = Data(pos=pos)
    data = RandomTranslate(0.1)(data)
    assert len(data) == 1
    assert data.pos.min().item() >= -0.1
    assert data.pos.max().item() <= 0.1

    data = Data(pos=pos)
    data = RandomTranslate([0.1, 1])(data)
    assert len(data) == 1
    assert data.pos[:, 0].min().item() >= -0.1
    assert data.pos[:, 0].max().item() <= 0.1
    assert data.pos[:, 1].min().item() >= -1
    assert data.pos[:, 1].max().item() <= 1 

def test_to_networkx():
    x = torch.Tensor([[1, 2], [3, 4]])
    pos = torch.Tensor([[0, 0], [1, 1]])
    edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
    edge_attr = torch.Tensor([1, 2, 3])
    data = Data(x=x, pos=pos, edge_index=edge_index, weight=edge_attr)

    for remove_self_loops in [True, False]:
        G = to_networkx(data, node_attrs=['x', 'pos'], edge_attrs=['weight'],
                        remove_self_loops=remove_self_loops)

        assert G.nodes[0]['x'] == [1, 2]
        assert G.nodes[1]['x'] == [3, 4]
        assert G.nodes[0]['pos'] == [0, 0]
        assert G.nodes[1]['pos'] == [1, 1]

        if remove_self_loops:
            assert nx.to_numpy_matrix(G).tolist() == [[0, 1], [2, 0]]
        else:
            assert nx.to_numpy_matrix(G).tolist() == [[3, 1], [2, 0]] 

def test_target_indegree():
    assert TargetIndegree().__repr__() == ('TargetIndegree(norm=True, '
                                           'max_value=None)')

    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
    edge_attr = torch.Tensor([1, 1, 1, 1])

    data = Data(edge_index=edge_index, num_nodes=3)
    data = TargetIndegree(norm=False)(data)
    assert len(data) == 2
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[2], [1], [1], [2]]

    data = Data(edge_index=edge_index, edge_attr=edge_attr, num_nodes=3)
    data = TargetIndegree(norm=True)(data)
    assert len(data) == 2
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.edge_attr.tolist() == [[1, 1], [1, 0.5], [1, 0.5], [1, 1]] 

def test_one_hot_degree():
    assert OneHotDegree(max_degree=3).__repr__() == 'OneHotDegree(3)'

    edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
    x = torch.Tensor([1, 1, 1, 1])

    data = Data(edge_index=edge_index, num_nodes=4)
    data = OneHotDegree(max_degree=3)(data)
    assert len(data) == 2
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.x.tolist() == [[0, 0, 0, 1], [0, 1, 0, 0], [0, 1, 0, 0],
                               [0, 1, 0, 0]]

    data = Data(edge_index=edge_index, x=x)
    data = OneHotDegree(max_degree=3)(data)
    assert len(data) == 2
    assert data.edge_index.tolist() == edge_index.tolist()
    assert data.x.tolist() == [[1, 0, 0, 0, 1], [1, 0, 1, 0, 0],
                               [1, 0, 1, 0, 0], [1, 0, 1, 0, 0]] 

def test_two_hop():
    assert TwoHop().__repr__() == 'TwoHop()'

    edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
    edge_attr = torch.tensor([1, 2, 3, 1, 2, 3], dtype=torch.float)

    data = Data(edge_index=edge_index, edge_attr=edge_attr, num_nodes=4)
    data = TwoHop()(data)
    assert len(data) == 2
    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3],
                                        [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2]]
    assert data.edge_attr.tolist() == [1, 2, 3, 1, 0, 0, 2, 0, 0, 3, 0, 0]

    data = Data(edge_index=edge_index, num_nodes=4)
    data = TwoHop()(data)
    assert len(data) == 1
    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3],
                                        [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2]] 

def test_random_scale():
    assert RandomScale([1, 2]).__repr__() == 'RandomScale([1, 2])'

    pos = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])

    data = Data(pos=pos)
    data = RandomScale([1, 1])(data)
    assert len(data) == 1
    assert data.pos.tolist() == pos.tolist()

    data = Data(pos=pos)
    data = RandomScale([2, 2])(data)
    assert len(data) == 1
    assert data.pos.tolist() == [[-2, -2], [-2, 2], [2, -2], [2, 2]] 

def test_normalize_rotation():
    assert NormalizeRotation().__repr__() == 'NormalizeRotation()'

    pos = torch.Tensor([[-2, -2], [-1, -1], [0, 0], [1, 1], [2, 2]])
    norm = torch.Tensor([[-1, 1], [-1, 1], [-1, 1], [-1, 1], [-1, 1]])
    data = Data(pos=pos)
    data.norm = norm
    data = NormalizeRotation()(data)
    assert len(data) == 2

    expected_pos = torch.Tensor([
        [-2 * sqrt(2), 0],
        [-sqrt(2), 0],
        [0, 0],
        [sqrt(2), 0],
        [2 * sqrt(2), 0],
    ])
    expected_norm = [[0, 1], [0, 1], [0, 1], [0, 1], [0, 1]]

    assert torch.allclose(data.pos, expected_pos, atol=1e-04)
    assert data.norm.tolist() == expected_norm

    data = Data(pos=pos)
    data.norm = norm
    data = NormalizeRotation(max_points=3)(data)
    assert len(data) == 2

    assert torch.allclose(data.pos, expected_pos, atol=1e-04)
    assert data.norm.tolist() == expected_norm 

def test_center():
    assert Center().__repr__() == 'Center()'

    pos = torch.Tensor([[0, 0], [2, 0], [4, 0]])

    data = Data(pos=pos)
    data = Center()(data)
    assert len(data) == 1
    assert data.pos.tolist() == [[-2, 0], [0, 0], [2, 0]] 

def test_face_to_edge():
    assert FaceToEdge().__repr__() == 'FaceToEdge()'

    face = torch.tensor([[0, 0], [1, 1], [2, 3]])

    data = Data(face=face, num_nodes=4)
    data = FaceToEdge()(data)
    assert len(data) == 1
    assert data.edge_index.tolist() == [[0, 0, 0, 1, 1, 1, 2, 2, 3, 3],
                                        [1, 2, 3, 0, 2, 3, 0, 1, 0, 1]] 

def test_dataloader():
    x = torch.Tensor([[1], [1], [1]])
    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
    face = torch.tensor([[0], [1], [2]])
    y = 2.
    z = torch.tensor(0.)
    name = 'data'

    data = Data(x=x, edge_index=edge_index, y=y, z=z, name=name)
    assert data.__repr__() == (
        'Data(edge_index=[2, 4], name="data", x=[3, 1], y=2.0, z=0.0)')
    data.face = face

    loader = DataLoader([data, data], batch_size=2, shuffle=False)

    for batch in loader:
        assert len(batch) == 7
        assert batch.batch.tolist() == [0, 0, 0, 1, 1, 1]
        assert batch.x.tolist() == [[1], [1], [1], [1], [1], [1]]
        assert batch.edge_index.tolist() == [[0, 1, 1, 2, 3, 4, 4, 5],
                                             [1, 0, 2, 1, 4, 3, 5, 4]]
        assert batch.y.tolist() == [2.0, 2.0]
        assert batch.z.tolist() == [0.0, 0.0]
        assert batch.name == ['data', 'data']
        assert batch.face.tolist() == [[0, 3], [1, 4], [2, 5]]

    loader = DataLoader([data, data], batch_size=2, shuffle=False,
                        follow_batch=['edge_index'])

    for batch in loader:
        assert len(batch) == 8
        assert batch.edge_index_batch.tolist() == [0, 0, 0, 0, 1, 1, 1, 1] 

def test_trimesh():
    pos = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0]],
                       dtype=torch.float)
    face = torch.tensor([[0, 1, 2], [1, 2, 3]]).t()

    data = Data(pos=pos, face=face)
    mesh = to_trimesh(data)
    data = from_trimesh(mesh)

    assert pos.tolist() == data.pos.tolist()
    assert face.tolist() == data.face.tolist()