Coverage for dynasor/modes/tools.py: 99%

1import numbers

2from fractions import Fraction

3from typing import Optional, Union

4import numba

5import numpy as np

6from numpy.typing import NDArray

9# Linalg

10def trace(A):

11 return A[0, 0] + A[1, 1] + A[2, 2]

14def det(A):

15 if len(A) == 2:

16 return A[0, 0] * A[1, 1] - A[0, 1] * A[1, 0]

17 d = 0

18 for i, B in enumerate(A[0]):

19 minor = np.hstack([A[1:, :i], A[1:, i+1:]])

20 d += (-1)**i * B * det(minor)

21 return d

24def inv(

25 A: Union[Fraction, NDArray[float]],

26 as_fraction: Optional[bool] = True,

27) -> Union[Fraction, NDArray[float]]:

28 """

29 Inverts the A matrix.

31 Parameters

32 ----------

33 A

34 Array to be inverted.

35 as_fraction

36 Boolean which determines if inverted matrix is returned

37 as :class:`Fraction` matrix or real matrix.

38 """

40 detx2 = det(A) * 2

42 rest = (trace(A)**2 - trace(A @ A)) * np.diag([1, 1, 1]) - 2 * A * trace(A) + 2 * A @ A

44 if detx2 < 0:

45 detx2 = -detx2

46 rest = -rest

48 A_inv = np.array([Fraction(n, detx2) for n in rest.flat]).reshape(3, 3)

50 assert np.all(A @ A_inv == np.eye(len(A)))

52 # return either real matrix or as fraction matrix

53 A_inv_real = np.reshape([float(x) for x in A_inv.flat], A_inv.shape)

54 assert np.allclose(A_inv_real, np.linalg.inv(A))

56 if as_fraction:

57 return A_inv

58 else:

59 return A_inv_real

62def symmetrize_eigenvectors(

63 eigenvectors: NDArray[float],

64 cell: Optional[NDArray[float]] = None,

65 max_iter: Optional[int] = 1000,

66 method: Optional[str] = 'varimax',

67 tol: float = 1e-12,

68) -> NDArray[float]:

69 """Takes a set of vectors and tries to make them nice

71 Parameters

72 ----------

73 eigenvectors

74 If there are n degenerate eigenvectors and m atoms in the basis the

75 array should be `(n,m,3)`.

76 cell

77 If default `None` nothing is done to the cartesian directions but a cell

78 can be provided so the directions are in scaled coordinates instead.

79 max_iter

80 Maximum number of iterations in the symmetrization procedure (safety cap).

81 method

82 Can be `'varimax'` or `'quartimax'` or a parameter between `0: quartimax` and

83 `1: varimax`. Depending on the choice one obtains, e.g., Equamax, Parsimax, etc.

84 tol

85 Convergence tolerance passed to :func:`factor_analysis`; iteration stops

86 early once the rotation matrix has converged. Set `tol=0` to always run

87 `max_iter` iterations.

89 """

91 if cell is None:

92 cell = np.eye(3)

94 # s = band, i = basis, a = axis

95 eigenvectors = np.einsum('sia,ab->sib', eigenvectors, np.linalg.inv(cell))

97 components = eigenvectors.reshape(len(eigenvectors), -1).T

99 rotation_matrix = factor_analysis(components, iterations=max_iter, method=method, tol=tol)

100

101 new_eigenvectors = np.dot(components, rotation_matrix).T

102

103 new_eigenvectors = new_eigenvectors.reshape(len(new_eigenvectors), -1, 3)

104

105 new_eigenvectors = np.einsum('sia,ab->sib', new_eigenvectors, cell)

106

107 return new_eigenvectors

108

109

110def factor_analysis(

111 L: NDArray[float],

112 iterations: Optional[int] = 1000,

113 method: Optional[str] = 'varimax',

114 tol: float = 1e-12,

115) -> NDArray[float]:

116 """Performs factor analysis on `L` finding rotation matrix `R` such that `L @ R = L'` is simple.

117

118 In the future consider using the scikit learn methods directly but beware

119 the changes need to accommodate complex numbers.

120

121 Parameters

122 ----------

123 L

124 Matrix whose columns are rotated.

125 iterations

126 Maximum number of iterations (safety cap).

127 method

128 Can be `'varimax'` or `'quartimax'` or a parameter between `0: quartimax`

129 and `1: varimax`.

130 tol

131 Convergence tolerance: iteration stops once the rotation matrix `R`

132 stops changing, i.e. once ``max|R - R_prev| < tol`` between two

133 successive sweeps, with `iterations` as a safety cap.

134

135 References:

136 * *Sparse Modeling of Landmark and Texture Variability using the Orthomax Criterion*

137 Mikkel B. Stegmann, Karl Sjöstrand, Rasmus Larsen

138 http://www2.imm.dtu.dk/pubdb/edoc/imm4041.pdf

139

140 * http://www.cs.ucl.ac.uk/staff/d.barber/brml

141

142 * https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.FactorAnalysis.html # noqa

143

144 * https://stats.stackexchange.com/questions/185216/factor-rotation-methods-varimax-quartimax-oblimin-etc-what-do-the-names # noqa

145

146 """

147

148 nrow, ncol = L.shape

149 R = np.eye(ncol)

150

151 if method == 'varimax':

152 gamma = 1

153 elif method == 'quartimax':

154 gamma = 0

155 else:

156 gamma = method

157

158 for _ in range(iterations):

159 LR = np.dot(L, R)

160 grad = LR * (np.abs(LR)**2 - gamma * np.mean(np.abs(LR)**2, axis=0))

161 G = L.T.conj() @ grad

162 u, s, vh = np.linalg.svd(G)

163 R_new = u @ vh

164 change = np.abs(R_new - R).max()

165 R = R_new

166 if change < tol:

167 break

168

169 return R

170

171

172def group_eigvals(

173 vals: NDArray[float],

174 tol: Optional[float] = 1e-6,

175) -> list[list[int]]:

176 assert sorted(vals) == list(vals), vals

177

178 groups = [[0]]

179 for i in range(1, len(vals)):

180 if np.abs(vals[i] - vals[i-1]) < tol:

181 groups[-1].append(i)

182 else:

183 groups.append([i])

184 return groups

185

186

187# misc

188def as_fraction(not_fraction):

189 if isinstance(not_fraction, numbers.Number):

190 return Fraction(not_fraction)

191

192 if isinstance(not_fraction, np.ndarray):

193 arr = np.array([Fraction(n) for n in not_fraction.flat])

194 return arr.reshape(not_fraction.shape)

195

196 if isinstance(not_fraction, tuple):

197 return tuple(Fraction(n) for n in not_fraction)

198

199 if isinstance(not_fraction, list): 199 ↛ exitline 199 didn't return from function 'as_fraction' because the condition on line 199 was always true

200 arr = [Fraction(n) for n in not_fraction]

201 return arr

202

203

204@numba.njit # pragma: no cover

205def get_dynamical_matrix(fc, offsets, indices, q):

206

207 n = indices.max() + 1

208 N = len(fc)

209 D = np.zeros(shape=(n, n, 3, 3), dtype=np.complex128)

210 for ia in range(n):

211 for a in range(N):

212 if ia != indices[a]:

213 continue

214 na = offsets[a]

215 for b in range(N):

216 ib = indices[b]

217 nb = offsets[b]

218 dn = (nb - na).astype(np.float64)

219 D[ia, ib] += fc[a, b] * np.exp(2j*np.pi * np.dot(dn, q))

220 break

221 return D

222

223

224# For debug, this is a slower but perhaps more accurate variant, if the fc

225# obeys translational invariance this should give the same result

226# @numba.njit

227# def get_dynamical_matrix_full(fc, offsets, indices, q):

228#

229# n = indices.max() + 1

230# N = len(fc)

231# D = np.zeros(shape=(n, n, 3, 3), dtype=np.complex128)

232# for I in range(N):

233# i = indices[I]

234# m = offsets[I]

235# for J in range(N):

236# j = indices[J]

237# n = offsets[J]

238#

239# off = (n - m).astype(np.float64)

240#

241# phase = np.exp(1j * 2*np.pi * np.dot(off, q))

242#

243# D[i, j] += fc[I, J] * phase

244#

245# D /= (N / D.shape[0])

246#

247# return D

248

249

250def make_table(M):

251 rows = []

252 for r in M:

253 rows.append(''.join(f'{e:<20.2f}' for e in r))

254 return '\n'.join(rows)