Test Failed
Push — master ( 9779af...19222a )
by Daniel
02:55
created

amd.utils.random_cell()   A

Complexity

Conditions 5

Size

Total Lines 27
Code Lines 18

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 18
dl 0
loc 27
rs 9.0333
c 0
b 0
f 0
cc 5
nop 3
1
"""General utility functions."""
2
3
from typing import Tuple
4
import sys
5
import time
6
import datetime
7
8
import numpy as np
9
import numba
10
from scipy.spatial.distance import squareform
11
12
13
def diameter(cell):
14
    """Diameter of a unit cell (as a square matrix in Cartesian/Orthogonal form)
15
    in 3 or fewer dimensions."""
16
17
    dims = cell.shape[0]
18
    if dims == 1:
19
        return cell[0][0]
20
    if dims == 2:
21
        d = np.amax(np.linalg.norm(np.array([cell[0] + cell[1], cell[0] - cell[1]]), axis=-1))
22
    elif dims == 3:
23
        diams = np.array([
24
            cell[0] + cell[1] + cell[2],
25
            cell[0] + cell[1] - cell[2],
26
            cell[0] - cell[1] + cell[2],
27
            - cell[0] + cell[1] + cell[2]
28
        ])
29
        d = np.amax(np.linalg.norm(diams, axis=-1))
30
    else:
31
        msg = f'diameter() not implemented for dims > 3 (passed cell shape {cell.shape}).'
32
        raise NotImplementedError(msg)
33
    return d
34
35
36
@numba.njit()
37
def cellpar_to_cell(a, b, c, alpha, beta, gamma):
0 ignored issues
show
best-practice introduced by
Too many arguments (6/5)
Loading history...
38
    """Simplified version of function from :mod:`ase.geometry` of the same name.
39
    3D unit cell parameters a,b,c,α,β,γ --> cell as 3x3 NumPy array.
40
    """
41
42
    eps = 2 * np.spacing(90.0)  # ~1.4e-14
43
44
    cos_alpha = 0. if abs(abs(alpha) - 90.) < eps else np.cos(alpha * np.pi / 180.)
45
    cos_beta = 0. if abs(abs(beta) - 90.) < eps else np.cos(beta * np.pi / 180.)
46
    cos_gamma = 0. if abs(abs(gamma) - 90.) < eps else np.cos(gamma * np.pi / 180.)
47
48
    if abs(gamma - 90) < eps:
49
        sin_gamma = 1.
50
    elif abs(gamma + 90) < eps:
51
        sin_gamma = -1.
52
    else:
53
        sin_gamma = np.sin(gamma * np.pi / 180.)
54
55
    cy = (cos_alpha - cos_beta * cos_gamma) / sin_gamma
56
    cz_sqr = 1. - cos_beta ** 2 - cy ** 2
57
    if cz_sqr < 0:
58
        raise RuntimeError('Could not create unit cell from given parameters.')
59
60
    cell = np.zeros((3, 3))
61
    cell[0, 0] = a
62
    cell[1, 0] = b * cos_gamma
63
    cell[1, 1] = b * sin_gamma
64
    cell[2, 0] = c * cos_beta
65
    cell[2, 1] = c * cy
66
    cell[2, 2] = c * np.sqrt(cz_sqr)
67
68
    return cell
69
70
71
@numba.njit()
72
def cellpar_to_cell_2D(a, b, alpha):
73
    """2D unit cell parameters a,b,α --> cell as 2x2 ndarray."""
74
75
    cell = np.zeros((2, 2))
76
    cell[0, 0] = a
77
    cell[1, 0] = b * np.cos(alpha * np.pi / 180.)
78
    cell[1, 1] = b * np.sin(alpha * np.pi / 180.)
79
80
    return cell
81
82
83
def cell_to_cellpar(cell):
84
    """Unit cell as a 3x3 NumPy array -> list of 6 lengths + angles."""
85
    lengths = np.linalg.norm(cell, axis=-1)
86
    angles = []
87
    for i, j in [(1, 2), (0, 2), (0, 1)]:
88
        ang_rad = np.arccos(np.dot(cell[i], cell[j]) / (lengths[i] * lengths[j]))
89
        angles.append(np.rad2deg(ang_rad))
90
    return np.concatenate((lengths, np.array(angles)))
91
92
93
def cell_to_cellpar_2D(cell):
94
    """Unit cell as a 2x2 NumPy array -> list of 2 lengths and an angle."""
95
    cellpar = np.zeros((3, ))
96
    lengths = np.linalg.norm(cell, axis=-1)
97
    ang_rad = np.arccos(np.dot(cell[0], cell[1]) / (lengths[0] * lengths[1]))
98
    cellpar[0] = lengths[0]
99
    cellpar[1] = lengths[1]
100
    cellpar[2] = np.rad2deg(ang_rad)
101
    return cellpar
102
103
104
def neighbours_from_distance_matrix(
105
        n: int,
106
        dm: np.ndarray
107
) -> Tuple[np.ndarray, np.ndarray]:
108
    """Given a distance matrix, find the n nearest neighbours of each item.
109
110
    Parameters
111
    ----------
112
    n : int
113
        Number of nearest neighbours to find for each item.
114
    dm : :class:`numpy.ndarray`
115
        2D distance matrix or 1D condensed distance matrix.
116
117
    Returns
118
    -------
119
    (nn_dm, inds) : Tuple[:class:`numpy.ndarray`, :class:`numpy.ndarray`]
120
        ``nn_dm[i][j]`` is the distance from item :math:`i` to its :math:`j+1` st
121
        nearest neighbour, and ``inds[i][j]`` is the index of this neighbour
122
        (:math:`j+1` since index 0 is the first nearest neighbour).
123
    """
124
125
    inds = None
126
127
    # 2D distance matrix
128
    if len(dm.shape) == 2:
129
        inds = np.array([np.argpartition(row, n)[:n] for row in dm])
130
131
    # 1D condensed distance vector
132
    elif len(dm.shape) == 1:
133
        dm = squareform(dm)
134
        inds = []
135
        for i, row in enumerate(dm):
136
            inds_row = np.argpartition(row, n+1)[:n+1]
137
            inds_row = inds_row[inds_row != i][:n]
138
            inds.append(inds_row)
139
        inds = np.array(inds)
140
141
    else:
142
        ValueError(
143
            'Input must be an ndarray, either a 2D distance matrix '
144
            'or a condensed distance matrix (returned by pdist).')
145
146
    # inds are the indexes of nns: inds[i,j] is the j-th nn to point i
147
    nn_dm = np.take_along_axis(dm, inds, axis=-1)
148
    sorted_inds = np.argsort(nn_dm, axis=-1)
149
    inds = np.take_along_axis(inds, sorted_inds, axis=-1)
150
    nn_dm = np.take_along_axis(nn_dm, sorted_inds, axis=-1)
151
    return nn_dm, inds
152
153
154
def random_cell(length_bounds=(1, 2), angle_bounds=(60, 120), dims=3):
155
    """Dimensions 2 and 3 only. Random unit cell with uniformally chosen length and
156
    angle parameters between bounds."""
157
158
    ll, lu = length_bounds
159
    al, au = angle_bounds
160
161
    if dims == 3:
162
        while True:
163
            lengths = [np.random.uniform(low=ll, high=lu) for _ in range(dims)]
164
            angles = [np.random.uniform(low=al, high=au) for _ in range(dims)]
165
166
            try:
167
                cell = cellpar_to_cell(*lengths, *angles)
168
                break
169
            except RuntimeError:
170
                continue
171
172
    elif dims == 2:
173
        lengths = [np.random.uniform(low=ll, high=lu) for _ in range(dims)]
174
        alpha = np.random.uniform(low=al, high=au)
175
        cell = cellpar_to_cell_2D(*lengths, alpha)
176
177
    else:
178
        raise ValueError(f'random_cell only implimented for dimensions 2 and 3 (passed {dims})')
179
180
    return cell
181
182
183
class _ETA:
0 ignored issues
show
best-practice introduced by
Too many instance attributes (8/7)
Loading history...
184
    """Pass the number to items to process, then call .update() on every loop.
185
    """
186
187
    # epochtime_{n+1} = factor * epochtime + (1-factor) * epochtime_{n}
188
    _moving_av_factor = 0.3
189
190
    def __init__(self, to_do, update_rate=1000):
191
        self.to_do = to_do
192
        self.update_rate = update_rate
193
        self.counter = 0
194
        self.this_epoch = 0
195
        self.start_time = time.perf_counter()
196
        self.time_per_epoch = None
197
        self.done = False
198
        self.tic = self.start_time
199
200
    def update(self):
201
        """Call on each loop."""
202
203
        self.counter += 1
204
205
        if self.counter >= self.to_do:
0 ignored issues
show
unused-code introduced by
Unnecessary "elif" after "return"
Loading history...
206
            if not self.done:
207
                self.finish()
208
            return
209
210
        elif not self.counter % self.update_rate:
211
            self.end_epoch()
212
213
    def finish(self):
214
        """Called when done."""
215
        total = time.perf_counter() - self.start_time
216
        self.done = True
217
        msg = f'Time ({self.counter} items): {round(total, 2)}s, ' \
218
              f'{round(self.to_do/total, 2)} items/s\r\n'
219
        sys.stdout.write(msg)
220
221
    def end_epoch(self):
222
        """Called at the end of every epoch, printing the ETA to stdout."""
223
        toc = time.perf_counter()
224
        epoch_time = toc - self.tic
225
        if self.time_per_epoch is None:
226
            self.time_per_epoch = epoch_time
227
        else:
228
            self.time_per_epoch = _ETA._moving_av_factor * epoch_time + \
229
                                  (1 - _ETA._moving_av_factor) * self.time_per_epoch
230
        percent = round(100 * self.counter / self.to_do, 2)
231
        remaining = int(((self.to_do - self.counter) / self.update_rate) * self.time_per_epoch)
232
        eta = str(datetime.timedelta(seconds=remaining))
233
        self.tic = time.perf_counter()
234
        sys.stdout.write(f'{percent}%, ETA {eta}' + ' ' * 20 + '\r')
235