Coverage for local_installation/dynasor/modes/mode

1from __future__ import annotations

2import numpy as np

3import warnings

4import pickle

5import itertools

7from ase.calculators.singlepoint import SinglePointCalculator

8from ase.units import fs

9from ase import Atoms

11from dynasor.units import radians_per_fs_to_THz

12from .tools import get_dynamical_matrix, group_eigvals, symmetrize_eigenvectors, get_displacements

13from .qpoint import QPoint

14from .atoms import Prim, Supercell

15from ..qpoints.tools import get_commensurate_lattice_points

18class ModeProjector:

19 """

20 The :class:`ModeProjector` maps between real atomic displacements `u` and

21 complex mode coordinates `Q`.

23 Some special python methods are implemented. The `__str__` and `__repr__`

24 provides useful info. :class:`QPoint` objects are representations of a

25 single q-point and associated information and can be accessed either by

26 call providing a reduced wavevector

28 >>> mp((1/2, 0, 0)) # doctest: +SKIP

30 or by index corresponding to reduced q-point accessible from

31 :attr:`ModeProjector.q_reduced`

33 >>> mp[2] # doctest: +SKIP

35 In addition to mode coordinates `Q` the class can also map the atomic

36 velocities `v` to mode momenta `P` as well as atomic forces `f` to mode

37 forces `F`. The class can also map back. This mapping is done using

38 getters and setters. Internally only Q, P and F are stored

40 >>> Q = mp.get_Q() # doctest: +SKIP

41 >>> mp.set_u(u) # doctest: +SKIP

43 In addition, the forces corresponding to the harmonic forces can be

44 accessed by :meth:`ModeProjector.get_f_harmonic()` and

45 :meth:`ModeProjector.get_F_harmonic()`. For ASE Atoms objects the

46 displacments etc. can be updated and applied by

48 >>> mp.update_from_atoms(atoms) # doctest: +SKIP

49 >>> atoms = mp.get_atoms(harmonic_forces=False) # doctest: +SKIP

51 The shapes for each property uses the follwoing varaibles

53 * N: Number of atoms in supercell

54 * Nu: unit cells (=N/Np)

55 * Np: primitive basis atoms (=N/Nu)

56 * Nb: bands (=Np*3)

57 * Nq: q-points (=Nu)

59 Please consult the documentation or the specific getters and setters to see

60 the exact transformations used.

62 Units

63 -----

64 The internal units in dynasor is Å, fs and eV. All frequencies are angular

65 (a.k.a. "physicist's" convention with 2π included). These are the units

66 dynasor will expect and return. In e.g. print functions conventional units

67 like fs, Å, THz, Da, meV will often be used.

69 Mass: The internal unit choice means that the mass unit is not Dalton

70 (a.k.a. amu or u) but instead 0.009648533290731906 Da. We denote this with

71 dmu (dynasor mass unit). In other words one Dalton is equal to

72 103.64269572045423 dmu units. This is typically not a problem since the

73 masses provided via e.g. ASE Atoms objects are converted internally.

75 Waves: dynasor will always communicate and expect spatial (angular)

76 frequencies in rad/Å and temporal (angular) frequencies in rad/fs. This

77 follows the "physicist's" convention as the factor of 2π is included in the

78 wave vectors. For instance the wavelength is given by λ=2π/q.

80 Mode amplitudes: Usually the mode amplitudes is communicated in Å√dmu

81 a.k.a. fs√eV.

83 Velocities: For modes the momenta (Å√dmu/fs or just √eV) is used but for

84 atoms the velocities (Å/fs) are used.

86 Mode forces: The force is just defined as the rate of change of the

87 canonical momenta so the unit would be Å√dmu/fs² (or √eV/fs)

89 Internal arrays

90 ---------------

91 For the curious, the internal data arrays are

93 * primitive, supercell, force_constants (input)

94 * _q, q_minus (reduced q-points and which q-points are related by inversion)

95 * _D, _w2, _W (dynamical matrices, frequencies (ev/Å²/Da), polarization vectors)

96 * _X (eigenmodes which are mass weighted "polarization vectors" in the supercell)

97 """

98 def __init__(self, primitive, supercell, force_constants):

99 """The mode projector is initialized by a primitive cell and a

100 supercell as well as harmonic force constants.

101

102 The force constants are assumed to be in units of eV/Å² as returned

103 from phonopy. Be careful about the permutations when working with force

104 constants and atoms object from different codes.

105

106 Parameters

107 ----------

108 primitive: Atoms

109 Primitive cell. Note that the masses are stored internally as

110 Dalton in ASE but will be converted to dynasor's internal

111 mass unit (dmu) internally in dynasor.

112 supercell: Atoms

113 Ideal supercell corresponding to the force constants.

114 force_constants

115 Force constants for the supercell in eV/Å² (N, N, 3, 3) wher N is

116 `len(supercell)`

117 """

118 if len(primitive) == len(supercell): 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true

119 warnings.warn('Primitive and supercell have the same size')

120 elif len(primitive) > len(supercell): 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true

121 raise ValueError('Primitive cell larger than supercell')

122 elif not (len(supercell) / len(primitive)).is_integer(): 122 ↛ 123line 122 didn't jump to line 123 because the condition on line 122 was never true

123 raise ValueError('supercell size is not multiple of primitive size')

124

125 if len(supercell) != len(force_constants): 125 ↛ 126line 125 didn't jump to line 126 because the condition on line 125 was never true

126 raise ValueError('force constants shape is not compatible with supercell size')

127

128 if force_constants.shape != (len(supercell), len(supercell), 3, 3): 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true

129 raise ValueError('force constants shape should be (N, N, 3, 3)')

130

131 self.primitive = Prim(primitive)

132 self.supercell = Supercell(supercell, primitive)

133 self.force_constants = force_constants

134

135 # Find q-points in reduced primitive cell coordinates

136 q_integer = get_commensurate_lattice_points(self.supercell.P.T)

137 q_reduced = np.dot(q_integer, self.supercell.P_inv.T)

138 self._q = np.array(sorted(tuple(q) for q in q_reduced))

139

140 # The equivalent q-point corresponding to -q

141 self._q_minus = [[tuple(q) for q in self._q].index(tuple((-q) % 1)) for q in self._q]

142

143 # Construct dynamical matrix and diagonalize at each q-point

144 self._D, self._w2, self._W = [], [], []

145 for qi, q in enumerate(self._q):

146

147 D = get_dynamical_matrix(

148 self.force_constants, self.supercell.offsets, self.supercell.indices,

149 q.astype(np.float64))

150

151 if qi == self.q_minus[qi]:

152 assert np.allclose(D.imag, 0)

153 D = D.real

154

155 D = np.einsum('ijab,i,j->ijab',

156 D, self.primitive.masses**-0.5, self.primitive.masses**-0.5)

157 D_matrix = D.transpose(0, 2, 1, 3).reshape(-1, self.primitive.n_atoms * 3)

158 assert np.allclose(D_matrix, D_matrix.T.conj())

159 w2, W = np.linalg.eigh(D_matrix)

160 W = W.T.reshape(-1, self.primitive.n_atoms, 3)

161

162 self._D.append(D)

163 self._w2.append(w2)

164 self._W.append(W)

165

166 self._D = np.array(self._D)

167 self._w2 = np.array(self._w2)

168 self._W = np.array(self._W)

169

170 # Post check basic symmetries, group eigenvalues and try to make degenerate modes nicer

171 for q, q_minus in enumerate(self.q_minus):

172 q_minus = self.q_minus[q]

173

174 assert np.allclose(self._D[q], self._D[q_minus].conj())

175 assert np.allclose(self._w2[q], self._w2[q_minus])

176

177 # tolerances for grouping and sorting eigenvalues and eigenvectors

178 round_decimals = 12

179 tolerance = 10**(-round_decimals)

180

181 for group in group_eigvals(self._w2[q], tolerance**0.5):

182 W = symmetrize_eigenvectors(self._W[q, group])

183

184 # Try to order them

185 W_sort = W.copy().transpose(0, 2, 1).reshape(len(W), -1)

186 # abs is because we want to consider the magnitude

187 # - (minus) basically reverts the sort order to place largest first

188 # T is just because how lexsort works, we want to consider each

189 # atom and direction as a key for the bands

190 # -1 is because we want to make the x-direction of the first

191 # atom the mist significant key

192 # At the end the first band should have the largest magnitude

193 # for the first atom in x

194 argsort = np.lexsort(np.round(-np.abs(W_sort).T[::-1], round_decimals))

195 self._W[q, group] = W[argsort]

196

197 self._W[q_minus] = self._W[q].conj()

198

199 # Construct supercell projection matrix

200 # q_ks = X_ksna u_na

201 self._X = np.zeros((len(self._q), self.primitive.n_atoms * 3, self.supercell.n_atoms, 3),

202 dtype=np.complex128)

203

204 for index in range(self.supercell.n_atoms):

205 i, N = self.supercell.indices[index], self.supercell.offsets[index]

206 for q, s, a in itertools.product(

207 range(len(self._q)), range(self.primitive.n_atoms * 3), range(3)):

208 phase = np.exp(-1j * 2*np.pi * self._q[q] @ N)

209 self._X[q, s, index, a] = (

210 self.primitive.masses[i]**0.5 * phase * self._W[q, s, i, a].conj())

211 self._X /= (self.supercell.n_atoms / self.primitive.n_atoms)**0.5

212

213 # Init arrays to hold Q, P and F

214 self._Q = np.zeros((len(self._q), self.primitive.n_atoms*3), dtype=np.complex128)

215 self._P = np.zeros_like(self._Q)

216 self._F = np.zeros_like(self._Q)

217

218 # Dunder functions

219 def __str__(self):

220 strings = ['### ModeProjector ###']

221 strings += [f'{self.supercell}']

222 strings += [f'{self.primitive}']

223 string = '\n'.join(strings)

224

225 # ASCII DOS!

226 width = 80

227 height = 24

228 dos = np.full((height, width), ' ')

229

230 THz = self.omegas * radians_per_fs_to_THz

231

232 hist, bins = np.histogram(THz.flat, bins=width)

233

234 for i, h in enumerate(hist):

235 dos[np.round(h * (height - 1) / hist.max()).astype(int), i] = '+' # '·' or 'x'

236

237 dos = dos[::-1]

238

239 dos[-1, dos[-1] == ' '] = '-'

240

241 dos = '\n'.join([''.join(d) for d in dos])

242

243 string += f'\n{dos}'

244

245 string += f'\n|{THz.min():<10.2f} THz' + ' '*(width - 26) + f'{THz.max():>10.2f}|'

246

247 return string

248

249 def __repr__(self):

250 return str(self)

251

252 def __getitem__(self, q) -> QPoint:

253 """Returns the q-point object based on its index"""

254 if q < 0 or q >= len(self._q):

255 raise IndexError

256 return QPoint(q, self)

257

258 def __call__(self, qpoint) -> QPoint:

259 """Tries to find a matching q-point based on reduced coordinate"""

260 qpoint = np.array(qpoint).astype(np.float64) % 1

261 for q, qpoint2 in enumerate(np.array(self._q).astype(np.float64)): 261 ↛ 264line 261 didn't jump to line 264 because the loop on line 261 didn't complete

262 if np.allclose(qpoint, qpoint2):

263 return QPoint(q, self)

264 raise ValueError('qpoint not compatible, check mp.q_reduced')

265

266 # Getters ans setters for internal mutable arrays

267 def get_Q(self):

268 """The mode coordinate amplitudes in Å√dmu"""

269 return self._Q.copy()

270

271 def get_P(self):

272 """The mode momentum amplitudes in √eV"""

273 return self._P.copy()

274

275 def get_F(self):

276 """The mode force amplitudes in eV/Å√dmu"""

277 return self._F.copy()

278

279 def get_F_harmonic(self):

280 """The harmonic mode forces, computed as -omega^2 * Q, in eV/Å√dmu"""

281 return -self._w2 * self.get_Q()

282

283 def set_Q(self, Q):

284 """Sets the internal mode coordinates Q

285

286 The functions ensures Q(-q)=Q*(q)

287 """

288 # This ensures that stuff like mp.set_Q(0) works while not updating the

289 # array until the assert

290 Q_new = self.get_Q()

291 Q_new[:] = Q

292 if not np.allclose(np.conjugate(Q_new), Q_new[self.q_minus]): 292 ↛ 293line 292 didn't jump to line 293 because the condition on line 292 was never true

293 raise ValueError('Supplied Q does not fulfill Q(-q) = Q(q)*')

294 self._Q[:] = Q_new

295

296 def set_P(self, P):

297 """Sets the internal mode momenta P.

298

299 The functions ensures :math:`P(-q)=P^*(q)`

300 """

301 P_new = self.get_P()

302 P_new[:] = P

303 if not np.allclose(np.conjugate(P_new), P_new[self.q_minus]): 303 ↛ 304line 303 didn't jump to line 304 because the condition on line 303 was never true

304 raise ValueError('Supplied P does not fulfill P(-q) = P(q)*')

305 self._P[:] = P_new

306

307 def set_F(self, F):

308 """Sets the internal mode forces F.

309

310 The functions ensures :math:`F(-q)=F^*(q)`

311 """

312 F_new = self.get_F()

313 F_new[:] = F

314 if not np.allclose(np.conjugate(F_new), F_new[self.q_minus]): 314 ↛ 315line 314 didn't jump to line 315 because the condition on line 314 was never true

315 raise ValueError('Supplied F does not fulfill F(-q) = F(q)*')

316 self._F[:] = F_new

317

318 def get_u(self):

319 """The atomic displacements in Å"""

320 u = np.einsum('ksna,ks,n->na', self._X.conj(), self._Q, 1 / self.supercell.masses)

321 assert np.allclose(u.imag, 0)

322 return u.real

323

324 def get_v(self):

325 """The atomic velocities in Å/fs"""

326 v = np.einsum('ksna,ks,n->na', self._X, self._P, 1 / self.supercell.masses)

327 assert np.allclose(v.imag, 0)

328 return v.real

329

330 def get_f(self):

331 """The atomic forces in eV/Å"""

332 f = np.einsum('ksna,ks->na', self._X, self._F)

333 assert np.allclose(f.imag, 0)

334 return f.real

335

336 def get_f_harmonic(self):

337 """The harmonic atomic forces for the current displacements"""

338 F_harmonic = self.get_F_harmonic()

339 f_harmonic = np.einsum('ksna,ks->na', self._X, F_harmonic)

340 assert np.allclose(f_harmonic.imag, 0)

341 return f_harmonic.real

342

343 def set_u(self, u):

344 """Sets the internal mode coordinates Q given the atomic displacements u

345

346 `Q = X * u`

347

348 Parameters

349 ----------

350 u

351 The atomic displacements in Å

352 """

353 Q = np.einsum('ksna,na->ks', self._X, u)

354 self.set_Q(Q)

355

356 def set_v(self, v):

357 """Sets the internal mode momenta P given the velocities v

358

359 `P = X.conj() * v`

360

361 Parameters

362 ----------

363 v

364 The atomic velocities in fs/Å

365 """

366 P = np.einsum('ksna,na->ks', self._X.conj(), v)

367 self.set_P(P)

368

369 def set_f(self, f):

370 """Sets the internal mode forces F given the forces f

371

372 `F = X.conj() * f / m`

373

374 Parameters

375 ----------

376 f

377 The atomic forces in eV/Å

378 """

379 F = np.einsum('ksna,na,n->ks', self._X.conj(), f, 1 / self.supercell.masses)

380 self.set_F(F)

381

382 # Convenience functions to handle ASE Atoms objects

383 def get_atoms(self, harmonic_forces=False) -> Atoms:

384 """Returns ase Atoms object with displacement, velocities, forces and harmonic energies set

385

386 Parameters

387 ----------

388 harmonic_forces

389 Whether the forces should be taken from the internal `F` or via `-w2 Q`

390 """

391 atoms = self.supercell.to_ase()

392 atoms.positions += self.get_u()

393 atoms.set_velocities(self.get_v() / fs)

394 E = self.potential_energies.sum()

395 f = self.get_f_harmonic() if harmonic_forces else self.get_f()

396

397 atoms.calc = SinglePointCalculator(

398 energy=E, forces=f, stress=None, magmoms=None, atoms=atoms)

399

400 return atoms

401

402 def update_from_atoms(self, atoms):

403 """Updates the ModeProjector with displacments, velocities and forces

404 from an ASE Atoms object

405

406 Checks for attached calculator in first place and next for forces array.

407

408 If no data sets corresponding array to zeros.

409

410 The masses and velocities are converted to dynasor units internally.

411 """

412

413 u = get_displacements(atoms, self.supercell)

414 if np.max(np.abs(u)) > 2.0: 414 ↛ 415line 414 didn't jump to line 415 because the condition on line 414 was never true

415 warnings.warn('Displacements larger than 2Å. Is the atoms object permuted?')

416 self.set_u(u)

417 self.set_v(atoms.get_velocities() * fs)

418 try:

419 self.set_f(atoms.get_forces())

420 except RuntimeError:

421 if 'forces' in atoms.arrays: 421 ↛ 424line 421 didn't jump to line 424 because the condition on line 421 was always true

422 self.set_f(atoms.arrays['forces'])

423 else:

424 self.set_f(np.zeros_like(atoms.positions))

425

426 # properties

427 @property

428 def q_minus(self):

429 """The index of the corresponding counter-propagating mode (-q)"""

430 return self._q_minus.copy()

431

432 @property

433 def q_reduced(self):

434 """The q-points in reduced coordinates.

435

436 A zone boundary mode would be e.g. (1/2, 0, 0)

437 """

438 return self._q.astype(float)

439

440 @property

441 def q_cartesian(self):

442 """The q-points in cartesian coordinates with unit of rad/Å (2π included)"""

443 return 2*np.pi * self.q_reduced @ self.primitive.inv_cell

444

445 @property

446 def omegas(self):

447 """The frequencies of each mode in (rad/fs).

448

449 Negative frequencies correspond to imaginary frequencies

450 """

451 return np.sign(self._w2) * np.sqrt(np.abs(self._w2))

452

453 @property

454 def polarizations(self):

455 """The polarization vectors for each mode (Nq, Nb, Np, 3)"""

456 return self._W

457

458 @property

459 def eigenmodes(self):

460 """The eigenmodes in the supercell as (Nq, Nb, N, 3)-array

461

462 The eigenmodes with masses included so that

463

464 `Q = X u`

465

466 where u is the supercell displacements

467 """

468 return self._X

469

470 @property

471 def potential_energies(self):

472 """Potential energy per mode as (Nq, Nb)-array

473

474 The potential energies are defined as :math:`1/2QQ^*` and should equal

475 :math:`1/2k_BT` in equilibrium for a harmonic system.`

476 """

477 return 1/2 * np.abs(self._Q)**2 * self._w2

478

479 @property

480 def kinetic_energies(self):

481 """Kinetic energy per mode as (Nq, Nb)-array

482

483 The kinetic energies are defined as :math:`1/2PP^*`. Should equal

484 :math:`1/2k_BT` in equilibrium.`

485 """

486 return 1/2 * np.abs(self._P)**2

487

488 @property

489 def virial_energies(self):

490 """The virial energies per mode as (Nq, Nb)-array

491

492 The virial energies are defined here as :math:`-1/2QF` which should have an

493 expectation value of :math:`1/2k_BT` per mode in equilibrium. For a harmonic

494 system this is simply equal to the potential energy. This means that

495 that the virial energy can be used to monitor the anharmonicity or

496 define a measure of the potential energy.

497 """

498 return -1/2 * self._Q * self._F

499

500 def write(self, file_name: str):

501 """Uses pickle to write mp to file"""

502 with open(file_name, 'wb') as f:

503 pickle.dump(self, f)

504

505 @classmethod

506 def read(cls, file_name: str) -> ModeProjector:

507 """Return mp instance from pickle file that was saved using mp.write"""

508 with open(file_name, 'rb') as f:

509 mp = pickle.load(f)

510 return mp

Coverage for local_installation/dynasor/modes/mode_projector.py: 87%

206 statements