Coverage for /builds/ase/ase/ase/io/vasp.py: 91.13%

496 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-08-02 00:12 +0000

1# fmt: off 

2 

3""" 

4This module contains functionality for reading and writing an ASE 

5Atoms object in VASP POSCAR format. 

6 

7""" 

8from __future__ import annotations 

9 

10import re 

11from pathlib import Path 

12from typing import List, Optional, TextIO, Tuple 

13 

14import numpy as np 

15 

16from ase import Atoms 

17from ase.constraints import FixAtoms, FixedLine, FixedPlane, FixScaled 

18from ase.io import ParseError 

19from ase.io.formats import string2index 

20from ase.io.utils import ImageIterator 

21from ase.symbols import Symbols 

22from ase.units import Ang, fs 

23from ase.utils import reader, writer 

24 

25from .vasp_parsers import vasp_outcar_parsers as vop 

26 

27__all__ = [ 

28 'read_vasp', 'read_vasp_out', 'iread_vasp_out', 'read_vasp_xdatcar', 

29 'read_vasp_xml', 'write_vasp', 'write_vasp_xdatcar' 

30] 

31 

32 

33def parse_poscar_scaling_factor(line: str) -> np.ndarray: 

34 """Parse scaling factor(s) in the second line in POSCAR/CONTCAR. 

35 

36 This can also be one negative number or three positive numbers. 

37 

38 https://www.vasp.at/wiki/index.php/POSCAR#Full_format_specification 

39 

40 """ 

41 scale = [] 

42 for _ in line.split()[:3]: 

43 try: 

44 scale.append(float(_)) 

45 except ValueError: 

46 break 

47 if len(scale) not in {1, 3}: 

48 raise RuntimeError('The number of scaling factors must be 1 or 3.') 

49 if len(scale) == 3 and any(_ < 0.0 for _ in scale): 

50 raise RuntimeError('All three scaling factors must be positive.') 

51 return np.array(scale) 

52 

53 

54def get_atomtypes(fname): 

55 """Given a file name, get the atomic symbols. 

56 

57 The function can get this information from OUTCAR and POTCAR 

58 format files. The files can also be compressed with gzip or 

59 bzip2. 

60 

61 """ 

62 fpath = Path(fname) 

63 

64 atomtypes = [] 

65 atomtypes_alt = [] 

66 if fpath.suffix == '.gz': 

67 import gzip 

68 opener = gzip.open 

69 elif fpath.suffix == '.bz2': 

70 import bz2 

71 opener = bz2.BZ2File 

72 else: 

73 opener = open 

74 with opener(fpath) as fd: 

75 for line in fd: 

76 if 'TITEL' in line: 

77 atomtypes.append(line.split()[3].split('_')[0].split('.')[0]) 

78 elif 'POTCAR:' in line: 

79 atomtypes_alt.append( 

80 line.split()[2].split('_')[0].split('.')[0]) 

81 

82 if len(atomtypes) == 0 and len(atomtypes_alt) > 0: 

83 # old VASP doesn't echo TITEL, but all versions print out species 

84 # lines preceded by "POTCAR:", twice 

85 if len(atomtypes_alt) % 2 != 0: 

86 raise ParseError( 

87 f'Tried to get atom types from {len(atomtypes_alt)}' 

88 '"POTCAR": lines in OUTCAR, but expected an even number' 

89 ) 

90 atomtypes = atomtypes_alt[0:len(atomtypes_alt) // 2] 

91 

92 return atomtypes 

93 

94 

95def atomtypes_outpot(posfname, numsyms): 

96 """Try to retrieve chemical symbols from OUTCAR or POTCAR 

97 

98 If getting atomtypes from the first line in POSCAR/CONTCAR fails, it might 

99 be possible to find the data in OUTCAR or POTCAR, if these files exist. 

100 

101 posfname -- The filename of the POSCAR/CONTCAR file we're trying to read 

102 

103 numsyms -- The number of symbols we must find 

104 

105 """ 

106 posfpath = Path(posfname) 

107 

108 # Check files with exactly same path except POTCAR/OUTCAR instead 

109 # of POSCAR/CONTCAR. 

110 fnames = [posfpath.with_name('POTCAR'), 

111 posfpath.with_name('OUTCAR')] 

112 # Try the same but with compressed files 

113 fsc = [] 

114 for fnpath in fnames: 

115 fsc.append(fnpath.parent / (fnpath.name + '.gz')) 

116 fsc.append(fnpath.parent / (fnpath.name + '.bz2')) 

117 for f in fsc: 

118 fnames.append(f) 

119 # Code used to try anything with POTCAR or OUTCAR in the name 

120 # but this is no longer supported 

121 

122 tried = [] 

123 for fn in fnames: 

124 if fn in posfpath.parent.iterdir(): 

125 tried.append(fn) 

126 at = get_atomtypes(fn) 

127 if len(at) == numsyms: 

128 return at 

129 

130 raise ParseError('Could not determine chemical symbols. Tried files ' + 

131 str(tried)) 

132 

133 

134def get_atomtypes_from_formula(formula): 

135 """Return atom types from chemical formula (optionally prepended 

136 with and underscore). 

137 """ 

138 from ase.symbols import string2symbols 

139 symbols = string2symbols(formula.split('_')[0]) 

140 atomtypes = [symbols[0]] 

141 for s in symbols[1:]: 

142 if s != atomtypes[-1]: 

143 atomtypes.append(s) 

144 return atomtypes 

145 

146 

147@reader 

148def read_vasp(fd): 

149 """Import POSCAR/CONTCAR type file. 

150 

151 Reads unitcell, atom positions and constraints from the POSCAR/CONTCAR 

152 file and tries to read atom types from POSCAR/CONTCAR header, if this 

153 fails the atom types are read from OUTCAR or POTCAR file. 

154 """ 

155 atoms = read_vasp_configuration(fd) 

156 velocities = read_velocities_if_present(fd, len(atoms)) 

157 if velocities is not None: 

158 atoms.set_velocities(velocities) 

159 return atoms 

160 

161 

162def read_vasp_configuration(fd): 

163 """Read common POSCAR/CONTCAR/CHGCAR/CHG quantities and return Atoms.""" 

164 from ase.data import chemical_symbols 

165 

166 # The first line is in principle a comment line, however in VASP 

167 # 4.x a common convention is to have it contain the atom symbols, 

168 # eg. "Ag Ge" in the same order as later in the file (and POTCAR 

169 # for the full vasp run). In the VASP 5.x format this information 

170 # is found on the fifth line. Thus we save the first line and use 

171 # it in case we later detect that we're reading a VASP 4.x format 

172 # file. 

173 line1 = fd.readline() 

174 

175 scale = parse_poscar_scaling_factor(fd.readline()) 

176 

177 # Now the lattice vectors 

178 cell = np.array([fd.readline().split()[:3] for _ in range(3)], dtype=float) 

179 # Negative scaling factor corresponds to the cell volume. 

180 if scale[0] < 0.0: 

181 scale = np.cbrt(-1.0 * scale / np.linalg.det(cell)) 

182 cell *= scale # This works for both one and three scaling factors. 

183 

184 # Number of atoms. Again this must be in the same order as 

185 # in the first line 

186 # or in the POTCAR or OUTCAR file 

187 atom_symbols = [] 

188 numofatoms = fd.readline().split() 

189 # Check whether we have a VASP 4.x or 5.x format file. If the 

190 # format is 5.x, use the fifth line to provide information about 

191 # the atomic symbols. 

192 vasp5 = False 

193 try: 

194 int(numofatoms[0]) 

195 except ValueError: 

196 vasp5 = True 

197 atomtypes = numofatoms 

198 numofatoms = fd.readline().split() 

199 

200 # check for comments in numofatoms line and get rid of them if necessary 

201 commentcheck = np.array(['!' in s for s in numofatoms]) 

202 if commentcheck.any(): 

203 # only keep the elements up to the first including a '!': 

204 numofatoms = numofatoms[:np.arange(len(numofatoms))[commentcheck][0]] 

205 

206 if not vasp5: 

207 # Split the comment line (first in the file) into words and 

208 # try to compose a list of chemical symbols 

209 from ase.formula import Formula 

210 atomtypes = [] 

211 for word in line1.split(): 

212 word_without_delims = re.sub(r"-|_|,|\.|=|[0-9]|^", "", word) 

213 if len(word_without_delims) < 1: 

214 continue 

215 try: 

216 atomtypes.extend(list(Formula(word_without_delims))) 

217 except ValueError: 

218 # print(atomtype, e, 'is comment') 

219 pass 

220 # Now the list of chemical symbols atomtypes must be formed. 

221 # For example: atomtypes = ['Pd', 'C', 'O'] 

222 

223 numsyms = len(numofatoms) 

224 if len(atomtypes) < numsyms: 

225 # First line in POSCAR/CONTCAR didn't contain enough symbols. 

226 

227 # Sometimes the first line in POSCAR/CONTCAR is of the form 

228 # "CoP3_In-3.pos". Check for this case and extract atom types 

229 if len(atomtypes) == 1 and '_' in atomtypes[0]: 

230 atomtypes = get_atomtypes_from_formula(atomtypes[0]) 

231 else: 

232 atomtypes = atomtypes_outpot(fd.name, numsyms) 

233 else: 

234 try: 

235 for atype in atomtypes[:numsyms]: 

236 if atype not in chemical_symbols: 

237 raise KeyError 

238 except KeyError: 

239 atomtypes = atomtypes_outpot(fd.name, numsyms) 

240 

241 for i, num in enumerate(numofatoms): 

242 numofatoms[i] = int(num) 

243 atom_symbols.extend(numofatoms[i] * [atomtypes[i]]) 

244 

245 # Check if Selective dynamics is switched on 

246 sdyn = fd.readline() 

247 selective_dynamics = sdyn[0].lower() == 's' 

248 

249 # Check if atom coordinates are cartesian or direct 

250 if selective_dynamics: 

251 ac_type = fd.readline() 

252 else: 

253 ac_type = sdyn 

254 cartesian = ac_type[0].lower() in ['c', 'k'] 

255 tot_natoms = sum(numofatoms) 

256 atoms_pos = np.empty((tot_natoms, 3)) 

257 if selective_dynamics: 

258 selective_flags = np.empty((tot_natoms, 3), dtype=bool) 

259 for atom in range(tot_natoms): 

260 ac = fd.readline().split() 

261 atoms_pos[atom] = [float(_) for _ in ac[0:3]] 

262 if selective_dynamics: 

263 selective_flags[atom] = [_ == 'F' for _ in ac[3:6]] 

264 

265 atoms = Atoms(symbols=atom_symbols, cell=cell, pbc=True) 

266 if cartesian: 

267 atoms_pos *= scale 

268 atoms.set_positions(atoms_pos) 

269 else: 

270 atoms.set_scaled_positions(atoms_pos) 

271 if selective_dynamics: 

272 set_constraints(atoms, selective_flags) 

273 

274 return atoms 

275 

276 

277def read_velocities_if_present(fd, natoms) -> np.ndarray | None: 

278 """Read velocities from POSCAR/CONTCAR if present, return in ASE units.""" 

279 ac_type = fd.readline().strip() 

280 

281 # Check if velocities are present 

282 if not ac_type: 

283 return None 

284 

285 atoms_vel = np.empty((natoms, 3)) 

286 for atom in range(natoms): 

287 words = fd.readline().split() 

288 assert len(words) == 3 

289 atoms_vel[atom] = (float(words[0]), float(words[1]), float(words[2])) 

290 

291 # unit conversion from Angstrom/fs to ASE units 

292 return atoms_vel * (Ang / fs) 

293 

294 

295def set_constraints(atoms: Atoms, selective_flags: np.ndarray): 

296 """Set constraints based on selective_flags""" 

297 from ase.constraints import FixAtoms, FixConstraint, FixScaled 

298 

299 constraints: List[FixConstraint] = [] 

300 indices = [] 

301 for ind, sflags in enumerate(selective_flags): 

302 if sflags.any() and not sflags.all(): 

303 constraints.append(FixScaled(ind, sflags, atoms.get_cell())) 

304 elif sflags.all(): 

305 indices.append(ind) 

306 if indices: 

307 constraints.append(FixAtoms(indices)) 

308 if constraints: 

309 atoms.set_constraint(constraints) 

310 

311 

312def iread_vasp_out(filename, index=-1): 

313 """Import OUTCAR type file, as a generator.""" 

314 it = ImageIterator(vop.outcarchunks) 

315 return it(filename, index=index) 

316 

317 

318@reader 

319def read_vasp_out(filename='OUTCAR', index=-1): 

320 """Import OUTCAR type file. 

321 

322 Reads unitcell, atom positions, energies, and forces from the OUTCAR file 

323 and attempts to read constraints (if any) from CONTCAR/POSCAR, if present. 

324 """ 

325 # "filename" is actually a file-descriptor thanks to @reader 

326 g = iread_vasp_out(filename, index=index) 

327 # Code borrowed from formats.py:read 

328 if isinstance(index, (slice, str)): 

329 # Return list of atoms 

330 return list(g) 

331 else: 

332 # Return single atoms object 

333 return next(g) 

334 

335 

336@reader 

337def read_vasp_xdatcar(filename='XDATCAR', index=-1): 

338 """Import XDATCAR file. 

339 

340 Parameters 

341 ---------- 

342 index : int or slice or str 

343 Which frame(s) to read. The default is -1 (last frame). 

344 See :func:`ase.io.read` for details. 

345 

346 Notes 

347 ----- 

348 Constraints ARE NOT stored in the XDATCAR, and as such, Atoms objects 

349 retrieved from the XDATCAR will not have constraints. 

350 """ 

351 fd = filename # @reader decorator ensures this is a file descriptor 

352 images = [] 

353 

354 cell = np.eye(3) 

355 atomic_formula = '' 

356 

357 while True: 

358 comment_line = fd.readline() 

359 if "Direct configuration=" not in comment_line: 

360 try: 

361 lattice_constant = float(fd.readline()) 

362 except Exception: 

363 # XXX: When would this happen? 

364 break 

365 

366 xx = [float(x) for x in fd.readline().split()] 

367 yy = [float(y) for y in fd.readline().split()] 

368 zz = [float(z) for z in fd.readline().split()] 

369 cell = np.array([xx, yy, zz]) * lattice_constant 

370 

371 symbols = fd.readline().split() 

372 numbers = [int(n) for n in fd.readline().split()] 

373 total = sum(numbers) 

374 

375 atomic_formula = ''.join(f'{sym:s}{numbers[n]:d}' 

376 for n, sym in enumerate(symbols)) 

377 

378 fd.readline() 

379 

380 coords = [np.array(fd.readline().split(), float) for _ in range(total)] 

381 

382 image = Atoms(atomic_formula, cell=cell, pbc=True) 

383 image.set_scaled_positions(np.array(coords)) 

384 images.append(image) 

385 

386 if index is None: 

387 index = -1 

388 

389 if isinstance(index, str): 

390 index = string2index(index) 

391 

392 return images[index] 

393 

394 

395def __get_xml_parameter(par): 

396 """An auxiliary function that enables convenient extraction of 

397 parameter values from a vasprun.xml file with proper type 

398 handling. 

399 

400 """ 

401 def to_bool(b): 

402 if b == 'T': 

403 return True 

404 else: 

405 return False 

406 

407 to_type = {'int': int, 'logical': to_bool, 'string': str, 'float': float} 

408 

409 text = par.text 

410 if text is None: 

411 text = '' 

412 

413 # Float parameters do not have a 'type' attrib 

414 var_type = to_type[par.attrib.get('type', 'float')] 

415 

416 try: 

417 if par.tag == 'v': 

418 return list(map(var_type, text.split())) 

419 else: 

420 return var_type(text.strip()) 

421 except ValueError: 

422 # Vasp can sometimes write "*****" due to overflow 

423 return None 

424 

425 

426def read_vasp_xml(filename='vasprun.xml', index=-1): 

427 """Parse vasprun.xml file. 

428 

429 Reads unit cell, atom positions, energies, forces, and constraints 

430 from vasprun.xml file 

431 

432 Examples: 

433 >>> import ase.io 

434 >>> ase.io.write("out.traj", ase.io.read("vasprun.xml", index=":")) 

435 """ 

436 

437 import xml.etree.ElementTree as ET 

438 from collections import OrderedDict 

439 

440 from ase.calculators.singlepoint import ( 

441 SinglePointDFTCalculator, 

442 SinglePointKPoint, 

443 ) 

444 from ase.units import GPa 

445 

446 tree = ET.iterparse(filename, events=['start', 'end']) 

447 

448 atoms_init = None 

449 calculation = [] 

450 ibz_kpts = None 

451 kpt_weights = None 

452 parameters = OrderedDict() 

453 

454 try: 

455 for event, elem in tree: 

456 

457 if event == 'end': 

458 if elem.tag == 'kpoints': 

459 for subelem in elem.iter(tag='generation'): 

460 kpts_params = OrderedDict() 

461 parameters['kpoints_generation'] = kpts_params 

462 for par in subelem.iter(): 

463 if par.tag in ['v', 'i'] and "name" in par.attrib: 

464 parname = par.attrib['name'].lower() 

465 kpts_params[parname] = __get_xml_parameter(par) 

466 

467 kpts = elem.findall("varray[@name='kpointlist']/v") 

468 ibz_kpts = np.zeros((len(kpts), 3)) 

469 

470 for i, kpt in enumerate(kpts): 

471 ibz_kpts[i] = [float(val) for val in kpt.text.split()] 

472 

473 kpt_weights = elem.findall('varray[@name="weights"]/v') 

474 kpt_weights = [float(val.text) for val in kpt_weights] 

475 

476 elif elem.tag == 'parameters': 

477 for par in elem.iter(): 

478 if par.tag in ['v', 'i']: 

479 parname = par.attrib['name'].lower() 

480 parameters[parname] = __get_xml_parameter(par) 

481 

482 elif elem.tag == 'atominfo': 

483 species = [] 

484 

485 for entry in elem.find("array[@name='atoms']/set"): 

486 species.append(entry[0].text.strip()) 

487 

488 natoms = len(species) 

489 

490 elif (elem.tag == 'structure' 

491 and elem.attrib.get('name') == 'initialpos'): 

492 cell_init = np.zeros((3, 3), dtype=float) 

493 

494 for i, v in enumerate( 

495 elem.find("crystal/varray[@name='basis']")): 

496 cell_init[i] = np.array( 

497 [float(val) for val in v.text.split()]) 

498 

499 scpos_init = np.zeros((natoms, 3), dtype=float) 

500 

501 for i, v in enumerate( 

502 elem.find("varray[@name='positions']")): 

503 scpos_init[i] = np.array( 

504 [float(val) for val in v.text.split()]) 

505 

506 constraints = [] 

507 fixed_indices = [] 

508 

509 for i, entry in enumerate( 

510 elem.findall("varray[@name='selective']/v")): 

511 flags = (np.array( 

512 entry.text.split() == np.array(['F', 'F', 'F']))) 

513 if flags.all(): 

514 fixed_indices.append(i) 

515 elif flags.any(): 

516 constraints.append(FixScaled(i, flags, cell_init)) 

517 

518 if fixed_indices: 

519 constraints.append(FixAtoms(fixed_indices)) 

520 

521 atoms_init = Atoms(species, 

522 cell=cell_init, 

523 scaled_positions=scpos_init, 

524 constraint=constraints, 

525 pbc=True) 

526 

527 elif elem.tag == 'dipole': 

528 dblock = elem.find('v[@name="dipole"]') 

529 if dblock is not None: 

530 dipole = np.array( 

531 [float(val) for val in dblock.text.split()]) 

532 

533 elif event == 'start' and elem.tag == 'calculation': 

534 calculation.append(elem) 

535 

536 except ET.ParseError as parse_error: 

537 if atoms_init is None: 

538 raise parse_error 

539 if calculation and calculation[-1].find("energy") is None: 

540 calculation = calculation[:-1] 

541 if not calculation: 

542 yield atoms_init 

543 

544 if calculation: 

545 if isinstance(index, int): 

546 steps = [calculation[index]] 

547 else: 

548 steps = calculation[index] 

549 else: 

550 steps = [] 

551 

552 for step in steps: 

553 # Workaround for VASP bug, e_0_energy contains the wrong value 

554 # in calculation/energy, but calculation/scstep/energy does not 

555 # include classical VDW corrections. So, first calculate 

556 # e_0_energy - e_fr_energy from calculation/scstep/energy, then 

557 # apply that correction to e_fr_energy from calculation/energy. 

558 lastscf = step.findall('scstep/energy')[-1] 

559 dipoles = step.findall('scstep/dipole') 

560 if dipoles: 

561 lastdipole = dipoles[-1] 

562 else: 

563 lastdipole = None 

564 

565 de = (float(lastscf.find('i[@name="e_0_energy"]').text) - 

566 float(lastscf.find('i[@name="e_fr_energy"]').text)) 

567 

568 free_energy = float(step.find('energy/i[@name="e_fr_energy"]').text) 

569 energy = free_energy + de 

570 

571 cell = np.zeros((3, 3), dtype=float) 

572 for i, vector in enumerate( 

573 step.find('structure/crystal/varray[@name="basis"]')): 

574 cell[i] = np.array([float(val) for val in vector.text.split()]) 

575 

576 scpos = np.zeros((natoms, 3), dtype=float) 

577 for i, vector in enumerate( 

578 step.find('structure/varray[@name="positions"]')): 

579 scpos[i] = np.array([float(val) for val in vector.text.split()]) 

580 

581 forces = None 

582 fblocks = step.find('varray[@name="forces"]') 

583 if fblocks is not None: 

584 forces = np.zeros((natoms, 3), dtype=float) 

585 for i, vector in enumerate(fblocks): 

586 forces[i] = np.array( 

587 [float(val) for val in vector.text.split()]) 

588 

589 stress = None 

590 sblocks = step.find('varray[@name="stress"]') 

591 if sblocks is not None: 

592 stress = np.zeros((3, 3), dtype=float) 

593 for i, vector in enumerate(sblocks): 

594 stress[i] = np.array( 

595 [float(val) for val in vector.text.split()]) 

596 stress *= -0.1 * GPa 

597 stress = stress.reshape(9)[[0, 4, 8, 5, 2, 1]] 

598 

599 dipole = None 

600 if lastdipole is not None: 

601 dblock = lastdipole.find('v[@name="dipole"]') 

602 if dblock is not None: 

603 dipole = np.zeros((1, 3), dtype=float) 

604 dipole = np.array([float(val) for val in dblock.text.split()]) 

605 

606 dblock = step.find('dipole/v[@name="dipole"]') 

607 if dblock is not None: 

608 dipole = np.zeros((1, 3), dtype=float) 

609 dipole = np.array([float(val) for val in dblock.text.split()]) 

610 

611 efermi = step.find('dos/i[@name="efermi"]') 

612 if efermi is not None: 

613 efermi = float(efermi.text) 

614 

615 kpoints = [] 

616 for ikpt in range(1, len(ibz_kpts) + 1): 

617 kblocks = step.findall( 

618 'eigenvalues/array/set/set/set[@comment="kpoint %d"]' % ikpt) 

619 if kblocks is not None: 

620 for spin, kpoint in enumerate(kblocks): 

621 eigenvals = kpoint.findall('r') 

622 eps_n = np.zeros(len(eigenvals)) 

623 f_n = np.zeros(len(eigenvals)) 

624 for j, val in enumerate(eigenvals): 

625 val = val.text.split() 

626 eps_n[j] = float(val[0]) 

627 f_n[j] = float(val[1]) 

628 if len(kblocks) == 1: 

629 f_n *= 2 

630 kpoints.append( 

631 SinglePointKPoint(kpt_weights[ikpt - 1], spin, ikpt, 

632 eps_n, f_n)) 

633 if len(kpoints) == 0: 

634 kpoints = None 

635 

636 # DFPT properties 

637 # dielectric tensor 

638 dielectric_tensor = None 

639 sblocks = step.find('varray[@name="dielectric_dft"]') 

640 if sblocks is not None: 

641 dielectric_tensor = np.zeros((3, 3), dtype=float) 

642 for ii, vector in enumerate(sblocks): 

643 dielectric_tensor[ii] = np.fromstring(vector.text, sep=' ') 

644 

645 # Born effective charges 

646 born_charges = None 

647 fblocks = step.find('array[@name="born_charges"]') 

648 if fblocks is not None: 

649 born_charges = np.zeros((natoms, 3, 3), dtype=float) 

650 for ii, block in enumerate(fblocks[1:]): # 1. element = dimension 

651 for jj, vector in enumerate(block): 

652 born_charges[ii, jj] = np.fromstring(vector.text, sep=' ') 

653 

654 atoms = atoms_init.copy() 

655 atoms.set_cell(cell) 

656 atoms.set_scaled_positions(scpos) 

657 atoms.calc = SinglePointDFTCalculator( 

658 atoms, 

659 energy=energy, 

660 forces=forces, 

661 stress=stress, 

662 free_energy=free_energy, 

663 ibzkpts=ibz_kpts, 

664 efermi=efermi, 

665 dipole=dipole, 

666 dielectric_tensor=dielectric_tensor, 

667 born_effective_charges=born_charges 

668 ) 

669 atoms.calc.name = 'vasp' 

670 atoms.calc.kpts = kpoints 

671 atoms.calc.parameters = parameters 

672 yield atoms 

673 

674 

675@writer 

676def write_vasp_xdatcar(fd, images, label=None): 

677 """Write VASP MD trajectory (XDATCAR) file 

678 

679 Only Vasp 5 format is supported (for consistency with read_vasp_xdatcar) 

680 

681 Args: 

682 fd (str, fp): Output file 

683 images (iterable of Atoms): Atoms images to write. These must have 

684 consistent atom order and lattice vectors - this will not be 

685 checked. 

686 label (str): Text for first line of file. If empty, default to list 

687 of elements. 

688 

689 """ 

690 

691 images = iter(images) 

692 image = next(images) 

693 

694 if not isinstance(image, Atoms): 

695 raise TypeError("images should be a sequence of Atoms objects.") 

696 

697 symbol_count = _symbol_count_from_symbols(image.get_chemical_symbols()) 

698 

699 if label is None: 

700 label = ' '.join([s for s, _ in symbol_count]) 

701 fd.write(label + '\n') 

702 

703 # Not using lattice constants, set it to 1 

704 fd.write(' 1\n') 

705 

706 # Lattice vectors; use first image 

707 float_string = '{:11.6f}' 

708 for row_i in range(3): 

709 fd.write(' ') 

710 fd.write(' '.join(float_string.format(x) for x in image.cell[row_i])) 

711 fd.write('\n') 

712 

713 fd.write(_symbol_count_string(symbol_count, vasp5=True)) 

714 _write_xdatcar_config(fd, image, index=1) 

715 for i, image in enumerate(images): 

716 # Index is off by 2: 1-indexed file vs 0-indexed Python; 

717 # and we already wrote the first block. 

718 _write_xdatcar_config(fd, image, i + 2) 

719 

720 

721def _write_xdatcar_config(fd, atoms, index): 

722 """Write a block of positions for XDATCAR file 

723 

724 Args: 

725 fd (fd): writeable Python file descriptor 

726 atoms (ase.Atoms): Atoms to write 

727 index (int): configuration number written to block header 

728 

729 """ 

730 fd.write(f"Direct configuration={index:6d}\n") 

731 float_string = '{:11.8f}' 

732 scaled_positions = atoms.get_scaled_positions() 

733 for row in scaled_positions: 

734 fd.write(' ') 

735 fd.write(' '.join([float_string.format(x) for x in row])) 

736 fd.write('\n') 

737 

738 

739def _symbol_count_from_symbols(symbols: Symbols) -> List[Tuple[str, int]]: 

740 """Reduce list of chemical symbols into compact VASP notation 

741 

742 Args: 

743 symbols (iterable of str) 

744 

745 Returns: 

746 list of pairs [(el1, c1), (el2, c2), ...] 

747 

748 Example: 

749 >>> s = Atoms('Ar3NeHe2ArNe').symbols 

750 >>> _symbols_count_from_symbols(s) 

751 [('Ar', 3), ('Ne', 1), ('He', 2), ('Ar', 1), ('Ne', 1)] 

752 """ 

753 sc = [] 

754 psym = str(symbols[0]) # we cast to str to appease mypy 

755 count = 0 

756 for sym in symbols: 

757 if sym != psym: 

758 sc.append((psym, count)) 

759 psym = sym 

760 count = 1 

761 else: 

762 count += 1 

763 

764 sc.append((psym, count)) 

765 return sc 

766 

767 

768@writer 

769def write_vasp( 

770 fd: TextIO, 

771 atoms: Atoms, 

772 direct: bool = False, 

773 sort: bool = False, 

774 symbol_count: Optional[List[Tuple[str, int]]] = None, 

775 vasp5: bool = True, 

776 vasp6: bool = False, 

777 ignore_constraints: bool = False, 

778 potential_mapping: Optional[dict] = None 

779) -> None: 

780 """Method to write VASP position (POSCAR/CONTCAR) files. 

781 

782 Writes label, scalefactor, unitcell, # of various kinds of atoms, 

783 positions in cartesian or scaled coordinates (Direct), and constraints 

784 to file. Cartesian coordinates is default and default label is the 

785 atomic species, e.g. 'C N H Cu'. 

786 

787 Args: 

788 fd (TextIO): writeable Python file descriptor 

789 atoms (ase.Atoms): Atoms to write 

790 direct (bool): Write scaled coordinates instead of cartesian 

791 sort (bool): Sort the atomic indices alphabetically by element 

792 symbol_count (list of tuples of str and int, optional): Use the given 

793 combination of symbols and counts instead of automatically compute 

794 them 

795 vasp5 (bool): Write to the VASP 5+ format, where the symbols are 

796 written to file 

797 vasp6 (bool): Write symbols in VASP 6 format (which allows for 

798 potential type and hash) 

799 ignore_constraints (bool): Ignore all constraints on `atoms` 

800 potential_mapping (dict, optional): Map of symbols to potential file 

801 (and hash). Only works if `vasp6=True`. See `_symbol_string_count` 

802 

803 Raises: 

804 RuntimeError: raised if any of these are true: 

805 

806 1. `atoms` is not a single `ase.Atoms` object. 

807 2. The cell dimensionality is lower than 3 (0D-2D) 

808 3. One FixedPlane normal is not parallel to a unit cell vector 

809 4. One FixedLine direction is not parallel to a unit cell vector 

810 """ 

811 if isinstance(atoms, (list, tuple)): 

812 if len(atoms) > 1: 

813 raise RuntimeError( 

814 'Only one atomic structure can be saved to VASP ' 

815 'POSCAR/CONTCAR. Several were given.' 

816 ) 

817 else: 

818 atoms = atoms[0] 

819 

820 # Check lattice vectors are finite 

821 if atoms.cell.rank < 3: 

822 raise RuntimeError( 

823 'Lattice vectors must be finite and non-parallel. At least ' 

824 'one lattice length or angle is zero.' 

825 ) 

826 

827 # Write atomic positions in scaled or cartesian coordinates 

828 if direct: 

829 coord = atoms.get_scaled_positions(wrap=False) 

830 else: 

831 coord = atoms.positions 

832 

833 # Convert ASE constraints to VASP POSCAR constraints 

834 constraints_present = atoms.constraints and not ignore_constraints 

835 if constraints_present: 

836 sflags = _handle_ase_constraints(atoms) 

837 

838 # Conditionally sort ordering of `atoms` alphabetically by symbols 

839 if sort: 

840 ind = np.argsort(atoms.symbols) 

841 symbols = atoms.symbols[ind] 

842 coord = coord[ind] 

843 if constraints_present: 

844 sflags = sflags[ind] 

845 else: 

846 symbols = atoms.symbols 

847 

848 # Set or create a list of (symbol, count) pairs 

849 sc = symbol_count or _symbol_count_from_symbols(symbols) 

850 

851 # Write header as atomic species in `symbol_count` order 

852 label = ' '.join(f'{sym:2s}' for sym, _ in sc) 

853 fd.write(label + '\n') 

854 

855 # For simplicity, we write the unitcell in real coordinates, so the 

856 # scaling factor is always set to 1.0. 

857 fd.write(f'{1.0:19.16f}\n') 

858 

859 for vec in atoms.cell: 

860 fd.write(' ' + ' '.join([f'{el:21.16f}' for el in vec]) + '\n') 

861 

862 # Write version-dependent species-and-count section 

863 sc_str = _symbol_count_string(sc, vasp5, vasp6, potential_mapping) 

864 fd.write(sc_str) 

865 

866 # Write POSCAR switches 

867 if constraints_present: 

868 fd.write('Selective dynamics\n') 

869 

870 fd.write('Direct\n' if direct else 'Cartesian\n') 

871 

872 # Write atomic positions and, if any, the cartesian constraints 

873 for iatom, atom in enumerate(coord): 

874 for dcoord in atom: 

875 fd.write(f' {dcoord:19.16f}') 

876 if constraints_present: 

877 flags = ['F' if flag else 'T' for flag in sflags[iatom]] 

878 fd.write(''.join([f'{f:>4s}' for f in flags])) 

879 fd.write('\n') 

880 

881 # if velocities in atoms object write velocities 

882 if atoms.has('momenta'): 

883 cform = 3 * ' {:19.16f}' + '\n' 

884 fd.write('Cartesian\n') 

885 # unit conversion to Angstrom / fs 

886 vel = atoms.get_velocities() / (Ang / fs) 

887 for vatom in vel: 

888 fd.write(cform.format(*vatom)) 

889 

890 

891def _handle_ase_constraints(atoms: Atoms) -> np.ndarray: 

892 """Convert the ASE constraints on `atoms` to VASP constraints 

893 

894 Returns a boolean array with dimensions Nx3, where N is the number of 

895 atoms. A value of `True` indicates that movement along the given lattice 

896 vector is disallowed for that atom. 

897 

898 Args: 

899 atoms (Atoms) 

900 

901 Returns: 

902 boolean numpy array with dimensions Nx3 

903 

904 Raises: 

905 RuntimeError: If there is a FixedPlane or FixedLine constraint, that 

906 is not parallel to a cell vector. 

907 """ 

908 sflags = np.zeros((len(atoms), 3), dtype=bool) 

909 for constr in atoms.constraints: 

910 if isinstance(constr, FixScaled): 

911 sflags[constr.index] = constr.mask 

912 elif isinstance(constr, FixAtoms): 

913 sflags[constr.index] = 3 * [True] 

914 elif isinstance(constr, FixedPlane): 

915 # Calculate if the plane normal is parallel to a cell vector 

916 mask = np.all( 

917 np.abs(np.cross(constr.dir, atoms.cell)) < 1e-5, axis=1 

918 ) 

919 if mask.sum() != 1: 

920 raise RuntimeError( 

921 'VASP requires that the direction of FixedPlane ' 

922 'constraints is parallel with one of the cell axis' 

923 ) 

924 sflags[constr.index] = mask 

925 elif isinstance(constr, FixedLine): 

926 # Calculate if line is parallel to a cell vector 

927 mask = np.all( 

928 np.abs(np.cross(constr.dir, atoms.cell)) < 1e-5, axis=1 

929 ) 

930 if mask.sum() != 1: 

931 raise RuntimeError( 

932 'VASP requires that the direction of FixedLine ' 

933 'constraints is parallel with one of the cell axis' 

934 ) 

935 sflags[constr.index] = ~mask 

936 

937 return sflags 

938 

939 

940def _symbol_count_string( 

941 symbol_count: List[Tuple[str, int]], vasp5: bool = True, 

942 vasp6: bool = True, symbol_mapping: Optional[dict] = None 

943) -> str: 

944 """Create the symbols-and-counts block for POSCAR or XDATCAR 

945 

946 Args: 

947 symbol_count (list of 2-tuple): list of paired elements and counts 

948 vasp5 (bool): if False, omit symbols and only write counts 

949 vasp6 (bool): if True, write symbols in VASP 6 format (allows for 

950 potential type and hash) 

951 symbol_mapping (dict): mapping of symbols to VASP 6 symbols 

952 

953 e.g. if sc is [(Sn, 4), (S, 6)] then write for vasp 5: 

954 Sn S 

955 4 6 

956 

957 and for vasp 6 with mapping {'Sn': 'Sn_d_GW', 'S': 'S_GW/357d'}: 

958 Sn_d_GW S_GW/357d 

959 4 6 

960 """ 

961 symbol_mapping = symbol_mapping or {} 

962 out_str = ' ' 

963 

964 # Allow for VASP 6 format, i.e., specifying the pseudopotential used 

965 if vasp6: 

966 out_str += ' '.join([ 

967 f"{symbol_mapping.get(s, s)[:14]:16s}" for s, _ in symbol_count 

968 ]) + "\n " 

969 out_str += ' '.join([f"{c:16d}" for _, c in symbol_count]) + '\n' 

970 return out_str 

971 

972 # Write the species for VASP 5+ 

973 if vasp5: 

974 out_str += ' '.join([f"{s:3s}" for s, _ in symbol_count]) + "\n " 

975 

976 # Write counts line 

977 out_str += ' '.join([f"{c:3d}" for _, c in symbol_count]) + '\n' 

978 

979 return out_str