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]