Coverage for /builds/hweiske/ase/ase/vibrations/vibrations.py: 89.01%

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, Tuple

9import numpy as np

11import ase.io

12import ase.units as units

13from ase.atoms import Atoms

14from ase.parallel import paropen, world

15from ase.utils.filecache import get_json_cache

17from .data import VibrationsData

20class AtomicDisplacements:

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

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

23 i = 'xyz'.index(i)

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

26 def _eq_disp(self):

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

29 @property

30 def ndof(self):

31 return 3 * len(self.indices)

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

35 'vib'])):

36 @property

37 def name(self):

38 if self.sign == 0:

39 return 'eq'

41 axisname = 'xyz'[self.i]

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

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

45 @property

46 def _cached(self):

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

49 def forces(self):

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

52 @property

53 def step(self):

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

56 # XXX dipole only valid for infrared

57 def dipole(self):

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

60 # XXX below stuff only valid for TDDFT excitation stuff

61 def save_ov_nn(self, ov_nn):

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

64 def load_ov_nn(self):

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

67 @property

68 def _exname(self):

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

71 def calculate_and_save_static_polarizability(self, atoms):

72 exobj = self.vib._new_exobj()

73 excitation_data = exobj(atoms)

74 np.savetxt(self._exname, excitation_data)

76 def load_static_polarizability(self):

77 return np.loadtxt(self._exname)

79 def read_exobj(self):

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

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

83 def calculate_and_save_exlist(self, atoms):

84 # exo = self.vib._new_exobj()

85 excalc = self.vib._new_exobj()

86 exlist = excalc.calculate(atoms)

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

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

91class Vibrations(AtomicDisplacements):

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

94 The vibrational modes are calculated from a finite difference

95 approximation of the Hessian matrix.

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

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

99 described in:

100

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

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

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

104

105 atoms: Atoms object

106 The atoms to work on.

107 indices: list of int

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

109 to vibrate all atoms.

110 name: str

111 Name to use for files.

112 delta: float

113 Magnitude of displacements.

114 nfree: int

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

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

117 -delta for each cartesian coordinate.

118

119 Example:

120

121 >>> from ase import Atoms

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

123 >>> from ase.optimize import BFGS

124 >>> from ase.vibrations import Vibrations

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

126 ... calculator=EMT())

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

128 BFGS: 0 16:01:21 0.440339 3.2518

129 BFGS: 1 16:01:21 0.271928 0.8211

130 BFGS: 2 16:01:21 0.263278 0.1994

131 BFGS: 3 16:01:21 0.262777 0.0088

132 >>> vib = Vibrations(n2)

133 >>> vib.run()

134 >>> vib.summary()

135 ---------------------

136 # meV cm^-1

137 ---------------------

138 0 0.0 0.0

139 1 0.0 0.0

140 2 0.0 0.0

141 3 1.4 11.5

142 4 1.4 11.5

143 5 152.7 1231.3

144 ---------------------

145 Zero-point energy: 0.078 eV

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

147

148 """

149

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

151 assert nfree in [2, 4]

152 self.atoms = atoms

153 self.calc = atoms.calc

154 if indices is None:

155 indices = range(len(atoms))

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

157 raise ValueError(

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

159 self.indices = np.asarray(indices)

160

161 self.delta = delta

162 self.nfree = nfree

163 self.H = None

164 self.ir = None

165 self._vibrations = None

166

167 self.cache = get_json_cache(name)

168

169 @property

170 def name(self):

171 return str(self.cache.directory)

172

173 def run(self):

174 """Run the vibration calculations.

175

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

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

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

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

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

181 displacement.

182

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

184 simultaneously by several independent processes. This feature relies

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

186 case it is not found.

187

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

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

190 forces on your own.

191 """

192

193 if not self.cache.writable:

194 raise RuntimeError(

195 'Cannot run calculation. '

196 'Cache must be removed or split in order '

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

198

199 self._check_old_pickles()

200

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

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

203 if handle is None:

204 continue

205

206 result = self.calculate(atoms, disp)

207

208 if world.rank == 0:

209 handle.save(result)

210

211 def _check_old_pickles(self):

212 from pathlib import Path

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

214 pickle2json_instructions = f"""\

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

216Please remove them and recalculate or run \

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

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

219 raise RuntimeError(pickle2json_instructions)

220

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

222 Iterator[Tuple[Displacement, Atoms]]:

223 """Iterate over initial and displaced structures.

224

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

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

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

228

229 Parameters:

230 ------------

231 inplace: bool

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

233 copy will be made.

234

235 Yields:

236 --------

237 disp: Displacement

238 Displacement object with information about the displacement.

239 atoms: Atoms

240 Atoms object with the displaced structure.

241 """

242 # XXX change of type of disp

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

244 displacements = self.displacements()

245 eq_disp = next(displacements)

246 assert eq_disp.name == 'eq'

247 yield eq_disp, atoms

248

249 for disp in displacements:

250 if not inplace:

251 atoms = self.atoms.copy()

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

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

254 yield disp, atoms

255

256 if inplace:

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

258

259 def iterimages(self):

260 """Yield initial and displaced structures."""

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

262 yield atoms

263

264 def _iter_ai(self):

265 for a in self.indices:

266 for i in range(3):

267 yield a, i

268

269 def displacements(self):

270 yield self._eq_disp()

271

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

273 for sign in [-1, 1]:

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

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

276

277 def calculate(self, atoms, disp):

278 results = {}

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

280

281 if self.ir:

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

283

284 return results

285

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

287 """Remove json-files.

288

289 Use empty_files=True to remove only empty files and

290 combined=False to not remove the combined file.

291

292 """

293

294 if world.rank != 0:

295 return 0

296

297 if empty_files:

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

299

300 nfiles = self.cache.filecount()

301 self.cache.clear()

302 return nfiles

303

304 def combine(self):

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

306

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

308 of data structure at a time.

309

310 """

311 nelements_before = self.cache.filecount()

312 self.cache = self.cache.combine()

313 return nelements_before

314

315 def split(self):

316 """Split combined json-file.

317

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

319 sort of data structure at a time.

320

321 """

322 count = self.cache.filecount()

323 self.cache = self.cache.split()

324 return count

325

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

327 self.method = method.lower()

328 self.direction = direction.lower()

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

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

331

332 n = 3 * len(self.indices)

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

334 r = 0

335

336 eq_disp = self._eq_disp()

337

338 if direction != 'central':

339 feq = eq_disp.forces()

340

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

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

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

344

345 fminus = disp_minus.forces()

346 fplus = disp_plus.forces()

347 if self.method == 'frederiksen':

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

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

350 if self.nfree == 4:

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

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

353 if self.method == 'frederiksen':

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

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

356 if self.direction == 'central':

357 if self.nfree == 2:

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

359 else:

360 assert self.nfree == 4

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

362 8 * fminus -

363 8 * fplus +

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

365 elif self.direction == 'forward':

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

367 else:

368 assert self.direction == 'backward'

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

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

371 r += 1

372 H += H.copy().T

373 self.H = H

374 masses = self.atoms.get_masses()

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

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

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

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

379

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

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

382 method=self.method,

383 direction=self.direction)

384

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

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

387

388 # Conversion factor:

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

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

391

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

393 read_cache=True, **kw):

394 """Get vibrations as VibrationsData object

395

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

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

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

399

400 Args:

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

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

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

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

405 the current atoms/Hessian/indices data.

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

407

408 Returns:

409 VibrationsData

410

411 """

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

413 return self._vibrations

414

415 else:

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

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

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

419

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

421 indices=self.indices)

422

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

424 """Get vibration energies in eV."""

425 return self.get_vibrations(method=method,

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

427

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

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

430 return self.get_vibrations(method=method,

431 direction=direction).get_frequencies()

432

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

434 log=sys.stdout):

435 if freq is not None:

436 energies = freq * units.invcm

437 else:

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

439

440 summary_lines = VibrationsData._tabulate_from_energies(energies)

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

442

443 if isinstance(log, str):

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

445 log_file.write(log_text)

446 else:

447 log.write(log_text)

448

449 def get_zero_point_energy(self, freq=None):

450 if freq:

451 raise NotImplementedError

452 return self.get_vibrations().get_zero_point_energy()

453

454 def get_mode(self, n):

455 """Get mode number ."""

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

457

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

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

460 writes all non-zero modes."""

461 if n is None:

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

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

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

465 return

466

467 else:

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

469

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

471 for image in (self.get_vibrations()

472 .iter_animated_mode(n,

473 temperature=kT, frames=nimages)):

474 traj.write(image)

475

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

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

478

479 def write_jmol(self):

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

481

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

483 self._write_jmol(fd)

484

485 def _write_jmol(self, fd):

486 symbols = self.atoms.get_chemical_symbols()

487 freq = self.get_frequencies()

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

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

490

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

492 c = 'i'

493 freq[n] = freq[n].imag

494

495 else:

496 freq[n] = freq[n].real

497 c = ' '

498

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

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

501

502 if self.ir:

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

504 else:

505 fd.write('.\n')

506

507 mode = self.get_mode(n)

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

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

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

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

512

513 def fold(self, frequencies, intensities,

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

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

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

517 and folding method (Gaussian/Lorentzian).

518 The energy unit is cm^-1.

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

520 intensity.

521 """

522

523 lctype = type.lower()

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

525 if not npts:

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

527 prefactor = 1

528 if lctype == 'lorentzian':

529 intensities = intensities * width * pi / 2.

530 if normalize:

531 prefactor = 2. / width / pi

532 else:

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

534 if normalize:

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

536

537 # Make array with spectrum data

538 spectrum = np.empty(npts)

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

540 for i, energy in enumerate(energies):

541 energies[i] = energy

542 if lctype == 'lorentzian':

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

544 ((frequencies - energy)**2 +

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

546 else:

547 spectrum[i] = (intensities *

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

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

550 return [energies, prefactor * spectrum]

551

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

553 npts=None, width=10,

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

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

556

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

558 folded vibrational density of states.

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

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

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

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

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

564 start, end, npts, width, type)

565

566 # Write out spectrum in file.

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

568 outdata.T[0] = energies

569 outdata.T[1] = spectrum

570

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

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

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

574 for row in outdata:

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

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