Neighbor Lists

Overview

The aenet-python package provides neighbor list functionality through the TorchNeighborList class exported by aenet.torch_featurize. The neighbor list is fully integrated with AtomicStructure objects and can be used both through high-level convenience methods and low-level direct access.

For a longer workflow-oriented walkthrough, including low-level edge access and optional GPU execution, see notebooks/example-07-neighbor-list.ipynb.

Note

The features described here make use of PyTorch. Make sure to install core torch support plus the matching torch-scatter and torch-cluster wheels as described in Installation & Set-up. For use without PyTorch, aenet-python also provides a (less efficient) pure-Python neighbor list implementation in aenet.geometry.nblist, which can be used with AtomicStructure objects.

Quick Start

Using with AtomicStructure

The simplest way to use the neighbor list is through AtomicStructure.get_neighbors():

>>> import numpy as np
>>> from aenet.geometry import AtomicStructure
>>> structure = AtomicStructure(
...     np.array([
...         [0.0, 0.0, 0.0],
...         [1.5, 0.0, 0.0],
...         [0.0, 1.5, 0.0],
...     ]),
...     ['O', 'H', 'H'],
... )
>>> neighbors = structure.get_neighbors(i=0, cutoff=2.0)
>>> neighbors.natoms
3
>>> neighbors.types.tolist()
['O', 'H', 'H']

Direct Usage

For more control, use TorchNeighborList directly:

>>> import numpy as np
>>> from aenet.torch_featurize import TorchNeighborList
>>> nbl = TorchNeighborList(cutoff=4.0, device='cpu')
>>> positions = np.array([
...     [0.0, 0.0, 0.0],
...     [1.5, 0.0, 0.0],
...     [3.0, 0.0, 0.0],
... ])
>>> result = nbl.get_neighbors_of_atom(0, positions)
>>> result['indices'].cpu().tolist()
[1, 2]
>>> [round(float(d), 1) for d in result['distances'].cpu().tolist()]
[1.5, 3.0]
>>> result['offsets'] is None
True

Configuration Options

Cutoff Radius

The primary parameter controlling neighbor search:

# Standard cutoff
nbl = TorchNeighborList(cutoff=4.0)

# Larger cutoff for long-range interactions
nbl = TorchNeighborList(cutoff=8.0)

Maximum Number of Neighbors

By default, the neighbor list can handle up to 256 neighbors per atom. For dense systems, increase this limit:

# Standard systems
nbl = TorchNeighborList(cutoff=4.0, max_num_neighbors=256)

# Dense systems (liquids, high-pressure solids)
nbl = TorchNeighborList(cutoff=4.0, max_num_neighbors=512)

# Very dense systems
nbl = TorchNeighborList(cutoff=4.0, max_num_neighbors=1024)

If the number of neighbors exceeds this limit during a call, the limit is automatically increased and a warning is issued. However, this should be avoided for performance reasons.

Periodic Boundary Conditions

For periodic systems, provide lattice vectors. When cell is given, positions are interpreted as fractional coordinates by default. Pass fractional=False when your coordinates are Cartesian:

import numpy as np
from aenet.torch_featurize import TorchNeighborList

cell = np.eye(3) * 5.0
cartesian_positions = np.array([
    [0.5, 2.5, 2.5],
    [4.5, 2.5, 2.5],
])

nbl = TorchNeighborList(cutoff=2.0, device='cpu')
result = nbl.get_neighbors_of_atom(
    0,
    cartesian_positions,
    cell=cell,
    fractional=False,
)

# Offsets show which periodic images each neighbor belongs to
print(result['offsets'])  # tensor([[-1, 0, 0]])

Advanced Features

Getting Neighbor Coordinates Directly

Instead of manually computing neighbor positions from indices and offsets, use return_coordinates=True:

import numpy as np
from aenet.torch_featurize import TorchNeighborList

cell = np.eye(3) * 5.0
cartesian_positions = np.array([
    [0.5, 2.5, 2.5],
    [4.5, 2.5, 2.5],
])
nbl = TorchNeighborList(cutoff=2.0, device='cpu')

# For isolated systems
isolated_positions = np.array([
    [0.0, 0.0, 0.0],
    [1.5, 0.0, 0.0],
    [3.0, 0.0, 0.0],
])
result = nbl.get_neighbors_of_atom(0, isolated_positions,
                                   return_coordinates=True)
neighbor_coords = result['coordinates']  # Actual 3D positions

# For periodic systems - offsets are automatically applied
result = nbl.get_neighbors_of_atom(
    0,
    cartesian_positions,
    cell=cell,
    fractional=False,
    return_coordinates=True,
)
neighbor_coords = result['coordinates']  # PBC offsets already applied

This is especially convenient for periodic systems, where the neighbor list automatically applies cell offsets to compute actual Cartesian coordinates.

Type-Dependent Cutoffs

For multi-component systems, different atom pairs may require different cutoffs:

import numpy as np
import torch
from aenet.torch_featurize import TorchNeighborList

positions = np.array([
    [0.00000, 0.00000, 0.11779],
    [0.00000, 0.75545, -0.47116],
    [0.00000, -0.75545, -0.47116],
])
atom_types = torch.tensor([8, 1, 1])  # O, H, H

# Define pair-specific cutoffs
cutoff_dict = {
    (1, 1): 1.0,
    (1, 8): 2.5,
    (8, 8): 3.0,
}

# Create neighbor list with type information
nbl = TorchNeighborList(
    cutoff=5.0,  # Must be >= every pair-specific cutoff
    atom_types=atom_types,
    cutoff_dict=cutoff_dict,
    device='cpu',
)

oxygen_neighbors = nbl.get_neighbors_of_atom(0, positions)
hydrogen_neighbors = nbl.get_neighbors_of_atom(1, positions)

print(oxygen_neighbors['indices'])   # tensor([1, 2])
print(hydrogen_neighbors['indices'])  # tensor([0])

Per-Atom Neighbor Access

Get neighbors for all atoms at once:

import torch
from aenet.torch_featurize import TorchNeighborList

positions = torch.tensor([
    [0.0, 0.0, 0.0],
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
], dtype=torch.float64)
nbl = TorchNeighborList(cutoff=1.5, device='cpu')

# Returns list of neighbor dicts, one per atom
all_neighbors = nbl.get_neighbors_by_atom(positions)

for i, atom_neighbors in enumerate(all_neighbors):
    print(f"Atom {i}: {len(atom_neighbors['indices'])} neighbors")

Factory Method from AtomicStructure

Create a neighbor list pre-configured for a structure:

>>> import numpy as np
>>> from aenet.geometry import AtomicStructure
>>> from aenet.torch_featurize import TorchNeighborList
>>> structure = AtomicStructure(
...     np.array([
...         [0.0, 0.0, 0.0],
...         [1.5, 0.0, 0.0],
...         [0.0, 1.5, 0.0],
...     ]),
...     ['O', 'H', 'H'],
... )
>>> nbl = TorchNeighborList.from_AtomicStructure(
...     structure,
...     cutoff=2.0,
...     device='cpu',
... )
>>> nbl.cutoff
2.0
>>> nbl.max_num_neighbors
256

GPU Acceleration

For large systems, use GPU acceleration:

# Create neighbor list on GPU
nbl = TorchNeighborList(cutoff=4.0, device='cuda')

# Input arrays are automatically moved to GPU
result = nbl.get_neighbors_of_atom(0, positions)

# Results are on GPU - convert back to numpy if needed
neighbor_indices = result['indices'].cpu().numpy()

Performance Considerations

Caching

The neighbor list automatically caches results:

import torch
from aenet.torch_featurize import TorchNeighborList

positions = torch.tensor([
    [0.0, 0.0, 0.0],
    [1.0, 0.0, 0.0],
    [0.0, 1.0, 0.0],
], dtype=torch.float64)
nbl = TorchNeighborList(cutoff=1.5, device='cpu')

# First call: computes neighbor list
result1 = nbl.get_neighbors_of_atom(0, positions)

# Subsequent calls with same positions: uses cache
result2 = nbl.get_neighbors_of_atom(1, positions)
result3 = nbl.get_neighbors_of_atom(2, positions)

# Cache invalidated when positions change
new_positions = positions + 0.1
result4 = nbl.get_neighbors_of_atom(0, new_positions)  # Recomputes

See Also

• PyTorch-Based Featurization - PyTorch-based featurization

• Structure conversion and manipulation - Working with AtomicStructure

• example-07-neighbor-list.ipynb - longer low-level, GPU, and type-dependent neighbor-list examples