Coverage for /builds/ase/ase/ase/vibrations/vibrations.py: 89.58%

1# fmt: off

3"""A class for computing vibrational modes"""

5import sys

6from collections import namedtuple

7from math import log, pi, sqrt

8from pathlib import Path

9from typing import Iterator, Tuple

11import numpy as np

13import ase.io

14import ase.units as units

15from ase.atoms import Atoms

16from ase.constraints import FixAtoms

17from ase.parallel import paropen, world

18from ase.utils.filecache import get_json_cache

20from .data import VibrationsData

23class AtomicDisplacements:

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

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

26 i = 'xyz'.index(i)

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

29 def _eq_disp(self):

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

32 @property

33 def ndof(self):

34 return 3 * len(self.indices)

37class Displacement(namedtuple('Displacement', ['a', 'i', 'sign', 'ndisp',

38 'vib'])):

39 @property

40 def name(self):

41 if self.sign == 0:

42 return 'eq'

44 axisname = 'xyz'[self.i]

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

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

48 @property

49 def _cached(self):

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

52 def forces(self):

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

55 @property

56 def step(self):

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

59 # XXX dipole only valid for infrared

60 def dipole(self):

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

63 # XXX below stuff only valid for TDDFT excitation stuff

64 def save_ov_nn(self, ov_nn):

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

67 def load_ov_nn(self):

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

70 @property

71 def _exname(self):

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

74 def calculate_and_save_static_polarizability(self, atoms):

75 exobj = self.vib._new_exobj()

76 excitation_data = exobj(atoms)

77 np.savetxt(self._exname, excitation_data)

79 def load_static_polarizability(self):

80 return np.loadtxt(self._exname)

82 def read_exobj(self):

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

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

86 def calculate_and_save_exlist(self, atoms):

87 # exo = self.vib._new_exobj()

88 excalc = self.vib._new_exobj()

89 exlist = excalc.calculate(atoms)

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

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

94class Vibrations(AtomicDisplacements):

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

97 The vibrational modes are calculated from a finite difference

98 approximation of the Hessian matrix.

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

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

102 described in:

103

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

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

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

107

108 atoms: Atoms object

109 The atoms to work on.

110 indices: list of int

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

112 to vibrate all atoms.

113 name: str

114 Name to use for files.

115 delta: float

116 Magnitude of displacements.

117 nfree: int

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

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

120 -delta for each cartesian coordinate.

121

122 Example:

123

124 >>> from ase import Atoms

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

126 >>> from ase.optimize import BFGS

127 >>> from ase.vibrations import Vibrations

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

129 ... calculator=EMT())

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

131 BFGS: 0 16:01:21 0.440339 3.2518

132 BFGS: 1 16:01:21 0.271928 0.8211

133 BFGS: 2 16:01:21 0.263278 0.1994

134 BFGS: 3 16:01:21 0.262777 0.0088

135 >>> vib = Vibrations(n2)

136 >>> vib.run()

137 >>> vib.summary()

138 ---------------------

139 # meV cm^-1

140 ---------------------

141 0 0.0 0.0

142 1 0.0 0.0

143 2 0.0 0.0

144 3 1.4 11.5

145 4 1.4 11.5

146 5 152.7 1231.3

147 ---------------------

148 Zero-point energy: 0.078 eV

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

150

151 """

152

153 def __init__(self, atoms, indices=None, name='vib', delta=0.01, nfree=2):

154 assert nfree in [2, 4]

155 self.atoms = atoms

156 self.calc = atoms.calc

157 if indices is None:

158 fixed_indices = []

159 for constr in atoms.constraints:

160 if isinstance(constr, FixAtoms):

161 fixed_indices.extend(constr.get_indices())

162 fixed_indices = list(set(fixed_indices))

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

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

165 raise ValueError(

166 'one (or more) indices included more than once')

167 self.indices = np.asarray(indices)

168

169 self.delta = delta

170 self.nfree = nfree

171 self.H = None

172 self.ir = None

173 self._vibrations = None

174

175 self.cache = get_json_cache(name)

176

177 @property

178 def name(self):

179 return str(self.cache.directory)

180

181 def run(self):

182 """Run the vibration calculations.

183

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

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

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

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

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

189 displacement.

190

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

192 simultaneously by several independent processes. This feature relies

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

194 case it is not found.

195

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

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

198 forces on your own.

199 """

200

201 if not self.cache.writable:

202 raise RuntimeError(

203 'Cannot run calculation. '

204 'Cache must be removed or split in order '

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

206

207 self._check_old_pickles()

208

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

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

211 if handle is None:

212 continue

213

214 result = self.calculate(atoms, disp)

215

216 if world.rank == 0:

217 handle.save(result)

218

219 def _check_old_pickles(self):

220 from pathlib import Path

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

222 pickle2json_instructions = f"""\

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

224Please remove them and recalculate or run \

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

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

227 raise RuntimeError(pickle2json_instructions)

228

229 def iterdisplace(self, inplace=False) -> \

230 Iterator[Tuple[Displacement, Atoms]]:

231 """Iterate over initial and displaced structures.

232

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

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

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

236

237 Parameters:

238 ------------

239 inplace: bool

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

241 copy will be made.

242

243 Yields:

244 --------

245 disp: Displacement

246 Displacement object with information about the displacement.

247 atoms: Atoms

248 Atoms object with the displaced structure.

249 """

250 # XXX change of type of disp

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

252 displacements = self.displacements()

253 eq_disp = next(displacements)

254 assert eq_disp.name == 'eq'

255 yield eq_disp, atoms

256

257 for disp in displacements:

258 if not inplace:

259 atoms = self.atoms.copy()

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

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

262 yield disp, atoms

263

264 if inplace:

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

266

267 def iterimages(self):

268 """Yield initial and displaced structures."""

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

270 yield atoms

271

272 def _iter_ai(self):

273 for a in self.indices:

274 for i in range(3):

275 yield a, i

276

277 def displacements(self):

278 yield self._eq_disp()

279

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

281 for sign in [-1, 1]:

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

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

284

285 def calculate(self, atoms, disp):

286 results = {}

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

288

289 if self.ir:

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

291

292 return results

293

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

295 """Remove json-files.

296

297 Use empty_files=True to remove only empty files and

298 combined=False to not remove the combined file.

299

300 """

301

302 if world.rank != 0:

303 return 0

304

305 if empty_files:

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

307

308 nfiles = self.cache.filecount()

309 self.cache.clear()

310 return nfiles

311

312 def combine(self):

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

314

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

316 of data structure at a time.

317

318 """

319 nelements_before = self.cache.filecount()

320 self.cache = self.cache.combine()

321 return nelements_before

322

323 def split(self):

324 """Split combined json-file.

325

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

327 sort of data structure at a time.

328

329 """

330 count = self.cache.filecount()

331 self.cache = self.cache.split()

332 return count

333

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

335 self.method = method.lower()

336 self.direction = direction.lower()

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

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

339

340 n = 3 * len(self.indices)

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

342 r = 0

343

344 eq_disp = self._eq_disp()

345

346 if direction != 'central':

347 feq = eq_disp.forces()

348

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

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

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

352

353 fminus = disp_minus.forces()

354 fplus = disp_plus.forces()

355 if self.method == 'frederiksen':

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

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

358 if self.nfree == 4:

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

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

361 if self.method == 'frederiksen':

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

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

364 if self.direction == 'central':

365 if self.nfree == 2:

366 H[r] = .5 * (fminus - fplus)[self.indices].ravel()

367 else:

368 assert self.nfree == 4

369 H[r] = H[r] = (-fminusminus +

370 8 * fminus -

371 8 * fplus +

372 fplusplus)[self.indices].ravel() / 12.0

373 elif self.direction == 'forward':

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

375 else:

376 assert self.direction == 'backward'

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

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

379 r += 1

380 H += H.copy().T

381 self.H = H

382 masses = self.atoms.get_masses()

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

384 raise RuntimeError('Zero mass encountered in one or more of '

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

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

387

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

389 self._vibrations = self.get_vibrations(read_cache=False,

390 method=self.method,

391 direction=self.direction)

392

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

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

395

396 # Conversion factor:

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

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

399

400 def get_vibrations(self, method='standard', direction='central',

401 read_cache=True, **kw):

402 """Get vibrations as VibrationsData object

403

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

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

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

407

408 Args:

409 method (str): Calculation method passed to read()

410 direction (str): Finite-difference scheme passed to read()

411 read_cache (bool): The VibrationsData object will be cached for

412 quick access. Set False to force regeneration of the cache with

413 the current atoms/Hessian/indices data.

414 **kw: Any remaining keyword arguments are passed to read()

415

416 Returns:

417 VibrationsData

418

419 """

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

421 return self._vibrations

422

423 else:

424 if (self.H is None or method.lower() != self.method or

425 direction.lower() != self.direction):

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

427

428 return VibrationsData.from_2d(self.atoms, self.H,

429 indices=self.indices)

430

431 def get_energies(self, method='standard', direction='central', **kw):

432 """Get vibration energies in eV."""

433 return self.get_vibrations(method=method,

434 direction=direction, **kw).get_energies()

435

436 def get_frequencies(self, method='standard', direction='central'):

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

438 return self.get_vibrations(method=method,

439 direction=direction).get_frequencies()

440

441 def summary(self, method='standard', direction='central', freq=None,

442 log=sys.stdout):

443 if freq is not None:

444 energies = freq * units.invcm

445 else:

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

447

448 summary_lines = VibrationsData._tabulate_from_energies(energies)

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

450

451 if isinstance(log, str):

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

453 log_file.write(log_text)

454 else:

455 log.write(log_text)

456

457 def get_zero_point_energy(self, freq=None):

458 if freq:

459 raise NotImplementedError

460 return self.get_vibrations().get_zero_point_energy()

461

462 def get_mode(self, n):

463 """Get mode number ."""

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

465

466 def write_mode(self, n=None, kT=units.kB * 300, nimages=30):

467 """Write mode number n to trajectory file. If n is not specified,

468 writes all non-zero modes."""

469 if n is None:

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

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

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

473 return

474

475 else:

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

477

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

479 for image in (self.get_vibrations()

480 .iter_animated_mode(n,

481 temperature=kT, frames=nimages)):

482 traj.write(image)

483

484 def show_as_force(self, n, scale=0.2, show=True):

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

486

487 def write_jmol(self):

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

489

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

491 self._write_jmol(fd)

492

493 def _write_jmol(self, fd):

494 symbols = self.atoms.get_chemical_symbols()

495 freq = self.get_frequencies()

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

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

498

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

500 c = 'i'

501 freq[n] = freq[n].imag

502

503 else:

504 freq[n] = freq[n].real

505 c = ' '

506

507 fd.write('Mode #%d, f = %.1f%s cm^-1'

508 % (n, float(freq[n].real), c))

509

510 if self.ir:

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

512 else:

513 fd.write('.\n')

514

515 mode = self.get_mode(n)

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

517 fd.write('%2s %12.5f %12.5f %12.5f %12.5f %12.5f %12.5f\n' %

518 (symbols[i], pos[0], pos[1], pos[2],

519 mode[i, 0], mode[i, 1], mode[i, 2]))

520

521 def fold(self, frequencies, intensities,

522 start=800.0, end=4000.0, npts=None, width=4.0,

523 type='Gaussian', normalize=False):

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

525 and folding method (Gaussian/Lorentzian).

526 The energy unit is cm^-1.

527 normalize=True ensures the integral over the peaks to give the

528 intensity.

529 """

530

531 lctype = type.lower()

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

533 if not npts:

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

535 prefactor = 1

536 if lctype == 'lorentzian':

537 intensities = intensities * width * pi / 2.

538 if normalize:

539 prefactor = 2. / width / pi

540 else:

541 sigma = width / 2. / sqrt(2. * log(2.))

542 if normalize:

543 prefactor = 1. / sigma / sqrt(2 * pi)

544

545 # Make array with spectrum data

546 spectrum = np.empty(npts)

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

548 for i, energy in enumerate(energies):

549 energies[i] = energy

550 if lctype == 'lorentzian':

551 spectrum[i] = (intensities * 0.5 * width / pi /

552 ((frequencies - energy)**2 +

553 0.25 * width**2)).sum()

554 else:

555 spectrum[i] = (intensities *

556 np.exp(-(frequencies - energy)**2 /

557 2. / sigma**2)).sum()

558 return [energies, prefactor * spectrum]

559

560 def write_dos(self, out='vib-dos.dat', start=800, end=4000,

561 npts=None, width=10,

562 type='Gaussian', method='standard', direction='central'):

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

564

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

566 folded vibrational density of states.

567 Start and end points, and width of the Gaussian/Lorentzian

568 should be given in cm^-1."""

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

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

571 energies, spectrum = self.fold(frequencies, intensities,

572 start, end, npts, width, type)

573

574 # Write out spectrum in file.

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

576 outdata.T[0] = energies

577 outdata.T[1] = spectrum

578

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

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

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

582 for row in outdata:

583 fd.write('%.3f %15.5e\n' %

584 (row[0], row[1]))