Passed
Push — master ( 4570af...abe023 )
by Daniel
03:02
created

amd.utils.random_cell()   B

Complexity

Conditions 5

Size

Total Lines 31
Code Lines 23

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
eloc 23
dl 0
loc 31
rs 8.8613
c 0
b 0
f 0
cc 5
nop 3
1
"""Helpful utility functions, e.g. unit cell diameter, converting
2
cell parameters to Cartesian form, and an ETA class."""
3
4
from typing import Tuple
5
6
import numpy as np
7
from scipy.spatial.distance import squareform
8
9
10
def diameter(cell):
11
    """Diameter of a unit cell (as a square matrix in Cartesian form)
12
    in 3 or fewer dimensions."""
13
14
    dims = cell.shape[0]
15
    if dims == 1:
16
        return cell[0][0]
17
    if dims == 2:
18
        d = np.amax(np.linalg.norm(np.array([cell[0] + cell[1], cell[0] - cell[1]]), axis=-1))
19
    elif dims == 3:
20
        diams = np.array([
21
            cell[0] + cell[1] + cell[2],
22
            cell[0] + cell[1] - cell[2],
23
            cell[0] - cell[1] + cell[2],
24
            -cell[0] + cell[1] + cell[2]
25
        ])
26
        d =  np.amax(np.linalg.norm(diams, axis=-1))
0 ignored issues
show
Coding Style introduced by
Exactly one space required after assignment
Loading history...
27
    else:
28
        raise ValueError(f'diameter only implimented for dimensions <= 3 (passed {dims})')
29
    return d
30
31
32
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...
33
    """Simplified version of function from :mod:`ase.geometry` of the same name.
34
    3D unit cell parameters a,b,c,α,β,γ --> cell as 3x3 NumPy array.
35
    """
36
37
    # Handle orthorhombic cells separately to avoid rounding errors
38
    eps = 2 * np.spacing(90.0, dtype=np.float64)  # around 1.4e-14
39
40
    cos_alpha = 0. if abs(abs(alpha) - 90.) < eps else np.cos(alpha * np.pi / 180.)
41
    cos_beta = 0. if abs(abs(beta) - 90.) < eps else np.cos(beta * np.pi / 180.)
42
    cos_gamma = 0. if abs(abs(gamma) - 90.) < eps else np.cos(gamma * np.pi / 180.)
43
44
    if abs(gamma - 90) < eps:
45
        sin_gamma = 1.
46
    elif abs(gamma + 90) < eps:
47
        sin_gamma = -1.
48
    else:
49
        sin_gamma = np.sin(gamma * np.pi / 180.)
50
51
    cy = (cos_alpha - cos_beta * cos_gamma) / sin_gamma
52
    cz_sqr = 1. - cos_beta ** 2 - cy ** 2
53
    if cz_sqr < 0:
54
        raise RuntimeError('Could not create unit cell from parameters ' + \
55
                           f'a={a},b={b},c={c},α={alpha},β={beta},γ={gamma}')
56
57
    return np.array([[a, 0, 0],
58
                     [b*cos_gamma, b*sin_gamma, 0],
59
                     [c*cos_beta, c*cy, c*np.sqrt(cz_sqr)]])
60
61
62
def cellpar_to_cell_2D(a, b, alpha):
63
    """2D unit cell parameters a,b,α --> cell as 2x2 ndarray."""
64
65
    return np.array([[a, 0],
66
                     [b * np.cos(alpha * np.pi / 180.), b * np.sin(alpha * np.pi / 180.)]])
67
68
69
def neighbours_from_distance_matrix(
70
        n: int,
71
        dm: np.ndarray
72
) -> Tuple[np.ndarray, np.ndarray]:
73
    """Given a distance matrix, find the n nearest neighbours of each item.
74
75
    Parameters
76
    ----------
77
    n : int
78
        Number of nearest neighbours to find for each item.
79
    dm : numpy.ndarray
80
        2D distance matrix or 1D condensed distance matrix.
81
82
    Returns
83
    -------
84
    nn_dm, inds : Tuple[numpy.ndarray, numpy.ndarray] 
0 ignored issues
show
Coding Style introduced by
Trailing whitespace
Loading history...
85
        ``nn_dm[i][j]`` is the distance from item ``i`` to its ``j+1`` st
86
        nearest neighbour, and ``inds[i][j]`` is the index of this neighbour 
0 ignored issues
show
Coding Style introduced by
Trailing whitespace
Loading history...
87
        (``j+1`` since index 0 is the first nearest neighbour).
88
    """
89
90
    inds = None
91
92
    # 2D distance matrix
93
    if len(dm.shape) == 2:
94
        inds = np.array([np.argpartition(row, n)[:n] for row in dm])
95
96
    # 1D condensed distance vector
97
    elif len(dm.shape) == 1:
98
        dm = squareform(dm)
99
        inds = []
100
        for i, row in enumerate(dm):
101
            inds_row = np.argpartition(row, n+1)[:n+1]
102
            inds_row = inds_row[inds_row != i][:n]
103
            inds.append(inds_row)
104
        inds = np.array(inds)
105
106
    else:
107
        ValueError(
108
            'Input must be an ndarray, either a 2D distance matrix '
109
            'or a condensed distance matrix (returned by pdist).')
110
111
    # inds are the indexes of nns: inds[i,j] is the j-th nn to point i
112
    nn_dm = np.take_along_axis(dm, inds, axis=-1)
113
    sorted_inds = np.argsort(nn_dm, axis=-1)
114
    inds = np.take_along_axis(inds, sorted_inds, axis=-1)
115
    nn_dm = np.take_along_axis(nn_dm, sorted_inds, axis=-1)
116
    return nn_dm, inds
117
118
119
def lattice_cubic(scale=1, dims=3):
120
    """Return a pair ``(motif, cell)`` representing a cubic lattice, passable to
121
    ``amd.AMD()`` or ``amd.PDD()``."""
122
123
    return (np.zeros((1, dims)), np.identity(dims) * scale)
124
125
126
def random_cell(length_bounds=(1, 2), angle_bounds=(60, 120), dims=3):
127
    """Random unit cell with uniformally chosen length and angle parameters
128
    between bounds."""
129
130
    if dims == 3:
131
        while True:
132
            lengths = [np.random.uniform(low=length_bounds[0],
133
                                         high=length_bounds[1])
134
                       for _ in range(dims)]
135
            angles = [np.random.uniform(low=angle_bounds[0],
136
                                        high=length_bounds[1])
137
                      for _ in range(dims)]
138
139
            try:
140
                cell = cellpar_to_cell(*lengths, *angles)
141
                break
142
            except RuntimeError:
143
                continue
144
145
    elif dims == 2:
146
        lengths = [np.random.uniform(low=length_bounds[0],
147
                                     high=length_bounds[1])
148
                       for _ in range(dims)]
0 ignored issues
show
Coding Style introduced by
Wrong continued indentation (remove 4 spaces).
Loading history...
149
        alpha = np.random.uniform(low=angle_bounds[0],
150
                                  high=length_bounds[1])
151
        cell = cellpar_to_cell_2D(*lengths, alpha)
152
153
    else:
154
        raise ValueError(f'random_cell only implimented for dimensions 2 and 3 (passed {dims})')
155
156
    return cell
157