def too_close_atom_pairs(structure: Structure, radius: float = 2, too_close_criterion_factor: float = 0.7) -> bool: """Check whether too close atomic pairs exist in the structure or not. Raise a ValueError if the number of atoms in structure is 1. Args: structure (Structure): Input structure """ distances = structure.get_all_neighbors(radius) for i, periodic_neighbors in enumerate(distances): elem1 = structure[i].species_string frac1 = structure[i].frac_coords for periodic_neighbor in periodic_neighbors: elem2 = periodic_neighbor.species_string frac2 = periodic_neighbor.frac_coords dist = periodic_neighbor.nn_distance rcore_sum = rcore[elem1] + rcore[elem2] if dist < rcore_sum * too_close_criterion_factor: logger.warning( f"Element {elem1} (index: {i}) at {frac1} & Element {elem2}" f" (index: {periodic_neighbor.index}) at {frac2} are too " f"close with distance {round(dist, 4)}.") return True return False
def test_get_all_neighbors_outside_cell(self): s = Structure(Lattice.cubic(2), ['Li', 'Li', 'Li', 'Si'], [[3.1] * 3, [0.11] * 3, [-1.91] * 3, [0.5] * 3]) all_nn = s.get_all_neighbors(0.2, True) for site, nns in zip(s, all_nn): for nn in nns: self.assertTrue(nn[0].is_periodic_image(s[nn[2]])) d = sum((site.coords - nn[0].coords) ** 2) ** 0.5 self.assertAlmostEqual(d, nn[1]) self.assertEqual(list(map(len, all_nn)), [2, 2, 2, 0])
def test_get_all_neighbors_outside_cell(self): s = Structure(Lattice.cubic(2), ['Li', 'Li', 'Li', 'Si'], [[3.1] * 3, [0.11] * 3, [-1.91] * 3, [0.5] * 3]) all_nn = s.get_all_neighbors(0.2, True) for site, nns in zip(s, all_nn): for nn in nns: self.assertTrue(nn[0].is_periodic_image(s[nn[2]])) d = sum((site.coords - nn[0].coords)**2)**0.5 self.assertAlmostEqual(d, nn[1]) self.assertEqual(list(map(len, all_nn)), [2, 2, 2, 0])
def neighbor_list_and_relative_vec_lattice(pos, lattice, r_max, self_interaction=True, r_min=1e-8): """ Create neighbor list (edge_index) and relative vectors (edge_attr) based on radial cutoff and periodic lattice. :param pos: torch.tensor of coordinates with shape (N, 3) :param r_max: float of radial cutoff :param self_interaction: whether or not to include self edge :return: list of edges [(2, num_edges)], Tensor of relative vectors [num_edges, 3] edges are given by the convention edge_list[0] = source (convolution center) edge_list[1] = target (neighbor index) Thus, the edge_list has the same convention vector notation for relative vectors \vec{r}_{source, target} Relative vectors are given for the different images of the neighbor atom within r_max. """ N, _ = pos.shape structure = Structure(lattice, ['H'] * N, pos, coords_are_cartesian=True) nei_list = [] geo_list = [] neighbors = structure.get_all_neighbors(r_max, include_index=True, include_image=True, numerical_tol=r_min) for i, (site, neis) in enumerate(zip(structure, neighbors)): indices, cart = zip(*[(n.index, n.coords) for n in neis]) cart = torch.tensor(cart) indices = torch.LongTensor([[i, target] for target in indices]) dist = cart - torch.tensor(site.coords) if self_interaction: self_index = torch.LongTensor([[i, i]]) indices = torch.cat([self_index, indices], dim=0) self_dist = torch.zeros(1, 3, dtype=dist.dtype) dist = torch.cat([self_dist, dist], dim=0) nei_list.append(indices) geo_list.append(dist) return torch.cat(nei_list, dim=0).transpose(1, 0), torch.cat(geo_list, dim=0)
def test_get_all_neighbors_and_get_neighbors(self): s = self.struct r = random.uniform(3, 6) all_nn = s.get_all_neighbors(r, True) for i in range(len(s)): self.assertEqual(len(all_nn[i]), len(s.get_neighbors(s[i], r))) for site, nns in zip(s, all_nn): for nn in nns: self.assertTrue(nn[0].is_periodic_image(s[nn[2]])) d = sum((site.coords - nn[0].coords)**2)**0.5 self.assertAlmostEqual(d, nn[1]) s = Structure(Lattice.cubic(1), ['Li'], [[0, 0, 0]]) s.make_supercell([2, 2, 2]) self.assertEqual(sum(map(len, s.get_all_neighbors(3))), 976)
def test_get_all_neighbors_and_get_neighbors(self): s = self.struct r = random.uniform(3, 6) all_nn = s.get_all_neighbors(r, True) for i in range(len(s)): self.assertEqual(len(all_nn[i]), len(s.get_neighbors(s[i], r))) for site, nns in zip(s, all_nn): for nn in nns: self.assertTrue(nn[0].is_periodic_image(s[nn[2]])) d = sum((site.coords - nn[0].coords) ** 2) ** 0.5 self.assertAlmostEqual(d, nn[1]) s = Structure(Lattice.cubic(1), ['Li'], [[0,0,0]]) s.make_supercell([2,2,2]) self.assertEqual(sum(map(len, s.get_all_neighbors(3))), 976)
def edge_features(self, structure: Structure, **kwargs): def expand_distance(distances, dmin=0, step=0.2, var=None): """ Parameters ---------- dmin: float Minimum interatomic distance dmax: float Maximum interatomic distance step: float Step size for the Gaussian filter """ filter_ = np.arange(dmin, self.radius + step, step) if var is None: var = step return np.exp(-(distances[..., np.newaxis] - filter_)**2 / var**2) all_nbrs = structure.get_all_neighbors(self.radius, include_index=True) all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs] nbr_fea_idx, nbr_fea = [], [] for nbr in all_nbrs: if len(nbr) < self.max_num_nbr: warnings.warn('can not find enough neighbors to build graph. ' 'If it happens frequently, consider increase ' 'radius.') nbr_fea_idx.append( list(map(lambda x: x[2], nbr)) + [0] * (self.max_num_nbr - len(nbr))) nbr_fea.append( list(map(lambda x: x[1], nbr)) + [self.radius + 1.] * (self.max_num_nbr - len(nbr))) else: nbr_fea_idx.append( list(map(lambda x: x[2], nbr[:self.max_num_nbr]))) nbr_fea.append( list(map(lambda x: x[1], nbr[:self.max_num_nbr]))) nbr_fea = np.array(nbr_fea) nbr_fea = expand_distance(nbr_fea) nbr_fea = torch.Tensor(nbr_fea) nbr_fea_idx = torch.LongTensor(nbr_fea_idx) return nbr_fea, nbr_fea_idx
def __getitem__(self, idx): # NOTE sites must be given in fractional co-ordinates cif_id, comp, target, cell, sites = self.df.iloc[idx] cif_id = str(cif_id) if self.use_cache: cache_path = os.path.join(self.cachedir, cif_id + ".pkl") if self.use_cache and os.path.exists(cache_path): with open(cache_path, "rb") as f: try: pkl_data = pickle.load(f) except EOFError: raise EOFError(f"Check {f} for issue") atom_fea = pkl_data[0] nbr_fea = pkl_data[1] self_fea_idx = pkl_data[2] nbr_fea_idx = pkl_data[3] else: cell, elems, coords = parse_cgcnn(cell, sites) # NOTE getting primative structure before constructing graph # significantly harms the performnace of this model. crystal = Structure( lattice=cell, species=elems, coords=coords, to_unit_cell=True ) # atom features atom_fea = [atom.specie.symbol for atom in crystal] # neighbours all_nbrs = crystal.get_all_neighbors(self.radius, include_index=True) all_nbrs = [sorted(nbrs, key=lambda x: x[1]) for nbrs in all_nbrs] self_fea_idx, nbr_fea_idx, nbr_fea = [], [], [] for i, nbr in enumerate(all_nbrs): # NOTE due to using a geometric learning library we do not # need to set a maximum number of neighbours but do so in # order to replicate the original code. if len(nbr) < self.max_num_nbr: nbr_fea_idx.extend(list(map(lambda x: x[2], nbr))) nbr_fea.extend(list(map(lambda x: x[1], nbr))) else: nbr_fea_idx.extend( list(map(lambda x: x[2], nbr[: self.max_num_nbr])) ) nbr_fea.extend(list(map(lambda x: x[1], nbr[: self.max_num_nbr]))) if len(nbr) == 0: raise ValueError( f"Isolated atom found in {cif_id} ({comp}) - " "increase maximum radius or remove structure" ) self_fea_idx.extend([i] * min(len(nbr), self.max_num_nbr)) nbr_fea = np.array(nbr_fea) if self.use_cache: with open(cache_path, "wb") as f: pickle.dump((atom_fea, nbr_fea, self_fea_idx, nbr_fea_idx), f) nbr_fea = self.gdf.expand(nbr_fea) atom_fea = np.vstack([self.ari.get_fea(atom) for atom in atom_fea]) atom_fea = torch.Tensor(atom_fea) nbr_fea = torch.Tensor(nbr_fea) self_fea_idx = torch.LongTensor(self_fea_idx) nbr_fea_idx = torch.LongTensor(nbr_fea_idx) if self.task == "regression": target = torch.Tensor([float(target)]) elif self.task == "classification": target = torch.LongTensor([target]) return (atom_fea, nbr_fea, self_fea_idx, nbr_fea_idx), target, comp, cif_id