1import numpy as np 

2 

3from ase.parallel import world, broadcast 

4from ase.geometry import get_distances 

5 

6 

7def random_unit_vector(rng): 

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

9 

10 Parameter 

11 --------- 

12 rng: random number generator object 

13 """ 

14 ct = -1 + 2 * rng.rand() 

15 phi = 2 * np.pi * rng.rand() 

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

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

18 

19 

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

21 """Return indices of nearest atoms""" 

22 p1 = atoms1.get_positions() 

23 p2 = atoms2.get_positions() 

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

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

26 return i1, i2, vd_aac[i1, i2] 

27 

28 

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

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

31 """Attach two structures 

32 

33 Parameters 

34 ---------- 

35 atoms1: Atoms 

36 cell and pbc of this object are used 

37 atoms2: Atoms 

38 distance: float 

39 minimal distance (Angstrom) 

40 direction: unit vector (3 floats) 

41 relative direction between center of masses 

42 maxiter: int 

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

44 accuracy: float 

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

46 """ 

47 atoms = atoms1.copy() 

48 atoms2 = atoms2.copy() 

49 

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

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

52 assert len(direction) == 3 

53 dist2 = distance**2 

54 

55 i1, i2, dv_c = nearest(atoms, atoms2, atoms.cell, atoms.pbc) 

56 

57 for i in range(maxiter): 

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

59 

60 vcost = np.dot(dv_c, direction) 

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

62 move = a - vcost 

63 if abs(move) < accuracy: 

64 atoms += atoms2 

65 return atoms 

66 

67 # we need to move 

68 atoms2.translate(direction * move) 

69 i1, i2, dv_c = nearest(atoms, atoms2, atoms.cell, atoms.pbc) 

70 

71 raise RuntimeError('attach did not converge') 

72 

73 

74def attach_randomly(atoms1, atoms2, distance, 

75 rng=np.random): 

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

77 

78 Parameters 

79 ---------- 

80 atoms1: Atoms object 

81 atoms2: Atoms object 

82 distance: float 

83 Required distance 

84 rng: random number generator object 

85 defaults to np.random.RandomState() 

86 

87 Returns 

88 ------- 

89 Joined structure as an atoms object. 

90 """ 

91 atoms2 = atoms2.copy() 

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

93 center=atoms2.get_center_of_mass()) 

94 return attach(atoms1, atoms2, distance, 

95 direction=random_unit_vector(rng)) 

96 

97 

98def attach_randomly_and_broadcast(atoms1, atoms2, distance, 

99 rng=np.random, 

100 comm=world): 

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

102 and ensure that these are distributed. 

103 

104 Parameters 

105 ---------- 

106 atoms1: Atoms object 

107 atoms2: Atoms object 

108 distance: float 

109 Required distance 

110 rng: random number generator object 

111 defaults to np.random.RandomState() 

112 comm: communicator to distribute 

113 Communicator to distribute the structure, default: world 

114 

115 Returns 

116 ------- 

117 Joined structure as an atoms object. 

118 """ 

119 if comm.rank == 0: 

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

121 broadcast(joined, 0, comm=comm) 

122 else: 

123 joined = broadcast(None, 0, comm) 

124 return joined