Coverage for /builds/kinetik161/ase/ase/io/vasp.py: 89.22%

464 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-12-10 11:04 +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 = [ 

335 np.array(fd.readline().split(), float) for ii in range(total) 

336 ] 

337 

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

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

340 images.append(image) 

341 

342 if index is None: 

343 index = -1 

344 

345 if isinstance(index, str): 

346 index = string2index(index) 

347 

348 return images[index] 

349 

350 

351def __get_xml_parameter(par): 

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

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

354 handling. 

355 

356 """ 

357 def to_bool(b): 

358 if b == 'T': 

359 return True 

360 else: 

361 return False 

362 

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

364 

365 text = par.text 

366 if text is None: 

367 text = '' 

368 

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

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

371 

372 try: 

373 if par.tag == 'v': 

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

375 else: 

376 return var_type(text.strip()) 

377 except ValueError: 

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

379 return None 

380 

381 

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

383 """Parse vasprun.xml file. 

384 

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

386 from vasprun.xml file 

387 

388 Examples: 

389 >>> import ase.io 

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

391 """ 

392 

393 import xml.etree.ElementTree as ET 

394 from collections import OrderedDict 

395 

396 from ase.calculators.singlepoint import ( 

397 SinglePointDFTCalculator, 

398 SinglePointKPoint, 

399 ) 

400 from ase.constraints import FixAtoms, FixScaled 

401 from ase.units import GPa 

402 

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

404 

405 atoms_init = None 

406 calculation = [] 

407 ibz_kpts = None 

408 kpt_weights = None 

409 parameters = OrderedDict() 

410 

411 try: 

412 for event, elem in tree: 

413 

414 if event == 'end': 

415 if elem.tag == 'kpoints': 

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

417 kpts_params = OrderedDict() 

418 parameters['kpoints_generation'] = kpts_params 

419 for par in subelem.iter(): 

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

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

422 kpts_params[parname] = __get_xml_parameter(par) 

423 

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

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

426 

427 for i, kpt in enumerate(kpts): 

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

429 

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

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

432 

433 elif elem.tag == 'parameters': 

434 for par in elem.iter(): 

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

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

437 parameters[parname] = __get_xml_parameter(par) 

438 

439 elif elem.tag == 'atominfo': 

440 species = [] 

441 

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

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

444 

445 natoms = len(species) 

446 

447 elif (elem.tag == 'structure' 

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

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

450 

451 for i, v in enumerate( 

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

453 cell_init[i] = np.array( 

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

455 

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

457 

458 for i, v in enumerate( 

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

460 scpos_init[i] = np.array( 

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

462 

463 constraints = [] 

464 fixed_indices = [] 

465 

466 for i, entry in enumerate( 

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

468 flags = (np.array( 

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

470 if flags.all(): 

471 fixed_indices.append(i) 

472 elif flags.any(): 

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

474 

475 if fixed_indices: 

476 constraints.append(FixAtoms(fixed_indices)) 

477 

478 atoms_init = Atoms(species, 

479 cell=cell_init, 

480 scaled_positions=scpos_init, 

481 constraint=constraints, 

482 pbc=True) 

483 

484 elif elem.tag == 'dipole': 

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

486 if dblock is not None: 

487 dipole = np.array( 

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

489 

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

491 calculation.append(elem) 

492 

493 except ET.ParseError as parse_error: 

494 if atoms_init is None: 

495 raise parse_error 

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

497 calculation = calculation[:-1] 

498 if not calculation: 

499 yield atoms_init 

500 

501 if calculation: 

502 if isinstance(index, int): 

503 steps = [calculation[index]] 

504 else: 

505 steps = calculation[index] 

506 else: 

507 steps = [] 

508 

509 for step in steps: 

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

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

512 # include classical VDW corrections. So, first calculate 

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

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

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

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

517 if dipoles: 

518 lastdipole = dipoles[-1] 

519 else: 

520 lastdipole = None 

521 

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

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

524 

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

526 energy = free_energy + de 

527 

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

529 for i, vector in enumerate( 

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

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

532 

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

534 for i, vector in enumerate( 

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

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

537 

538 forces = None 

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

540 if fblocks is not None: 

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

542 for i, vector in enumerate(fblocks): 

543 forces[i] = np.array( 

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

545 

546 stress = None 

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

548 if sblocks is not None: 

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

550 for i, vector in enumerate(sblocks): 

551 stress[i] = np.array( 

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

553 stress *= -0.1 * GPa 

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

555 

556 dipole = None 

557 if lastdipole is not None: 

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

559 if dblock is not None: 

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

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

562 

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

564 if dblock is not None: 

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

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

567 

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

569 if efermi is not None: 

570 efermi = float(efermi.text) 

571 

572 kpoints = [] 

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

574 kblocks = step.findall( 

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

576 if kblocks is not None: 

577 for spin, kpoint in enumerate(kblocks): 

578 eigenvals = kpoint.findall('r') 

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

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

581 for j, val in enumerate(eigenvals): 

582 val = val.text.split() 

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

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

585 if len(kblocks) == 1: 

586 f_n *= 2 

587 kpoints.append( 

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

589 eps_n, f_n)) 

590 if len(kpoints) == 0: 

591 kpoints = None 

592 

593 # DFPT properties 

594 # dielectric tensor 

595 dielectric_tensor = None 

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

597 if sblocks is not None: 

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

599 for ii, vector in enumerate(sblocks): 

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

601 

602 # Born effective charges 

603 born_charges = None 

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

605 if fblocks is not None: 

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

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

608 for jj, vector in enumerate(block): 

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

610 

611 atoms = atoms_init.copy() 

612 atoms.set_cell(cell) 

613 atoms.set_scaled_positions(scpos) 

614 atoms.calc = SinglePointDFTCalculator( 

615 atoms, 

616 energy=energy, 

617 forces=forces, 

618 stress=stress, 

619 free_energy=free_energy, 

620 ibzkpts=ibz_kpts, 

621 efermi=efermi, 

622 dipole=dipole, 

623 dielectric_tensor=dielectric_tensor, 

624 born_effective_charges=born_charges 

625 ) 

626 atoms.calc.name = 'vasp' 

627 atoms.calc.kpts = kpoints 

628 atoms.calc.parameters = parameters 

629 yield atoms 

630 

631 

632@writer 

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

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

635 

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

637 

638 Args: 

639 fd (str, fp): Output file 

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

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

642 checked. 

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

644 of elements. 

645 

646 """ 

647 

648 images = iter(images) 

649 image = next(images) 

650 

651 if not isinstance(image, Atoms): 

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

653 

654 symbol_count = _symbol_count_from_symbols(image.get_chemical_symbols()) 

655 

656 if label is None: 

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

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

659 

660 # Not using lattice constants, set it to 1 

661 fd.write(' 1\n') 

662 

663 # Lattice vectors; use first image 

664 float_string = '{:11.6f}' 

665 for row_i in range(3): 

666 fd.write(' ') 

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

668 fd.write('\n') 

669 

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

671 _write_xdatcar_config(fd, image, index=1) 

672 for i, image in enumerate(images): 

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

674 # and we already wrote the first block. 

675 _write_xdatcar_config(fd, image, i + 2) 

676 

677 

678def _write_xdatcar_config(fd, atoms, index): 

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

680 

681 Args: 

682 fd (fd): writeable Python file descriptor 

683 atoms (ase.Atoms): Atoms to write 

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

685 

686 """ 

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

688 float_string = '{:11.8f}' 

689 scaled_positions = atoms.get_scaled_positions() 

690 for row in scaled_positions: 

691 fd.write(' ') 

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

693 fd.write('\n') 

694 

695 

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

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

698 

699 Args: 

700 symbols (iterable of str) 

701 

702 Returns: 

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

704 

705 Example: 

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

707 >>> _symbols_count_from_symbols(s) 

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

709 """ 

710 sc = [] 

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

712 count = 0 

713 for sym in symbols: 

714 if sym != psym: 

715 sc.append((psym, count)) 

716 psym = sym 

717 count = 1 

718 else: 

719 count += 1 

720 

721 sc.append((psym, count)) 

722 return sc 

723 

724 

725@writer 

726def write_vasp( 

727 fd: TextIO, 

728 atoms: Atoms, 

729 direct: bool = False, 

730 sort: bool = False, 

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

732 vasp5: bool = True, 

733 vasp6: bool = False, 

734 ignore_constraints: bool = False, 

735 potential_mapping: Optional[dict] = None 

736) -> None: 

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

738 

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

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

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

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

743 

744 Args: 

745 fd (TextIO): writeable Python file descriptor 

746 atoms (ase.Atoms): Atoms to write 

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

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

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

750 combination of symbols and counts instead of automatically compute 

751 them 

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

753 written to file 

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

755 potential type and hash) 

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

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

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

759 

760 Raises: 

761 RuntimeError: raised if any of these are true: 

762 

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

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

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

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

767 """ 

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

769 if len(atoms) > 1: 

770 raise RuntimeError( 

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

772 'POSCAR/CONTCAR. Several were given.' 

773 ) 

774 else: 

775 atoms = atoms[0] 

776 

777 # Check lattice vectors are finite 

778 if atoms.cell.rank < 3: 

779 raise RuntimeError( 

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

781 'one lattice length or angle is zero.' 

782 ) 

783 

784 # Write atomic positions in scaled or cartesian coordinates 

785 coord = atoms.get_scaled_positions() if direct else atoms.positions 

786 

787 # Convert ASE constraints to VASP POSCAR constraints 

788 constraints_present = atoms.constraints and not ignore_constraints 

789 if constraints_present: 

790 sflags = _handle_ase_constraints(atoms) 

791 

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

793 if sort: 

794 ind = np.argsort(atoms.symbols) 

795 symbols = atoms.symbols[ind] 

796 coord = coord[ind] 

797 if constraints_present: 

798 sflags = sflags[ind] 

799 else: 

800 symbols = atoms.symbols 

801 

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

803 sc = symbol_count or _symbol_count_from_symbols(symbols) 

804 

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

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

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

808 

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

810 # scaling factor is always set to 1.0. 

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

812 

813 for vec in atoms.cell: 

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

815 

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

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

818 fd.write(sc_str) 

819 

820 # Write POSCAR switches 

821 if constraints_present: 

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

823 

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

825 

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

827 for iatom, atom in enumerate(coord): 

828 for dcoord in atom: 

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

830 if constraints_present: 

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

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

833 fd.write('\n') 

834 

835 

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

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

838 

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

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

841 vector is disallowed for that atom. 

842 

843 Args: 

844 atoms (Atoms) 

845 

846 Returns: 

847 boolean numpy array with dimensions Nx3 

848 

849 Raises: 

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

851 is not parallel to a cell vector. 

852 """ 

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

854 for constr in atoms.constraints: 

855 if isinstance(constr, FixScaled): 

856 sflags[constr.index] = constr.mask 

857 elif isinstance(constr, FixAtoms): 

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

859 elif isinstance(constr, FixedPlane): 

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

861 mask = np.all( 

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

863 ) 

864 if sum(mask) != 1: 

865 raise RuntimeError( 

866 'VASP requires that the direction of FixedPlane ' 

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

868 ) 

869 sflags[constr.index] = mask 

870 elif isinstance(constr, FixedLine): 

871 # Calculate if line is parallel to a cell vector 

872 mask = np.all( 

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

874 ) 

875 if sum(mask) != 1: 

876 raise RuntimeError( 

877 'VASP requires that the direction of FixedLine ' 

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

879 ) 

880 sflags[constr.index] = ~mask 

881 

882 return sflags 

883 

884 

885def _symbol_count_string( 

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

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

888) -> str: 

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

890 

891 Args: 

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

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

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

895 potential type and hash) 

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

897 

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

899 Sn S 

900 4 6 

901 

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

903 Sn_d_GW S_GW/357d 

904 4 6 

905 """ 

906 symbol_mapping = symbol_mapping or dict() 

907 out_str = ' ' 

908 

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

910 if vasp6: 

911 out_str += ' '.join([ 

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

913 ]) + "\n " 

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

915 return out_str 

916 

917 # Write the species for VASP 5+ 

918 if vasp5: 

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

920 

921 # Write counts line 

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

923 

924 return out_str