find k nearest neighbors within r_cutoff

[yzchen08]

Can vesin find k nearest neighbors within r_cutoff and output idx_i, idx_j, shifts? For now I can only calculate all neighbors within r_cutoff.

calculator = vesin.NeighborList(cutoff=self.r_cutoff, full_list=True)
idx_i, idx_j, shifts, distances = calculator.compute(
    points=atoms.positions, 
    box=atoms.cell.array, 
    periodic=True, 
    quantities="ijSd"
    )

[Luthaf]

In the current state, you'll need to do some post-processing to get the k nearest neighbors: sorting the distances array; sorting the (i, j, S) arrays in the same way; and then taking the first k elements.

It could make sense to add an option on vesin side to do the sorting according to distances (the current sorted argument sorts the pair indices); and then you would only have to slice the arrays to get the data you need.

[dembart]

Here is one possible solution. I'm not sure if it's the fastest one. Also, It would benefit from flexible handling of "quantities".

import numpy as np
from vesin import NeighborList
from ase import Atoms



atoms = Atoms(positions = [[0., 0., 0.]],
              numbers = [1], cell = [[3.0, 0.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 3.0]],
              pbc = True)

r_cut = 10.0

calculator = NeighborList(cutoff=r_cut, full_list=True)
i, j, ij_length, offset, ij_vector = calculator.compute(
    points=atoms.positions,
    box=atoms.cell,
    periodic=True in atoms.pbc, 
    quantities="ijdSD"
)

n_max_neighbors = 12
if n_max_neighbors is not None:
    filtered_indices = []
    for atom_idx in np.unique(i):
        mask = (i == atom_idx)
        indices = np.where(mask)[0]
        sorted_indices = indices[np.argsort(ij_length[mask])]
        filtered_indices.extend(sorted_indices[:n_max_neighbors])
    i = i[filtered_indices]
    j = j[filtered_indices]
    ij_length = ij_length[filtered_indices]
    offset = offset[filtered_indices]
    ij_vector = ij_vector[filtered_indices]