Coverage for icet/core/neighbor_list.py: 93%

1""" 

2This module provides a simple wrapper of the ASE Neighbor List class, 

3returning a list of Lattice Sites. 

4""" 

5from typing import List, Union 

6from ase import Atoms 

7from ase.neighborlist import NeighborList as ASENeighborList 

8from .lattice_site import LatticeSite 

9from .structure import Structure 

10 

11 

12def get_neighbor_lists(structure: Union[Atoms, Structure], 

13 cutoffs: List, 

14 position_tolerance: float = 1e-5) -> List[List]: 

15 """ 

16 Returns a list of :program:`icet` neighbor lists given a configuration and cutoffs. 

17 

18 Parameters 

19 ---------- 

20 structure 

21 Atomic configuration. 

22 cutoffs 

23 Positive floats indicating the cutoffs for the various clusters. 

24 position_tolerance 

25 Tolerance applied when comparing positions in Cartesian coordinates. 

26 """ 

27 # deal with different types of structure objects 

28 if isinstance(structure, Structure): 

29 structure = Structure.to_atoms(structure) 

30 elif not isinstance(structure, Atoms): 30 ↛ 31line 30 didn't jump to line 31, because the condition on line 30 was never true

31 msg = ['Unknown structure format'] 

32 msg += [f'{type(structure)} (get_neighbor_list)'] 

33 raise Exception(' '.join(msg)) 

34 

35 neighbor_lists = [] 

36 for cutoff in cutoffs: 

37 neighbor_list = [] 

38 ase_nl = ASENeighborList(len(structure)*[cutoff/2.0], 

39 skin=2*position_tolerance, bothways=True, 

40 self_interaction=False) 

41 ase_nl.update(structure) 

42 

43 for i in range(len(structure)): 

44 ase_indices, ase_offsets = ase_nl.get_neighbors(i) 

45 site = [] 

46 # Update the final list of sites with LatticeSite objects 

47 for index, offset in zip(ase_indices, ase_offsets): 

48 element = LatticeSite(index, offset) 

49 site.append(element) 

50 # sort by unitcell_offset first 

51 # have to cast to a list, since numpy array isn't iterable 

52 site = sorted(site, key=lambda s: list(s.unitcell_offset)) 

53 # now sort by index 

54 site = sorted(site, key=lambda s: s.index) 

55 

56 neighbor_list.append(site) 

57 neighbor_lists.append(neighbor_list) 

58 

59 # This returns list[list[list[LatticeSite]]], which corresponds to 

60 # <std::vector<std::vector<std::vector<LatticeSite>>>> 

61 return neighbor_lists