Statistiques
| Branche: | Tag: | Révision :

dockonsurf / modules / ASANN.py @ b75bf97d

Historique | Voir | Annoter | Télécharger (41,34 ko)

1 9dcee788 Carles Marti
#!/usr/bin/env python3
2 9dcee788 Carles Marti
# -*- coding: utf-8 -*-
3 9dcee788 Carles Marti
4 9dcee788 Carles Marti
"""Python3 implementation of the ASANN algorithm (Anisotropically corrected
5 9dcee788 Carles Marti
Solid-Angle based Nearest-Neighbors) """
6 9dcee788 Carles Marti
7 9dcee788 Carles Marti
# EXTERNAL MODULES IMPORTS ###
8 9dcee788 Carles Marti
import numpy as np
9 9dcee788 Carles Marti
import math
10 9dcee788 Carles Marti
11 9dcee788 Carles Marti
12 9dcee788 Carles Marti
# FUNCTIONS DEFINITION ###
13 9dcee788 Carles Marti
14 9dcee788 Carles Marti
# PRE-PROCESSING FUNCTIONS ##
15 9dcee788 Carles Marti
def add_periodic_images(coords, cell, mode):
16 9dcee788 Carles Marti
    """Explicitely add adjacent or surrounding periodic images.
17 9dcee788 Carles Marti

18 9dcee788 Carles Marti
    Parameters: coords (2D array): List of atoms coordinates to consider.
19 9dcee788 Carles Marti
    Important: direct coordinates are expected (torus topology of side 1) if
20 9dcee788 Carles Marti
    `pbc` is set to True. Shape expected: (nb_atoms, nb_dim), where nb_atoms
21 9dcee788 Carles Marti
    is the number of atoms considered, and nb_dim is the dimensionnality of
22 9dcee788 Carles Marti
    the system.
23 9dcee788 Carles Marti

24 9dcee788 Carles Marti
        cell (numpy 2D array): List of cell vectors to consider for periodic
25 9dcee788 Carles Marti
        boundaries conditions. Shape expected: (nb_dim, nb_dim), where nb_dim
26 9dcee788 Carles Marti
        is the dimensionnality of the system considered. Exemple:
27 9dcee788 Carles Marti
        cell_vectors=[[v1_x, v1_y, v1_z], [v2_x, v2_y, v2_z], [v3_x, v3_y,
28 9dcee788 Carles Marti
        v3_z]]
29 9dcee788 Carles Marti

30 9dcee788 Carles Marti
        mode (str): Determines which periodic images should be included. If
31 9dcee788 Carles Marti
        'adjacent', all adjacent periodic images are included (all periodic
32 9dcee788 Carles Marti
        images sharing a face), increasing the number of coordinates 9-fold.
33 9dcee788 Carles Marti
        If 'full', all surrounding periodic images are included (all periodic
34 9dcee788 Carles Marti
        images sharing a face, an edge or a point), increasing the number of
35 9dcee788 Carles Marti
        coordinates 27-fold.
36 9dcee788 Carles Marti

37 9dcee788 Carles Marti
    Returns: (nb_atoms, new_coords) nb_atoms: number of real atom coordinates
38 9dcee788 Carles Marti
    new_coords: numpy 2D array containing updated coordinates (initial +
39 9dcee788 Carles Marti
    periodic images) Shape: (nb_coords, nb_dim)
40 9dcee788 Carles Marti
    """
41 9dcee788 Carles Marti
    # Initialize new coordinates
42 9dcee788 Carles Marti
    new_coords = np.mod(coords, 1)  # Project coords inside initial cell
43 9dcee788 Carles Marti
44 9dcee788 Carles Marti
    # Iterate over cell vectors
45 9dcee788 Carles Marti
    for vect in np.eye(cell.shape[0]):
46 9dcee788 Carles Marti
        # Add periodic images
47 9dcee788 Carles Marti
        if mode == 'adjacent':
48 9dcee788 Carles Marti
            new_coords = np.vstack((new_coords, coords + vect, coords - vect))
49 9dcee788 Carles Marti
        elif mode == 'full':
50 9dcee788 Carles Marti
            new_coords = np.vstack((new_coords, new_coords + vect,
51 9dcee788 Carles Marti
                                    new_coords - vect))  # Recursive process
52 9dcee788 Carles Marti
            # to include all dimension combinaisons
53 9dcee788 Carles Marti
        else:
54 9dcee788 Carles Marti
            raise NotImplementedError
55 9dcee788 Carles Marti
56 9dcee788 Carles Marti
    return coords.shape[0], new_coords
57 9dcee788 Carles Marti
58 9dcee788 Carles Marti
59 9dcee788 Carles Marti
def get_pbc_vectors(coords, pbc, nb_atoms=None):
60 9dcee788 Carles Marti
    """Compute pairwise vectors with or without periodic boundaries conditions.
61 9dcee788 Carles Marti

62 9dcee788 Carles Marti
    Parameters:
63 9dcee788 Carles Marti
        coords (numpy 2D array): List of atoms coordinates to
64 9dcee788 Carles Marti
        consider. Important: direct coordinates are expected (torus topology of
65 9dcee788 Carles Marti
        side 1) if `pbc` is set to True. Shape expected: (nb_coords, nb_dim),
66 9dcee788 Carles Marti
        where nb_coords is the number of coordinates considered, and nb_dim is
67 9dcee788 Carles Marti
        the dimensionnality of the system.
68 9dcee788 Carles Marti

69 9dcee788 Carles Marti
        pbc (bool): Determines if periodic boundaries conditions should be
70 9dcee788 Carles Marti
        applied. Default to True. If True, coordinates are interpreted as
71 9dcee788 Carles Marti
        direct coordinates and the distance between points is computed as the
72 9dcee788 Carles Marti
        minimum euclidian distance between all duplicates (due to periodic
73 9dcee788 Carles Marti
        boundaries conditions) of these points. If false, coordinates are
74 9dcee788 Carles Marti
        interpreted as cartesian coordinates and the metric used is simply
75 9dcee788 Carles Marti
        the euclidian distance. Note: the minimum image convention is not
76 9dcee788 Carles Marti
        applied if periodic images are already explicitely included.
77 9dcee788 Carles Marti

78 9dcee788 Carles Marti
        nb_atoms (int, optional): Number of real atoms coordinates (i.e. for
79 9dcee788 Carles Marti
        which distances must be computed). This is particularly useful for
80 9dcee788 Carles Marti
        excluding periodic images coordinates as central atoms. The real
81 9dcee788 Carles Marti
        atoms coordinates are supposed to be the first coordinates in `coords`
82 9dcee788 Carles Marti

83 9dcee788 Carles Marti
    Returns: numpy 2D array containing pairwise vectors from `coords`.
84 9dcee788 Carles Marti
    Important: if coords are direct coordinates (i.e. `pbc` is set to True),
85 9dcee788 Carles Marti
    the vectors are in direct coordinates. If coords are in cartesian
86 9dcee788 Carles Marti
    coordinates (i.e. `pbc` is set to False), the vectors are in cartesian
87 9dcee788 Carles Marti
    coordinates. vectors[i,j] = (v_ij_x, v_ij_y, v_ij_z) = (r_j_x - r_i_x,
88 9dcee788 Carles Marti
    r_j_y - r_i_y, r_j_z - r_i_z) Shape : (nb_atoms, nb_coords, nb_dim)
89 9dcee788 Carles Marti
    """
90 9dcee788 Carles Marti
    # Retrieve number of non-virtual atoms (from which distances are computed)
91 9dcee788 Carles Marti
    if nb_atoms is None:
92 9dcee788 Carles Marti
        nb_atoms = coords.shape[0]
93 9dcee788 Carles Marti
94 9dcee788 Carles Marti
    # Compute cartesian vectors
95 9dcee788 Carles Marti
    vectors = coords[np.newaxis, :, :] - coords[:nb_atoms, np.newaxis, :]
96 9dcee788 Carles Marti
97 9dcee788 Carles Marti
    # Applying PBC (if minimum image convention is required)
98 9dcee788 Carles Marti
    if pbc and pbc not in ('adjacent', 'full'):
99 9dcee788 Carles Marti
        # vectors = np.mod(vectors + 0.5, 1) - 0.5 # modulo operation is a
100 9dcee788 Carles Marti
        # bit slower than floor operation...
101 9dcee788 Carles Marti
        vectors -= np.floor(vectors + 0.5)
102 9dcee788 Carles Marti
103 9dcee788 Carles Marti
    return vectors
104 9dcee788 Carles Marti
105 9dcee788 Carles Marti
106 9dcee788 Carles Marti
def get_sorted_distances(coords, pbc, nb_atoms=None, cell=np.eye(3)):
107 9dcee788 Carles Marti
    """Compute distances-sorted pairwise distances and vectors with or
108 9dcee788 Carles Marti
    without periodic boundaries conditions.
109 9dcee788 Carles Marti

110 9dcee788 Carles Marti
    Parameters: coords (numpy 2D array): List of atoms coordinates to
111 9dcee788 Carles Marti
    consider. Important: direct coordinates are expected (torus topology of
112 9dcee788 Carles Marti
    side 1) if `pbc` is set to True. Shape expected: (nb_atoms, nb_dim),
113 9dcee788 Carles Marti
    where nb_atoms is the number of atoms considered, and nb_dim is the
114 9dcee788 Carles Marti
    dimensionnality of the system.
115 9dcee788 Carles Marti

116 9dcee788 Carles Marti
        pbc (bool): Determines if periodic boundaries conditions should be
117 9dcee788 Carles Marti
        applied. Default to True. If True, coordinates are interpreted as
118 9dcee788 Carles Marti
        direct coordinates and the distance between points is computed as the
119 9dcee788 Carles Marti
        minimum euclidian distance between all duplicates (due to periodic
120 9dcee788 Carles Marti
        boundaries conditions) of these points. If false, coordinates are
121 9dcee788 Carles Marti
        interpreted as cartesian coordinates and the metric used is simply
122 9dcee788 Carles Marti
        the euclidian distance.
123 9dcee788 Carles Marti

124 9dcee788 Carles Marti
        nb_atoms (int, optional): Number of real atoms coordinates (i.e. for
125 9dcee788 Carles Marti
        which distances must be computed). This is particularly useful for
126 9dcee788 Carles Marti
        excluding periodic images coordinates as central atoms. If None,
127 9dcee788 Carles Marti
        all `coords` coordinates are considered. Default: None
128 9dcee788 Carles Marti

129 9dcee788 Carles Marti
        cell (2D array, optional): List of cell vectors to consider when
130 9dcee788 Carles Marti
        periodic boundaries conditions are considered. Shape expected: (
131 9dcee788 Carles Marti
        nb_dim, nb_dim), where nb_dim is the dimensionnality of the system
132 9dcee788 Carles Marti
        considered. Important: For this parameter to be taken into account,
133 9dcee788 Carles Marti
        `pbc` must be set to True. Exemple: cell_vectors=[[v1_x, v1_y, v1_z],
134 9dcee788 Carles Marti
        [v2_x, v2_y, v2_z], [v3_x, v3_y, v3_z]] Default: numpy.eye(3) (i.e. a
135 9dcee788 Carles Marti
        cubic cell of side 1).
136 9dcee788 Carles Marti

137 9dcee788 Carles Marti
    Returns: (sorted_distances, sorted_vectors, sorted_indexes)
138 9dcee788 Carles Marti
    sorted_vectors: numpy 3D array containing pairwise vectors from `coords`,
139 9dcee788 Carles Marti
    where each i-th row is sorted by increasing distance with respect to the
140 9dcee788 Carles Marti
    i-th coordinates. Important: The vectors are in cartesian coordinates.
141 9dcee788 Carles Marti
    sorted_vectors[i,j] = (v_ij_x, v_ij_y, v_ij_z) = (r_j_x - r_i_x, r_j_y -
142 9dcee788 Carles Marti
    r_i_y, r_j_z - r_i_z) Shape : (nb_atoms, nb_coords, nb_dim)
143 9dcee788 Carles Marti
    sorted_distances: numpy 2D array containing pairwise distances from
144 9dcee788 Carles Marti
    `coords`, where each i-th row is sorted by increasing distance with
145 9dcee788 Carles Marti
    respect to the i-th coordinates. Important: Cartesian euclidian distance
146 9dcee788 Carles Marti
    is used here. sorted_distances[i,j-1] <= sorted_distances[i,
147 9dcee788 Carles Marti
    j] <= sorted_distances[i,j+1] Shape : (nb_atoms, nb_coords)
148 9dcee788 Carles Marti
    sorted_indexes: numpy 2D array containing for each row the
149 9dcee788 Carles Marti
    distances-sorted indexes of neighbors. In other words, the atom indexed
150 9dcee788 Carles Marti
    sorted_indexes[i,j] is the j-th nearest neighbor of atom i. Shape : (
151 9dcee788 Carles Marti
    nb_atoms, nb_coords)
152 9dcee788 Carles Marti
    """
153 9dcee788 Carles Marti
    # Retrieve number of atoms if not given
154 9dcee788 Carles Marti
    if nb_atoms is None:
155 9dcee788 Carles Marti
        nb_atoms = coords.shape[0]
156 9dcee788 Carles Marti
157 9dcee788 Carles Marti
    # Computes pairwise vectors
158 9dcee788 Carles Marti
    vectors = get_pbc_vectors(coords, pbc, nb_atoms=nb_atoms)  # vector.shape =
159 9dcee788 Carles Marti
    # (nb_atoms, nb_coords, nb_dim)
160 9dcee788 Carles Marti
161 9dcee788 Carles Marti
    # Convert into cartesian coordinates if pbc
162 9dcee788 Carles Marti
    if pbc:
163 9dcee788 Carles Marti
        vectors = np.dot(vectors, cell)  # dot product with cell vectors to
164 9dcee788 Carles Marti
        # have cartesian coordinates (for distance computation)
165 9dcee788 Carles Marti
166 9dcee788 Carles Marti
    # Computes pairwise distances
167 9dcee788 Carles Marti
    distances = np.sqrt(np.sum(vectors ** 2, axis=-1))  # simply the square
168 9dcee788 Carles Marti
    # root of the sum of squared components for each pairwise vector
169 9dcee788 Carles Marti
170 9dcee788 Carles Marti
    # Sorting vectors and distances (with respect to distance) for improved
171 9dcee788 Carles Marti
    # performance of the (CA)SANN algorithm Getting sorted indexes to apply to
172 9dcee788 Carles Marti
    # distances and vectors
173 9dcee788 Carles Marti
    sorted_index_axis1 = np.argsort(distances, axis=-1)  # sorting columns
174 9dcee788 Carles Marti
    sorted_index_axis0 = np.arange(nb_atoms)[:, None]  # keeping rows
175 9dcee788 Carles Marti
    # Rearranging distances and vectors so that each row is sorted by
176 9dcee788 Carles Marti
    # increasing distance. (i.e. distances[i, j-1] <= distances[i,
177 9dcee788 Carles Marti
    # j] <= distances[i, j+1])
178 9dcee788 Carles Marti
    distances = distances[sorted_index_axis0, sorted_index_axis1]
179 9dcee788 Carles Marti
    vectors = vectors[sorted_index_axis0, sorted_index_axis1]
180 9dcee788 Carles Marti
181 9dcee788 Carles Marti
    return distances, vectors, sorted_index_axis1
182 9dcee788 Carles Marti
183 9dcee788 Carles Marti
184 9dcee788 Carles Marti
# SANN IMPLEMENTATION ##
185 9dcee788 Carles Marti
def get_SANN(all_distances):
186 9dcee788 Carles Marti
    """Computes coordination numbers according to the SANN algorithm,
187 9dcee788 Carles Marti
    from all pairwise distances.
188 9dcee788 Carles Marti

189 9dcee788 Carles Marti
    Parameters: all_distances: numpy 2D array containing pairwise distances
190 9dcee788 Carles Marti
    from `coords`, where each i-th row is sorted by increasing distance with
191 9dcee788 Carles Marti
    respect to the i-th coordinates. Important: Cartesian euclidian distance
192 9dcee788 Carles Marti
    is used here. sorted_distances[i,j-1] <= sorted_distances[i,
193 9dcee788 Carles Marti
    j] <= sorted_distances[i,j+1] Shape : (nb_atoms, nb_coords)
194 9dcee788 Carles Marti

195 9dcee788 Carles Marti
    Returns: list_sann_CN : Numpy 1D array containing SANN-based coordination
196 9dcee788 Carles Marti
    numbers (i.e. number of neighbors).
197 9dcee788 Carles Marti

198 9dcee788 Carles Marti
        list_sann_radius : Numpy 1D array containing SANN-based coordination
199 9dcee788 Carles Marti
        radius (i.e. coordination sphere radius).
200 9dcee788 Carles Marti
    """
201 9dcee788 Carles Marti
202 9dcee788 Carles Marti
    # Retrieve number of atoms
203 9dcee788 Carles Marti
    nb_coords = all_distances.shape[1]
204 9dcee788 Carles Marti
205 9dcee788 Carles Marti
    # Initialize coordination number vector (i-th element is the coordination
206 9dcee788 Carles Marti
    # number of the i-th atom) and coordination radius
207 9dcee788 Carles Marti
    list_sann_CN = list()
208 9dcee788 Carles Marti
    list_sann_radius = list()
209 9dcee788 Carles Marti
    # Initialize sum of distances vector (i-th element is meant to temporarly
210 9dcee788 Carles Marti
    # store the sum of the i-th atom's 3 nearest neighbors distances)
211 9dcee788 Carles Marti
    list_dist_sum = all_distances[:, 1:4].sum(axis=1)
212 9dcee788 Carles Marti
213 9dcee788 Carles Marti
    # Treat each atom separately (since CN can be different for each atom,
214 9dcee788 Carles Marti
    # Numpy methods are unsuited here)
215 9dcee788 Carles Marti
    for (dist_sum, atom_distances) in zip(list_dist_sum, all_distances):
216 9dcee788 Carles Marti
        sann_CN = 3  # Set CN to 3 (i.e. the minimum CN value for the SANN
217 9dcee788 Carles Marti
        # algorithm) SANN algorithm (i.e. while SANN radius sum(r_ij,1,
218 9dcee788 Carles Marti
        # m)/(m-2) > r_i(m+1), increase m by 1)
219 9dcee788 Carles Marti
        while (sann_CN + 1 < nb_coords) and (
220 9dcee788 Carles Marti
                dist_sum / (sann_CN - 2) >= atom_distances[sann_CN + 1]):
221 9dcee788 Carles Marti
            dist_sum += atom_distances[sann_CN + 1]
222 9dcee788 Carles Marti
            sann_CN += 1
223 9dcee788 Carles Marti
        # Store SANN CN found
224 9dcee788 Carles Marti
        list_sann_CN.append(sann_CN)
225 9dcee788 Carles Marti
        list_sann_radius.append(dist_sum / (sann_CN - 2))
226 9dcee788 Carles Marti
227 9dcee788 Carles Marti
    return (np.array(list_sann_CN), np.array(list_sann_radius))
228 9dcee788 Carles Marti
229 9dcee788 Carles Marti
230 9dcee788 Carles Marti
## ANISOTROPY HANDLING ##
231 9dcee788 Carles Marti
def dist_to_barycenter(nearest_neighbors, nearest_distances, radius):
232 9dcee788 Carles Marti
    """Compute the distance from the central atom to the barycenter of
233 9dcee788 Carles Marti
    nearest neighbors.
234 9dcee788 Carles Marti

235 9dcee788 Carles Marti
    Parameters: nearest_neighbors: numpy 2D array containing nearest
236 9dcee788 Carles Marti
    neighbors vectors from the central atom, sorted by increasing distance
237 9dcee788 Carles Marti
    with respect to the central atom. Important: The vectors must be in
238 9dcee788 Carles Marti
    cartesian coordinates.
239 9dcee788 Carles Marti

240 9dcee788 Carles Marti
        nearest_distances: numpy 1D array containing nearest neighbors
241 9dcee788 Carles Marti
        distances from the central atom, sorted by increasing distance with
242 9dcee788 Carles Marti
        respect to the central atom.
243 9dcee788 Carles Marti

244 9dcee788 Carles Marti
        radius (float): radius R_i_m of the sphere of coordination.
245 9dcee788 Carles Marti

246 9dcee788 Carles Marti
    Returns: dist_bary (float): distance from the central atom to the
247 9dcee788 Carles Marti
    barycenter of nearest neighbors, weighted by relative solid angles
248 9dcee788 Carles Marti

249 9dcee788 Carles Marti
    Computational details: The barycenter is computed using a solid angle
250 9dcee788 Carles Marti
    weight (i.e. the solid angle associated with the corresponding neighbor).
251 9dcee788 Carles Marti
    """
252 9dcee788 Carles Marti
253 9dcee788 Carles Marti
    # Compute solid angles (SA) of neighbors
254 9dcee788 Carles Marti
    list_SA = 1 - nearest_distances / radius
255 9dcee788 Carles Marti
256 9dcee788 Carles Marti
    # Compute SA-based barycenter
257 9dcee788 Carles Marti
    bary_vector = np.sum(nearest_neighbors * list_SA[:, np.newaxis],
258 9dcee788 Carles Marti
                         axis=0) / np.sum(list_SA)
259 9dcee788 Carles Marti
260 9dcee788 Carles Marti
    # Returns distance from the central atom to the barycenter
261 9dcee788 Carles Marti
    return (math.sqrt(np.sum(bary_vector ** 2)), bary_vector)
262 9dcee788 Carles Marti
263 9dcee788 Carles Marti
264 9dcee788 Carles Marti
def angular_correction(nearest_neighbors, nearest_distances, radius):
265 9dcee788 Carles Marti
    """Compute the angular correction `ang_corr`, such that R_i_m = sum(r_ij,
266 9dcee788 Carles Marti
    j=1..m)/(m-2(1-ang_corr)).
267 9dcee788 Carles Marti

268 9dcee788 Carles Marti
    Parameters: nearest_neighbors: numpy 2D array containing the m nearest
269 9dcee788 Carles Marti
    neighbors vectors from the central atom, sorted by increasing distance
270 9dcee788 Carles Marti
    with respect to the central atom. Important: The vectors must be in
271 9dcee788 Carles Marti
    cartesian coordinates.
272 9dcee788 Carles Marti

273 9dcee788 Carles Marti
        nearest_distances: numpy 1D array containing the m nearest neighbors
274 9dcee788 Carles Marti
        distances from the central atom, sorted by increasing distance with
275 9dcee788 Carles Marti
        respect to the central atom.
276 9dcee788 Carles Marti

277 9dcee788 Carles Marti
        radius (float): radius R_i_m of the sphere of coordination.
278 9dcee788 Carles Marti

279 9dcee788 Carles Marti
    Returns:
280 9dcee788 Carles Marti
        ang_corr (float): angular correction.
281 9dcee788 Carles Marti

282 9dcee788 Carles Marti
    Computational details: The angular correction is computed from the
283 9dcee788 Carles Marti
    distance between the nearest neighbor barycenter and the central atom
284 9dcee788 Carles Marti
    `dist_bary`.
285 9dcee788 Carles Marti

286 9dcee788 Carles Marti
        Let us define `alpha` such that: dist_bary = alpha * radius Then,
287 9dcee788 Carles Marti
        mathematical derivations show that: ang_corr = (alpha + sqrt(alpha**2
288 9dcee788 Carles Marti
        + 3*alpha))/3
289 9dcee788 Carles Marti
    """
290 9dcee788 Carles Marti
291 9dcee788 Carles Marti
    # Computing the ratio between the distance to the nearest neighbors
292 9dcee788 Carles Marti
        # barycenter and the radius
293 9dcee788 Carles Marti
    alpha = dist_to_barycenter(nearest_neighbors, nearest_distances, radius)[
294 9dcee788 Carles Marti
                0] / radius
295 9dcee788 Carles Marti
    vector = dist_to_barycenter(nearest_neighbors, nearest_distances, radius)[1]
296 9dcee788 Carles Marti
297 9dcee788 Carles Marti
    # Computing angular correction
298 9dcee788 Carles Marti
    return ((alpha + math.sqrt(alpha ** 2 + 3 * alpha)) / 3, vector)
299 9dcee788 Carles Marti
300 9dcee788 Carles Marti
301 9dcee788 Carles Marti
## ASANN IMPLEMENTATION ##
302 9dcee788 Carles Marti
def get_ASANN(sorted_distances, sorted_vectors, sann_CNs, sann_radii):
303 9dcee788 Carles Marti
    """Update ASANN-based coordination numbers using an angular correction term.
304 9dcee788 Carles Marti

305 9dcee788 Carles Marti
    Parameters: sorted_vectors: numpy 3D array containing pairwise vectors
306 9dcee788 Carles Marti
    from `coords`, where each i-th row is sorted by increasing distance with
307 9dcee788 Carles Marti
    respect to the i-th coordinates. Important: The vectors must be in
308 9dcee788 Carles Marti
    cartesian coordinates. sorted_vectors[i,j] = (v_ij_x, v_ij_y, v_ij_z) = (
309 9dcee788 Carles Marti
    r_j_x - r_i_x, r_j_y - r_i_y, r_j_z - r_i_z) Shape : (nb_atoms,
310 9dcee788 Carles Marti
    nb_coords, nb_dim)
311 9dcee788 Carles Marti

312 9dcee788 Carles Marti
        sorted_distances: numpy 2D array containing pairwise distances from
313 9dcee788 Carles Marti
        `coords`, where each i-th row is sorted by increasing distance with
314 9dcee788 Carles Marti
        respect to the i-th coordinates. Important: Cartesian euclidian
315 9dcee788 Carles Marti
        distance is used here. sorted_distances[i,j-1] <= sorted_distances[i,
316 9dcee788 Carles Marti
        j] <= sorted_distances[i,j+1] Shape : (nb_atoms, nb_coords)
317 9dcee788 Carles Marti

318 9dcee788 Carles Marti
        sann_CNs: Numpy 1D array containing SANN-based coordination numbers (
319 9dcee788 Carles Marti
        i.e. number of neighbors).
320 9dcee788 Carles Marti

321 9dcee788 Carles Marti
        sann_radii: Numpy 1D array containing SANN-based coordination radius
322 9dcee788 Carles Marti
        (i.e. radius of the coordination spheres).
323 9dcee788 Carles Marti

324 9dcee788 Carles Marti
    Returns: list_asann_CN : Numpy 1D array containing ASANN-based
325 9dcee788 Carles Marti
    coordination numbers (i.e. number of neighbors, with an angular
326 9dcee788 Carles Marti
    correction term).
327 9dcee788 Carles Marti

328 9dcee788 Carles Marti
        list_asann_radius : Numpy 1D array containing ASANN-based
329 9dcee788 Carles Marti
        coordination radius (i.e. coordination sphere radius).
330 9dcee788 Carles Marti

331 9dcee788 Carles Marti
    Computational details: ASANN-based coordination number is defined as the
332 9dcee788 Carles Marti
    maximum coordination number m' such that forall m>=m', R_ang_i_m = sum(
333 9dcee788 Carles Marti
    r_ij, j=1..m)/(m-2(1-ang_corr)) < r_i(m+1)
334 9dcee788 Carles Marti

335 9dcee788 Carles Marti
        It is easy to show that R_ang_i_m <= R_i_m, and therefore, m'<=m.
336 9dcee788 Carles Marti
        Consequently, the ASANN algorithm is sure to converge.
337 9dcee788 Carles Marti

338 9dcee788 Carles Marti
        Unlike SC-ASANN algorithm, the angular correction is solely computed
339 9dcee788 Carles Marti
        using the SANN radius (instead of a self-coherent approach, where the
340 9dcee788 Carles Marti
        angular term is defined by (and defines) the ASANN radius itself)
341 9dcee788 Carles Marti
    """
342 9dcee788 Carles Marti
343 9dcee788 Carles Marti
    # Retrieve number of atoms
344 9dcee788 Carles Marti
    nb_coords = sorted_distances.shape[1]
345 9dcee788 Carles Marti
346 9dcee788 Carles Marti
    # Create coordination number vector (i-th element is the coordination
347 9dcee788 Carles Marti
    # number of the i-th atom) and coordination radius
348 9dcee788 Carles Marti
    list_asann_CN = list()
349 9dcee788 Carles Marti
    list_asann_radius = list()
350 9dcee788 Carles Marti
    list_bary_vector = list()
351 9dcee788 Carles Marti
352 9dcee788 Carles Marti
    # Treat each atom separately (since CN can be different for each atom,
353 9dcee788 Carles Marti
    # Numpy methods are unsuited here)
354 9dcee788 Carles Marti
    for (atom_distances, atom_neighbors, sann_CN, sann_radius) in zip(
355 9dcee788 Carles Marti
            sorted_distances, sorted_vectors, sann_CNs, sann_radii):
356 9dcee788 Carles Marti
357 9dcee788 Carles Marti
        # Computes angular correction
358 9dcee788 Carles Marti
        nearest_distances = atom_distances[1:sann_CN + 1]
359 9dcee788 Carles Marti
        nearest_neighbors = atom_neighbors[1:sann_CN + 1]
360 9dcee788 Carles Marti
        ang_corr, vec = angular_correction(nearest_neighbors, nearest_distances,
361 9dcee788 Carles Marti
                                           sann_radius)
362 9dcee788 Carles Marti
        beta = 2 * (1 - ang_corr)
363 9dcee788 Carles Marti
364 9dcee788 Carles Marti
        # ASANN algorithm (i.e. while ASANN radius sum(r_ij, j=1..m)/(m-2*(
365 9dcee788 Carles Marti
        # 1-ang_corr)) >= r_i(m+1), increase m by 1)
366 9dcee788 Carles Marti
        asann_CN = int(
367 9dcee788 Carles Marti
            beta) + 1  # Set CN to floor(2*(1-ang_corr)) + 1 (i.e. the
368 9dcee788 Carles Marti
        # minimum CN value for the ASANN algorithm)
369 9dcee788 Carles Marti
        dist_sum = atom_distances[
370 9dcee788 Carles Marti
                   1:asann_CN + 1].sum()  # Initialize sum of distances
371 9dcee788 Carles Marti
        while (asann_CN + 1 < nb_coords) and (
372 9dcee788 Carles Marti
                dist_sum / (asann_CN - beta) >= atom_distances[asann_CN + 1]):
373 9dcee788 Carles Marti
            dist_sum += atom_distances[asann_CN + 1]
374 9dcee788 Carles Marti
            asann_CN += 1
375 9dcee788 Carles Marti
376 9dcee788 Carles Marti
        # Store ASANN CN and radius found
377 9dcee788 Carles Marti
        list_asann_CN.append(asann_CN)
378 9dcee788 Carles Marti
        list_asann_radius.append(dist_sum / (asann_CN - beta))
379 9dcee788 Carles Marti
        list_bary_vector.append(vec)
380 9dcee788 Carles Marti
381 9dcee788 Carles Marti
    return (np.array(list_asann_CN), np.array(list_asann_radius),
382 9dcee788 Carles Marti
            np.array(list_bary_vector))
383 9dcee788 Carles Marti
384 9dcee788 Carles Marti
385 9dcee788 Carles Marti
# VARIANTS DEFINITIONS ##
386 9dcee788 Carles Marti
def get_self_consistent_ASANN(sorted_distances, sorted_vectors, sann_CNs,
387 9dcee788 Carles Marti
                              radius_eps=1e-2):
388 9dcee788 Carles Marti
    """Update ASANN-based coordination numbers using an angular correction
389 9dcee788 Carles Marti
    term computed in a self-consistent manner.
390 9dcee788 Carles Marti

391 9dcee788 Carles Marti
    Parameters: sorted_vectors: numpy 3D array containing pairwise vectors
392 9dcee788 Carles Marti
    from `coords`, where each i-th row is sorted by increasing distance with
393 9dcee788 Carles Marti
    respect to the i-th coordinates. Important: The vectors must be in
394 9dcee788 Carles Marti
    cartesian coordinates. sorted_vectors[i,j] = (v_ij_x, v_ij_y, v_ij_z) = (
395 9dcee788 Carles Marti
    r_j_x - r_i_x, r_j_y - r_i_y, r_j_z - r_i_z) Shape: (nb_atoms, nb_atoms,
396 9dcee788 Carles Marti
    nb_dim)
397 9dcee788 Carles Marti

398 9dcee788 Carles Marti
        sorted_distances: numpy 2D array containing pairwise distances from
399 9dcee788 Carles Marti
        `coords`, where each i-th row is sorted by increasing distance with
400 9dcee788 Carles Marti
        respect to the i-th coordinates. Important: Cartesian euclidian
401 9dcee788 Carles Marti
        distance is used here. sorted_distances[i,j-1] <= sorted_distances[i,
402 9dcee788 Carles Marti
        j] <= sorted_distances[i,j+1] Shape: (nb_atoms, nb_atoms)
403 9dcee788 Carles Marti

404 9dcee788 Carles Marti
        sann_CNs: Numpy 1D array containing SANN-based coordination numbers (
405 9dcee788 Carles Marti
        i.e. number of neighbors).
406 9dcee788 Carles Marti

407 9dcee788 Carles Marti
        radius_eps: Convergence threshold used for stopping the
408 9dcee788 Carles Marti
        self-consistent radius computation. Default: 1e-2
409 9dcee788 Carles Marti

410 9dcee788 Carles Marti
    Returns: list_asann_CN : Numpy 1D array containing ASANN-based
411 9dcee788 Carles Marti
    coordination numbers (i.e. number of neighbors, with an angular
412 9dcee788 Carles Marti
    correction term).
413 9dcee788 Carles Marti

414 9dcee788 Carles Marti
        list_asann_radius : Numpy 1D array containing ASANN-based
415 9dcee788 Carles Marti
        coordination radius (i.e. coordination sphere radius).
416 9dcee788 Carles Marti

417 9dcee788 Carles Marti
    Computational details: SC-ASANN-based coordination number is defined as
418 9dcee788 Carles Marti
    the maximum coordination number m' such that forall m>=m', R_ang_i_m =
419 9dcee788 Carles Marti
    sum(r_ij, j=1..m)/(m-2(1-ang_corr)) < r_i(m+1)
420 9dcee788 Carles Marti

421 9dcee788 Carles Marti
        Note that the angular correction term is computed here using the ASANN radius (defined itself from the angular correction term). In this approach, the angular correction term is computed in a self-consistent fashion.
422 9dcee788 Carles Marti
    """
423 9dcee788 Carles Marti
424 9dcee788 Carles Marti
    # Create coordination number vector (i-th element is the coordination
425 9dcee788 Carles Marti
    # number of the i-th atom) and coordination radius
426 9dcee788 Carles Marti
    list_asann_CN = list()
427 9dcee788 Carles Marti
    list_asann_radius = list()
428 9dcee788 Carles Marti
    list_vectors = list()
429 9dcee788 Carles Marti
430 9dcee788 Carles Marti
    # Treat each atom separately (since CN can be different for each atom,
431 9dcee788 Carles Marti
    # Numpy methods are unsuited here)
432 9dcee788 Carles Marti
    for (atom_distances, atom_neighbors, sann_CN) in zip(sorted_distances,
433 9dcee788 Carles Marti
                                                         sorted_vectors,
434 9dcee788 Carles Marti
                                                         sann_CNs):
435 9dcee788 Carles Marti
        asann_CN = sann_CN  # Set initial CN to 1 above the maximum CN that
436 9dcee788 Carles Marti
        # can break ASANN relation (i.e. sann_CN-1)
437 9dcee788 Carles Marti
        radius = 0
438 9dcee788 Carles Marti
        prev_radius = 0
439 9dcee788 Carles Marti
440 9dcee788 Carles Marti
        # ASANN update algorithm (i.e. while ASANN radius sum(r_ij, j=1..m)/(
441 9dcee788 Carles Marti
        # m-2(1-ang_corr)) < r_i(m+1), decrease m by 1)
442 9dcee788 Carles Marti
        while True:
443 9dcee788 Carles Marti
            # Extract properties of nearest neighbors
444 9dcee788 Carles Marti
            nearest_distances = atom_distances[1:asann_CN + 1]
445 9dcee788 Carles Marti
            nearest_neighbors = atom_neighbors[1:asann_CN + 1]
446 9dcee788 Carles Marti
447 9dcee788 Carles Marti
            # Computes radius iteratively
448 9dcee788 Carles Marti
            sum_nearest = np.sum(
449 9dcee788 Carles Marti
                nearest_distances)  # store sum of nearest distances
450 9dcee788 Carles Marti
            radius = sum_nearest / (
451 9dcee788 Carles Marti
                    asann_CN - 2)  # set initial radius to SANN value
452 9dcee788 Carles Marti
            delta_radius = math.inf
453 9dcee788 Carles Marti
            # Update radius until convergence
454 9dcee788 Carles Marti
            while delta_radius > radius_eps:
455 9dcee788 Carles Marti
                delta_radius = sum_nearest / (asann_CN - 2 * (
456 9dcee788 Carles Marti
                        1 - angular_correction(nearest_neighbors,
457 9dcee788 Carles Marti
                                               nearest_distances,
458 9dcee788 Carles Marti
                                               radius)[0])) - radius  #
459 9dcee788 Carles Marti
                # new_radius(old_radius) - old_radius
460 9dcee788 Carles Marti
                vec = angular_correction(nearest_neighbors, nearest_distances,
461 9dcee788 Carles Marti
                                         radius)[1]
462 9dcee788 Carles Marti
                radius += delta_radius
463 9dcee788 Carles Marti
464 9dcee788 Carles Marti
            # Check if ASANN relation is broken
465 9dcee788 Carles Marti
            if radius >= atom_distances[asann_CN + 1]:
466 9dcee788 Carles Marti
                break
467 9dcee788 Carles Marti
            asann_CN -= 1
468 9dcee788 Carles Marti
            prev_radius = radius
469 9dcee788 Carles Marti
            if asann_CN < 1:  # ASANN CN is not defined for less than 1
470 9dcee788 Carles Marti
                # neighbor
471 9dcee788 Carles Marti
                break
472 9dcee788 Carles Marti
473 9dcee788 Carles Marti
        # Store ASANN CN and radius found (before breaking the ASANN relation)
474 9dcee788 Carles Marti
        list_asann_CN.append(asann_CN + 1)
475 9dcee788 Carles Marti
        list_asann_radius.append(prev_radius)
476 9dcee788 Carles Marti
        list_vectors.append(vec)
477 9dcee788 Carles Marti
478 9dcee788 Carles Marti
    return (np.array(list_asann_CN), np.array(list_asann_radius),
479 9dcee788 Carles Marti
            np.array(list_vectors))
480 9dcee788 Carles Marti
481 9dcee788 Carles Marti
482 9dcee788 Carles Marti
def convert_continuous(list_CN, list_radius, sorted_distances):
483 9dcee788 Carles Marti
    """Convert integer coordination numbers into continuous decimal values.
484 9dcee788 Carles Marti

485 9dcee788 Carles Marti
    Parameters: list_CN: Numpy 1D array containing discretized coordination
486 9dcee788 Carles Marti
    numbers (i.e. number of neighbors).
487 9dcee788 Carles Marti

488 9dcee788 Carles Marti
        list_radius: Numpy 1D array containing coordination radius (i.e.
489 9dcee788 Carles Marti
        radius of the coordination spheres).
490 9dcee788 Carles Marti

491 9dcee788 Carles Marti
        sorted_distances: Numpy 2D array containing pairwise distances
492 9dcee788 Carles Marti
        between coordinates, where each i-th row is sorted by increasing
493 9dcee788 Carles Marti
        distance with respect to the i-th coordinates.
494 9dcee788 Carles Marti

495 9dcee788 Carles Marti
    Returns: list_continuous_CN: Numpy 1D array containing continuous
496 9dcee788 Carles Marti
    coordination numbers.
497 9dcee788 Carles Marti

498 9dcee788 Carles Marti
        list_weights: 2D array containing the distance-related weights to
499 9dcee788 Carles Marti
        apply to each neighbors of each atom.
500 9dcee788 Carles Marti

501 9dcee788 Carles Marti
    Computational details: In the discretized version, each neighbor is
502 9dcee788 Carles Marti
    comptabilized exactly once. In the continuous version, the contribution
503 9dcee788 Carles Marti
    of each neighbor is scaled by the following weight : SA/SA_max (where SA
504 9dcee788 Carles Marti
    is the solid angle associated with the corresponding neighbor, and SA_max
505 9dcee788 Carles Marti
    is the maximum solid angle (i.e. the solid angle associated with the
506 9dcee788 Carles Marti
    nearest neighbor)).
507 9dcee788 Carles Marti

508 9dcee788 Carles Marti
        The continuous coordination number is then the sum of all weights.
509 9dcee788 Carles Marti
    """
510 9dcee788 Carles Marti
511 9dcee788 Carles Marti
    # Create array to store continuous coordination numbers
512 9dcee788 Carles Marti
    list_continuous_CN = list()
513 9dcee788 Carles Marti
    list_weights = list()
514 9dcee788 Carles Marti
515 9dcee788 Carles Marti
    # Treat each atom separately (since CN can be different for each atom,
516 9dcee788 Carles Marti
    # Numpy methods are unsuited here)
517 9dcee788 Carles Marti
    for (atom_CN, atom_radius, atom_distances) in zip(list_CN, list_radius,
518 9dcee788 Carles Marti
                                                      sorted_distances):
519 9dcee788 Carles Marti
        # Extract distances of nearest neighbors
520 9dcee788 Carles Marti
        nearest_distances = atom_distances[1:atom_CN + 1]
521 9dcee788 Carles Marti
522 9dcee788 Carles Marti
        # Compute solid angles of nearest neighbors
523 9dcee788 Carles Marti
        nearest_SA = 1 - nearest_distances / atom_radius
524 9dcee788 Carles Marti
525 9dcee788 Carles Marti
        # Retrieve maximum solid angle (whose contribution will be 1)
526 9dcee788 Carles Marti
        max_SA = nearest_SA[0]
527 9dcee788 Carles Marti
528 9dcee788 Carles Marti
        # Compute and store weights
529 9dcee788 Carles Marti
        weights = nearest_SA / max_SA
530 9dcee788 Carles Marti
        list_weights.append(weights)
531 9dcee788 Carles Marti
532 9dcee788 Carles Marti
        # Compute sum of contributions (i.e. continuous CN)
533 9dcee788 Carles Marti
        continuous_CN = np.sum(weights)
534 9dcee788 Carles Marti
        list_continuous_CN.append(continuous_CN)
535 9dcee788 Carles Marti
536 9dcee788 Carles Marti
    return np.array(list_continuous_CN), list_weights
537 9dcee788 Carles Marti
538 9dcee788 Carles Marti
539 9dcee788 Carles Marti
def get_generalized(list_CN, list_weights, sorted_indexes, max_CN=None):
540 9dcee788 Carles Marti
    """Convert coordination numbers into generalized coordination numbers.
541 9dcee788 Carles Marti

542 9dcee788 Carles Marti
    Parameters:
543 9dcee788 Carles Marti
        list_CN: Numpy 1D array containing the coordination numbers.
544 9dcee788 Carles Marti

545 9dcee788 Carles Marti
        list_weights: 2D list containing the distance-related weights to
546 9dcee788 Carles Marti
        apply to each neighbors of each atom. Note: In the case of
547 9dcee788 Carles Marti
        discretized version, each weight is equal to 1.
548 9dcee788 Carles Marti

549 9dcee788 Carles Marti
        sorted_indexes: Numpy 2D array containing for each row the
550 9dcee788 Carles Marti
        distances-sorted indexes of neighbors. In other words, the atom
551 9dcee788 Carles Marti
        indexed sorted_indexes[i,j] is the j-th nearest neighbor of atom i.
552 9dcee788 Carles Marti
        Shape = (nb_atoms, nb_coords)
553 9dcee788 Carles Marti

554 9dcee788 Carles Marti
        max_CN (int, optional): Value to use for the maximum coordination
555 9dcee788 Carles Marti
        number, i.e. bulk coordination. If None, the maximum/bulk
556 9dcee788 Carles Marti
        coordination number will be guessed. Default: None
557 9dcee788 Carles Marti

558 9dcee788 Carles Marti
    Returns: list_generalized_CN: Numpy 1D array containing generalized
559 9dcee788 Carles Marti
    coordination numbers.
560 9dcee788 Carles Marti

561 9dcee788 Carles Marti
    Computational details: The generalized coordination number algorithm
562 9dcee788 Carles Marti
    scales the contributions of each neighbor by the following weight :
563 9dcee788 Carles Marti
    CN/CN_max (where CN is the coordination number associated with the
564 9dcee788 Carles Marti
    corresponding neighbor, and CN_max is the maximum coordination achievable
565 9dcee788 Carles Marti
    (i.e. bulk coordination)) The generalized coordination number is then the
566 9dcee788 Carles Marti
    sum of all weights (GCN = sum(CN/max_CN, neighbors))
567 9dcee788 Carles Marti

568 9dcee788 Carles Marti
        This sum can be weighted futhermore (GCN = sum(CN/max_CN*SA/max_SA,
569 9dcee788 Carles Marti
        neighbors)), using the continuous algorithm if requested (see `help(
570 9dcee788 Carles Marti
        convert_continuous)` for more details on this algorithm).
571 9dcee788 Carles Marti
    """
572 9dcee788 Carles Marti
573 9dcee788 Carles Marti
    # Create array to store generalized coordination numbers
574 9dcee788 Carles Marti
    list_generalized_CN = list()
575 9dcee788 Carles Marti
576 9dcee788 Carles Marti
    # Retrieve maximal coordination number
577 9dcee788 Carles Marti
    if max_CN is None:
578 9dcee788 Carles Marti
        max_CN = max(list_CN)
579 9dcee788 Carles Marti
580 9dcee788 Carles Marti
    # Treat each atom separately (since CN can be different for each atom,
581 9dcee788 Carles Marti
    # Numpy methods are unsuited here)
582 9dcee788 Carles Marti
    for (weights, indexes) in zip(list_weights, sorted_indexes):
583 9dcee788 Carles Marti
        # Initialize atom coordination number, and maximal coordination number
584 9dcee788 Carles Marti
        atom_CN = 0
585 9dcee788 Carles Marti
586 9dcee788 Carles Marti
        # Loop over all neighbors, compute and add the corresponding weights
587 9dcee788 Carles Marti
        for (weight, index) in zip(weights, indexes[1:]):
588 9dcee788 Carles Marti
            try:
589 9dcee788 Carles Marti
                neighbor_CN = list_CN[index]
590 9dcee788 Carles Marti
            except IndexError:
591 9dcee788 Carles Marti
                # if neighbor is virtual, use corresponding original one instead
592 9dcee788 Carles Marti
                neighbor_CN = list_CN[index % sorted_indexes.shape[0]]
593 9dcee788 Carles Marti
            atom_CN += weight * neighbor_CN
594 9dcee788 Carles Marti
595 9dcee788 Carles Marti
        # Divide by maximal coordination number
596 9dcee788 Carles Marti
        list_generalized_CN.append(atom_CN / max_CN)
597 9dcee788 Carles Marti
598 9dcee788 Carles Marti
    return (np.array(list_generalized_CN))
599 9dcee788 Carles Marti
600 9dcee788 Carles Marti
601 9dcee788 Carles Marti
def get_edges(list_CN, sorted_indexes, reduce_mode=None, nb_atoms=None):
602 9dcee788 Carles Marti
    """Compute all edges corresponding with the connectivity graph of the
603 9dcee788 Carles Marti
    given structure, based on discretized coordination numbers.
604 9dcee788 Carles Marti

605 9dcee788 Carles Marti
    Parameters: list_CN: Numpy 1D array containing the discretized
606 9dcee788 Carles Marti
    coordination numbers (i.e. number of neighbors).
607 9dcee788 Carles Marti

608 9dcee788 Carles Marti
        sorted_indexes: Numpy 2D array containing for each row the
609 9dcee788 Carles Marti
        distances-sorted indexes of neighbors. In other words, the atom
610 9dcee788 Carles Marti
        indexed sorted_indexes[i,j] is the j-th nearest neighbor of atom i.
611 9dcee788 Carles Marti
        Shape: (nb_atoms, nb_coords)
612 9dcee788 Carles Marti

613 9dcee788 Carles Marti
        reduce_mode: Edges counting mode. The ASANN/SANN algorithm can only
614 9dcee788 Carles Marti
        find directed edges (i.e. find i->j but not j->i). This parameter
615 9dcee788 Carles Marti
        defines the conversion mode from directed to undirected edges.
616 9dcee788 Carles Marti
        Possible values: None: All directed edges are given, no reduction is
617 9dcee788 Carles Marti
        performed. 'both': An undirected edge (i,j) is present only if both
618 9dcee788 Carles Marti
        related directed edges (i->j and j->i) are found. 'any': An
619 9dcee788 Carles Marti
        undirected edge (i,j) is present if any related directed edge (i->j
620 9dcee788 Carles Marti
        or j->i) is found.
621 9dcee788 Carles Marti

622 9dcee788 Carles Marti
        nb_atoms (int, optional): Number of real atoms coordinates (i.e. for
623 9dcee788 Carles Marti
        which distances must be computed). This is particularly useful for
624 9dcee788 Carles Marti
        excluding periodic images coordinates as central atoms. If None,
625 9dcee788 Carles Marti
        all coordinates are considered. Default: None
626 9dcee788 Carles Marti

627 9dcee788 Carles Marti
    Returns:
628 9dcee788 Carles Marti
        list_edges: List containing all edges of the connectivity graph.
629 9dcee788 Carles Marti
            The edges are in the form of a couple (index_node_i, index_node_j)
630 9dcee788 Carles Marti
            Shape: (nb_bonds_found, 2)
631 9dcee788 Carles Marti
    """
632 9dcee788 Carles Marti
633 9dcee788 Carles Marti
    # Create array to store all edges
634 9dcee788 Carles Marti
    list_edges = list()
635 9dcee788 Carles Marti
636 9dcee788 Carles Marti
    # Treat each atom separately (since CN can be different for each atom,
637 9dcee788 Carles Marti
    # Numpy methods are unsuited here)
638 9dcee788 Carles Marti
    for (atom_CN, indexes) in zip(list_CN, sorted_indexes):
639 9dcee788 Carles Marti
        # Retrieve current atom index
640 9dcee788 Carles Marti
        index_i = indexes[0]
641 9dcee788 Carles Marti
        # Loop over all neighbors, and add the corresponding edges
642 9dcee788 Carles Marti
        for index_j in indexes[1:atom_CN + 1]:
643 9dcee788 Carles Marti
            list_edges.append(
644 9dcee788 Carles Marti
                (index_i, index_j) if reduce_mode is None else tuple(sorted((
645 9dcee788 Carles Marti
                    index_i,
646 9dcee788 Carles Marti
                    index_j))))  # add sorted edge instead (representing an
647 9dcee788 Carles Marti
            # undirected edge) if conversion is required (reduce_mode not None)
648 9dcee788 Carles Marti
649 9dcee788 Carles Marti
    # Re-map to correct atom index if explicit periodic images are included
650 9dcee788 Carles Marti
    if nb_atoms is not None:
651 9dcee788 Carles Marti
        list_edges = [(index_i % nb_atoms, index_j % nb_atoms) for
652 9dcee788 Carles Marti
                      (index_i, index_j) in list_edges]
653 9dcee788 Carles Marti
654 9dcee788 Carles Marti
    # Conversion of directed edges set to undirected edges set
655 9dcee788 Carles Marti
    if reduce_mode == 'both':
656 9dcee788 Carles Marti
        # Extract only edges that are present multiple times
657 9dcee788 Carles Marti
        seen = dict()
658 9dcee788 Carles Marti
        duplicates = []
659 9dcee788 Carles Marti
        for edge in list_edges:
660 9dcee788 Carles Marti
            if edge in seen:
661 9dcee788 Carles Marti
                duplicates.append(edge)
662 9dcee788 Carles Marti
            else:
663 9dcee788 Carles Marti
                seen[edge] = None
664 9dcee788 Carles Marti
        list_edges = duplicates
665 9dcee788 Carles Marti
    elif reduce_mode == 'any':
666 9dcee788 Carles Marti
        # Retrieve all unique undirected edges
667 9dcee788 Carles Marti
        list_edges = list(set(list_edges))
668 9dcee788 Carles Marti
669 9dcee788 Carles Marti
    return (list_edges)
670 9dcee788 Carles Marti
671 9dcee788 Carles Marti
672 9dcee788 Carles Marti
# FULL WRAPPER FUNCTION ##
673 9dcee788 Carles Marti
def coordination_numbers(list_coords, pbc=True, cell_vectors=np.eye(3),
674 9dcee788 Carles Marti
                         continuous=False, generalized=False, edges=True,
675 9dcee788 Carles Marti
                         correction='ASANN', parallel=False, reduce_mode=None):
676 9dcee788 Carles Marti
    """Computes coordination numbers according to the CASANN algorithm.
677 9dcee788 Carles Marti

678 9dcee788 Carles Marti
    Parameters:
679 9dcee788 Carles Marti
        list_coords (2D array): List of atoms coordinates to
680 9dcee788 Carles Marti
        consider. Important: direct coordinates are expected (torus topology of
681 9dcee788 Carles Marti
        side 1), unless `pbc` is set to False. Note: This will be converted into
682 9dcee788 Carles Marti
        a numpy.ndarray. Shape expected: (nb_atoms, nb_dim), where nb_atoms is
683 9dcee788 Carles Marti
        the number of atoms considered, and nb_dim is the dimensionnality of the
684 9dcee788 Carles Marti
        system.
685 9dcee788 Carles Marti

686 9dcee788 Carles Marti
        pbc (bool or str, optional): Determines if periodic boundaries
687 9dcee788 Carles Marti
        conditions should be applied. Default to True. If True, coordinates
688 9dcee788 Carles Marti
        are interpreted as direct coordinates and the distance between points
689 9dcee788 Carles Marti
        is computed as the minimum euclidian distance between all duplicates
690 9dcee788 Carles Marti
        (due to periodic boundaries conditions) of these points. If False,
691 9dcee788 Carles Marti
        coordinates are interpreted as cartesian coordinates and the metric
692 9dcee788 Carles Marti
        used is simply the euclidian distance. To explicitely include all
693 9dcee788 Carles Marti
        adjacent periodic images (not only the minimum image convention) set
694 9dcee788 Carles Marti
        `pbc` to 'adjacent'. This mode is particularly pertinent for small
695 9dcee788 Carles Marti
        enough cells, but increases 9-fold the number of atoms. To
696 9dcee788 Carles Marti
        explicitely include all surrounding periodic images set `pbc` to
697 9dcee788 Carles Marti
        'full'. This mode is particularly pertinent for very small cells,
698 9dcee788 Carles Marti
        but increases 27-fold the number of atoms. Note: This option implies
699 9dcee788 Carles Marti
        the use of cell vectors (see `cell` parameter) for the computation of
700 9dcee788 Carles Marti
        distance. Default: True.
701 9dcee788 Carles Marti

702 9dcee788 Carles Marti
        cell_vectors (2D array, optional): List of cell vectors to consider
703 9dcee788 Carles Marti
        when periodic boundaries conditions are considered. Note: This will
704 9dcee788 Carles Marti
        be converted into a numpy.ndarray. Shape expected: (nb_dim, nb_dim),
705 9dcee788 Carles Marti
        where nb_dim is the dimensionnality of the system considered.
706 9dcee788 Carles Marti
        Important: For this parameter to be taken into account, `pbc` must be
707 9dcee788 Carles Marti
        set to True. Exemple: cell_vectors=[[v1_x, v1_y, v1_z], [v2_x, v2_y,
708 9dcee788 Carles Marti
        v2_z], [v3_x, v3_y, v3_z]] Default: numpy.eye(3) (i.e. a cubic cell
709 9dcee788 Carles Marti
        of side 1).
710 9dcee788 Carles Marti

711 9dcee788 Carles Marti
        continuous (bool, optional): If True, computes continuous
712 9dcee788 Carles Marti
        coordination numbers. If False, computes discretized coordination
713 9dcee788 Carles Marti
        numbers. Default to True. In the discretized version,
714 9dcee788 Carles Marti
        the coordination number is equal to the number of detected neighbors.
715 9dcee788 Carles Marti
        In the continuous version, each neighbors' contribution to the
716 9dcee788 Carles Marti
        coordination number is 1 weighted by SA/SA_max, where SA is the solid
717 9dcee788 Carles Marti
        angle corresponding to this neighbor and SA_max is the solid angle
718 9dcee788 Carles Marti
        corresponding to the nearest neighbor (i.e. the maximum solid angle
719 9dcee788 Carles Marti
        amongs all neighbors detected). Default: False.
720 9dcee788 Carles Marti

721 9dcee788 Carles Marti
        generalized (bool, optional): If True, computes generalized
722 9dcee788 Carles Marti
        coordination numbers, where each neighbor is weighted by its own
723 9dcee788 Carles Marti
        coordination number Default: False.
724 9dcee788 Carles Marti

725 9dcee788 Carles Marti
        edges (bool, optional): If True, computes edges of the connectivity
726 9dcee788 Carles Marti
        graph defined by the discretized coordination numbers computed.
727 9dcee788 Carles Marti
        Default: True.
728 9dcee788 Carles Marti

729 9dcee788 Carles Marti
        correction (str, optional): Determines if a correction term should be
730 9dcee788 Carles Marti
        used. Default to 'ASANN'. The SANN algorithm suffers
731 9dcee788 Carles Marti
        overdetermination of the coordination numbers at interfaces (i.e.
732 9dcee788 Carles Marti
        high density gradient) basically because neighbors are always
733 9dcee788 Carles Marti
        expected to fill the whole neighboring space (i.e. the sum of solid
734 9dcee788 Carles Marti
        angles must be 4π), but at interfaces, neighbors are only expected to
735 9dcee788 Carles Marti
        fill a portion of that space depending on the geometry of the
736 9dcee788 Carles Marti
        interface. Possible values: - 'ASANN': If `correction` is set to
737 9dcee788 Carles Marti
        'ASANN', the total sum of solid angles is rescaled by the ASANN
738 9dcee788 Carles Marti
        angular correction term. - 'SC-ASANN': If `correction` is set to
739 9dcee788 Carles Marti
        'SC-ASANN', the total sum of solid angles is rescaled by the ASANN
740 9dcee788 Carles Marti
        angular correction term, computed in a self-consistent manner.
741 9dcee788 Carles Marti
        Arguably better CNs, but more time consuming and less regularized
742 9dcee788 Carles Marti
        continuous CNs. - None (or anything else): No correction term
743 9dcee788 Carles Marti
        applied. This is equivalent to the SANN algorithm.
744 9dcee788 Carles Marti

745 9dcee788 Carles Marti
        reduce_mode: Edges counting mode. The ASANN/SANN algorithm can only
746 9dcee788 Carles Marti
        find directed edges (e.g. find i->j but not j->i). This parameter
747 9dcee788 Carles Marti
        defines the conversion mode from directed to undirected edges.
748 9dcee788 Carles Marti
        Possible values: - 'both': An undirected edge (i,j) is present only
749 9dcee788 Carles Marti
        if both related directed edges (i->j and j->i) are found. - 'any': An
750 9dcee788 Carles Marti
        undirected edge (i,j) is present if any related directed edge (i->j
751 9dcee788 Carles Marti
        or j->i) is found. - None: All directed edges are given, no reduction
752 9dcee788 Carles Marti
        is performed. Default: None.
753 9dcee788 Carles Marti

754 9dcee788 Carles Marti
    Returns: (asann_CNs, asann_radius, asann_edges) asann_CNs (numpy array):
755 9dcee788 Carles Marti
    list of coordination numbers computed. The order is the same as provided
756 9dcee788 Carles Marti
    in list_coords. asann_radius (numpy array): list of coordination radii
757 9dcee788 Carles Marti
    computed. The order is the same as provided in list_coords. asann_edges (
758 9dcee788 Carles Marti
    list of tuples): list of edges found. Each edge is represented as a tuple
759 9dcee788 Carles Marti
    (i,j), meaning that j was found as a neighbor of i. Note: if `edges` is
760 9dcee788 Carles Marti
    set to False, asann_edges is set to None.
761 9dcee788 Carles Marti

762 9dcee788 Carles Marti
    Computational details:
763 9dcee788 Carles Marti
        Correction:
764 9dcee788 Carles Marti

765 9dcee788 Carles Marti
            The SANN algorithm computes the coordination number (i.e. CN) m of
766 9dcee788 Carles Marti
            atom i, as the minimum integer m>=3 satisfying R_i_m = sum(r_ij,
767 9dcee788 Carles Marti
            j=1..m)/(m-2) < r_i(m+1) (assuming r_ij are sorted: r_i(j-1) <=
768 9dcee788 Carles Marti
            r_ij <= r_i(j+1))
769 9dcee788 Carles Marti

770 9dcee788 Carles Marti
            This formula is equivalent to determining a sphere of radius
771 9dcee788 Carles Marti
            R_i_m, such that the sum of all solid angles of inner atoms is 4π.
772 9dcee788 Carles Marti
            The solid angle of atom j with respect to atom i, is the solid
773 9dcee788 Carles Marti
            angle of the spherical cap of height r_ij on the sphere of radius
774 9dcee788 Carles Marti
            R_i_m (i.e. SA_ij = 2π(1-r_ij/R_i_m)). However, at interfaces,
775 9dcee788 Carles Marti
            it is expected that neighbors do not fill the whole space (i.e.
776 9dcee788 Carles Marti
            the sum of solid angles should not be 4π)
777 9dcee788 Carles Marti

778 9dcee788 Carles Marti
            Hence we propose here an angular correction which is exact in the
779 9dcee788 Carles Marti
            case of infinitely evenly distributed neighbors along a spherical
780 9dcee788 Carles Marti
            cap. Indeed, assuming that a normal vector can be meaningfully
781 9dcee788 Carles Marti
            defined at the interface, a correction is needed when the
782 9dcee788 Carles Marti
            neighbors barycenter is far from the central atom.
783 9dcee788 Carles Marti

784 9dcee788 Carles Marti
            Therefore, the neighbors barycenter is computed. From this
785 9dcee788 Carles Marti
            barycenter, one deduces the corresponding spherical cap on the
786 9dcee788 Carles Marti
            sphere of radius R_i_m. Its solid angle is then taken instead of
787 9dcee788 Carles Marti
            4π.
788 9dcee788 Carles Marti

789 9dcee788 Carles Marti
            Consequently, this angular correction assumes a spherical
790 9dcee788 Carles Marti
            cap-like distribution of the nearest neighbors.
791 9dcee788 Carles Marti

792 9dcee788 Carles Marti
        Continuous:
793 9dcee788 Carles Marti

794 9dcee788 Carles Marti
            The continuous coordination number algorithm scales the
795 9dcee788 Carles Marti
            contributions of each neighbor by the following weight :
796 9dcee788 Carles Marti
            SA/SA_max (where SA is the solid angle associated with the
797 9dcee788 Carles Marti
            corresponding neighbor, and SA_max is the maximum solid angle (
798 9dcee788 Carles Marti
            i.e. the solid angle associated with the nearest neighbor)) The
799 9dcee788 Carles Marti
            continuous coordination number is then the sum of all weights.
800 9dcee788 Carles Marti

801 9dcee788 Carles Marti
        Generalized:
802 9dcee788 Carles Marti

803 9dcee788 Carles Marti
            The generalized coordination number algorithm scales the
804 9dcee788 Carles Marti
            contributions of each neighbor by the following weight :
805 9dcee788 Carles Marti
            CN/CN_max (where CN is the coordination number associated with
806 9dcee788 Carles Marti
            the corresponding neighbor, and CN_max is the maximum
807 9dcee788 Carles Marti
            coordination number (i.e. the coordination number associated with
808 9dcee788 Carles Marti
            the most coordinated neighbor)) The generalized coordination
809 9dcee788 Carles Marti
            number is then the sum of all weights (weighted futhermore,
810 9dcee788 Carles Marti
            or not, by the continuous algorithm).
811 9dcee788 Carles Marti
    """
812 9dcee788 Carles Marti
813 9dcee788 Carles Marti
    # Conversion to numpy.ndarray
814 9dcee788 Carles Marti
    coords = np.array(list_coords)
815 9dcee788 Carles Marti
    if pbc:
816 9dcee788 Carles Marti
        cell = np.array(cell_vectors)
817 9dcee788 Carles Marti
    else:
818 9dcee788 Carles Marti
        cell = None
819 9dcee788 Carles Marti
    nb_atoms = None
820 9dcee788 Carles Marti
821 9dcee788 Carles Marti
    # Retrieve parameters and check dimension consistency
822 9dcee788 Carles Marti
    assert ((not pbc) or (coords.shape[1] == cell.shape[0] == cell.shape[1]))
823 9dcee788 Carles Marti
824 9dcee788 Carles Marti
    # Explicitely add adjacent periodic images if requested
825 9dcee788 Carles Marti
    if pbc in ('adjacent', 'full'):
826 9dcee788 Carles Marti
        nb_atoms, coords = add_periodic_images(coords, cell, pbc)
827 9dcee788 Carles Marti
828 9dcee788 Carles Marti
    # Retrieve distance-sorted pairwise distances, vectors and indexes
829 9dcee788 Carles Marti
    sorted_distances, sorted_vectors, sorted_indexes = get_sorted_distances(
830 9dcee788 Carles Marti
        coords, pbc, nb_atoms=nb_atoms, cell=cell)
831 9dcee788 Carles Marti
832 9dcee788 Carles Marti
    # Retrieve number of nearest neighbors and coordination radius with SANN
833 9dcee788 Carles Marti
    # algorithm
834 9dcee788 Carles Marti
    asann_CNs, asann_radius = get_SANN(sorted_distances)
835 9dcee788 Carles Marti
836 9dcee788 Carles Marti
    # Apply angular correction if requested
837 9dcee788 Carles Marti
    if correction == 'ASANN':
838 9dcee788 Carles Marti
        asann_CNs, asann_radius, vectors = get_ASANN(sorted_distances,
839 9dcee788 Carles Marti
                                                     sorted_vectors, asann_CNs,
840 9dcee788 Carles Marti
                                                     asann_radius)
841 9dcee788 Carles Marti
    elif correction == 'SC-ASANN':
842 9dcee788 Carles Marti
        asann_CNs, asann_radius, vectors = get_self_consistent_ASANN(
843 9dcee788 Carles Marti
            sorted_distances, sorted_vectors, asann_CNs)
844 9dcee788 Carles Marti
845 9dcee788 Carles Marti
    # Compute edges
846 9dcee788 Carles Marti
    if edges:
847 9dcee788 Carles Marti
        asann_edges = get_edges(asann_CNs, sorted_indexes,
848 9dcee788 Carles Marti
                                reduce_mode=reduce_mode, nb_atoms=nb_atoms)
849 9dcee788 Carles Marti
    else:
850 9dcee788 Carles Marti
        asann_edges = None
851 9dcee788 Carles Marti
852 9dcee788 Carles Marti
    # Convert coordination numbers into continuous values if requested
853 9dcee788 Carles Marti
    if continuous:
854 9dcee788 Carles Marti
        # Compute continuous CN by weighting each neighbor contribution
855 9dcee788 Carles Marti
        asann_CNs, list_weights = convert_continuous(asann_CNs, asann_radius,
856 9dcee788 Carles Marti
                                                     sorted_distances)
857 9dcee788 Carles Marti
    elif generalized:
858 9dcee788 Carles Marti
        list_weights = [[1] * asann_CN for asann_CN in asann_CNs]
859 9dcee788 Carles Marti
860 9dcee788 Carles Marti
    # Compute generalized coordination numbers if requested
861 9dcee788 Carles Marti
    if generalized:
862 9dcee788 Carles Marti
        asann_CNs = get_generalized(asann_CNs, list_weights, sorted_indexes)
863 9dcee788 Carles Marti
864 9dcee788 Carles Marti
    return asann_CNs, asann_radius, asann_edges, vectors
865 9dcee788 Carles Marti
866 9dcee788 Carles Marti
867 9dcee788 Carles Marti
# Program being executed when used as a script (instead of a module)
868 9dcee788 Carles Marti
if __name__ == '__main__':
869 9dcee788 Carles Marti
    # Imports
870 9dcee788 Carles Marti
    import sys
871 9dcee788 Carles Marti
872 9dcee788 Carles Marti
    # Import local structure reader sub-module
873 9dcee788 Carles Marti
    try:
874 9dcee788 Carles Marti
        from structure_reader import structure_from_file
875 9dcee788 Carles Marti
    except ImportError as err:
876 9dcee788 Carles Marti
        print(
877 9dcee788 Carles Marti
            'Unable to find file structure_reader.py, which allows reading '
878 9dcee788 Carles Marti
            'molecular structures. Aborting.',
879 9dcee788 Carles Marti
            file=sys.stderr)
880 9dcee788 Carles Marti
        print(
881 9dcee788 Carles Marti
            'Plase copy structure_reader.py in either the same folder '
882 9dcee788 Carles Marti
            'containing this script ({}), or in your working '
883 9dcee788 Carles Marti
            'directory'.format(
884 9dcee788 Carles Marti
                sys.argv[0]), file=sys.stderr)
885 9dcee788 Carles Marti
886 9dcee788 Carles Marti
        raise err
887 9dcee788 Carles Marti
888 9dcee788 Carles Marti
    # Retrieve filename of molecular structure to treat
889 9dcee788 Carles Marti
    try:
890 9dcee788 Carles Marti
        filename = sys.argv[1]
891 9dcee788 Carles Marti
    except IndexError as err:
892 9dcee788 Carles Marti
        print(
893 9dcee788 Carles Marti
            'Unable to parse a filename to treat (containing coordinates in '
894 9dcee788 Carles Marti
            'supported format: XYZ, POSCAR, CIF, ...)',
895 9dcee788 Carles Marti
            file=sys.stderr)
896 9dcee788 Carles Marti
        print('Usage: {} filename'.format(sys.argv[0]), file=sys.stderr)
897 9dcee788 Carles Marti
        raise err
898 9dcee788 Carles Marti
899 9dcee788 Carles Marti
    # Read file structure
900 9dcee788 Carles Marti
    file_structure = structure_from_file(filename)
901 9dcee788 Carles Marti
    coordinates = file_structure.get_coords()
902 9dcee788 Carles Marti
    cell_vectors = file_structure.get_cell_matrix()
903 9dcee788 Carles Marti
    pbc_mode = file_structure.pbc_enabled
904 9dcee788 Carles Marti
905 9dcee788 Carles Marti
    # Explicitely add next periodic images if number of atoms is too low. A
906 9dcee788 Carles Marti
    # simple modulo operation is not enough when a single point is found as
907 9dcee788 Carles Marti
    # neighbor more than once (as part of periodic images)
908 9dcee788 Carles Marti
    if pbc_mode and len(coordinates) < 100:
909 9dcee788 Carles Marti
        pbc_mode = 'full'  # Explictely include all 27 next periodic cells
910 9dcee788 Carles Marti
    elif pbc_mode and len(coordinates) < 1000:
911 9dcee788 Carles Marti
        pbc_mode = 'adjacent'  # Explicitely include all 9 adjacent cells
912 9dcee788 Carles Marti
913 9dcee788 Carles Marti
    asann_CNs, asann_radii, asann_edges, vectors = coordination_numbers(
914 9dcee788 Carles Marti
        coordinates, pbc=pbc_mode, cell_vectors=cell_vectors, continuous=False,
915 9dcee788 Carles Marti
        generalized=False, edges=True, correction='ASANN')
916 9dcee788 Carles Marti
917 9dcee788 Carles Marti
    print('ASANN vectors')
918 9dcee788 Carles Marti
    for l in vectors:
919 9dcee788 Carles Marti
        print(-l[0], -l[1], -l[2])