Coverage for /builds/hweiske/ase/ase/io/vasp.py: 89.27%

466 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2024-04-22 11:22 +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 get_atomtypes(fname): 

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

30 

31 The function can get this information from OUTCAR and POTCAR 

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

33 bzip2. 

34 

35 """ 

36 fpath = Path(fname) 

37 

38 atomtypes = [] 

39 atomtypes_alt = [] 

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

41 import gzip 

42 opener = gzip.open 

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

44 import bz2 

45 opener = bz2.BZ2File 

46 else: 

47 opener = open 

48 with opener(fpath) as fd: 

49 for line in fd: 

50 if 'TITEL' in line: 

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

52 elif 'POTCAR:' in line: 

53 atomtypes_alt.append( 

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

55 

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

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

58 # lines preceded by "POTCAR:", twice 

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

60 raise ParseError( 

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

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

63 ) 

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

65 

66 return atomtypes 

67 

68 

69def atomtypes_outpot(posfname, numsyms): 

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

71 

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

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

74 

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

76 

77 numsyms -- The number of symbols we must find 

78 

79 """ 

80 posfpath = Path(posfname) 

81 

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

83 # of POSCAR/CONTCAR. 

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

85 posfpath.with_name('OUTCAR')] 

86 # Try the same but with compressed files 

87 fsc = [] 

88 for fnpath in fnames: 

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

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

91 for f in fsc: 

92 fnames.append(f) 

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

94 # but this is no longer supported 

95 

96 tried = [] 

97 for fn in fnames: 

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

99 tried.append(fn) 

100 at = get_atomtypes(fn) 

101 if len(at) == numsyms: 

102 return at 

103 

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

105 str(tried)) 

106 

107 

108def get_atomtypes_from_formula(formula): 

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

110 with and underscore). 

111 """ 

112 from ase.symbols import string2symbols 

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

114 atomtypes = [symbols[0]] 

115 for s in symbols[1:]: 

116 if s != atomtypes[-1]: 

117 atomtypes.append(s) 

118 return atomtypes 

119 

120 

121@reader 

122def read_vasp(filename='CONTCAR'): 

123 """Import POSCAR/CONTCAR type file. 

124 

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

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

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

128 """ 

129 

130 from ase.data import chemical_symbols 

131 

132 fd = filename 

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

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

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

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

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

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

139 # file. 

140 line1 = fd.readline() 

141 

142 # Scaling factor 

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

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

145 scale = np.array(fd.readline().split()[:3], dtype=float) 

146 if len(scale) not in [1, 3]: 

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

148 if len(scale) == 3 and np.any(scale < 0.0): 

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

150 

151 # Now the lattice vectors 

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

153 # Negative scaling factor corresponds to the cell volume. 

154 if scale[0] < 0.0: 

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

156 cell *= scale 

157 

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

159 # in the first line 

160 # or in the POTCAR or OUTCAR file 

161 atom_symbols = [] 

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

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

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

165 # the atomic symbols. 

166 vasp5 = False 

167 try: 

168 int(numofatoms[0]) 

169 except ValueError: 

170 vasp5 = True 

171 atomtypes = numofatoms 

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

173 

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

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

176 if commentcheck.any(): 

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

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

179 

180 if not vasp5: 

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

182 # try to compose a list of chemical symbols 

183 from ase.formula import Formula 

184 atomtypes = [] 

185 for word in line1.split(): 

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

187 if len(word_without_delims) < 1: 

188 continue 

189 try: 

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

191 except ValueError: 

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

193 pass 

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

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

196 

197 numsyms = len(numofatoms) 

198 if len(atomtypes) < numsyms: 

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

200 

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

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

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

204 atomtypes = get_atomtypes_from_formula(atomtypes[0]) 

205 else: 

206 atomtypes = atomtypes_outpot(fd.name, numsyms) 

207 else: 

208 try: 

209 for atype in atomtypes[:numsyms]: 

210 if atype not in chemical_symbols: 

211 raise KeyError 

212 except KeyError: 

213 atomtypes = atomtypes_outpot(fd.name, numsyms) 

214 

215 for i, num in enumerate(numofatoms): 

216 numofatoms[i] = int(num) 

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

218 

219 # Check if Selective dynamics is switched on 

220 sdyn = fd.readline() 

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

222 

223 # Check if atom coordinates are cartesian or direct 

224 if selective_dynamics: 

225 ac_type = fd.readline() 

226 else: 

227 ac_type = sdyn 

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

229 tot_natoms = sum(numofatoms) 

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

231 if selective_dynamics: 

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

233 for atom in range(tot_natoms): 

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

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

236 if selective_dynamics: 

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

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

239 if cartesian: 

240 atoms_pos *= scale 

241 atoms.set_positions(atoms_pos) 

242 else: 

243 atoms.set_scaled_positions(atoms_pos) 

244 if selective_dynamics: 

245 set_constraints(atoms, selective_flags) 

246 return atoms 

247 

248 

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

250 """Set constraints based on selective_flags""" 

251 from ase.constraints import FixAtoms, FixConstraint, FixScaled 

252 

253 constraints: List[FixConstraint] = [] 

254 indices = [] 

255 for ind, sflags in enumerate(selective_flags): 

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

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

258 elif sflags.all(): 

259 indices.append(ind) 

260 if indices: 

261 constraints.append(FixAtoms(indices)) 

262 if constraints: 

263 atoms.set_constraint(constraints) 

264 

265 

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

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

268 it = ImageIterator(vop.outcarchunks) 

269 return it(filename, index=index) 

270 

271 

272@reader 

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

274 """Import OUTCAR type file. 

275 

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

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

278 """ 

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

280 g = iread_vasp_out(filename, index=index) 

281 # Code borrowed from formats.py:read 

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

283 # Return list of atoms 

284 return list(g) 

285 else: 

286 # Return single atoms object 

287 return next(g) 

288 

289 

290@reader 

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

292 """Import XDATCAR file. 

293 

294 Parameters 

295 ---------- 

296 index : int or slice or str 

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

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

299 

300 Notes 

301 ----- 

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

303 retrieved from the XDATCAR will not have constraints. 

304 """ 

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

306 images = [] 

307 

308 cell = np.eye(3) 

309 atomic_formula = '' 

310 

311 while True: 

312 comment_line = fd.readline() 

313 if "Direct configuration=" not in comment_line: 

314 try: 

315 lattice_constant = float(fd.readline()) 

316 except Exception: 

317 # XXX: When would this happen? 

318 break 

319 

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

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

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

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

324 

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

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

327 total = sum(numbers) 

328 

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

330 for n, sym in enumerate(symbols)) 

331 

332 fd.readline() 

333 

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

335 

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

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

338 images.append(image) 

339 

340 if index is None: 

341 index = -1 

342 

343 if isinstance(index, str): 

344 index = string2index(index) 

345 

346 return images[index] 

347 

348 

349def __get_xml_parameter(par): 

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

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

352 handling. 

353 

354 """ 

355 def to_bool(b): 

356 if b == 'T': 

357 return True 

358 else: 

359 return False 

360 

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

362 

363 text = par.text 

364 if text is None: 

365 text = '' 

366 

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

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

369 

370 try: 

371 if par.tag == 'v': 

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

373 else: 

374 return var_type(text.strip()) 

375 except ValueError: 

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

377 return None 

378 

379 

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

381 """Parse vasprun.xml file. 

382 

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

384 from vasprun.xml file 

385 

386 Examples: 

387 >>> import ase.io 

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

389 """ 

390 

391 import xml.etree.ElementTree as ET 

392 from collections import OrderedDict 

393 

394 from ase.calculators.singlepoint import (SinglePointDFTCalculator, 

395 SinglePointKPoint) 

396 from ase.constraints import FixAtoms, FixScaled 

397 from ase.units import GPa 

398 

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

400 

401 atoms_init = None 

402 calculation = [] 

403 ibz_kpts = None 

404 kpt_weights = None 

405 parameters = OrderedDict() 

406 

407 try: 

408 for event, elem in tree: 

409 

410 if event == 'end': 

411 if elem.tag == 'kpoints': 

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

413 kpts_params = OrderedDict() 

414 parameters['kpoints_generation'] = kpts_params 

415 for par in subelem.iter(): 

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

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

418 kpts_params[parname] = __get_xml_parameter(par) 

419 

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

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

422 

423 for i, kpt in enumerate(kpts): 

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

425 

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

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

428 

429 elif elem.tag == 'parameters': 

430 for par in elem.iter(): 

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

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

433 parameters[parname] = __get_xml_parameter(par) 

434 

435 elif elem.tag == 'atominfo': 

436 species = [] 

437 

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

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

440 

441 natoms = len(species) 

442 

443 elif (elem.tag == 'structure' 

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

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

446 

447 for i, v in enumerate( 

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

449 cell_init[i] = np.array( 

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

451 

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

453 

454 for i, v in enumerate( 

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

456 scpos_init[i] = np.array( 

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

458 

459 constraints = [] 

460 fixed_indices = [] 

461 

462 for i, entry in enumerate( 

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

464 flags = (np.array( 

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

466 if flags.all(): 

467 fixed_indices.append(i) 

468 elif flags.any(): 

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

470 

471 if fixed_indices: 

472 constraints.append(FixAtoms(fixed_indices)) 

473 

474 atoms_init = Atoms(species, 

475 cell=cell_init, 

476 scaled_positions=scpos_init, 

477 constraint=constraints, 

478 pbc=True) 

479 

480 elif elem.tag == 'dipole': 

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

482 if dblock is not None: 

483 dipole = np.array( 

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

485 

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

487 calculation.append(elem) 

488 

489 except ET.ParseError as parse_error: 

490 if atoms_init is None: 

491 raise parse_error 

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

493 calculation = calculation[:-1] 

494 if not calculation: 

495 yield atoms_init 

496 

497 if calculation: 

498 if isinstance(index, int): 

499 steps = [calculation[index]] 

500 else: 

501 steps = calculation[index] 

502 else: 

503 steps = [] 

504 

505 for step in steps: 

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

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

508 # include classical VDW corrections. So, first calculate 

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

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

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

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

513 if dipoles: 

514 lastdipole = dipoles[-1] 

515 else: 

516 lastdipole = None 

517 

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

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

520 

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

522 energy = free_energy + de 

523 

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

525 for i, vector in enumerate( 

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

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

528 

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

530 for i, vector in enumerate( 

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

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

533 

534 forces = None 

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

536 if fblocks is not None: 

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

538 for i, vector in enumerate(fblocks): 

539 forces[i] = np.array( 

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

541 

542 stress = None 

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

544 if sblocks is not None: 

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

546 for i, vector in enumerate(sblocks): 

547 stress[i] = np.array( 

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

549 stress *= -0.1 * GPa 

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

551 

552 dipole = None 

553 if lastdipole is not None: 

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

555 if dblock is not None: 

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

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

558 

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

560 if dblock is not None: 

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

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

563 

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

565 if efermi is not None: 

566 efermi = float(efermi.text) 

567 

568 kpoints = [] 

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

570 kblocks = step.findall( 

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

572 if kblocks is not None: 

573 for spin, kpoint in enumerate(kblocks): 

574 eigenvals = kpoint.findall('r') 

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

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

577 for j, val in enumerate(eigenvals): 

578 val = val.text.split() 

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

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

581 if len(kblocks) == 1: 

582 f_n *= 2 

583 kpoints.append( 

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

585 eps_n, f_n)) 

586 if len(kpoints) == 0: 

587 kpoints = None 

588 

589 # DFPT properties 

590 # dielectric tensor 

591 dielectric_tensor = None 

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

593 if sblocks is not None: 

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

595 for ii, vector in enumerate(sblocks): 

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

597 

598 # Born effective charges 

599 born_charges = None 

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

601 if fblocks is not None: 

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

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

604 for jj, vector in enumerate(block): 

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

606 

607 atoms = atoms_init.copy() 

608 atoms.set_cell(cell) 

609 atoms.set_scaled_positions(scpos) 

610 atoms.calc = SinglePointDFTCalculator( 

611 atoms, 

612 energy=energy, 

613 forces=forces, 

614 stress=stress, 

615 free_energy=free_energy, 

616 ibzkpts=ibz_kpts, 

617 efermi=efermi, 

618 dipole=dipole, 

619 dielectric_tensor=dielectric_tensor, 

620 born_effective_charges=born_charges 

621 ) 

622 atoms.calc.name = 'vasp' 

623 atoms.calc.kpts = kpoints 

624 atoms.calc.parameters = parameters 

625 yield atoms 

626 

627 

628@writer 

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

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

631 

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

633 

634 Args: 

635 fd (str, fp): Output file 

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

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

638 checked. 

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

640 of elements. 

641 

642 """ 

643 

644 images = iter(images) 

645 image = next(images) 

646 

647 if not isinstance(image, Atoms): 

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

649 

650 symbol_count = _symbol_count_from_symbols(image.get_chemical_symbols()) 

651 

652 if label is None: 

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

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

655 

656 # Not using lattice constants, set it to 1 

657 fd.write(' 1\n') 

658 

659 # Lattice vectors; use first image 

660 float_string = '{:11.6f}' 

661 for row_i in range(3): 

662 fd.write(' ') 

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

664 fd.write('\n') 

665 

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

667 _write_xdatcar_config(fd, image, index=1) 

668 for i, image in enumerate(images): 

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

670 # and we already wrote the first block. 

671 _write_xdatcar_config(fd, image, i + 2) 

672 

673 

674def _write_xdatcar_config(fd, atoms, index): 

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

676 

677 Args: 

678 fd (fd): writeable Python file descriptor 

679 atoms (ase.Atoms): Atoms to write 

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

681 

682 """ 

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

684 float_string = '{:11.8f}' 

685 scaled_positions = atoms.get_scaled_positions() 

686 for row in scaled_positions: 

687 fd.write(' ') 

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

689 fd.write('\n') 

690 

691 

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

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

694 

695 Args: 

696 symbols (iterable of str) 

697 

698 Returns: 

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

700 

701 Example: 

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

703 >>> _symbols_count_from_symbols(s) 

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

705 """ 

706 sc = [] 

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

708 count = 0 

709 for sym in symbols: 

710 if sym != psym: 

711 sc.append((psym, count)) 

712 psym = sym 

713 count = 1 

714 else: 

715 count += 1 

716 

717 sc.append((psym, count)) 

718 return sc 

719 

720 

721@writer 

722def write_vasp( 

723 fd: TextIO, 

724 atoms: Atoms, 

725 direct: bool = False, 

726 sort: bool = False, 

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

728 vasp5: bool = True, 

729 vasp6: bool = False, 

730 ignore_constraints: bool = False, 

731 potential_mapping: Optional[dict] = None 

732) -> None: 

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

734 

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

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

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

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

739 

740 Args: 

741 fd (TextIO): writeable Python file descriptor 

742 atoms (ase.Atoms): Atoms to write 

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

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

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

746 combination of symbols and counts instead of automatically compute 

747 them 

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

749 written to file 

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

751 potential type and hash) 

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

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

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

755 

756 Raises: 

757 RuntimeError: raised if any of these are true: 

758 

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

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

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

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

763 """ 

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

765 if len(atoms) > 1: 

766 raise RuntimeError( 

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

768 'POSCAR/CONTCAR. Several were given.' 

769 ) 

770 else: 

771 atoms = atoms[0] 

772 

773 # Check lattice vectors are finite 

774 if atoms.cell.rank < 3: 

775 raise RuntimeError( 

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

777 'one lattice length or angle is zero.' 

778 ) 

779 

780 # Write atomic positions in scaled or cartesian coordinates 

781 if direct: 

782 coord = atoms.get_scaled_positions(wrap=False) 

783 else: 

784 coord = atoms.positions 

785 

786 # Convert ASE constraints to VASP POSCAR constraints 

787 constraints_present = atoms.constraints and not ignore_constraints 

788 if constraints_present: 

789 sflags = _handle_ase_constraints(atoms) 

790 

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

792 if sort: 

793 ind = np.argsort(atoms.symbols) 

794 symbols = atoms.symbols[ind] 

795 coord = coord[ind] 

796 if constraints_present: 

797 sflags = sflags[ind] 

798 else: 

799 symbols = atoms.symbols 

800 

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

802 sc = symbol_count or _symbol_count_from_symbols(symbols) 

803 

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

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

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

807 

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

809 # scaling factor is always set to 1.0. 

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

811 

812 for vec in atoms.cell: 

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

814 

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

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

817 fd.write(sc_str) 

818 

819 # Write POSCAR switches 

820 if constraints_present: 

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

822 

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

824 

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

826 for iatom, atom in enumerate(coord): 

827 for dcoord in atom: 

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

829 if constraints_present: 

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

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

832 fd.write('\n') 

833 

834 

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

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

837 

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

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

840 vector is disallowed for that atom. 

841 

842 Args: 

843 atoms (Atoms) 

844 

845 Returns: 

846 boolean numpy array with dimensions Nx3 

847 

848 Raises: 

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

850 is not parallel to a cell vector. 

851 """ 

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

853 for constr in atoms.constraints: 

854 if isinstance(constr, FixScaled): 

855 sflags[constr.index] = constr.mask 

856 elif isinstance(constr, FixAtoms): 

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

858 elif isinstance(constr, FixedPlane): 

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

860 mask = np.all( 

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

862 ) 

863 if sum(mask) != 1: 

864 raise RuntimeError( 

865 'VASP requires that the direction of FixedPlane ' 

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

867 ) 

868 sflags[constr.index] = mask 

869 elif isinstance(constr, FixedLine): 

870 # Calculate if line is parallel to a cell vector 

871 mask = np.all( 

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

873 ) 

874 if sum(mask) != 1: 

875 raise RuntimeError( 

876 'VASP requires that the direction of FixedLine ' 

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

878 ) 

879 sflags[constr.index] = ~mask 

880 

881 return sflags 

882 

883 

884def _symbol_count_string( 

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

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

887) -> str: 

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

889 

890 Args: 

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

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

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

894 potential type and hash) 

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

896 

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

898 Sn S 

899 4 6 

900 

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

902 Sn_d_GW S_GW/357d 

903 4 6 

904 """ 

905 symbol_mapping = symbol_mapping or {} 

906 out_str = ' ' 

907 

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

909 if vasp6: 

910 out_str += ' '.join([ 

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

912 ]) + "\n " 

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

914 return out_str 

915 

916 # Write the species for VASP 5+ 

917 if vasp5: 

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

919 

920 # Write counts line 

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

922 

923 return out_str