Coverage for /builds/kinetik161/ase/ase/vibrations/data.py: 98.32%

1"""Storage and analysis for vibrational data"""

3import collections

4from math import pi, sin, sqrt

5from numbers import Integral, Real

6from typing import Any, Dict, Iterator, List, Sequence, Tuple, TypeVar, Union

8import numpy as np

10import ase.io

11import ase.units as units

12from ase.atoms import Atoms

13from ase.calculators.singlepoint import SinglePointCalculator

14from ase.constraints import FixAtoms, FixCartesian, constrained_indices

15from ase.spectrum.doscollection import DOSCollection

16from ase.spectrum.dosdata import RawDOSData

17from ase.utils import jsonable, lazymethod

19RealSequence4D = Sequence[Sequence[Sequence[Sequence[Real]]]]

20VD = TypeVar('VD', bound='VibrationsData')

23@jsonable('vibrationsdata')

24class VibrationsData:

25 """Class for storing and analyzing vibrational data (i.e. Atoms + Hessian)

27 This class is not responsible for calculating Hessians; the Hessian should

28 be computed by a Calculator or some other algorithm. Once the

29 VibrationsData has been constructed, this class provides some common

30 processing options; frequency calculation, mode animation, DOS etc.

32 If the Atoms object is a periodic supercell, VibrationsData may be

33 converted to a PhononData using the VibrationsData.to_phonondata() method.

34 This provides access to q-point-dependent analyses such as phonon

35 dispersion plotting.

37 If the Atoms object has FixedAtoms or FixedCartesian constraints, these

38 will be respected and the Hessian should be sized accordingly.

40 Args:

41 atoms:

42 Equilibrium geometry of vibrating system. This will be stored as a

43 full copy.

45 hessian: Second-derivative in energy with respect to

46 Cartesian nuclear movements as an (N, 3, N, 3) array.

48 indices: indices of atoms which are included

49 in Hessian. Default value (None) includes all freely

50 moving atoms (i.e. not fixed ones). Leave at None if

51 constraints should be determined automatically from the

52 atoms object.

54 """

56 def __init__(self,

57 atoms: Atoms,

58 hessian: Union[RealSequence4D, np.ndarray],

59 indices: Union[Sequence[int], np.ndarray] = None,

60 ) -> None:

62 if indices is None:

63 indices = np.asarray(self.indices_from_constraints(atoms),

64 dtype=int)

66 self._indices = np.array(indices, dtype=int)

68 n_atoms = self._check_dimensions(atoms, np.asarray(hessian),

69 indices=self._indices)

70 self._atoms = atoms.copy()

72 self._hessian2d = (np.asarray(hessian)

73 .reshape(3 * n_atoms, 3 * n_atoms).copy())

75 _setter_error = ("VibrationsData properties cannot be modified: construct "

76 "a new VibrationsData with consistent atoms, Hessian and "

77 "(optionally) indices/mask.")

79 @classmethod

80 def from_2d(cls, atoms: Atoms,

81 hessian_2d: Union[Sequence[Sequence[Real]], np.ndarray],

82 indices: Sequence[int] = None) -> 'VibrationsData':

83 """Instantiate VibrationsData when the Hessian is in a 3Nx3N format

85 Args:

86 atoms: Equilibrium geometry of vibrating system

88 hessian: Second-derivative in energy with respect to

89 Cartesian nuclear movements as a (3N, 3N) array.

91 indices: Indices of (non-frozen) atoms included in Hessian

93 """

94 if indices is None:

95 indices = range(len(atoms))

96 assert indices is not None # Show Mypy that indices is now a sequence

98 hessian_2d_array = np.asarray(hessian_2d)

99 n_atoms = cls._check_dimensions(atoms, hessian_2d_array,

100 indices=indices, two_d=True)

101

102 return cls(atoms, hessian_2d_array.reshape(n_atoms, 3, n_atoms, 3),

103 indices=indices)

104

105 @staticmethod

106 def indices_from_constraints(atoms: Atoms) -> List[int]:

107 """Indices corresponding to Atoms Constraints

108

109 Deduces the freely moving atoms from the constraints set on the

110 atoms object. VibrationsData only supports FixCartesian and

111 FixAtoms. All others are neglected.

112

113 Args:

114 atoms: Atoms object.

115

116 Retruns:

117 indices of free atoms.

118

119 """

120 # Only fully fixed atoms supported by VibrationsData

121 const_indices = constrained_indices(

122 atoms, only_include=(FixCartesian, FixAtoms))

123 # Invert the selection to get free atoms

124 indices = np.setdiff1d(

125 np.array(

126 range(

127 len(atoms))),

128 const_indices).astype(int)

129 return indices.tolist()

130

131 @staticmethod

132 def indices_from_mask(mask: Union[Sequence[bool], np.ndarray]

133 ) -> List[int]:

134 """Indices corresponding to boolean mask

135

136 This is provided as a convenience for instantiating VibrationsData with

137 a boolean mask. For example, if the Hessian data includes only the H

138 atoms in a structure::

139

140 h_mask = atoms.get_chemical_symbols() == 'H'

141 vib_data = VibrationsData(atoms, hessian,

142 VibrationsData.indices_from_mask(h_mask))

143

144 Take care to ensure that the length of the mask corresponds to the full

145 number of atoms; this function is only aware of the mask it has been

146 given.

147

148 Args:

149 mask: a sequence of True, False values

150

151 Returns:

152 indices of True elements

153

154 """

155 return np.where(mask)[0].tolist()

156

157 @staticmethod

158 def _check_dimensions(atoms: Atoms,

159 hessian: np.ndarray,

160 indices: Union[np.ndarray, Sequence[int]],

161 two_d: bool = False) -> int:

162 """Sanity check on array shapes from input data

163

164 Args:

165 atoms: Structure

166 indices: Indices of atoms used in Hessian

167 hessian: Proposed Hessian array

168

169 Returns:

170 Number of atoms contributing to Hessian

171

172 Raises:

173 ValueError if Hessian dimensions are not (N, 3, N, 3)

174

175 """

176

177 n_atoms = len(atoms[indices])

178

179 if two_d:

180 ref_shape = [n_atoms * 3, n_atoms * 3]

181 ref_shape_txt = '{n:d}x{n:d}'.format(n=(n_atoms * 3))

182

183 else:

184 ref_shape = [n_atoms, 3, n_atoms, 3]

185 ref_shape_txt = '{n:d}x3x{n:d}x3'.format(n=n_atoms)

186

187 if (isinstance(hessian, np.ndarray)

188 and hessian.shape == tuple(ref_shape)):

189 return n_atoms

190 else:

191 raise ValueError("Hessian for these atoms should be a "

192 "{} numpy array.".format(ref_shape_txt))

193

194 def get_atoms(self) -> Atoms:

195 return self._atoms.copy()

196

197 def get_indices(self) -> np.ndarray:

198 return self._indices.copy()

199

200 def get_mask(self) -> np.ndarray:

201 """Boolean mask of atoms selected by indices"""

202 return self._mask_from_indices(self._atoms, self.get_indices())

203

204 @staticmethod

205 def _mask_from_indices(atoms: Atoms,

206 indices: Union[None, Sequence[int], np.ndarray]

207 ) -> np.ndarray:

208 """Boolean mask of atoms selected by indices"""

209 natoms = len(atoms)

210

211 # Wrap indices to allow negative values

212 indices = np.asarray(indices) % natoms

213

214 mask = np.full(natoms, False, dtype=bool)

215 mask[indices] = True

216 return mask

217

218 def get_hessian(self) -> np.ndarray:

219 """The Hessian; second derivative of energy wrt positions

220

221 This format is preferred for iteration over atoms and when

222 addressing specific elements of the Hessian.

223

224 Returns:

225 array with shape (n_atoms, 3, n_atoms, 3) where

226

227 - the first and third indices identify atoms in self.get_atoms()

228

229 - the second and fourth indices cover the corresponding

230 Cartesian movements in x, y, z

231

232 e.g. the element h[0, 2, 1, 0] gives a harmonic force exerted on

233 atoms[1] in the x-direction in response to a movement in the

234 z-direction of atoms[0]

235 """

236 n_atoms = int(self._hessian2d.shape[0] / 3)

237 return self._hessian2d.reshape(n_atoms, 3, n_atoms, 3).copy()

238

239 def get_hessian_2d(self) -> np.ndarray:

240 """Get the Hessian as a 2-D array

241

242 This format may be preferred for use with standard linear algebra

243 functions

244

245 Returns:

246 array with shape (n_atoms * 3, n_atoms * 3) where the elements are

247 ordered by atom and Cartesian direction::

248

249 >> [[at1x_at1x, at1x_at1y, at1x_at1z, at1x_at2x, ...],

250 >> [at1y_at1x, at1y_at1y, at1y_at1z, at1y_at2x, ...],

251 >> [at1z_at1x, at1z_at1y, at1z_at1z, at1z_at2x, ...],

252 >> [at2x_at1x, at2x_at1y, at2x_at1z, at2x_at2x, ...],

253 >> ...]

254

255

256 e.g. the element h[2, 3] gives a harmonic force exerted on

257 atoms[1] in the x-direction in response to a movement in the

258 z-direction of atoms[0]

259 """

260 return self._hessian2d.copy()

261

262 def todict(self) -> Dict[str, Any]:

263 if np.allclose(self._indices, range(len(self._atoms))):

264 indices = None

265 else:

266 indices = self.get_indices()

267

268 return {'atoms': self.get_atoms(),

269 'hessian': self.get_hessian(),

270 'indices': indices}

271

272 @classmethod

273 def fromdict(cls, data: Dict[str, Any]) -> 'VibrationsData':

274 # mypy is understandably suspicious of data coming from a dict that

275 # holds mixed types, but it can see if we sanity-check with 'assert'

276 assert isinstance(data['atoms'], Atoms)

277 assert isinstance(data['hessian'], (collections.abc.Sequence,

278 np.ndarray))

279 if data['indices'] is not None:

280 assert isinstance(data['indices'], (collections.abc.Sequence,

281 np.ndarray))

282 for index in data['indices']:

283 assert isinstance(index, Integral)

284

285 return cls(data['atoms'], data['hessian'], indices=data['indices'])

286

287 @lazymethod

288 def _energies_and_modes(self) -> Tuple[np.ndarray, np.ndarray]:

289 """Diagonalise the Hessian to obtain harmonic modes

290

291 This method is an internal implementation of get_energies_and_modes(),

292 see the docstring of that method for more information.

293

294 """

295 active_atoms = self._atoms[self.get_mask()]

296 n_atoms = len(active_atoms)

297 masses = active_atoms.get_masses()

298

299 if not np.all(masses):

300 raise ValueError('Zero mass encountered in one or more of '

301 'the vibrated atoms. Use Atoms.set_masses()'

302 ' to set all masses to non-zero values.')

303 mass_weights = np.repeat(masses**-0.5, 3)

304

305 omega2, vectors = np.linalg.eigh(mass_weights

306 * self.get_hessian_2d()

307 * mass_weights[:, np.newaxis])

308

309 unit_conversion = units._hbar * units.m / sqrt(units._e * units._amu)

310 energies = unit_conversion * omega2.astype(complex)**0.5

311

312 modes = vectors.T.reshape(n_atoms * 3, n_atoms, 3)

313 modes = modes * masses[np.newaxis, :, np.newaxis]**-0.5

314

315 return (energies, modes)

316

317 def get_energies_and_modes(self, all_atoms: bool = False

318 ) -> Tuple[np.ndarray, np.ndarray]:

319 """Diagonalise the Hessian to obtain harmonic modes

320

321 Results are cached so diagonalization will only be performed once for

322 this object instance.

323

324 Args:

325 all_atoms:

326 If True, return modes as (3N, [N + N_frozen], 3) array where

327 the second axis corresponds to the full list of atoms in the

328 attached atoms object. Atoms that were not included in the

329 Hessian will have displacement vectors of (0, 0, 0).

330

331 Returns:

332 tuple (energies, modes)

333

334 Energies are given in units of eV. (To convert these to frequencies

335 in cm-1, divide by ase.units.invcm.)

336

337 Modes are given in Cartesian coordinates as a (3N, N, 3) array

338 where indices correspond to the (mode_index, atom, direction).

339

340 """

341

342 energies, modes_from_hessian = self._energies_and_modes()

343

344 if all_atoms:

345 n_active_atoms = len(self.get_indices())

346 n_all_atoms = len(self._atoms)

347 modes = np.zeros((3 * n_active_atoms, n_all_atoms, 3))

348 modes[:, self.get_mask(), :] = modes_from_hessian

349 else:

350 modes = modes_from_hessian.copy()

351

352 return (energies.copy(), modes)

353

354 def get_modes(self, all_atoms: bool = False) -> np.ndarray:

355 """Diagonalise the Hessian to obtain harmonic modes

356

357 Results are cached so diagonalization will only be performed once for

358 this object instance.

359

360 all_atoms:

361 If True, return modes as (3N, [N + N_frozen], 3) array where

362 the second axis corresponds to the full list of atoms in the

363 attached atoms object. Atoms that were not included in the

364 Hessian will have displacement vectors of (0, 0, 0).

365

366 Returns:

367 Modes in Cartesian coordinates as a (3N, N, 3) array where indices

368 correspond to the (mode_index, atom, direction).

369

370 """

371 return self.get_energies_and_modes(all_atoms=all_atoms)[1]

372

373 def get_energies(self) -> np.ndarray:

374 """Diagonalise the Hessian to obtain eigenvalues

375

376 Results are cached so diagonalization will only be performed once for

377 this object instance.

378

379 Returns:

380 Harmonic mode energies in units of eV

381

382 """

383 return self.get_energies_and_modes()[0]

384

385 def get_frequencies(self) -> np.ndarray:

386 """Diagonalise the Hessian to obtain frequencies in cm^-1

387

388 Results are cached so diagonalization will only be performed once for

389 this object instance.

390

391 Returns:

392 Harmonic mode frequencies in units of cm^-1

393

394 """

395

396 return self.get_energies() / units.invcm

397

398 def get_zero_point_energy(self) -> float:

399 """Diagonalise the Hessian and sum hw/2 to obtain zero-point energy

400

401 Args:

402 energies:

403 Pre-computed energy eigenvalues. Use if available to avoid

404 re-calculating these from the Hessian.

405

406 Returns:

407 zero-point energy in eV

408 """

409 return self._calculate_zero_point_energy(self.get_energies())

410

411 @staticmethod

412 def _calculate_zero_point_energy(energies: Union[Sequence[complex],

413 np.ndarray]) -> float:

414 return 0.5 * np.asarray(energies).real.sum()

415

416 def tabulate(self, im_tol: float = 1e-8) -> str:

417 """Get a summary of the vibrational frequencies.

418

419 Args:

420 im_tol:

421 Tolerance for imaginary frequency in eV. If frequency has a

422 larger imaginary component than im_tol, the imaginary component

423 is shown in the summary table.

424

425 Returns:

426 Summary table as formatted text

427 """

428

429 energies = self.get_energies()

430

431 return ('\n'.join(self._tabulate_from_energies(energies,

432 im_tol=im_tol))

433 + '\n')

434

435 @classmethod

436 def _tabulate_from_energies(cls,

437 energies: Union[Sequence[complex], np.ndarray],

438 im_tol: float = 1e-8) -> List[str]:

439 summary_lines = ['---------------------',

440 ' # meV cm^-1',

441 '---------------------']

442

443 for n, e in enumerate(energies):

444 if abs(e.imag) > im_tol:

445 c = 'i'

446 e = e.imag

447 else:

448 c = ''

449 e = e.real

450

451 summary_lines.append('{index:3d} {mev:6.1f}{im:1s} {cm:7.1f}{im}'

452 .format(index=n, mev=(e * 1e3),

453 cm=(e / units.invcm), im=c))

454 summary_lines.append('---------------------')

455 summary_lines.append('Zero-point energy: {:.3f} eV'.format(

456 cls._calculate_zero_point_energy(energies=energies)))

457

458 return summary_lines

459

460 def iter_animated_mode(self, mode_index: int,

461 temperature: float = units.kB * 300,

462 frames: int = 30) -> Iterator[Atoms]:

463 """Obtain animated mode as a series of Atoms

464

465 Args:

466 mode_index: Selection of mode to animate

467 temperature: In energy units - use units.kB * T_IN_KELVIN

468 frames: number of image frames in animation

469

470 Yields:

471 Displaced atoms following vibrational mode

472

473 """

474

475 mode = (self.get_modes(all_atoms=True)[mode_index]

476 * sqrt(temperature / abs(self.get_energies()[mode_index])))

477

478 for phase in np.linspace(0, 2 * pi, frames, endpoint=False):

479 atoms = self.get_atoms()

480 atoms.positions += sin(phase) * mode

481

482 yield atoms

483

484 def show_as_force(self,

485 mode: int,

486 scale: float = 0.2,

487 show: bool = True) -> Atoms:

488 """Illustrate mode as "forces" on atoms

489

490 Args:

491 mode: mode index

492 scale: scale factor

493 show: if True, open the ASE GUI and show atoms

494

495 Returns:

496 Atoms with scaled forces corresponding to mode eigenvectors (using

497 attached SinglePointCalculator).

498

499 """

500

501 atoms = self.get_atoms()

502 mode = self.get_modes(all_atoms=True)[mode] * len(atoms) * 3 * scale

503 atoms.calc = SinglePointCalculator(atoms, forces=mode)

504 if show:

505 atoms.edit()

506

507 return atoms

508

509 def write_jmol(self,

510 filename: str = 'vib.xyz',

511 ir_intensities: Union[Sequence[float], np.ndarray] = None

512 ) -> None:

513 """Writes file for viewing of the modes with jmol.

514

515 This is an extended XYZ file with eigenvectors given as extra columns

516 and metadata given in the label/comment line for each image. The format

517 is not quite human-friendly, but has the advantage that it can be

518 imported back into ASE with ase.io.read.

519

520 Args:

521 filename: Path for output file

522 ir_intensities: If available, IR intensities can be included in the

523 header lines. This does not affect the visualisation, but may

524 be convenient when comparing to experimental data.

525 """

526

527 all_images = list(self._get_jmol_images(atoms=self.get_atoms(),

528 energies=self.get_energies(),

529 modes=self.get_modes(

530 all_atoms=True),

531 ir_intensities=ir_intensities))

532 ase.io.write(filename, all_images, format='extxyz')

533

534 @staticmethod

535 def _get_jmol_images(atoms: Atoms,

536 energies: np.ndarray,

537 modes: np.ndarray,

538 ir_intensities:

539 Union[Sequence[float], np.ndarray] = None

540 ) -> Iterator[Atoms]:

541 """Get vibrational modes as a series of Atoms with attached data

542

543 For each image (Atoms object):

544

545 - eigenvalues are attached to image.arrays['mode']

546 - "mode#" and "frequency_cm-1" are set in image.info

547 - "IR_intensity" is set if provided in ir_intensities

548 - "masses" is removed

549

550 This is intended to set up the object for JMOL-compatible export using

551 ase.io.extxyz.

552

553

554 Args:

555 atoms: The base atoms object; all images have the same positions

556 energies: Complex vibrational energies in eV

557 modes: Eigenvectors array corresponding to atoms and energies. This

558 should cover the full set of atoms (i.e. modes =

559 vib.get_modes(all_atoms=True)).

560 ir_intensities: If available, IR intensities can be included in the

561 header lines. This does not affect the visualisation, but may

562 be convenient when comparing to experimental data.

563 Returns:

564 Iterator of Atoms objects

565

566 """

567 for i, (energy, mode) in enumerate(zip(energies, modes)):

568 # write imaginary frequencies as negative numbers

569 if energy.imag > energy.real:

570 energy = float(-energy.imag)

571 else:

572 energy = energy.real

573

574 image = atoms.copy()

575 image.info.update({'mode#': str(i),

576 'frequency_cm-1': energy / units.invcm,

577 })

578 image.arrays['mode'] = mode

579

580 # Custom masses are quite useful in vibration analysis, but will

581 # show up in the xyz file unless we remove them

582 if image.has('masses'):

583 del image.arrays['masses']

584

585 if ir_intensities is not None:

586 image.info['IR_intensity'] = float(ir_intensities[i])

587

588 yield image

589

590 def get_dos(self) -> RawDOSData:

591 """Total phonon DOS"""

592 energies = self.get_energies()

593 return RawDOSData(energies, np.ones_like(energies))

594

595 def get_pdos(self) -> DOSCollection:

596 """Phonon DOS, including atomic contributions"""

597 energies = self.get_energies()

598 masses = self._atoms[self.get_mask()].get_masses()

599

600 # Get weights as N_moving_atoms x N_modes array

601 vectors = self.get_modes() / masses[np.newaxis, :, np.newaxis]**-0.5

602 all_weights = (np.linalg.norm(vectors, axis=-1)**2).T

603

604 mask = self.get_mask()

605 all_info = [{'index': i, 'symbol': a.symbol}

606 for i, a in enumerate(self._atoms) if mask[i]]

607

608 return DOSCollection([RawDOSData(energies, weights, info=info)

609 for weights, info in zip(all_weights, all_info)])

610

611 def with_new_masses(self: VD, masses: Union[Sequence[float], np.ndarray]

612 ) -> VD:

613 """Get a copy of vibrations with modified masses and the same Hessian

614

615 Args:

616 masses:

617 New sequence of masses corresponding to the atom order in

618 self.get_atoms()

619 Returns:

620 A copy of the data with new masses for the same Hessian

621 """

622

623 new_atoms = self.get_atoms()

624 new_atoms.set_masses(masses)

625 return self.__class__(new_atoms, self.get_hessian(),

626 indices=self.get_indices())