def construct_graph(self, traj_coords, lattices, atom_types, target_index):
if self.backend == 'kdtree':
nbr_lists, diag_lattices = [], []
for coord, lattice in tqdm(zip(traj_coords, lattices),
total=len(traj_coords),
disable=not self.verbose):
# take the diagonal part of the lattice matrix
lattice = np.diagonal(lattice)
diag_lattices.append(lattice)
pkdt = PeriodicCKDTree(lattice, coord)
all_nbrs_idxes = pkdt.query_ball_point(coord, self.radius)
nbr_list = []
for idx, nbr_idxes in enumerate(all_nbrs_idxes):
nbr_dists = distance_pbc(coord[idx], coord[nbr_idxes],
lattice)
nbr_idx_dist = sorted(zip(nbr_idxes, nbr_dists),
key=lambda x: x[1])
assert nbr_idx_dist[0][1] == 0 and\
nbr_idx_dist[0][0] == idx and\
len(nbr_idx_dist) >= self.n_nbrs + 1
nbr_list.append(
[idx for idx, dist in nbr_idx_dist[1:self.n_nbrs + 1]])
nbr_lists.append(np.stack(np.array(nbr_list, dtype='int32')))
nbr_lists = np.stack(nbr_lists)
diag_lattices = np.stack(diag_lattices)
return {
'traj_coords': traj_coords,
'lattices': diag_lattices,
'atom_types': atom_types,
'target_index': target_index,
'nbr_lists': nbr_lists
}
elif self.backend == 'direct':
nbr_lists, nbr_dists = [], []
for coord, lattice in tqdm(zip(traj_coords, lattices),
total=len(traj_coords),
disable=not self.verbose):
crystal = IStructure(lattice=lattice,
species=atom_types,
coords=coord,
coords_are_cartesian=True)
all_nbrs = crystal.get_all_neighbors(r=self.radius,
include_index=True)
all_nbrs = [
sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs
]
nbr_list, nbr_dist = [], []
for nbr in all_nbrs:
assert len(nbr) >= self.n_nbrs, 'not find enough neighbors'
nbr_list.append(
list(map(lambda x: x[2], nbr[:self.n_nbrs])))
nbr_dist.append(
list(map(lambda x: x[1], nbr[:self.n_nbrs])))
nbr_lists.append(np.array(nbr_list, dtype='int32'))
nbr_dists.append(np.array(nbr_dist, dtype='float32'))
nbr_lists, nbr_dists = np.stack(nbr_lists), np.stack(nbr_dists)
return {
'traj_coords': traj_coords,
'atom_types': atom_types,
'target_index': target_index,
'nbr_lists': nbr_lists,
'nbr_dists': nbr_dists
}
elif self.backend == 'ndirect':
stcs = [
Structure(lattice=lattices[i],
species=atom_types,
coords=traj_coords[i],
coords_are_cartesian=True)
for i in tqdm(range(len(traj_coords)),
desc='Generating structure...',
disable=not self.verbose)
]
a, b, c = [
np.ceil(2 * self.radius / d).astype('int')
for d in stcs[0].lattice.abc
]
if [a, b, c] != [1, 1, 1]:
_ = [
stc.make_supercell([
np.ceil(2 * self.radius / d).astype('int')
for d in stc.lattice.abc
]) for stc in tqdm(stcs,
desc='Building supercell...',
disable=not self.verbose)
]
nbr_lists = np.array([
stc.distance_matrix.argsort()[:, 1:1 + self.n_nbrs]
for stc in tqdm(stcs,
desc='Generating neighbor index...',
disable=not self.verbose)
],
dtype='int32')
nbr_dists = np.array([
np.sort(stc.distance_matrix)[:, 1:1 + self.n_nbrs]
for stc in tqdm(stcs,
desc='Generating neighbor distance...',
disable=not self.verbose)
],
dtype='float32')
nbr_lists, nbr_dists = np.stack(nbr_lists), np.stack(nbr_dists)
if not np.all((nbr_dists < self.radius) & (nbr_dists > 0)):
raise ('not find enough neighbors')
return {
'traj_coords': traj_coords,
'atom_types': atom_types,
'target_index': target_index,
'nbr_lists': nbr_lists,
'nbr_dists': nbr_dists
}
def construct_graph(self, traj_coords, lattices, atom_types, target_index):
if self.backend == 'kdtree':
nbr_lists, diag_lattices = [], []
for coord, lattice in tqdm(zip(traj_coords, lattices),
total=len(traj_coords),
disable=not self.verbose):
# take the diagonal part of the lattice matrix
lattice = np.diagonal(lattice)
diag_lattices.append(lattice)
pkdt = PeriodicCKDTree(lattice, coord)
all_nbrs_idxes = pkdt.query_ball_point(coord, self.radius)
nbr_list = []
for idx, nbr_idxes in enumerate(all_nbrs_idxes):
nbr_dists = distance_pbc(coord[idx], coord[nbr_idxes],
lattice)
nbr_idx_dist = sorted(zip(nbr_idxes, nbr_dists),
key=lambda x: x[1])
assert nbr_idx_dist[0][1] == 0 and\
nbr_idx_dist[0][0] == idx and\
len(nbr_idx_dist) >= self.n_nbrs + 1
nbr_list.append(
[idx for idx, dist in nbr_idx_dist[1:self.n_nbrs + 1]])
nbr_lists.append(np.stack(np.array(nbr_list, dtype='int32')))
nbr_lists = np.stack(nbr_lists)
diag_lattices = np.stack(diag_lattices)
return {
'traj_coords': traj_coords,
'lattices': diag_lattices,
'atom_types': atom_types,
'target_index': target_index,
'nbr_lists': nbr_lists
}
elif self.backend == 'direct':
nbr_lists, nbr_dists = [], []
for coord, lattice in tqdm(zip(traj_coords, lattices),
total=len(traj_coords),
disable=not self.verbose):
crystal = IStructure(lattice=lattice,
species=atom_types,
coords=coord,
coords_are_cartesian=True)
all_nbrs = crystal.get_all_neighbors(r=self.radius,
include_index=True)
all_nbrs = [
sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs
]
nbr_list, nbr_dist = [], []
for nbr in all_nbrs:
assert len(nbr) >= self.n_nbrs, 'not find enough neighbors'
nbr_list.append(
list(map(lambda x: x[2], nbr[:self.n_nbrs])))
nbr_dist.append(
list(map(lambda x: x[1], nbr[:self.n_nbrs])))
nbr_lists.append(np.array(nbr_list, dtype='int32'))
nbr_dists.append(np.array(nbr_dist, dtype='float32'))
nbr_lists, nbr_dists = np.stack(nbr_lists), np.stack(nbr_dists)
return {
'traj_coords': traj_coords,
'atom_types': atom_types,
'target_index': target_index,
'nbr_lists': nbr_lists,
'nbr_dists': nbr_dists
}
