enerzyme.models.leftnet.core.get_max_neighbors_mask#

enerzyme.models.leftnet.core.get_max_neighbors_mask(natoms, index, atom_distance, max_num_neighbors_threshold)[source]#

Give a mask that filters out edges so that each atom has at most max_num_neighbors_threshold neighbors. Assumes that index is sorted.