def test_points_in_spheres(self): points = [[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]] lattice = Lattice.cubic(3) center_points = [[1.5, 1.5, 1.5]] nns = find_points_in_spheres( all_coords=np.array(points), center_coords=np.array(center_points), r=3, pbc=np.array([0, 0, 0], dtype=int), lattice=np.array(lattice.matrix), tol=1e-8, ) self.assertEqual(len(nns[0]), 2) # two neighbors nns = find_points_in_spheres( all_coords=np.array(points), center_coords=np.array(center_points), r=3, pbc=np.array([1, 1, 1], dtype=int), lattice=np.array(lattice.matrix), ) self.assertEqual(len(nns[0]), 12) nns = find_points_in_spheres( all_coords=np.array(points), center_coords=np.array(center_points), r=3, pbc=np.array([True, False, False], dtype=int), lattice=np.array(lattice.matrix), ) self.assertEqual(len(nns[0]), 4)
def get_neighbors_within_cutoff(structure: Structure, cutoff: float = 5.0, tolerence: float = 1e-8) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: period_boundary_condition = np.array([1, 1, 1]) lattice_matrix = np.array(structure.lattice.matrix) cart_coords = structure.cart_coords center_indices, neighbor_indices, _, distances = find_points_in_spheres( cart_coords, cart_coords, cutoff, period_boundary_condition, lattice_matrix, tolerence) exclude_self = (center_indices != neighbor_indices) | ( distances > tolerence) return center_indices[exclude_self], neighbor_indices[exclude_self], distances[exclude_self]
def get_radius_in_spheres( structure: StructureOrMolecule, nn_strategy=None, cutoff: float = 5.0, numerical_tol: float = 1e-6, pbc=True, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Get graph representations from structure within cutoff. Args: structure (pymatgen Structure or molecule) cutoff (float): cutoff radius numerical_tol (float): numerical tolerance nn_strategy(str):not used Returns: center_indices, neighbor_indices, images, distances, center_prop """ _ = nn_strategy if isinstance(structure, Structure): lattice_matrix = np.ascontiguousarray(np.array( structure.lattice.matrix), dtype=float) if pbc is not False: pbc = re_pbc(pbc, return_type="int") else: pbc = np.array([0, 0, 0]) elif isinstance(structure, Molecule): lattice_matrix = np.array( [[1000.0, 0.0, 0.0], [0.0, 1000.0, 0.0], [0.0, 0.0, 1000.0]], dtype=float) pbc = np.array([0, 0, 0]) else: raise ValueError("structure type not supported") r = float(cutoff) cart_coords = np.ascontiguousarray(np.array(structure.cart_coords), dtype=float) center_indices, neighbor_indices, images, distances = find_points_in_spheres( cart_coords, cart_coords, r=r, pbc=pbc, lattice=lattice_matrix, tol=numerical_tol) center_indices = center_indices.astype(np.int64) neighbor_indices = neighbor_indices.astype(np.int64) images = images.astype(np.int64) distances = distances.astype(np.float32) exclude_self = (center_indices != neighbor_indices) | (distances > numerical_tol) return center_indices[exclude_self], neighbor_indices[exclude_self], images[exclude_self], \ distances[exclude_self], np.array(None)
def get_graphs_within_cutoff( structure: StructureOrMolecule, cutoff: float = 5.0, numerical_tol: float = 1e-8 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Get graph representations from structure within cutoff Args: structure (pymatgen Structure or molecule) cutoff (float): cutoff radius numerical_tol (float): numerical tolerance Returns: center_indices, neighbor_indices, images, distances """ if isinstance(structure, Structure): lattice_matrix = np.ascontiguousarray(np.array( structure.lattice.matrix), dtype=float) pbc = np.array([1, 1, 1], dtype=int) elif isinstance(structure, Molecule): lattice_matrix = np.array( [[1000.0, 0.0, 0.0], [0.0, 1000.0, 0.0], [0.0, 0.0, 1000.0]], dtype=float) pbc = np.array([0, 0, 0], dtype=int) else: raise ValueError("structure type not supported") r = float(cutoff) cart_coords = np.ascontiguousarray(np.array(structure.cart_coords), dtype=float) center_indices, neighbor_indices, images, distances = find_points_in_spheres( cart_coords, cart_coords, r=r, pbc=pbc, lattice=lattice_matrix, tol=numerical_tol) center_indices = center_indices.astype(DataType.np_int) neighbor_indices = neighbor_indices.astype(DataType.np_int) images = images.astype(DataType.np_int) distances = distances.astype(DataType.np_float) exclude_self = (center_indices != neighbor_indices) | (distances > numerical_tol) return center_indices[exclude_self], neighbor_indices[ exclude_self], images[exclude_self], distances[exclude_self]
def get_graphs_within_cutoff(structure: Union[Structure, MEGNetMolecule, Molecule], cutoff: float = 5.0, numerical_tol: float = 1e-8): """ Get graph representations from structure within cutoff Args: structure: (pymatgen Structure) cutoff: (float) cutoff radius numerical_tol: (float) numerical tolerance Returns: center_indices, neighbor_indices, images, distances """ if isinstance(structure, Structure): lattice_matrix = np.ascontiguousarray(np.array( structure.lattice.matrix), dtype=float) pbc = np.array([1, 1, 1], dtype=int) elif isinstance(structure, MEGNetMolecule) or isinstance( structure, Molecule): lattice_matrix = np.array( [[1000.0, 0., 0.], [0., 1000., 0.], [0., 0., 1000.]], dtype=float) pbc = np.array([0, 0, 0], dtype=int) else: raise ValueError('structure type not supported') r = float(cutoff) cart_coords = np.ascontiguousarray(np.array(structure.cart_coords), dtype=float) center_indices, neighbor_indices, images, distances = \ find_points_in_spheres(cart_coords, cart_coords, r=r, pbc=pbc, lattice=lattice_matrix, tol=numerical_tol) exclude_self = (center_indices != neighbor_indices) | (distances > numerical_tol) return center_indices[exclude_self], neighbor_indices[ exclude_self], images[exclude_self], distances[exclude_self]