def read_off(self, file):
# Some files may contain a bug and do not have a carriage return after OFF.
with open(file, 'r') as f:
src = f.read().split('\n')[:-1]
if src[0] == 'OFF':
src = src[1:]
else:
src[0] = src[0][3:]
num_nodes, num_faces = [int(item) for item in src[0].split()[:2]]
pos = parse_txt_array(src[1:1 + num_nodes])
# pos = pos - pos.mean(dim=-2, keepdim=True)
scale = (1 / pos.abs().max()) * 0.999999
pos = pos * scale
face = src[1 + num_nodes:1 + num_nodes + num_faces]
face = face_to_tri(face)
edge_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
edge_index = to_undirected(edge_index, num_nodes=num_nodes)
data = Data(edge_index=edge_index, face=face, pos=pos)
return data
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 parse_off(src):
# Some files may contain a bug and do not have a carriage return after OFF.
if src[0] == 'OFF':
src = src[1:]
else:
src[0] = src[0][3:]
num_nodes, num_faces = [int(item) for item in src[0].split()[:2]]
pos = parse_txt_array(src[1:1 + num_nodes])
face = src[1 + num_nodes:1 + num_nodes + num_faces]
face = face_to_tri(face)
data = Data(pos=pos)
data.face = face
return data