# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
from __future__ import annotations
from functools import cached_property
from numbers import Integral
from typing import Optional
import numpy as np
from sisl import Geometry
from sisl._internal import set_module
__all__ = [
"UniqueNeighborList",
"FullNeighborList",
"PartialNeighborList",
"AtomNeighborList",
"PointsNeighborList",
"PointNeighborList",
]
@set_module("sisl.geom")
class Neighbors:
def __init__(
self,
geometry: Geometry,
finder_results: np.ndarray,
split_indices: Optional[np.ndarray] = None,
):
self.geometry = geometry
self._finder_results = finder_results
self._split_indices = split_indices
@property
def i(self) -> np.ndarray:
"""For each neighbor pair (i, j), the first index."""
return self._finder_results[:, 0]
@property
def j(self) -> np.ndarray:
"""For each neighbor pair (i, j), the second index."""
return self._finder_results[:, 1]
@property
def isc(self) -> np.ndarray:
"""For each neighbor pair (i, j), the supercell indices of `j`."""
return self._finder_results[:, 2:]
@cached_property
def min_nsc(self) -> np.ndarray:
"""The minimum nsc for an auxiliary supercell to contain all neighbors in this object."""
if len(self.isc) == 0:
return np.ones(3, dtype=int)
else:
return np.max(np.abs(self.isc), axis=0) * 2 + 1
@property
def n_neighbors(self):
raise NotImplementedError
@cached_property
def split_indices(self) -> np.ndarray:
"""Indices to split the interactions of each atom."""
if self._split_indices is None:
return np.cumsum(self.n_neighbors)
else:
return self._split_indices
def __len__(self):
return len(self.split_indices)
@set_module("sisl.geom")
class AtomsNeighborList(Neighbors):
"""Base class for interactions between atoms."""
@cached_property
def n_neighbors(self) -> np.ndarray:
"""Number of neighbors that each atom has."""
if self._split_indices is None:
n_neighbors = np.zeros(self.geometry.na, dtype=int)
index, counts = np.unique(self.i, return_counts=True)
n_neighbors[index] = counts
return n_neighbors
else:
return np.diff(self._split_indices, prepend=0)
@set_module("sisl.geom")
class UniqueNeighborList(AtomsNeighborList):
"""Full neighbors list of a system, but **containing only the upper triangle of the adjacency matrix**.
What this means, is that the the object only contains one direction of each interaction.
This is only possible if the interaction is symmetric (there is no directionality and
thresholds for interaction do not depend on direction).
Examples
--------
You can get a unique neighbors list from the `find_unique_pairs` method of a `NeighborFinder` object.
Then, you can retreive the neighbors from it:
.. code-block:: python
import sisl
# Build a graphene supercell with a vacancy
graphene = sisl.geom.graphene().tile(2, 0).tile(2, 1)
graphene = graphene.remove(2).translate2uc()
# Initialize a finder for neighbors that are within 1.5 Angstrom
finder = sisl.geom.NeighborFinder(graphene, R=1.5)
# Get the list of unique neighbor pairs
neighbors = finder.find_unique_pairs()
# You can get the neighbor pairs (i,j) from the i and j attributes
# The supercell index of atom J is in the isc attribute.
print("ATOM I", neighbors.i)
print("ATOM J (NEIGHBOR)", neighbors.j)
print("NEIGHBORS ISC:", neighbors.isc)
# Notice that I is always smaller than J. Each connection is only
# stored once.
# You can convert to a full neighbors list, which will contain both
# directions.
full_neighbors = neighbors.to_full()
"""
@cached_property
def n_neighbors(self) -> np.ndarray:
"""Number of neighbors that each atom has."""
if self._split_indices is None:
n_neighbors = np.zeros(self.geometry.na, dtype=int)
index, counts = np.unique([self.i, self.j], return_counts=True)
n_neighbors[index] = counts
return n_neighbors
else:
return np.diff(self._split_indices, prepend=0)
[docs]
def to_full(self) -> FullNeighborList:
"""Converts the unique neighbors list to a full neighbors list."""
upper_tri = self._finder_results
lower_tri = np.column_stack(
[self.j, self.i, -self.isc[:, 0], -self.isc[:, 1], -self.isc[:, 2]]
)
self_interactions = (self.i == self.j) & np.all(self.isc == 0, axis=1)
lower_tri = lower_tri[~self_interactions]
# Concatenate the lower triangular with the upper triangular part
all_finder_results = np.concatenate([upper_tri, lower_tri], axis=0)
# Sort by i and then by j
sorted_indices = np.lexsort(all_finder_results[:, [1, 0]].T)
all_finder_results = all_finder_results[sorted_indices]
return FullNeighborList(
self.geometry,
all_finder_results,
)
@set_module("sisl.geom")
class FullNeighborList(AtomsNeighborList):
"""Full neighbors list of a system.
This class, contrary to `UniqueNeighborList`, (possibly) contains the two directions
of an interaction between two given atoms. Notice that it is possible that there is
a connection from atom `i` to atom `j` but not the other way around.
Examples
--------
You can get a full neighbors list from the `find_neighbors` method of a `NeighborFinder` object.
Then, you can retreive the neighbors from it:
.. code-block:: python
import sisl
# Build a graphene supercell with a vacancy
graphene = sisl.geom.graphene().tile(2, 0).tile(2, 1)
graphene = graphene.remove(2).translate2uc()
# Initialize a finder for neighbors that are within 1.5 Angstrom
finder = sisl.geom.NeighborFinder(graphene, R=1.5)
# Get the full neighbors list
neighbors = finder.find_neighbors()
# You can loop through atoms to get their neighbors
for at_neighs in neighbors:
print()
print(f"NEIGHBORS OF ATOM {at_neighs.atom} ({at_neighs.n_neighbors} neighbors): ")
print("J", at_neighs.j)
print("ISC", at_neighs.isc)
# Or get the neighbors of a particular atom:
neighbors[0].j
See Also
--------
AtomNeighborList
The object returned by this list when iterating or indexing.
"""
def __getitem__(self, item) -> AtomNeighborList:
"""Returns the interactions of a given atom."""
if isinstance(item, Integral):
start = 0 if item == 0 else self.split_indices[item - 1]
end = self.split_indices[item]
return AtomNeighborList(
self.geometry,
self._finder_results[start:end],
atom=item,
)
else:
raise ValueError("Only integer indexing is supported.")
[docs]
def to_unique(self) -> UniqueNeighborList:
"""Converts the full neighbors list to a unique neighbors list."""
full_finder_results = self._finder_results
unique_finder_results = full_finder_results[self.i <= self.j]
# Concatenate the uc connections with the rest of the connections.
return UniqueNeighborList(
geometry=self.geometry, finder_results=unique_finder_results
)
@set_module("sisl.geom")
class PartialNeighborList(AtomsNeighborList):
"""Neighbors list containing only the neighbors of some atoms.
Examples
--------
You can get a partial neighbors list from the `find_neighbors` method of a
`NeighborFinder` object if you pass the `atoms` argument. Then, you can
retreive the neighbors from it:
.. code-block:: python
import sisl
# Build a graphene supercell with a vacancy
graphene = sisl.geom.graphene().tile(2, 0).tile(2, 1)
graphene = graphene.remove(2).translate2uc()
# Initialize a finder for neighbors that are within 1.5 Angstrom
finder = sisl.geom.NeighborFinder(graphene, R=1.5)
# Get a partial neighbors list
neighbors = finder.find_neighbors(atoms=[2, 4])
# You can loop through atoms to get their neighbors
# In this case, the loop will go through atoms 2 and 4
for at_neighs in neighbors:
print()
print(f"NEIGHBORS OF ATOM {at_neighs.atom} ({at_neighs.n_neighbors} neighbors): ")
print("J", at_neighs.j)
print("ISC", at_neighs.isc)
# Or get the neighbors of a particular atom
neighbors[0].atom # This will be 2
neighbors[0].j
See Also
--------
AtomNeighborList
The object returned by this list when iterating or indexing.
"""
#: The atoms for which the neighbors are stored.
atoms: np.ndarray
def __init__(
self, geometry: Geometry, finder_results, atoms: np.ndarray, split_indices=None
):
self.atoms = atoms
super().__init__(geometry, finder_results, split_indices)
@cached_property
def n_neighbors(self) -> np.ndarray:
"""Number of neighbors that each atom has."""
if self._split_indices is None:
return np.array([np.sum(self.i == at) for at in self.atoms])
else:
return np.diff(self._split_indices, prepend=0)
def __getitem__(self, item) -> AtomNeighborList:
"""Returns the interactions of a given atom."""
if isinstance(item, Integral):
start = 0 if item == 0 else self.split_indices[item - 1]
end = self.split_indices[item]
return AtomNeighborList(
self.geometry,
self._finder_results[start:end],
atom=self.atoms[item],
)
else:
raise ValueError("Only integer indexing is supported.")
@set_module("sisl.geom")
class AtomNeighborList(Neighbors):
"""List of atoms that are close to a given atom.
The usual way to get an `AtomNeighborList` object is by iterating over a
`FullNeighborList` or a `PartialNeighborList`. See their documentation for
examples.
See Also
--------
FullNeighborList, PartialNeighborList
The lists that, when iterated, return `AtomNeighborList` objects.
"""
#: The atom for which the neighbors are stored.
atom: int
def __init__(self, geometry, finder_results, atom: int):
self.atom = atom
super().__init__(geometry, finder_results)
@cached_property
def n_neighbors(self) -> int:
"""Number of neighbors of the atom."""
return len(self._finder_results)
@set_module("sisl.geom")
class PointsNeighborList(Neighbors):
"""List of atoms that are close to a set of points in space.
Examples
--------
You can get a points neighbors list from the `find_close` method of a
`NeighborFinder` object. Then, you can retreive the neighbors from it:
.. code-block:: python
import sisl
# Build a graphene supercell with a vacancy
graphene = sisl.geom.graphene().tile(2, 0).tile(2, 1)
graphene = graphene.remove(2).translate2uc()
# Initialize a finder for neighbors that are within 1.5 Angstrom
finder = sisl.geom.NeighborFinder(graphene, R=1.5)
# Get the full neighbors list
points = [[0, 0, 0], [2, 0, 0]]
neighbors = finder.find_close(points)
# You can loop through points to get their neighbors
for point_neighs in neighbors:
print()
print(f"NEIGHBORS OF POINT {point_neighs.point} ({point_neighs.n_neighbors} neighbors): ")
print("J", point_neighs.j)
print("ISC", point_neighs.isc)
# Or get the neighbors of a particular point:
neighbors[0].j
See Also
--------
PointNeighborList
The object returned by this list when iterating or indexing.
"""
def __init__(
self, geometry, points: np.ndarray, finder_results, split_indices=None
):
self.points = points
super().__init__(geometry, finder_results, split_indices)
@cached_property
def n_neighbors(self) -> np.ndarray:
"""Number of atoms that are close to each point"""
if self._split_indices is None:
n_neighbors = np.zeros(len(self.points), dtype=int)
index, counts = np.unique(self.i, return_counts=True)
n_neighbors[index] = counts
return n_neighbors
else:
return np.diff(self._split_indices, prepend=0)
def __getitem__(self, item):
"""Returns the interactions of a given point."""
if isinstance(item, Integral):
start = 0 if item == 0 else self.split_indices[item - 1]
end = self.split_indices[item]
return PointNeighborList(
self.geometry,
self.points[item],
self._finder_results[start:end],
)
else:
raise ValueError("Only integer indexing is supported.")
@set_module("sisl.geom")
class PointNeighborList(Neighbors):
"""List of atoms that are close to a point in space.
The usual way to get a `PointNeighborList` object is by iterating over a
`PointsNeighborList`. See its documentation for examples.
See Also
--------
PointsNeighborList
The list that, when iterated, returns `PointNeighborList` objects.
"""
#: The point for which the neighbors are stored.
point: np.ndarray
def __init__(self, geometry, point: np.ndarray, finder_results: np.ndarray):
self.point = point
super().__init__(geometry, finder_results)
@cached_property
def n_neighbors(self) -> int:
"""Number of neighbors of the point."""
return len(self._finder_results)