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

45 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-12-10 11:04 +0000

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 List of Atoms object that connected_indices calls connected. 

54 """ 

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

56 

57 separated = [] 

58 while indices: 

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

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

61 for i in my_indcs: 

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

63 del indices[indices.index(i)] 

64 

65 return separated 

66 

67 

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

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

70 

71 index1: index of first atom 

72 index2: index of second atom 

73 kwargs: kwargs transferred to connected_atoms 

74 

75 Returns two Atoms objects 

76 """ 

77 assert index1 != index2 

78 if index2 > index1: 

79 shift = 0, 1 

80 else: 

81 shift = 1, 0 

82 

83 atoms_copy = atoms.copy() 

84 del atoms_copy[index2] 

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

86 

87 atoms_copy = atoms.copy() 

88 del atoms_copy[index1] 

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

90 

91 return atoms1, atoms2