1""" 

2This module contains functionality for reading and writing an ASE 

3Atoms object in VASP POSCAR format. 

4 

5""" 

6 

7import re 

8 

9import numpy as np 

10 

11from ase import Atoms 

12from ase.utils import reader, writer 

13from ase.io.utils import ImageIterator 

14from ase.io import ParseError 

15from .vasp_parsers import vasp_outcar_parsers as vop 

16from pathlib import Path 

17 

18__all__ = [ 

19 'read_vasp', 'read_vasp_out', 'iread_vasp_out', 'read_vasp_xdatcar', 

20 'read_vasp_xml', 'write_vasp', 'write_vasp_xdatcar' 

21] 

22 

23 

24def get_atomtypes(fname): 

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

26 

27 The function can get this information from OUTCAR and POTCAR 

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

29 bzip2. 

30 

31 """ 

32 fpath = Path(fname) 

33 

34 atomtypes = [] 

35 atomtypes_alt = [] 

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

37 import gzip 

38 opener = gzip.open 

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

40 import bz2 

41 opener = bz2.BZ2File 

42 else: 

43 opener = open 

44 with opener(fpath) as fd: 

45 for line in fd: 

46 if 'TITEL' in line: 

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

48 elif 'POTCAR:' in line: 

49 atomtypes_alt.append( 

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

51 

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

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

54 # preceded by "POTCAR:", twice 

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

56 raise ParseError( 

57 f'Tried to get atom types from {len(atomtypes_alt)} "POTCAR": ' 

58 'lines in OUTCAR, but expected an even number') 

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

60 

61 return atomtypes 

62 

63 

64def atomtypes_outpot(posfname, numsyms): 

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

66 

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

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

69 

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

71 

72 numsyms -- The number of symbols we must find 

73 

74 """ 

75 posfpath = Path(posfname) 

76 

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

78 # of POSCAR/CONTCAR. 

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

80 posfpath.with_name('OUTCAR')] 

81 # Try the same but with compressed files 

82 fsc = [] 

83 for fnpath in fnames: 

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

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

86 for f in fsc: 

87 fnames.append(f) 

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

89 # but this is no longer supported 

90 

91 tried = [] 

92 for fn in fnames: 

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

94 tried.append(fn) 

95 at = get_atomtypes(fn) 

96 if len(at) == numsyms: 

97 return at 

98 

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

100 str(tried)) 

101 

102 

103def get_atomtypes_from_formula(formula): 

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

105 with and underscore). 

106 """ 

107 from ase.symbols import string2symbols 

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

109 atomtypes = [symbols[0]] 

110 for s in symbols[1:]: 

111 if s != atomtypes[-1]: 

112 atomtypes.append(s) 

113 return atomtypes 

114 

115 

116@reader 

117def read_vasp(filename='CONTCAR'): 

118 """Import POSCAR/CONTCAR type file. 

119 

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

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

122 the atom types are read from OUTCAR or POTCAR file. 

123 """ 

124 

125 from ase.constraints import FixAtoms, FixScaled 

126 from ase.data import chemical_symbols 

127 

128 fd = filename 

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

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

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

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

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

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

135 # file. 

136 line1 = fd.readline() 

137 

138 lattice_constant = float(fd.readline().split()[0]) 

139 

140 # Now the lattice vectors 

141 a = [] 

142 for _ in range(3): 

143 s = fd.readline().split() 

144 floatvect = float(s[0]), float(s[1]), float(s[2]) 

145 a.append(floatvect) 

146 

147 basis_vectors = np.array(a) * lattice_constant 

148 

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

150 # in the first line 

151 # or in the POTCAR or OUTCAR file 

152 atom_symbols = [] 

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

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

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

156 # the atomic symbols. 

157 vasp5 = False 

158 try: 

159 int(numofatoms[0]) 

160 except ValueError: 

161 vasp5 = True 

162 atomtypes = numofatoms 

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

164 

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

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

167 if commentcheck.any(): 

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

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

170 

171 if not vasp5: 

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

173 # try to compose a list of chemical symbols 

174 from ase.formula import Formula 

175 atomtypes = [] 

176 for word in line1.split(): 

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

178 if len(word_without_delims) < 1: 

179 continue 

180 try: 

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

182 except ValueError: 

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

184 pass 

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

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

187 

188 numsyms = len(numofatoms) 

189 if len(atomtypes) < numsyms: 

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

191 

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

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

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

195 atomtypes = get_atomtypes_from_formula(atomtypes[0]) 

196 else: 

197 atomtypes = atomtypes_outpot(fd.name, numsyms) 

198 else: 

199 try: 

200 for atype in atomtypes[:numsyms]: 

201 if atype not in chemical_symbols: 

202 raise KeyError 

203 except KeyError: 

204 atomtypes = atomtypes_outpot(fd.name, numsyms) 

205 

206 for i, num in enumerate(numofatoms): 

207 numofatoms[i] = int(num) 

208 [atom_symbols.append(atomtypes[i]) for na in range(numofatoms[i])] 

209 

210 # Check if Selective dynamics is switched on 

211 sdyn = fd.readline() 

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

213 

214 # Check if atom coordinates are cartesian or direct 

215 if selective_dynamics: 

216 ac_type = fd.readline() 

217 else: 

218 ac_type = sdyn 

219 cartesian = ac_type[0].lower() == 'c' or ac_type[0].lower() == 'k' 

220 tot_natoms = sum(numofatoms) 

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

222 if selective_dynamics: 

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

224 for atom in range(tot_natoms): 

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

226 atoms_pos[atom] = (float(ac[0]), float(ac[1]), float(ac[2])) 

227 if selective_dynamics: 

228 curflag = [] 

229 for flag in ac[3:6]: 

230 curflag.append(flag == 'F') 

231 selective_flags[atom] = curflag 

232 if cartesian: 

233 atoms_pos *= lattice_constant 

234 atoms = Atoms(symbols=atom_symbols, cell=basis_vectors, pbc=True) 

235 if cartesian: 

236 atoms.set_positions(atoms_pos) 

237 else: 

238 atoms.set_scaled_positions(atoms_pos) 

239 if selective_dynamics: 

240 constraints = [] 

241 indices = [] 

242 for ind, sflags in enumerate(selective_flags): 

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

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

245 elif sflags.all(): 

246 indices.append(ind) 

247 if indices: 

248 constraints.append(FixAtoms(indices)) 

249 if constraints: 

250 atoms.set_constraint(constraints) 

251 return atoms 

252 

253 

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

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

256 it = ImageIterator(vop.outcarchunks) 

257 return it(filename, index=index) 

258 

259 

260@reader 

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

262 """Import OUTCAR type file. 

263 

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

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

266 """ 

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

268 g = iread_vasp_out(filename, index=index) 

269 # Code borrowed from formats.py:read 

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

271 # Return list of atoms 

272 return list(g) 

273 else: 

274 # Return single atoms object 

275 return next(g) 

276 

277 

278@reader 

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

280 """Import XDATCAR file 

281 

282 Reads all positions from the XDATCAR and returns a list of 

283 Atoms objects. Useful for viewing optimizations runs 

284 from VASP5.x 

285 

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

287 objects retrieved from the XDATCAR will not have constraints set. 

288 """ 

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

290 images = [] 

291 

292 cell = np.eye(3) 

293 atomic_formula = str() 

294 

295 while True: 

296 comment_line = fd.readline() 

297 if "Direct configuration=" not in comment_line: 

298 try: 

299 lattice_constant = float(fd.readline()) 

300 except Exception: 

301 # XXX: When would this happen? 

302 break 

303 

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

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

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

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

308 

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

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

311 total = sum(numbers) 

312 

313 atomic_formula = ''.join('{:s}{:d}'.format(sym, numbers[n]) 

314 for n, sym in enumerate(symbols)) 

315 

316 fd.readline() 

317 

318 coords = [ 

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

320 ] 

321 

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

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

324 images.append(image) 

325 

326 if not index: 

327 return images 

328 else: 

329 return images[index] 

330 

331 

332def __get_xml_parameter(par): 

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

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

335 handling. 

336 

337 """ 

338 def to_bool(b): 

339 if b == 'T': 

340 return True 

341 else: 

342 return False 

343 

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

345 

346 text = par.text 

347 if text is None: 

348 text = '' 

349 

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

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

352 

353 try: 

354 if par.tag == 'v': 

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

356 else: 

357 return var_type(text.strip()) 

358 except ValueError: 

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

360 return None 

361 

362 

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

364 """Parse vasprun.xml file. 

365 

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

367 from vasprun.xml file 

368 """ 

369 

370 import xml.etree.ElementTree as ET 

371 from ase.constraints import FixAtoms, FixScaled 

372 from ase.calculators.singlepoint import (SinglePointDFTCalculator, 

373 SinglePointKPoint) 

374 from ase.units import GPa 

375 from collections import OrderedDict 

376 

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

378 

379 atoms_init = None 

380 calculation = [] 

381 ibz_kpts = None 

382 kpt_weights = None 

383 parameters = OrderedDict() 

384 

385 try: 

386 for event, elem in tree: 

387 

388 if event == 'end': 

389 if elem.tag == 'kpoints': 

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

391 kpts_params = OrderedDict() 

392 parameters['kpoints_generation'] = kpts_params 

393 for par in subelem.iter(): 

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

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

396 kpts_params[parname] = __get_xml_parameter(par) 

397 

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

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

400 

401 for i, kpt in enumerate(kpts): 

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

403 

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

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

406 

407 elif elem.tag == 'parameters': 

408 for par in elem.iter(): 

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

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

411 parameters[parname] = __get_xml_parameter(par) 

412 

413 elif elem.tag == 'atominfo': 

414 species = [] 

415 

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

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

418 

419 natoms = len(species) 

420 

421 elif (elem.tag == 'structure' 

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

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

424 

425 for i, v in enumerate( 

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

427 cell_init[i] = np.array( 

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

429 

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

431 

432 for i, v in enumerate( 

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

434 scpos_init[i] = np.array( 

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

436 

437 constraints = [] 

438 fixed_indices = [] 

439 

440 for i, entry in enumerate( 

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

442 flags = (np.array( 

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

444 if flags.all(): 

445 fixed_indices.append(i) 

446 elif flags.any(): 

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

448 

449 if fixed_indices: 

450 constraints.append(FixAtoms(fixed_indices)) 

451 

452 atoms_init = Atoms(species, 

453 cell=cell_init, 

454 scaled_positions=scpos_init, 

455 constraint=constraints, 

456 pbc=True) 

457 

458 elif elem.tag == 'dipole': 

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

460 if dblock is not None: 

461 dipole = np.array( 

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

463 

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

465 calculation.append(elem) 

466 

467 except ET.ParseError as parse_error: 

468 if atoms_init is None: 

469 raise parse_error 

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

471 calculation = calculation[:-1] 

472 if not calculation: 

473 yield atoms_init 

474 

475 if calculation: 

476 if isinstance(index, int): 

477 steps = [calculation[index]] 

478 else: 

479 steps = calculation[index] 

480 else: 

481 steps = [] 

482 

483 for step in steps: 

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

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

486 # include classical VDW corrections. So, first calculate 

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

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

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

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

491 if dipoles: 

492 lastdipole = dipoles[-1] 

493 else: 

494 lastdipole = None 

495 

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

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

498 

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

500 energy = free_energy + de 

501 

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

503 for i, vector in enumerate( 

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

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

506 

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

508 for i, vector in enumerate( 

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

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

511 

512 forces = None 

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

514 if fblocks is not None: 

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

516 for i, vector in enumerate(fblocks): 

517 forces[i] = np.array( 

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

519 

520 stress = None 

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

522 if sblocks is not None: 

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

524 for i, vector in enumerate(sblocks): 

525 stress[i] = np.array( 

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

527 stress *= -0.1 * GPa 

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

529 

530 dipole = None 

531 if lastdipole is not None: 

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

533 if dblock is not None: 

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

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

536 

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

538 if dblock is not None: 

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

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

541 

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

543 if efermi is not None: 

544 efermi = float(efermi.text) 

545 

546 kpoints = [] 

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

548 kblocks = step.findall( 

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

550 if kblocks is not None: 

551 for spin, kpoint in enumerate(kblocks): 

552 eigenvals = kpoint.findall('r') 

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

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

555 for j, val in enumerate(eigenvals): 

556 val = val.text.split() 

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

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

559 if len(kblocks) == 1: 

560 f_n *= 2 

561 kpoints.append( 

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

563 eps_n, f_n)) 

564 if len(kpoints) == 0: 

565 kpoints = None 

566 

567 # DFPT properties 

568 # dielectric tensor 

569 dielectric_tensor = None 

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

571 if sblocks is not None: 

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

573 for ii, vector in enumerate(sblocks): 

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

575 

576 # Born effective charges 

577 born_charges = None 

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

579 if fblocks is not None: 

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

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

582 for jj, vector in enumerate(block): 

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

584 

585 atoms = atoms_init.copy() 

586 atoms.set_cell(cell) 

587 atoms.set_scaled_positions(scpos) 

588 atoms.calc = SinglePointDFTCalculator( 

589 atoms, 

590 energy=energy, 

591 forces=forces, 

592 stress=stress, 

593 free_energy=free_energy, 

594 ibzkpts=ibz_kpts, 

595 efermi=efermi, 

596 dipole=dipole, 

597 dielectric_tensor=dielectric_tensor, 

598 born_effective_charges=born_charges 

599 ) 

600 atoms.calc.name = 'vasp' 

601 atoms.calc.kpts = kpoints 

602 atoms.calc.parameters = parameters 

603 yield atoms 

604 

605 

606@writer 

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

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

609 

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

611 

612 Args: 

613 fd (str, fp): Output file 

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

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

616 checked. 

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

618 elements. 

619 

620 """ 

621 

622 images = iter(images) 

623 image = next(images) 

624 

625 if not isinstance(image, Atoms): 

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

627 

628 symbol_count = _symbol_count_from_symbols(image.get_chemical_symbols()) 

629 

630 if label is None: 

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

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

633 

634 # Not using lattice constants, set it to 1 

635 fd.write(' 1\n') 

636 

637 # Lattice vectors; use first image 

638 float_string = '{:11.6f}' 

639 for row_i in range(3): 

640 fd.write(' ') 

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

642 fd.write('\n') 

643 

644 _write_symbol_count(fd, symbol_count) 

645 _write_xdatcar_config(fd, image, index=1) 

646 for i, image in enumerate(images): 

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

648 # and we already wrote the first block. 

649 _write_xdatcar_config(fd, image, i + 2) 

650 

651 

652def _write_xdatcar_config(fd, atoms, index): 

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

654 

655 Args: 

656 fd (fd): writeable Python file descriptor 

657 atoms (ase.Atoms): Atoms to write 

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

659 

660 """ 

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

662 float_string = '{:11.8f}' 

663 scaled_positions = atoms.get_scaled_positions() 

664 for row in scaled_positions: 

665 fd.write(' ') 

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

667 fd.write('\n') 

668 

669 

670def _symbol_count_from_symbols(symbols): 

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

672 

673 args: 

674 symbols (iterable of str) 

675 

676 returns: 

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

678 """ 

679 sc = [] 

680 psym = symbols[0] 

681 count = 0 

682 for sym in symbols: 

683 if sym != psym: 

684 sc.append((psym, count)) 

685 psym = sym 

686 count = 1 

687 else: 

688 count += 1 

689 sc.append((psym, count)) 

690 return sc 

691 

692 

693def _write_symbol_count(fd, sc, vasp5=True): 

694 """Write the symbols and numbers block for POSCAR or XDATCAR 

695 

696 Args: 

697 f (fd): Descriptor for writable file 

698 sc (list of 2-tuple): list of paired elements and counts 

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

700 

701 e.g. if sc is [(Sn, 4), (S, 6)] then write:: 

702 

703 Sn S 

704 4 6 

705 

706 """ 

707 if vasp5: 

708 for sym, _ in sc: 

709 fd.write(' {:3s}'.format(sym)) 

710 fd.write('\n') 

711 

712 for _, count in sc: 

713 fd.write(' {:3d}'.format(count)) 

714 fd.write('\n') 

715 

716 

717@writer 

718def write_vasp(filename, 

719 atoms, 

720 label=None, 

721 direct=False, 

722 sort=None, 

723 symbol_count=None, 

724 long_format=True, 

725 vasp5=True, 

726 ignore_constraints=False, 

727 wrap=False): 

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

729 

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

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

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

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

734 """ 

735 

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

737 

738 fd = filename # @writer decorator ensures this arg is a file descriptor 

739 

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

741 if len(atoms) > 1: 

742 raise RuntimeError('Don\'t know how to save more than ' + 

743 'one image to VASP input') 

744 else: 

745 atoms = atoms[0] 

746 

747 # Check lattice vectors are finite 

748 if np.any(atoms.cell.cellpar() == 0.): 

749 raise RuntimeError( 

750 'Lattice vectors must be finite and not coincident. ' 

751 'At least one lattice length or angle is zero.') 

752 

753 # Write atom positions in scaled or cartesian coordinates 

754 if direct: 

755 coord = atoms.get_scaled_positions(wrap=wrap) 

756 else: 

757 coord = atoms.get_positions(wrap=wrap) 

758 

759 constraints = atoms.constraints and not ignore_constraints 

760 

761 if constraints: 

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

763 for constr in atoms.constraints: 

764 if isinstance(constr, FixScaled): 

765 sflags[constr.a] = constr.mask 

766 elif isinstance(constr, FixAtoms): 

767 sflags[constr.index] = [True, True, True] 

768 elif isinstance(constr, FixedPlane): 

769 mask = np.all(np.abs(np.cross(constr.dir, atoms.cell)) < 1e-5, 

770 axis=1) 

771 if sum(mask) != 1: 

772 raise RuntimeError( 

773 'VASP requires that the direction of FixedPlane ' 

774 'constraints is parallel with one of the cell axis') 

775 sflags[constr.a] = mask 

776 elif isinstance(constr, FixedLine): 

777 mask = np.all(np.abs(np.cross(constr.dir, atoms.cell)) < 1e-5, 

778 axis=1) 

779 if sum(mask) != 1: 

780 raise RuntimeError( 

781 'VASP requires that the direction of FixedLine ' 

782 'constraints is parallel with one of the cell axis') 

783 sflags[constr.a] = ~mask 

784 

785 if sort: 

786 ind = np.argsort(atoms.get_chemical_symbols()) 

787 symbols = np.array(atoms.get_chemical_symbols())[ind] 

788 coord = coord[ind] 

789 if constraints: 

790 sflags = sflags[ind] 

791 else: 

792 symbols = atoms.get_chemical_symbols() 

793 

794 # Create a list sc of (symbol, count) pairs 

795 if symbol_count: 

796 sc = symbol_count 

797 else: 

798 sc = _symbol_count_from_symbols(symbols) 

799 

800 # Create the label 

801 if label is None: 

802 label = '' 

803 for sym, c in sc: 

804 label += '%2s ' % sym 

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

806 

807 # Write unitcell in real coordinates and adapt to VASP convention 

808 # for unit cell 

809 # ase Atoms doesn't store the lattice constant separately, so always 

810 # write 1.0. 

811 fd.write('%19.16f\n' % 1.0) 

812 if long_format: 

813 latt_form = ' %21.16f' 

814 else: 

815 latt_form = ' %11.6f' 

816 for vec in atoms.get_cell(): 

817 fd.write(' ') 

818 for el in vec: 

819 fd.write(latt_form % el) 

820 fd.write('\n') 

821 

822 # Write out symbols (if VASP 5.x) and counts of atoms 

823 _write_symbol_count(fd, sc, vasp5=vasp5) 

824 

825 if constraints: 

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

827 

828 if direct: 

829 fd.write('Direct\n') 

830 else: 

831 fd.write('Cartesian\n') 

832 

833 if long_format: 

834 cform = ' %19.16f' 

835 else: 

836 cform = ' %9.6f' 

837 for iatom, atom in enumerate(coord): 

838 for dcoord in atom: 

839 fd.write(cform % dcoord) 

840 if constraints: 

841 for flag in sflags[iatom]: 

842 if flag: 

843 s = 'F' 

844 else: 

845 s = 'T' 

846 fd.write('%4s' % s) 

847 fd.write('\n')