Coverage for dynasor/modes/tools.py: 99%
88 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-22 07:00 +0000
1import numbers
2from fractions import Fraction
3from typing import Optional, Union
4import numba
5import numpy as np
6from numpy.typing import NDArray
9# Mode projections
10def mode_coordinates(X: NDArray[complex], u: NDArray[float]) -> NDArray[complex]:
11 r"""Projects atomic displacements onto modes, :math:`Q = X u`.
13 Parameters
14 ----------
15 X
16 Modes with shape ``(..., N, 3)``, where the leading dimensions define the output shape.
17 u
18 Atomic displacements in Å.
20 Returns
21 -------
22 The mode coordinates in Å√dmu, with shape ``X.shape[:-2]``.
23 """
24 return np.einsum('...na,na->...', X, u, optimize=True)
27def mode_momenta(X: NDArray[complex], v: NDArray[float]) -> NDArray[complex]:
28 r"""Projects atomic velocities onto modes, :math:`P = X^* v`.
30 Parameters
31 ----------
32 X
33 Modes with shape ``(..., N, 3)``, where the leading dimensions define the output shape.
34 v
35 Atomic velocities in Å/fs.
37 Returns
38 -------
39 The mode momenta in √eV, with shape ``X.shape[:-2]``.
40 """
41 return np.einsum('...na,na->...', X.conj(), v, optimize=True)
44def mode_forces(
45 X: NDArray[complex],
46 f: NDArray[float],
47 masses: NDArray[float],
48) -> NDArray[complex]:
49 r"""Projects atomic forces onto modes, :math:`F = X^* f / m`.
51 Parameters
52 ----------
53 X
54 Modes with shape ``(..., N, 3)``, where the leading dimensions define the output shape.
55 f
56 Atomic forces in eV/Å.
57 masses
58 Atomic masses in dmu, one per atom.
60 Returns
61 -------
62 The mode forces in eV/Å√dmu, with shape ``X.shape[:-2]``.
63 """
64 # The mass divides the forces rather than the modes, since f is the smaller array.
65 return np.einsum('...na,na->...', X.conj(), f / masses[:, None], optimize=True)
68# Linalg
69def trace(A):
70 return A[0, 0] + A[1, 1] + A[2, 2]
73def det(A):
74 if len(A) == 2:
75 return A[0, 0] * A[1, 1] - A[0, 1] * A[1, 0]
76 d = 0
77 for i, B in enumerate(A[0]):
78 minor = np.hstack([A[1:, :i], A[1:, i+1:]])
79 d += (-1)**i * B * det(minor)
80 return d
83def inv(
84 A: Union[Fraction, NDArray[float]],
85 as_fraction: Optional[bool] = True,
86) -> Union[Fraction, NDArray[float]]:
87 """
88 Inverts the A matrix.
90 Parameters
91 ----------
92 A
93 Array to be inverted.
94 as_fraction
95 Boolean which determines if inverted matrix is returned
96 as :class:`Fraction` matrix or real matrix.
97 """
99 detx2 = det(A) * 2
101 rest = (trace(A)**2 - trace(A @ A)) * np.diag([1, 1, 1]) - 2 * A * trace(A) + 2 * A @ A
103 if detx2 < 0:
104 detx2 = -detx2
105 rest = -rest
107 A_inv = np.array([Fraction(n, detx2) for n in rest.flat]).reshape(3, 3)
109 assert np.all(A @ A_inv == np.eye(len(A)))
111 # return either real matrix or as fraction matrix
112 A_inv_real = np.reshape([float(x) for x in A_inv.flat], A_inv.shape)
113 assert np.allclose(A_inv_real, np.linalg.inv(A))
115 if as_fraction:
116 return A_inv
117 else:
118 return A_inv_real
121def symmetrize_eigenvectors(
122 eigenvectors: NDArray[float],
123 cell: Optional[NDArray[float]] = None,
124 max_iter: Optional[int] = 1000,
125 method: Optional[str] = 'varimax',
126 tol: float = 1e-12,
127) -> NDArray[float]:
128 """Takes a set of vectors and tries to make them nice
130 Parameters
131 ----------
132 eigenvectors
133 If there are n degenerate eigenvectors and m atoms in the basis the
134 array should be `(n,m,3)`.
135 cell
136 If default `None` nothing is done to the cartesian directions but a cell
137 can be provided so the directions are in scaled coordinates instead.
138 max_iter
139 Maximum number of iterations in the symmetrization procedure (safety cap).
140 method
141 Can be `'varimax'` or `'quartimax'` or a parameter between `0: quartimax` and
142 `1: varimax`. Depending on the choice one obtains, e.g., Equamax, Parsimax, etc.
143 tol
144 Convergence tolerance passed to :func:`factor_analysis`; iteration stops
145 early once the rotation matrix has converged. Set `tol=0` to always run
146 `max_iter` iterations.
148 """
150 if cell is None:
151 cell = np.eye(3)
153 # s = band, i = basis, a = axis
154 eigenvectors = np.einsum('sia,ab->sib', eigenvectors, np.linalg.inv(cell))
156 components = eigenvectors.reshape(len(eigenvectors), -1).T
158 rotation_matrix = factor_analysis(components, iterations=max_iter, method=method, tol=tol)
160 new_eigenvectors = np.dot(components, rotation_matrix).T
162 new_eigenvectors = new_eigenvectors.reshape(len(new_eigenvectors), -1, 3)
164 new_eigenvectors = np.einsum('sia,ab->sib', new_eigenvectors, cell)
166 return new_eigenvectors
169def factor_analysis(
170 L: NDArray[float],
171 iterations: Optional[int] = 1000,
172 method: Optional[str] = 'varimax',
173 tol: float = 1e-12,
174) -> NDArray[float]:
175 """Performs factor analysis on `L` finding rotation matrix `R` such that `L @ R = L'` is simple.
177 In the future consider using the scikit learn methods directly but beware
178 the changes need to accommodate complex numbers.
180 Parameters
181 ----------
182 L
183 Matrix whose columns are rotated.
184 iterations
185 Maximum number of iterations (safety cap).
186 method
187 Can be `'varimax'` or `'quartimax'` or a parameter between `0: quartimax`
188 and `1: varimax`.
189 tol
190 Convergence tolerance: iteration stops once the rotation matrix `R`
191 stops changing, i.e. once ``max|R - R_prev| < tol`` between two
192 successive sweeps, with `iterations` as a safety cap.
194 References:
195 * *Sparse Modeling of Landmark and Texture Variability using the Orthomax Criterion*
196 Mikkel B. Stegmann, Karl Sjöstrand, Rasmus Larsen
197 http://www2.imm.dtu.dk/pubdb/edoc/imm4041.pdf
199 * http://www.cs.ucl.ac.uk/staff/d.barber/brml
201 * https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.FactorAnalysis.html # noqa
203 * https://stats.stackexchange.com/questions/185216/factor-rotation-methods-varimax-quartimax-oblimin-etc-what-do-the-names # noqa
205 """
207 nrow, ncol = L.shape
208 R = np.eye(ncol)
210 if method == 'varimax':
211 gamma = 1
212 elif method == 'quartimax':
213 gamma = 0
214 else:
215 gamma = method
217 for _ in range(iterations):
218 LR = np.dot(L, R)
219 grad = LR * (np.abs(LR)**2 - gamma * np.mean(np.abs(LR)**2, axis=0))
220 G = L.T.conj() @ grad
221 u, s, vh = np.linalg.svd(G)
222 R_new = u @ vh
223 change = np.abs(R_new - R).max()
224 R = R_new
225 if change < tol:
226 break
228 return R
231def group_eigvals(
232 vals: NDArray[float],
233 tol: Optional[float] = 1e-6,
234) -> list[list[int]]:
235 assert sorted(vals) == list(vals), vals
237 groups = [[0]]
238 for i in range(1, len(vals)):
239 if np.abs(vals[i] - vals[i-1]) < tol:
240 groups[-1].append(i)
241 else:
242 groups.append([i])
243 return groups
246# misc
247def as_fraction(not_fraction):
248 if isinstance(not_fraction, numbers.Number):
249 return Fraction(not_fraction)
251 if isinstance(not_fraction, np.ndarray):
252 arr = np.array([Fraction(n) for n in not_fraction.flat])
253 return arr.reshape(not_fraction.shape)
255 if isinstance(not_fraction, tuple):
256 return tuple(Fraction(n) for n in not_fraction)
258 if isinstance(not_fraction, list): 258 ↛ exitline 258 didn't return from function 'as_fraction' because the condition on line 258 was always true
259 arr = [Fraction(n) for n in not_fraction]
260 return arr
263@numba.njit # pragma: no cover
264def get_dynamical_matrix(fc, offsets, indices, q):
266 n = indices.max() + 1
267 N = len(fc)
268 D = np.zeros(shape=(n, n, 3, 3), dtype=np.complex128)
269 for ia in range(n):
270 for a in range(N):
271 if ia != indices[a]:
272 continue
273 na = offsets[a]
274 for b in range(N):
275 ib = indices[b]
276 nb = offsets[b]
277 dn = (nb - na).astype(np.float64)
278 D[ia, ib] += fc[a, b] * np.exp(2j*np.pi * np.dot(dn, q))
279 break
280 return D
283# For debug, this is a slower but perhaps more accurate variant, if the fc
284# obeys translational invariance this should give the same result
285# @numba.njit
286# def get_dynamical_matrix_full(fc, offsets, indices, q):
287#
288# n = indices.max() + 1
289# N = len(fc)
290# D = np.zeros(shape=(n, n, 3, 3), dtype=np.complex128)
291# for I in range(N):
292# i = indices[I]
293# m = offsets[I]
294# for J in range(N):
295# j = indices[J]
296# n = offsets[J]
297#
298# off = (n - m).astype(np.float64)
299#
300# phase = np.exp(1j * 2*np.pi * np.dot(off, q))
301#
302# D[i, j] += fc[I, J] * phase
303#
304# D /= (N / D.shape[0])
305#
306# return D
309def make_table(M):
310 rows = []
311 for r in M:
312 rows.append(''.join(f'{e:<20.2f}' for e in r))
313 return '\n'.join(rows)