Coverage for /builds/debichem-team/python-ase/ase/vibrations/vibrations.py: 89.58%

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.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(namedtuple('Displacement', ['a', 'i', 'sign', 'ndisp',

36 'vib'])):

37 @property

38 def name(self):

39 if self.sign == 0:

40 return 'eq'

42 axisname = 'xyz'[self.i]

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

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

46 @property

47 def _cached(self):

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

50 def forces(self):

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

53 @property

54 def step(self):

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

57 # XXX dipole only valid for infrared

58 def dipole(self):

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

61 # XXX below stuff only valid for TDDFT excitation stuff

62 def save_ov_nn(self, ov_nn):

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

65 def load_ov_nn(self):

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

68 @property

69 def _exname(self):

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

72 def calculate_and_save_static_polarizability(self, atoms):

73 exobj = self.vib._new_exobj()

74 excitation_data = exobj(atoms)

75 np.savetxt(self._exname, excitation_data)

77 def load_static_polarizability(self):

78 return np.loadtxt(self._exname)

80 def read_exobj(self):

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

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

84 def calculate_and_save_exlist(self, atoms):

85 # exo = self.vib._new_exobj()

86 excalc = self.vib._new_exobj()

87 exlist = excalc.calculate(atoms)

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

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

92class Vibrations(AtomicDisplacements):

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

95 The vibrational modes are calculated from a finite difference

96 approximation of the Hessian matrix.

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

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

100 described in:

101

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

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

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

105

106 atoms: Atoms object

107 The atoms to work on.

108 indices: list of int

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

110 to vibrate all atoms.

111 name: str

112 Name to use for files.

113 delta: float

114 Magnitude of displacements.

115 nfree: int

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

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

118 -delta for each cartesian coordinate.

119

120 Example:

121

122 >>> from ase import Atoms

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

124 >>> from ase.optimize import BFGS

125 >>> from ase.vibrations import Vibrations

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

127 ... calculator=EMT())

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

129 BFGS: 0 16:01:21 0.440339 3.2518

130 BFGS: 1 16:01:21 0.271928 0.8211

131 BFGS: 2 16:01:21 0.263278 0.1994

132 BFGS: 3 16:01:21 0.262777 0.0088

133 >>> vib = Vibrations(n2)

134 >>> vib.run()

135 >>> vib.summary()

136 ---------------------

137 # meV cm^-1

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

139 0 0.0 0.0

140 1 0.0 0.0

141 2 0.0 0.0

142 3 1.4 11.5

143 4 1.4 11.5

144 5 152.7 1231.3

145 ---------------------

146 Zero-point energy: 0.078 eV

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

148

149 """

150

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

152 assert nfree in [2, 4]

153 self.atoms = atoms

154 self.calc = atoms.calc

155 if indices is None:

156 fixed_indices = []

157 for constr in atoms.constraints:

158 if isinstance(constr, FixAtoms):

159 fixed_indices.extend(constr.get_indices())

160 fixed_indices = list(set(fixed_indices))

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

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

163 raise ValueError(

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

165 self.indices = np.asarray(indices)

166

167 self.delta = delta

168 self.nfree = nfree

169 self.H = None

170 self.ir = None

171 self._vibrations = None

172

173 self.cache = get_json_cache(name)

174

175 @property

176 def name(self):

177 return str(self.cache.directory)

178

179 def run(self):

180 """Run the vibration calculations.

181

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

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

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

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

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

187 displacement.

188

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

190 simultaneously by several independent processes. This feature relies

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

192 case it is not found.

193

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

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

196 forces on your own.

197 """

198

199 if not self.cache.writable:

200 raise RuntimeError(

201 'Cannot run calculation. '

202 'Cache must be removed or split in order '

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

204

205 self._check_old_pickles()

206

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

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

209 if handle is None:

210 continue

211

212 result = self.calculate(atoms, disp)

213

214 if world.rank == 0:

215 handle.save(result)

216

217 def _check_old_pickles(self):

218 from pathlib import Path

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

220 pickle2json_instructions = f"""\

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

222Please remove them and recalculate or run \

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

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

225 raise RuntimeError(pickle2json_instructions)

226

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

228 Iterator[Tuple[Displacement, Atoms]]:

229 """Iterate over initial and displaced structures.

230

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

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

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

234

235 Parameters:

236 ------------

237 inplace: bool

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

239 copy will be made.

240

241 Yields:

242 --------

243 disp: Displacement

244 Displacement object with information about the displacement.

245 atoms: Atoms

246 Atoms object with the displaced structure.

247 """

248 # XXX change of type of disp

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

250 displacements = self.displacements()

251 eq_disp = next(displacements)

252 assert eq_disp.name == 'eq'

253 yield eq_disp, atoms

254

255 for disp in displacements:

256 if not inplace:

257 atoms = self.atoms.copy()

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

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

260 yield disp, atoms

261

262 if inplace:

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

264

265 def iterimages(self):

266 """Yield initial and displaced structures."""

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

268 yield atoms

269

270 def _iter_ai(self):

271 for a in self.indices:

272 for i in range(3):

273 yield a, i

274

275 def displacements(self):

276 yield self._eq_disp()

277

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

279 for sign in [-1, 1]:

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

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

282

283 def calculate(self, atoms, disp):

284 results = {}

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

286

287 if self.ir:

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

289

290 return results

291

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

293 """Remove json-files.

294

295 Use empty_files=True to remove only empty files and

296 combined=False to not remove the combined file.

297

298 """

299

300 if world.rank != 0:

301 return 0

302

303 if empty_files:

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

305

306 nfiles = self.cache.filecount()

307 self.cache.clear()

308 return nfiles

309

310 def combine(self):

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

312

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

314 of data structure at a time.

315

316 """

317 nelements_before = self.cache.filecount()

318 self.cache = self.cache.combine()

319 return nelements_before

320

321 def split(self):

322 """Split combined json-file.

323

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

325 sort of data structure at a time.

326

327 """

328 count = self.cache.filecount()

329 self.cache = self.cache.split()

330 return count

331

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

333 self.method = method.lower()

334 self.direction = direction.lower()

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

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

337

338 n = 3 * len(self.indices)

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

340 r = 0

341

342 eq_disp = self._eq_disp()

343

344 if direction != 'central':

345 feq = eq_disp.forces()

346

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

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

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

350

351 fminus = disp_minus.forces()

352 fplus = disp_plus.forces()

353 if self.method == 'frederiksen':

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

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

356 if self.nfree == 4:

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

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

359 if self.method == 'frederiksen':

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

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

362 if self.direction == 'central':

363 if self.nfree == 2:

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

365 else:

366 assert self.nfree == 4

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

368 8 * fminus -

369 8 * fplus +

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

371 elif self.direction == 'forward':

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

373 else:

374 assert self.direction == 'backward'

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

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

377 r += 1

378 H += H.copy().T

379 self.H = H

380 masses = self.atoms.get_masses()

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

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

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

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

385

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

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

388 method=self.method,

389 direction=self.direction)

390

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

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

393

394 # Conversion factor:

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

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

397

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

399 read_cache=True, **kw):

400 """Get vibrations as VibrationsData object

401

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

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

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

405

406 Args:

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

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

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

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

411 the current atoms/Hessian/indices data.

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

413

414 Returns:

415 VibrationsData

416

417 """

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

419 return self._vibrations

420

421 else:

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

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

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

425

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

427 indices=self.indices)

428

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

430 """Get vibration energies in eV."""

431 return self.get_vibrations(method=method,

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

433

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

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

436 return self.get_vibrations(method=method,

437 direction=direction).get_frequencies()

438

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

440 log=sys.stdout):

441 if freq is not None:

442 energies = freq * units.invcm

443 else:

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

445

446 summary_lines = VibrationsData._tabulate_from_energies(energies)

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

448

449 if isinstance(log, str):

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

451 log_file.write(log_text)

452 else:

453 log.write(log_text)

454

455 def get_zero_point_energy(self, freq=None):

456 if freq:

457 raise NotImplementedError

458 return self.get_vibrations().get_zero_point_energy()

459

460 def get_mode(self, n):

461 """Get mode number ."""

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

463

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

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

466 writes all non-zero modes."""

467 if n is None:

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

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

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

471 return

472

473 else:

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

475

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

477 for image in (self.get_vibrations()

478 .iter_animated_mode(n,

479 temperature=kT, frames=nimages)):

480 traj.write(image)

481

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

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

484

485 def write_jmol(self):

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

487

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

489 self._write_jmol(fd)

490

491 def _write_jmol(self, fd):

492 symbols = self.atoms.get_chemical_symbols()

493 freq = self.get_frequencies()

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

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

496

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

498 c = 'i'

499 freq[n] = freq[n].imag

500

501 else:

502 freq[n] = freq[n].real

503 c = ' '

504

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

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

507

508 if self.ir:

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

510 else:

511 fd.write('.\n')

512

513 mode = self.get_mode(n)

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

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

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

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

518

519 def fold(self, frequencies, intensities,

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

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

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

523 and folding method (Gaussian/Lorentzian).

524 The energy unit is cm^-1.

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

526 intensity.

527 """

528

529 lctype = type.lower()

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

531 if not npts:

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

533 prefactor = 1

534 if lctype == 'lorentzian':

535 intensities = intensities * width * pi / 2.

536 if normalize:

537 prefactor = 2. / width / pi

538 else:

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

540 if normalize:

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

542

543 # Make array with spectrum data

544 spectrum = np.empty(npts)

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

546 for i, energy in enumerate(energies):

547 energies[i] = energy

548 if lctype == 'lorentzian':

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

550 ((frequencies - energy)**2 +

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

552 else:

553 spectrum[i] = (intensities *

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

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

556 return [energies, prefactor * spectrum]

557

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

559 npts=None, width=10,

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

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

562

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

564 folded vibrational density of states.

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

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

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

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

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

570 start, end, npts, width, type)

571

572 # Write out spectrum in file.

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

574 outdata.T[0] = energies

575 outdata.T[1] = spectrum

576

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

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

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

580 for row in outdata:

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

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