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 125