Python data.Data方法代码示例

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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))) 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]) 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]] 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]] 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]] 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]] 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]] 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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) 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]] 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]] 

# 需要导入模块: from torch_geometric import data [as 别名]
# 或者: from torch_geometric.data import Data [as 别名]
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]]