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