Coverage for ase/vibrations/vibrations.py: 89.55%

1"""A class for computing vibrational modes"""

3import sys

4from collections import namedtuple

5from math import log, pi, sqrt

6from pathlib import Path

7from typing import Iterator

9import numpy as np

11import ase.io

12import ase.units as units

13from ase.atoms import Atoms

14from ase.constraints import FixAtoms

15from ase.parallel import paropen, world

16from ase.utils.filecache import get_json_cache

18from .data import VibrationsData

21class AtomicDisplacements:

22 def _disp(self, a, i, step):

23 if isinstance(i, str): # XXX Simplify by removing this.

24 i = 'xyz'.index(i)

25 return Displacement(a, i, np.sign(step), abs(step), self)

27 def _eq_disp(self):

28 return self._disp(0, 0, 0)

30 @property

31 def ndof(self):

32 return 3 * len(self.indices)

35class Displacement(

36 namedtuple('Displacement', ['a', 'i', 'sign', 'ndisp', 'vib'])

37):

38 @property

39 def name(self):

40 if self.sign == 0:

41 return 'eq'

43 axisname = 'xyz'[self.i]

44 dispname = self.ndisp * ' +-'[self.sign]

45 return f'{self.a}{axisname}{dispname}'

47 @property

48 def _cached(self):

49 return self.vib.cache[self.name]

51 def forces(self):

52 return self._cached['forces'].copy()

54 @property

55 def step(self):

56 return self.ndisp * self.sign * self.vib.delta

58 # XXX dipole only valid for infrared

59 def dipole(self):

60 return self._cached['dipole'].copy()

62 # XXX below stuff only valid for TDDFT excitation stuff

63 def save_ov_nn(self, ov_nn):

64 np.save(Path(self.vib.exname) / (self.name + '.ov'), ov_nn)

66 def load_ov_nn(self):

67 return np.load(Path(self.vib.exname) / (self.name + '.ov.npy'))

69 @property

70 def _exname(self):

71 return Path(self.vib.exname) / f'ex.{self.name}{self.vib.exext}'

73 def calculate_and_save_static_polarizability(self, atoms):

74 exobj = self.vib._new_exobj()

75 excitation_data = exobj(atoms)

76 np.savetxt(self._exname, excitation_data)

78 def load_static_polarizability(self):

79 return np.loadtxt(self._exname)

81 def read_exobj(self):

82 # XXX each exobj should allow for self._exname as Path

83 return self.vib.read_exobj(str(self._exname))

85 def calculate_and_save_exlist(self, atoms):

86 # exo = self.vib._new_exobj()

87 excalc = self.vib._new_exobj()

88 exlist = excalc.calculate(atoms)

89 # XXX each exobj should allow for self._exname as Path

90 exlist.write(str(self._exname))

93class Vibrations(AtomicDisplacements):

94 """Class for calculating vibrational modes using finite difference.

96 The vibrational modes are calculated from a finite difference

97 approximation of the Hessian matrix.

99 The *summary()*, *get_energies()* and *get_frequencies()* methods all take

100 an optional *method* keyword. Use method='Frederiksen' to use the method

101 described in:

102

103 T. Frederiksen, M. Paulsson, M. Brandbyge, A. P. Jauho:

104 "Inelastic transport theory from first-principles: methodology and

105 applications for nanoscale devices", Phys. Rev. B 75, 205413 (2007)

106

107 Parameters

108 ----------

109 atoms: :class:`~ase.Atoms`

110 The atoms to work on.

111 indices: list[int] | :obj:`None`, default: None

112 List of indices of atoms to vibrate. Default behavior is

113 to vibrate all atoms.

114 name: str

115 Name to use for files.

116 delta: float

117 Magnitude of displacements.

118 nfree: int

119 Number of displacements per atom and cartesian coordinate, 2 and 4 are

120 supported. Default is 2 which will displace each atom +delta and

121 -delta for each cartesian coordinate.

122

123 Examples

124 --------

125 >>> from ase import Atoms

126 >>> from ase.calculators.emt import EMT

127 >>> from ase.optimize import BFGS

128 >>> from ase.vibrations import Vibrations

129 >>> n2 = Atoms('N2', [(0, 0, 0), (0, 0, 1.1)],

130 ... calculator=EMT())

131 >>> BFGS(n2).run(fmax=0.01)

132 BFGS: 0 16:01:21 0.440339 3.2518

133 BFGS: 1 16:01:21 0.271928 0.8211

134 BFGS: 2 16:01:21 0.263278 0.1994

135 BFGS: 3 16:01:21 0.262777 0.0088

136 >>> vib = Vibrations(n2)

137 >>> vib.run()

138 >>> vib.summary()

139 ---------------------

140 # meV cm^-1

141 ---------------------

142 0 0.0 0.0

143 1 0.0 0.0

144 2 0.0 0.0

145 3 1.4 11.5

146 4 1.4 11.5

147 5 152.7 1231.3

148 ---------------------

149 Zero-point energy: 0.078 eV

150 >>> vib.write_mode(-1) # write last mode to trajectory file

151

152 """

153

154 def __init__(

155 self,

156 atoms: Atoms,

157 indices: list[int] | None = None,

158 name: str = 'vib',

159 delta: float = 0.01,

160 nfree: int = 2,

161 ):

162 assert nfree in [2, 4]

163 self.atoms = atoms

164 self.calc = atoms.calc

165 if indices is None:

166 fixed_indices = []

167 for constr in atoms.constraints:

168 if isinstance(constr, FixAtoms):

169 fixed_indices.extend(constr.get_indices())

170 fixed_indices = list(set(fixed_indices))

171 indices = [i for i in range(len(atoms)) if i not in fixed_indices]

172 if len(indices) != len(set(indices)):

173 raise ValueError('one (or more) indices included more than once')

174 self.indices = np.asarray(indices)

175

176 self.delta = delta

177 self.nfree = nfree

178 self.H = None

179 self.ir = None

180 self._vibrations = None

181

182 self.cache = get_json_cache(name)

183

184 @property

185 def name(self):

186 return str(self.cache.directory)

187

188 def run(self):

189 """Run the vibration calculations.

190

191 This will calculate the forces for 6 displacements per atom +/-x,

192 +/-y, +/-z. Only those calculations that are not already done will be

193 started. Be aware that an interrupted calculation may produce an empty

194 file (ending with .json), which must be deleted before restarting the

195 job. Otherwise the forces will not be calculated for that

196 displacement.

197

198 Note that the calculations for the different displacements can be done

199 simultaneously by several independent processes. This feature relies

200 on the existence of files and the subsequent creation of the file in

201 case it is not found.

202

203 If the program you want to use does not have a calculator in ASE, use

204 ``iterdisplace`` to get all displaced structures and calculate the

205 forces on your own.

206 """

207

208 if not self.cache.writable:

209 raise RuntimeError(

210 'Cannot run calculation. '

211 'Cache must be removed or split in order '

212 'to have only one sort of data structure at a time.'

213 )

214

215 self._check_old_pickles()

216

217 for disp, atoms in self.iterdisplace(inplace=True):

218 with self.cache.lock(disp.name) as handle:

219 if handle is None:

220 continue

221

222 result = self.calculate(atoms, disp)

223

224 if world.rank == 0:

225 handle.save(result)

226

227 def _check_old_pickles(self):

228 eq_pickle_path = Path(f'{self.name}.eq.pckl')

229 pickle2json_instructions = f"""\

230Found old pickle files such as {eq_pickle_path}. \

231Please remove them and recalculate or run \

232"python -m ase.vibrations.pickle2json --help"."""

233 if len(self.cache) == 0 and eq_pickle_path.exists():

234 raise RuntimeError(pickle2json_instructions)

235

236 def iterdisplace(

237 self,

238 inplace: bool = False,

239 ) -> Iterator[tuple[Displacement, Atoms]]:

240 """Iterate over initial and displaced structures.

241

242 Use this to export the structures for each single-point calculation

243 to an external program instead of using ``run()``. Then save the

244 calculated gradients to <name>.json and continue using this instance.

245

246 Parameters

247 ----------

248 inplace: bool, default: :obj:`False`

249 If True, the atoms object will be modified in-place. Otherwise a

250 copy will be made.

251

252 Yields

253 ------

254 disp: :class:`Displacement`

255 Displacement object with information about the displacement.

256 atoms: :class:`~ase.Atoms`

257 Atoms object with the displaced structure.

258

259 """

260 # XXX change of type of disp

261 atoms = self.atoms if inplace else self.atoms.copy()

262 displacements = self.displacements()

263 eq_disp = next(displacements)

264 assert eq_disp.name == 'eq'

265 yield eq_disp, atoms

266

267 for disp in displacements:

268 if not inplace:

269 atoms = self.atoms.copy()

270 pos0 = atoms.positions[disp.a, disp.i]

271 atoms.positions[disp.a, disp.i] += disp.step

272 yield disp, atoms

273

274 if inplace:

275 atoms.positions[disp.a, disp.i] = pos0

276

277 def iterimages(self):

278 """Yield initial and displaced structures."""

279 for name, atoms in self.iterdisplace():

280 yield atoms

281

282 def _iter_ai(self):

283 for a in self.indices:

284 for i in range(3):

285 yield a, i

286

287 def displacements(self):

288 yield self._eq_disp()

289

290 for a, i in self._iter_ai():

291 for sign in [-1, 1]:

292 for ndisp in range(1, self.nfree // 2 + 1):

293 yield self._disp(a, i, sign * ndisp)

294

295 def calculate(self, atoms, disp):

296 results = {}

297 results['forces'] = self.calc.get_forces(atoms)

298

299 if self.ir:

300 results['dipole'] = self.calc.get_dipole_moment(atoms)

301

302 return results

303

304 def clean(self, empty_files=False, combined=True):

305 """Remove json-files.

306

307 Use empty_files=True to remove only empty files and

308 combined=False to not remove the combined file.

309

310 """

311

312 if world.rank != 0:

313 return 0

314

315 if empty_files:

316 return self.cache.strip_empties() # XXX Fails on combined cache

317

318 nfiles = self.cache.filecount()

319 self.cache.clear()

320 return nfiles

321

322 def combine(self):

323 """Combine json-files to one file ending with '.all.json'.

324

325 The other json-files will be removed in order to have only one sort

326 of data structure at a time.

327

328 """

329 nelements_before = self.cache.filecount()

330 self.cache = self.cache.combine()

331 return nelements_before

332

333 def split(self):

334 """Split combined json-file.

335

336 The combined json-file will be removed in order to have only one

337 sort of data structure at a time.

338

339 """

340 count = self.cache.filecount()

341 self.cache = self.cache.split()

342 return count

343

344 def read(self, method='standard', direction='central'):

345 self.method = method.lower()

346 self.direction = direction.lower()

347 assert self.method in ['standard', 'frederiksen']

348 assert self.direction in ['central', 'forward', 'backward']

349

350 n = 3 * len(self.indices)

351 H = np.empty((n, n))

352 r = 0

353

354 eq_disp = self._eq_disp()

355

356 if direction != 'central':

357 feq = eq_disp.forces()

358

359 for a, i in self._iter_ai():

360 disp_minus = self._disp(a, i, -1)

361 disp_plus = self._disp(a, i, 1)

362

363 fminus = disp_minus.forces()

364 fplus = disp_plus.forces()

365 if self.method == 'frederiksen':

366 fminus[a] -= fminus.sum(0)

367 fplus[a] -= fplus.sum(0)

368 if self.nfree == 4:

369 fminusminus = self._disp(a, i, -2).forces()

370 fplusplus = self._disp(a, i, 2).forces()

371 if self.method == 'frederiksen':

372 fminusminus[a] -= fminusminus.sum(0)

373 fplusplus[a] -= fplusplus.sum(0)

374 if self.direction == 'central':

375 if self.nfree == 2:

376 H[r] = 0.5 * (fminus - fplus)[self.indices].ravel()

377 else:

378 assert self.nfree == 4

379 H[r] = H[r] = (

380 -fminusminus + 8 * fminus - 8 * fplus + fplusplus

381 )[self.indices].ravel() / 12.0

382 elif self.direction == 'forward':

383 H[r] = (feq - fplus)[self.indices].ravel()

384 else:

385 assert self.direction == 'backward'

386 H[r] = (fminus - feq)[self.indices].ravel()

387 H[r] /= 2 * self.delta

388 r += 1

389 H += H.copy().T

390 self.H = H

391 masses = self.atoms.get_masses()

392 if any(masses[self.indices] == 0):

393 raise RuntimeError(

394 'Zero mass encountered in one or more of '

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

396 ' to set all masses to non-zero values.'

397 )

398

399 self.im = np.repeat(masses[self.indices] ** -0.5, 3)

400 self._vibrations = self.get_vibrations(

401 read_cache=False, method=self.method, direction=self.direction

402 )

403

404 omega2, modes = np.linalg.eigh(self.im[:, None] * H * self.im)

405 self.modes = modes.T.copy()

406

407 # Conversion factor:

408 s = units._hbar * 1e10 / sqrt(units._e * units._amu)

409 self.hnu = s * omega2.astype(complex) ** 0.5

410

411 def get_vibrations(

412 self,

413 method='standard',

414 direction='central',

415 read_cache=True,

416 **kw,

417 ):

418 """Get vibrations as VibrationsData object

419

420 If read() has not yet been called, this will be called to assemble data

421 from the outputs of run(). Most of the arguments to this function are

422 options to be passed to read() in this case.

423

424 Parameters

425 ----------

426 method : str

427 Calculation method passed to read()

428 direction : str

429 Finite-difference scheme passed to read()

430 read_cache : bool

431 The :class:`VibrationsData` object will be cached for quick access.

432 Set :obj:`False` to force regeneration of the cache with the

433 current atoms/Hessian/indices data.

434 **kw : dict

435 Any remaining keyword arguments are passed to read()

436

437 Returns

438 -------

439 :class:`VibrationsData`

440

441 """

442 if read_cache and (self._vibrations is not None):

443 return self._vibrations

444

445 else:

446 if (

447 self.H is None

448 or method.lower() != self.method

449 or direction.lower() != self.direction

450 ):

451 self.read(method, direction, **kw)

452

453 return VibrationsData.from_2d(

454 self.atoms,

455 self.H,

456 indices=self.indices,

457 )

458

459 def get_energies(

460 self,

461 method: str = 'standard',

462 direction: str = 'central',

463 **kw,

464 ):

465 """Get vibration energies in eV."""

466 return self.get_vibrations(

467 method=method, direction=direction, **kw

468 ).get_energies()

469

470 def get_frequencies(

471 self,

472 method: str = 'standard',

473 direction: str = 'central',

474 ):

475 """Get vibration frequencies in cm^-1."""

476 return self.get_vibrations(

477 method=method, direction=direction

478 ).get_frequencies()

479

480 def summary(

481 self,

482 method: str = 'standard',

483 direction: str = 'central',

484 freq=None,

485 log=sys.stdout,

486 ):

487 if freq is not None:

488 energies = freq * units.invcm

489 else:

490 energies = self.get_energies(method=method, direction=direction)

491

492 summary_lines = VibrationsData._tabulate_from_energies(energies)

493 log_text = '\n'.join(summary_lines) + '\n'

494

495 if isinstance(log, str):

496 with paropen(log, 'a') as log_file:

497 log_file.write(log_text)

498 else:

499 log.write(log_text)

500

501 def get_zero_point_energy(self, freq=None):

502 if freq:

503 raise NotImplementedError

504 return self.get_vibrations().get_zero_point_energy()

505

506 def get_mode(self, n):

507 """Get mode number ."""

508 return self.get_vibrations().get_modes(all_atoms=True)[n]

509

510 def write_mode(

511 self,

512 n: int | None = None,

513 kT: float = units.kB * 300.0,

514 nimages: int = 30,

515 ) -> None:

516 """Write mode number n to trajectory file.

517

518 If n is not specified, writes all non-zero modes.

519 """

520 if n is None:

521 for index, energy in enumerate(self.get_energies()):

522 if abs(energy) > 1e-5:

523 self.write_mode(n=index, kT=kT, nimages=nimages)

524 return

525

526 else:

527 n %= len(self.get_energies())

528

529 with ase.io.Trajectory('%s.%d.traj' % (self.name, n), 'w') as traj:

530 for image in self.get_vibrations().iter_animated_mode(

531 n, temperature=kT, frames=nimages

532 ):

533 traj.write(image)

534

535 def show_as_force(self, n: int, scale: float = 0.2, show: bool = True):

536 return self.get_vibrations().show_as_force(n, scale=scale, show=show)

537

538 def write_jmol(self):

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

540

541 with open(self.name + '.xyz', 'w') as fd:

542 self._write_jmol(fd)

543

544 def _write_jmol(self, fd):

545 symbols = self.atoms.get_chemical_symbols()

546 freq = self.get_frequencies()

547 for n in range(3 * len(self.indices)):

548 fd.write('%6d\n' % len(self.atoms))

549

550 if freq[n].imag != 0:

551 c = 'i'

552 freq[n] = freq[n].imag

553

554 else:

555 freq[n] = freq[n].real

556 c = ' '

557

558 fd.write('Mode #%d, f = %.1f%s cm^-1' % (n, float(freq[n].real), c))

559

560 if self.ir:

561 fd.write(', I = %.4f (D/Å)^2 amu^-1.\n' % self.intensities[n])

562 else:

563 fd.write('.\n')

564

565 mode = self.get_mode(n)

566 for i, pos in enumerate(self.atoms.positions):

567 fd.write(

568 '%2s %12.5f %12.5f %12.5f %12.5f %12.5f %12.5f\n'

569 % (

570 symbols[i],

571 pos[0],

572 pos[1],

573 pos[2],

574 mode[i, 0],

575 mode[i, 1],

576 mode[i, 2],

577 )

578 )

579

580 def fold(

581 self,

582 frequencies,

583 intensities,

584 start: float = 800.0,

585 end: float = 4000.0,

586 npts: int | None = None,

587 width: float = 4.0,

588 type: str = 'Gaussian',

589 normalize: bool = False,

590 ):

591 """Fold frequencies and intensities within the given range

592 and folding method (Gaussian/Lorentzian).

593 The energy unit is cm^-1.

594

595 Paramters

596 ---------

597 normalize : bool, default: False

598 If :obj:`True`, the integral over the peaks gives the intensity.

599

600 """

601

602 lctype = type.lower()

603 assert lctype in ['gaussian', 'lorentzian']

604 if not npts:

605 npts = int((end - start) / width * 10 + 1)

606 prefactor = 1.0

607 if lctype == 'lorentzian':

608 intensities = intensities * width * pi / 2.0

609 if normalize:

610 prefactor = 2.0 / width / pi

611 else:

612 sigma = width / 2.0 / sqrt(2.0 * log(2.0))

613 if normalize:

614 prefactor = 1.0 / sigma / sqrt(2 * pi)

615

616 # Make array with spectrum data

617 spectrum = np.empty(npts)

618 energies = np.linspace(start, end, npts)

619 for i, energy in enumerate(energies):

620 energies[i] = energy

621 if lctype == 'lorentzian':

622 spectrum[i] = (

623 intensities

624 * 0.5

625 * width

626 / pi

627 / ((frequencies - energy) ** 2 + 0.25 * width**2)

628 ).sum()

629 else:

630 spectrum[i] = (

631 intensities

632 * np.exp(-((frequencies - energy) ** 2) / 2.0 / sigma**2)

633 ).sum()

634 return [energies, prefactor * spectrum]

635

636 def write_dos(

637 self,

638 out: str = 'vib-dos.dat',

639 start: float = 800.0,

640 end: float = 4000.0,

641 npts: int | None = None,

642 width: float = 10.0,

643 type: str = 'Gaussian',

644 method: str = 'standard',

645 direction: str = 'central',

646 ) -> None:

647 """Write out the vibrational density of states to file.

648

649 First column is the wavenumber in cm^-1, the second column the

650 folded vibrational density of states.

651

652 Parameters

653 ----------

654 start : float

655 Start points in cm^-1.

656 end : float

657 End points in cm^-1.

658 width : float

659 Width of the Gaussian/Lorentzian in cm^-1.

660

661 """

662 frequencies = self.get_frequencies(method, direction).real

663 intensities = np.ones(len(frequencies))

664 energies, spectrum = self.fold(

665 frequencies, intensities, start, end, npts, width, type

666 )

667

668 # Write out spectrum in file.

669 outdata = np.empty([len(energies), 2])

670 outdata.T[0] = energies

671 outdata.T[1] = spectrum

672

673 with open(out, 'w') as fd:

674 fd.write(f'# {type.title()} folded, width={width:g} cm^-1\n')

675 fd.write('# [cm^-1] arbitrary\n')

676 for row in outdata:

677 fd.write('%.3f %15.5e\n' % (row[0], row[1]))

Coverage for ase / vibrations / vibrations.py: 89.55%

287 statements