Coverage for /builds/ase/ase/ase/build/attach.py: 86.05%

1# fmt: off 

2 

3import numpy as np 

4 

5from ase.geometry import get_distances 

6from ase.parallel import broadcast, world 

7 

8 

9def random_unit_vector(rng): 

10 """Random unit vector equally distributed on the sphere 

11 

12 Parameter 

13 --------- 

14 rng: random number generator object 

15 """ 

16 ct = -1 + 2 * rng.random() 

17 phi = 2 * np.pi * rng.random() 

18 st = np.sqrt(1 - ct**2) 

19 return np.array([st * np.cos(phi), st * np.sin(phi), ct]) 

20 

21 

22def nearest(atoms1, atoms2, cell=None, pbc=None): 

23 """Return indices of nearest atoms""" 

24 p1 = atoms1.get_positions() 

25 p2 = atoms2.get_positions() 

26 vd_aac, d2_aa = get_distances(p1, p2, cell, pbc) 

27 i1, i2 = np.argwhere(d2_aa == d2_aa.min())[0] 

28 return i1, i2, vd_aac[i1, i2] 

29 

30 

31def attach(atoms1, atoms2, distance, direction=(1, 0, 0), 

32 maxiter=50, accuracy=1e-5): 

33 """Attach two structures 

34 

35 Parameters 

36 ---------- 

37 atoms1: Atoms 

38 cell and pbc of this object are used 

39 atoms2: Atoms 

40 distance: float 

41 minimal distance (Angstrom) 

42 direction: unit vector (3 floats) 

43 relative direction between center of masses 

44 maxiter: int 

45 maximal number of iterations to get required distance, default 100 

46 accuracy: float 

47 required accuracy for minimal distance (Angstrom), default 1e-5 

48 

49 Returns 

50 ------- 

51 Joined structure as an atoms object. 

52 """ 

53 atoms = atoms1.copy() 

54 atoms2 = atoms2.copy() 

55 

56 direction = np.array(direction, dtype=float) 

57 direction /= np.linalg.norm(direction) 

58 assert len(direction) == 3 

59 dist2 = distance**2 

60 

61 _i1, _i2, dv_c = nearest(atoms, atoms2, atoms.cell, atoms.pbc) 

62 

63 for _ in range(maxiter): 

64 dv2 = (dv_c**2).sum() 

65 

66 vcost = np.dot(dv_c, direction) 

67 a = np.sqrt(max(0, dist2 - dv2 + vcost**2)) 

68 move = a - vcost 

69 if abs(move) < accuracy: 

70 atoms += atoms2 

71 return atoms 

72 

73 # we need to move 

74 atoms2.translate(direction * move) 

75 _i1, _i2, dv_c = nearest(atoms, atoms2, atoms.cell, atoms.pbc) 

76 

77 raise RuntimeError('attach did not converge') 

78 

79 

80def attach_randomly(atoms1, atoms2, distance, 

81 rng=np.random): 

82 """Randomly attach two structures with a given minimal distance 

83 

84 Parameters 

85 ---------- 

86 atoms1: Atoms object 

87 atoms2: Atoms object 

88 distance: float 

89 Required distance 

90 rng: random number generator object 

91 defaults to np.random.RandomState() 

92 

93 Returns 

94 ------- 

95 Joined structure as an atoms object. 

96 """ 

97 atoms2 = atoms2.copy() 

98 atoms2.rotate('x', random_unit_vector(rng), 

99 center=atoms2.get_center_of_mass()) 

100 return attach(atoms1, atoms2, distance, 

101 direction=random_unit_vector(rng)) 

102 

103 

104def attach_randomly_and_broadcast(atoms1, atoms2, distance, 

105 rng=np.random, 

106 comm=world): 

107 """Randomly attach two structures with a given minimal distance 

108 and ensure that these are distributed. 

109 

110 Parameters 

111 ---------- 

112 atoms1: Atoms object 

113 atoms2: Atoms object 

114 distance: float 

115 Required distance 

116 rng: random number generator object 

117 defaults to np.random.RandomState() 

118 comm: communicator to distribute 

119 Communicator to distribute the structure, default: world 

120 

121 Returns 

122 ------- 

123 Joined structure as an atoms object. 

124 """ 

125 if comm.rank == 0: 

126 joined = attach_randomly(atoms1, atoms2, distance, rng) 

127 broadcast(joined, 0, comm=comm) 

128 else: 

129 joined = broadcast(None, 0, comm) 

130 return joined