Coverage for /builds/debichem-team/python-ase/ase/io/vasp.py: 90.70%

473 statements  

« prev     ^ index     » next       coverage.py v7.5.3, created at 2025-03-06 04:00 +0000

1""" 

2This module contains functionality for reading and writing an ASE 

3Atoms object in VASP POSCAR format. 

4 

5""" 

6import re 

7from pathlib import Path 

8from typing import List, Optional, TextIO, Tuple 

9 

10import numpy as np 

11 

12from ase import Atoms 

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

14from ase.io import ParseError 

15from ase.io.formats import string2index 

16from ase.io.utils import ImageIterator 

17from ase.symbols import Symbols 

18from ase.utils import reader, writer 

19 

20from .vasp_parsers import vasp_outcar_parsers as vop 

21 

22__all__ = [ 

23 'read_vasp', 'read_vasp_out', 'iread_vasp_out', 'read_vasp_xdatcar', 

24 'read_vasp_xml', 'write_vasp', 'write_vasp_xdatcar' 

25] 

26 

27 

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

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

30 

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

32 

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

34 

35 """ 

36 scale = [] 

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

38 try: 

39 scale.append(float(_)) 

40 except ValueError: 

41 break 

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

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

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

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

46 return np.array(scale) 

47 

48 

49def get_atomtypes(fname): 

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

51 

52 The function can get this information from OUTCAR and POTCAR 

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

54 bzip2. 

55 

56 """ 

57 fpath = Path(fname) 

58 

59 atomtypes = [] 

60 atomtypes_alt = [] 

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

62 import gzip 

63 opener = gzip.open 

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

65 import bz2 

66 opener = bz2.BZ2File 

67 else: 

68 opener = open 

69 with opener(fpath) as fd: 

70 for line in fd: 

71 if 'TITEL' in line: 

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

73 elif 'POTCAR:' in line: 

74 atomtypes_alt.append( 

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

76 

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

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

79 # lines preceded by "POTCAR:", twice 

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

81 raise ParseError( 

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

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

84 ) 

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

86 

87 return atomtypes 

88 

89 

90def atomtypes_outpot(posfname, numsyms): 

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

92 

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

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

95 

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

97 

98 numsyms -- The number of symbols we must find 

99 

100 """ 

101 posfpath = Path(posfname) 

102 

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

104 # of POSCAR/CONTCAR. 

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

106 posfpath.with_name('OUTCAR')] 

107 # Try the same but with compressed files 

108 fsc = [] 

109 for fnpath in fnames: 

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

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

112 for f in fsc: 

113 fnames.append(f) 

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

115 # but this is no longer supported 

116 

117 tried = [] 

118 for fn in fnames: 

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

120 tried.append(fn) 

121 at = get_atomtypes(fn) 

122 if len(at) == numsyms: 

123 return at 

124 

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

126 str(tried)) 

127 

128 

129def get_atomtypes_from_formula(formula): 

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

131 with and underscore). 

132 """ 

133 from ase.symbols import string2symbols 

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

135 atomtypes = [symbols[0]] 

136 for s in symbols[1:]: 

137 if s != atomtypes[-1]: 

138 atomtypes.append(s) 

139 return atomtypes 

140 

141 

142@reader 

143def read_vasp(filename='CONTCAR'): 

144 """Import POSCAR/CONTCAR type file. 

145 

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

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

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

149 """ 

150 

151 from ase.data import chemical_symbols 

152 

153 fd = filename 

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

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

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

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

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

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

160 # file. 

161 line1 = fd.readline() 

162 

163 scale = parse_poscar_scaling_factor(fd.readline()) 

164 

165 # Now the lattice vectors 

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

167 # Negative scaling factor corresponds to the cell volume. 

168 if scale[0] < 0.0: 

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

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

171 

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

173 # in the first line 

174 # or in the POTCAR or OUTCAR file 

175 atom_symbols = [] 

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

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

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

179 # the atomic symbols. 

180 vasp5 = False 

181 try: 

182 int(numofatoms[0]) 

183 except ValueError: 

184 vasp5 = True 

185 atomtypes = numofatoms 

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

187 

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

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

190 if commentcheck.any(): 

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

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

193 

194 if not vasp5: 

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

196 # try to compose a list of chemical symbols 

197 from ase.formula import Formula 

198 atomtypes = [] 

199 for word in line1.split(): 

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

201 if len(word_without_delims) < 1: 

202 continue 

203 try: 

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

205 except ValueError: 

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

207 pass 

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

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

210 

211 numsyms = len(numofatoms) 

212 if len(atomtypes) < numsyms: 

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

214 

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

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

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

218 atomtypes = get_atomtypes_from_formula(atomtypes[0]) 

219 else: 

220 atomtypes = atomtypes_outpot(fd.name, numsyms) 

221 else: 

222 try: 

223 for atype in atomtypes[:numsyms]: 

224 if atype not in chemical_symbols: 

225 raise KeyError 

226 except KeyError: 

227 atomtypes = atomtypes_outpot(fd.name, numsyms) 

228 

229 for i, num in enumerate(numofatoms): 

230 numofatoms[i] = int(num) 

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

232 

233 # Check if Selective dynamics is switched on 

234 sdyn = fd.readline() 

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

236 

237 # Check if atom coordinates are cartesian or direct 

238 if selective_dynamics: 

239 ac_type = fd.readline() 

240 else: 

241 ac_type = sdyn 

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

243 tot_natoms = sum(numofatoms) 

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

245 if selective_dynamics: 

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

247 for atom in range(tot_natoms): 

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

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

250 if selective_dynamics: 

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

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

253 if cartesian: 

254 atoms_pos *= scale 

255 atoms.set_positions(atoms_pos) 

256 else: 

257 atoms.set_scaled_positions(atoms_pos) 

258 if selective_dynamics: 

259 set_constraints(atoms, selective_flags) 

260 return atoms 

261 

262 

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

264 """Set constraints based on selective_flags""" 

265 from ase.constraints import FixAtoms, FixConstraint, FixScaled 

266 

267 constraints: List[FixConstraint] = [] 

268 indices = [] 

269 for ind, sflags in enumerate(selective_flags): 

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

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

272 elif sflags.all(): 

273 indices.append(ind) 

274 if indices: 

275 constraints.append(FixAtoms(indices)) 

276 if constraints: 

277 atoms.set_constraint(constraints) 

278 

279 

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

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

282 it = ImageIterator(vop.outcarchunks) 

283 return it(filename, index=index) 

284 

285 

286@reader 

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

288 """Import OUTCAR type file. 

289 

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

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

292 """ 

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

294 g = iread_vasp_out(filename, index=index) 

295 # Code borrowed from formats.py:read 

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

297 # Return list of atoms 

298 return list(g) 

299 else: 

300 # Return single atoms object 

301 return next(g) 

302 

303 

304@reader 

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

306 """Import XDATCAR file. 

307 

308 Parameters 

309 ---------- 

310 index : int or slice or str 

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

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

313 

314 Notes 

315 ----- 

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

317 retrieved from the XDATCAR will not have constraints. 

318 """ 

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

320 images = [] 

321 

322 cell = np.eye(3) 

323 atomic_formula = '' 

324 

325 while True: 

326 comment_line = fd.readline() 

327 if "Direct configuration=" not in comment_line: 

328 try: 

329 lattice_constant = float(fd.readline()) 

330 except Exception: 

331 # XXX: When would this happen? 

332 break 

333 

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

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

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

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

338 

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

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

341 total = sum(numbers) 

342 

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

344 for n, sym in enumerate(symbols)) 

345 

346 fd.readline() 

347 

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

349 

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

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

352 images.append(image) 

353 

354 if index is None: 

355 index = -1 

356 

357 if isinstance(index, str): 

358 index = string2index(index) 

359 

360 return images[index] 

361 

362 

363def __get_xml_parameter(par): 

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

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

366 handling. 

367 

368 """ 

369 def to_bool(b): 

370 if b == 'T': 

371 return True 

372 else: 

373 return False 

374 

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

376 

377 text = par.text 

378 if text is None: 

379 text = '' 

380 

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

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

383 

384 try: 

385 if par.tag == 'v': 

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

387 else: 

388 return var_type(text.strip()) 

389 except ValueError: 

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

391 return None 

392 

393 

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

395 """Parse vasprun.xml file. 

396 

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

398 from vasprun.xml file 

399 

400 Examples: 

401 >>> import ase.io 

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

403 """ 

404 

405 import xml.etree.ElementTree as ET 

406 from collections import OrderedDict 

407 

408 from ase.calculators.singlepoint import ( 

409 SinglePointDFTCalculator, 

410 SinglePointKPoint, 

411 ) 

412 from ase.units import GPa 

413 

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

415 

416 atoms_init = None 

417 calculation = [] 

418 ibz_kpts = None 

419 kpt_weights = None 

420 parameters = OrderedDict() 

421 

422 try: 

423 for event, elem in tree: 

424 

425 if event == 'end': 

426 if elem.tag == 'kpoints': 

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

428 kpts_params = OrderedDict() 

429 parameters['kpoints_generation'] = kpts_params 

430 for par in subelem.iter(): 

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

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

433 kpts_params[parname] = __get_xml_parameter(par) 

434 

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

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

437 

438 for i, kpt in enumerate(kpts): 

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

440 

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

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

443 

444 elif elem.tag == 'parameters': 

445 for par in elem.iter(): 

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

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

448 parameters[parname] = __get_xml_parameter(par) 

449 

450 elif elem.tag == 'atominfo': 

451 species = [] 

452 

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

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

455 

456 natoms = len(species) 

457 

458 elif (elem.tag == 'structure' 

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

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

461 

462 for i, v in enumerate( 

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

464 cell_init[i] = np.array( 

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

466 

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

468 

469 for i, v in enumerate( 

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

471 scpos_init[i] = np.array( 

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

473 

474 constraints = [] 

475 fixed_indices = [] 

476 

477 for i, entry in enumerate( 

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

479 flags = (np.array( 

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

481 if flags.all(): 

482 fixed_indices.append(i) 

483 elif flags.any(): 

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

485 

486 if fixed_indices: 

487 constraints.append(FixAtoms(fixed_indices)) 

488 

489 atoms_init = Atoms(species, 

490 cell=cell_init, 

491 scaled_positions=scpos_init, 

492 constraint=constraints, 

493 pbc=True) 

494 

495 elif elem.tag == 'dipole': 

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

497 if dblock is not None: 

498 dipole = np.array( 

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

500 

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

502 calculation.append(elem) 

503 

504 except ET.ParseError as parse_error: 

505 if atoms_init is None: 

506 raise parse_error 

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

508 calculation = calculation[:-1] 

509 if not calculation: 

510 yield atoms_init 

511 

512 if calculation: 

513 if isinstance(index, int): 

514 steps = [calculation[index]] 

515 else: 

516 steps = calculation[index] 

517 else: 

518 steps = [] 

519 

520 for step in steps: 

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

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

523 # include classical VDW corrections. So, first calculate 

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

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

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

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

528 if dipoles: 

529 lastdipole = dipoles[-1] 

530 else: 

531 lastdipole = None 

532 

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

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

535 

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

537 energy = free_energy + de 

538 

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

540 for i, vector in enumerate( 

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

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

543 

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

545 for i, vector in enumerate( 

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

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

548 

549 forces = None 

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

551 if fblocks is not None: 

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

553 for i, vector in enumerate(fblocks): 

554 forces[i] = np.array( 

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

556 

557 stress = None 

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

559 if sblocks is not None: 

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

561 for i, vector in enumerate(sblocks): 

562 stress[i] = np.array( 

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

564 stress *= -0.1 * GPa 

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

566 

567 dipole = None 

568 if lastdipole is not None: 

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

570 if dblock is not None: 

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

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

573 

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

575 if dblock is not None: 

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

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

578 

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

580 if efermi is not None: 

581 efermi = float(efermi.text) 

582 

583 kpoints = [] 

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

585 kblocks = step.findall( 

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

587 if kblocks is not None: 

588 for spin, kpoint in enumerate(kblocks): 

589 eigenvals = kpoint.findall('r') 

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

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

592 for j, val in enumerate(eigenvals): 

593 val = val.text.split() 

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

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

596 if len(kblocks) == 1: 

597 f_n *= 2 

598 kpoints.append( 

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

600 eps_n, f_n)) 

601 if len(kpoints) == 0: 

602 kpoints = None 

603 

604 # DFPT properties 

605 # dielectric tensor 

606 dielectric_tensor = None 

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

608 if sblocks is not None: 

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

610 for ii, vector in enumerate(sblocks): 

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

612 

613 # Born effective charges 

614 born_charges = None 

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

616 if fblocks is not None: 

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

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

619 for jj, vector in enumerate(block): 

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

621 

622 atoms = atoms_init.copy() 

623 atoms.set_cell(cell) 

624 atoms.set_scaled_positions(scpos) 

625 atoms.calc = SinglePointDFTCalculator( 

626 atoms, 

627 energy=energy, 

628 forces=forces, 

629 stress=stress, 

630 free_energy=free_energy, 

631 ibzkpts=ibz_kpts, 

632 efermi=efermi, 

633 dipole=dipole, 

634 dielectric_tensor=dielectric_tensor, 

635 born_effective_charges=born_charges 

636 ) 

637 atoms.calc.name = 'vasp' 

638 atoms.calc.kpts = kpoints 

639 atoms.calc.parameters = parameters 

640 yield atoms 

641 

642 

643@writer 

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

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

646 

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

648 

649 Args: 

650 fd (str, fp): Output file 

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

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

653 checked. 

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

655 of elements. 

656 

657 """ 

658 

659 images = iter(images) 

660 image = next(images) 

661 

662 if not isinstance(image, Atoms): 

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

664 

665 symbol_count = _symbol_count_from_symbols(image.get_chemical_symbols()) 

666 

667 if label is None: 

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

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

670 

671 # Not using lattice constants, set it to 1 

672 fd.write(' 1\n') 

673 

674 # Lattice vectors; use first image 

675 float_string = '{:11.6f}' 

676 for row_i in range(3): 

677 fd.write(' ') 

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

679 fd.write('\n') 

680 

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

682 _write_xdatcar_config(fd, image, index=1) 

683 for i, image in enumerate(images): 

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

685 # and we already wrote the first block. 

686 _write_xdatcar_config(fd, image, i + 2) 

687 

688 

689def _write_xdatcar_config(fd, atoms, index): 

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

691 

692 Args: 

693 fd (fd): writeable Python file descriptor 

694 atoms (ase.Atoms): Atoms to write 

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

696 

697 """ 

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

699 float_string = '{:11.8f}' 

700 scaled_positions = atoms.get_scaled_positions() 

701 for row in scaled_positions: 

702 fd.write(' ') 

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

704 fd.write('\n') 

705 

706 

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

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

709 

710 Args: 

711 symbols (iterable of str) 

712 

713 Returns: 

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

715 

716 Example: 

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

718 >>> _symbols_count_from_symbols(s) 

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

720 """ 

721 sc = [] 

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

723 count = 0 

724 for sym in symbols: 

725 if sym != psym: 

726 sc.append((psym, count)) 

727 psym = sym 

728 count = 1 

729 else: 

730 count += 1 

731 

732 sc.append((psym, count)) 

733 return sc 

734 

735 

736@writer 

737def write_vasp( 

738 fd: TextIO, 

739 atoms: Atoms, 

740 direct: bool = False, 

741 sort: bool = False, 

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

743 vasp5: bool = True, 

744 vasp6: bool = False, 

745 ignore_constraints: bool = False, 

746 potential_mapping: Optional[dict] = None 

747) -> None: 

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

749 

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

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

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

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

754 

755 Args: 

756 fd (TextIO): writeable Python file descriptor 

757 atoms (ase.Atoms): Atoms to write 

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

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

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

761 combination of symbols and counts instead of automatically compute 

762 them 

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

764 written to file 

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

766 potential type and hash) 

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

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

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

770 

771 Raises: 

772 RuntimeError: raised if any of these are true: 

773 

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

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

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

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

778 """ 

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

780 if len(atoms) > 1: 

781 raise RuntimeError( 

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

783 'POSCAR/CONTCAR. Several were given.' 

784 ) 

785 else: 

786 atoms = atoms[0] 

787 

788 # Check lattice vectors are finite 

789 if atoms.cell.rank < 3: 

790 raise RuntimeError( 

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

792 'one lattice length or angle is zero.' 

793 ) 

794 

795 # Write atomic positions in scaled or cartesian coordinates 

796 if direct: 

797 coord = atoms.get_scaled_positions(wrap=False) 

798 else: 

799 coord = atoms.positions 

800 

801 # Convert ASE constraints to VASP POSCAR constraints 

802 constraints_present = atoms.constraints and not ignore_constraints 

803 if constraints_present: 

804 sflags = _handle_ase_constraints(atoms) 

805 

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

807 if sort: 

808 ind = np.argsort(atoms.symbols) 

809 symbols = atoms.symbols[ind] 

810 coord = coord[ind] 

811 if constraints_present: 

812 sflags = sflags[ind] 

813 else: 

814 symbols = atoms.symbols 

815 

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

817 sc = symbol_count or _symbol_count_from_symbols(symbols) 

818 

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

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

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

822 

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

824 # scaling factor is always set to 1.0. 

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

826 

827 for vec in atoms.cell: 

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

829 

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

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

832 fd.write(sc_str) 

833 

834 # Write POSCAR switches 

835 if constraints_present: 

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

837 

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

839 

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

841 for iatom, atom in enumerate(coord): 

842 for dcoord in atom: 

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

844 if constraints_present: 

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

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

847 fd.write('\n') 

848 

849 

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

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

852 

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

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

855 vector is disallowed for that atom. 

856 

857 Args: 

858 atoms (Atoms) 

859 

860 Returns: 

861 boolean numpy array with dimensions Nx3 

862 

863 Raises: 

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

865 is not parallel to a cell vector. 

866 """ 

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

868 for constr in atoms.constraints: 

869 if isinstance(constr, FixScaled): 

870 sflags[constr.index] = constr.mask 

871 elif isinstance(constr, FixAtoms): 

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

873 elif isinstance(constr, FixedPlane): 

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

875 mask = np.all( 

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

877 ) 

878 if sum(mask) != 1: 

879 raise RuntimeError( 

880 'VASP requires that the direction of FixedPlane ' 

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

882 ) 

883 sflags[constr.index] = mask 

884 elif isinstance(constr, FixedLine): 

885 # Calculate if line is parallel to a cell vector 

886 mask = np.all( 

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

888 ) 

889 if sum(mask) != 1: 

890 raise RuntimeError( 

891 'VASP requires that the direction of FixedLine ' 

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

893 ) 

894 sflags[constr.index] = ~mask 

895 

896 return sflags 

897 

898 

899def _symbol_count_string( 

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

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

902) -> str: 

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

904 

905 Args: 

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

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

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

909 potential type and hash) 

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

911 

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

913 Sn S 

914 4 6 

915 

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

917 Sn_d_GW S_GW/357d 

918 4 6 

919 """ 

920 symbol_mapping = symbol_mapping or {} 

921 out_str = ' ' 

922 

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

924 if vasp6: 

925 out_str += ' '.join([ 

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

927 ]) + "\n " 

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

929 return out_str 

930 

931 # Write the species for VASP 5+ 

932 if vasp5: 

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

934 

935 # Write counts line 

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

937 

938 return out_str