aenet.torch_featurize.TorchNeighborList

class aenet.torch_featurize.TorchNeighborList(cutoff: float, atom_types: torch.Tensor | None = None, cutoff_dict: Dict[Tuple[int, int], float] | None = None, device: str = 'cpu', dtype: torch.dtype = torch.float64, max_num_neighbors: int = 256, truncation_handling: str = 'auto', auto_multiplier: int = 2, auto_max_neighbors: int = 65536, pbc_backend: str = 'ghost')[source]

PyTorch-based neighbor list for atomic structures.

Supports: - Periodic boundary conditions (PBC) - Isolated systems - GPU acceleration - Double precision

Example

>>> nbl = TorchNeighborList(cutoff=4.0, device='cpu')
>>> positions = torch.randn(10, 3, dtype=torch.float64)
>>> result = nbl.get_neighbors(positions)
>>> edge_index = result['edge_index']  # (2, num_edges)
>>> distances = result['distances']    # (num_edges,)

Notes

  • Default PBC backend is ‘ghost’ which constructs a support set of ghost images near periodic boundaries and performs a single bipartite radius search (support -> central).

  • Offsets returned under PBC are integer lattice offsets defined relative to wrapped fractional coordinates (positions wrapped to [0,1) before offsetting). Downstream code should reconstruct displacements as: r_ij = ((frac[j] + offsets) @ cell) - (frac[i] @ cell) where frac = remainder(positions @ inv(cell), 1.0).

  • Fully differentiable and GPU-compatible. Truncation is handled by auto-growing max_num_neighbors up to auto_max_neighbors with warnings.

__init__(cutoff: float, atom_types: torch.Tensor | None = None, cutoff_dict: Dict[Tuple[int, int], float] | None = None, device: str = 'cpu', dtype: torch.dtype = torch.float64, max_num_neighbors: int = 256, truncation_handling: str = 'auto', auto_multiplier: int = 2, auto_max_neighbors: int = 65536, pbc_backend: str = 'ghost')[source]

Initialize neighbor list.

Parameters:

  • cutoff – Maximum interaction cutoff radius in Angstroms

  • atom_types – (N,) tensor of atom types (e.g., atomic numbers)

  • cutoff_dict – Dict mapping (type_i, type_j) tuples to cutoff distances. Keys should be sorted tuples: (min, max)

  • device – ‘cpu’ or ‘cuda’

  • dtype – torch.float32 or torch.float64 (recommended: float64)

  • max_num_neighbors – Maximum number of neighbors per atom to consider (default: 256). Increase if you encounter systems with very dense neighbor environments.

Raises:

  • ValueError – If cutoff_dict contains types not in atom_types:

  • ValueError – If cutoff_dict values exceed maximum cutoff:

Methods

__init__(cutoff[, atom_types, cutoff_dict, ...])

Initialize neighbor list.

from_AtomicStructure(structure, cutoff[, ...])

Factory method: create neighbor list from AtomicStructure.

get_neighbors(positions[, cell, pbc, fractional])

Unified interface for neighbor finding.

get_neighbors_by_atom(positions[, cell, ...])

Get neighbors for all atoms in structured format.

get_neighbors_isolated(positions)

Find neighbors for isolated system (no PBC).

get_neighbors_of_atom(atom_idx, positions[, ...])

Get neighbors of a specific atom.

get_neighbors_pbc(positions, cell[, pbc, ...])

Find neighbors with periodic boundary conditions.