Coverage for ase / build / connected.py: 100.00%

1from ase.atoms import Atoms 

2from ase.data import covalent_radii 

3from ase.neighborlist import NeighborList 

4 

5 

6def connected_atoms(atoms, index, dmax=None, scale=1.5): 

7 """Find all atoms connected to atoms[index] and return them.""" 

8 return atoms[connected_indices(atoms, index, dmax, scale)] 

9 

10 

11def connected_indices(atoms, index, dmax=None, scale=1.5): 

12 """Find atoms connected to atoms[index] and return their indices. 

13 

14 If dmax is not None: 

15 Atoms are defined to be connected if they are nearer than dmax 

16 to each other. 

17 

18 If dmax is None: 

19 Atoms are defined to be connected if they are nearer than the 

20 sum of their covalent radii * scale to each other. 

21 

22 """ 

23 if index < 0: 

24 index = len(atoms) + index 

25 

26 # set neighbor lists 

27 if dmax is None: 

28 # define neighbors according to covalent radii 

29 radii = scale * covalent_radii[atoms.get_atomic_numbers()] 

30 else: 

31 # define neighbors according to distance 

32 radii = [0.5 * dmax] * len(atoms) 

33 nl = NeighborList(radii, skin=0, self_interaction=False, bothways=True) 

34 nl.update(atoms) 

35 

36 connected = [index] + list(nl.get_neighbors(index)[0]) 

37 isolated = False 

38 while not isolated: 

39 isolated = True 

40 for i in connected: 

41 for j in nl.get_neighbors(i)[0]: 

42 if j not in connected: 

43 connected.append(j) 

44 isolated = False 

45 

46 return connected 

47 

48 

49def separate(atoms, **kwargs): 

50 """Split atoms into separated entities 

51 

52 Returns 

53 ------- 

54 List of Atoms object that connected_indices calls connected. 

55 """ 

56 indices = list(range(len(atoms))) 

57 

58 separated = [] 

59 while indices: 

60 my_indcs = connected_indices(atoms, indices[0], **kwargs) 

61 separated.append(Atoms(cell=atoms.cell, pbc=atoms.pbc)) 

62 for i in my_indcs: 

63 separated[-1].append(atoms[i]) 

64 del indices[indices.index(i)] 

65 

66 return separated 

67 

68 

69def split_bond(atoms, index1, index2, **kwargs): 

70 """Split atoms by a bond specified by indices 

71 

72 index1: index of first atom 

73 index2: index of second atom 

74 kwargs: kwargs transferred to connected_atoms 

75 

76 Returns two Atoms objects 

77 """ 

78 assert index1 != index2 

79 if index2 > index1: 

80 shift = 0, 1 

81 else: 

82 shift = 1, 0 

83 

84 atoms_copy = atoms.copy() 

85 del atoms_copy[index2] 

86 atoms1 = connected_atoms(atoms_copy, index1 - shift[0], **kwargs) 

87 

88 atoms_copy = atoms.copy() 

89 del atoms_copy[index1] 

90 atoms2 = connected_atoms(atoms_copy, index2 - shift[1], **kwargs) 

91 

92 return atoms1, atoms2