Coverage for /builds/ase/ase/ase/io/vasp.py: 91.13%
496 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3"""
4This module contains functionality for reading and writing an ASE
5Atoms object in VASP POSCAR format.
7"""
8from __future__ import annotations
10import re
11from pathlib import Path
12from typing import List, Optional, TextIO, Tuple
14import numpy as np
16from ase import Atoms
17from ase.constraints import FixAtoms, FixedLine, FixedPlane, FixScaled
18from ase.io import ParseError
19from ase.io.formats import string2index
20from ase.io.utils import ImageIterator
21from ase.symbols import Symbols
22from ase.units import Ang, fs
23from ase.utils import reader, writer
25from .vasp_parsers import vasp_outcar_parsers as vop
27__all__ = [
28 'read_vasp', 'read_vasp_out', 'iread_vasp_out', 'read_vasp_xdatcar',
29 'read_vasp_xml', 'write_vasp', 'write_vasp_xdatcar'
30]
33def parse_poscar_scaling_factor(line: str) -> np.ndarray:
34 """Parse scaling factor(s) in the second line in POSCAR/CONTCAR.
36 This can also be one negative number or three positive numbers.
38 https://www.vasp.at/wiki/index.php/POSCAR#Full_format_specification
40 """
41 scale = []
42 for _ in line.split()[:3]:
43 try:
44 scale.append(float(_))
45 except ValueError:
46 break
47 if len(scale) not in {1, 3}:
48 raise RuntimeError('The number of scaling factors must be 1 or 3.')
49 if len(scale) == 3 and any(_ < 0.0 for _ in scale):
50 raise RuntimeError('All three scaling factors must be positive.')
51 return np.array(scale)
54def get_atomtypes(fname):
55 """Given a file name, get the atomic symbols.
57 The function can get this information from OUTCAR and POTCAR
58 format files. The files can also be compressed with gzip or
59 bzip2.
61 """
62 fpath = Path(fname)
64 atomtypes = []
65 atomtypes_alt = []
66 if fpath.suffix == '.gz':
67 import gzip
68 opener = gzip.open
69 elif fpath.suffix == '.bz2':
70 import bz2
71 opener = bz2.BZ2File
72 else:
73 opener = open
74 with opener(fpath) as fd:
75 for line in fd:
76 if 'TITEL' in line:
77 atomtypes.append(line.split()[3].split('_')[0].split('.')[0])
78 elif 'POTCAR:' in line:
79 atomtypes_alt.append(
80 line.split()[2].split('_')[0].split('.')[0])
82 if len(atomtypes) == 0 and len(atomtypes_alt) > 0:
83 # old VASP doesn't echo TITEL, but all versions print out species
84 # lines preceded by "POTCAR:", twice
85 if len(atomtypes_alt) % 2 != 0:
86 raise ParseError(
87 f'Tried to get atom types from {len(atomtypes_alt)}'
88 '"POTCAR": lines in OUTCAR, but expected an even number'
89 )
90 atomtypes = atomtypes_alt[0:len(atomtypes_alt) // 2]
92 return atomtypes
95def atomtypes_outpot(posfname, numsyms):
96 """Try to retrieve chemical symbols from OUTCAR or POTCAR
98 If getting atomtypes from the first line in POSCAR/CONTCAR fails, it might
99 be possible to find the data in OUTCAR or POTCAR, if these files exist.
101 posfname -- The filename of the POSCAR/CONTCAR file we're trying to read
103 numsyms -- The number of symbols we must find
105 """
106 posfpath = Path(posfname)
108 # Check files with exactly same path except POTCAR/OUTCAR instead
109 # of POSCAR/CONTCAR.
110 fnames = [posfpath.with_name('POTCAR'),
111 posfpath.with_name('OUTCAR')]
112 # Try the same but with compressed files
113 fsc = []
114 for fnpath in fnames:
115 fsc.append(fnpath.parent / (fnpath.name + '.gz'))
116 fsc.append(fnpath.parent / (fnpath.name + '.bz2'))
117 for f in fsc:
118 fnames.append(f)
119 # Code used to try anything with POTCAR or OUTCAR in the name
120 # but this is no longer supported
122 tried = []
123 for fn in fnames:
124 if fn in posfpath.parent.iterdir():
125 tried.append(fn)
126 at = get_atomtypes(fn)
127 if len(at) == numsyms:
128 return at
130 raise ParseError('Could not determine chemical symbols. Tried files ' +
131 str(tried))
134def get_atomtypes_from_formula(formula):
135 """Return atom types from chemical formula (optionally prepended
136 with and underscore).
137 """
138 from ase.symbols import string2symbols
139 symbols = string2symbols(formula.split('_')[0])
140 atomtypes = [symbols[0]]
141 for s in symbols[1:]:
142 if s != atomtypes[-1]:
143 atomtypes.append(s)
144 return atomtypes
147@reader
148def read_vasp(fd):
149 """Import POSCAR/CONTCAR type file.
151 Reads unitcell, atom positions and constraints from the POSCAR/CONTCAR
152 file and tries to read atom types from POSCAR/CONTCAR header, if this
153 fails the atom types are read from OUTCAR or POTCAR file.
154 """
155 atoms = read_vasp_configuration(fd)
156 velocities = read_velocities_if_present(fd, len(atoms))
157 if velocities is not None:
158 atoms.set_velocities(velocities)
159 return atoms
162def read_vasp_configuration(fd):
163 """Read common POSCAR/CONTCAR/CHGCAR/CHG quantities and return Atoms."""
164 from ase.data import chemical_symbols
166 # The first line is in principle a comment line, however in VASP
167 # 4.x a common convention is to have it contain the atom symbols,
168 # eg. "Ag Ge" in the same order as later in the file (and POTCAR
169 # for the full vasp run). In the VASP 5.x format this information
170 # is found on the fifth line. Thus we save the first line and use
171 # it in case we later detect that we're reading a VASP 4.x format
172 # file.
173 line1 = fd.readline()
175 scale = parse_poscar_scaling_factor(fd.readline())
177 # Now the lattice vectors
178 cell = np.array([fd.readline().split()[:3] for _ in range(3)], dtype=float)
179 # Negative scaling factor corresponds to the cell volume.
180 if scale[0] < 0.0:
181 scale = np.cbrt(-1.0 * scale / np.linalg.det(cell))
182 cell *= scale # This works for both one and three scaling factors.
184 # Number of atoms. Again this must be in the same order as
185 # in the first line
186 # or in the POTCAR or OUTCAR file
187 atom_symbols = []
188 numofatoms = fd.readline().split()
189 # Check whether we have a VASP 4.x or 5.x format file. If the
190 # format is 5.x, use the fifth line to provide information about
191 # the atomic symbols.
192 vasp5 = False
193 try:
194 int(numofatoms[0])
195 except ValueError:
196 vasp5 = True
197 atomtypes = numofatoms
198 numofatoms = fd.readline().split()
200 # check for comments in numofatoms line and get rid of them if necessary
201 commentcheck = np.array(['!' in s for s in numofatoms])
202 if commentcheck.any():
203 # only keep the elements up to the first including a '!':
204 numofatoms = numofatoms[:np.arange(len(numofatoms))[commentcheck][0]]
206 if not vasp5:
207 # Split the comment line (first in the file) into words and
208 # try to compose a list of chemical symbols
209 from ase.formula import Formula
210 atomtypes = []
211 for word in line1.split():
212 word_without_delims = re.sub(r"-|_|,|\.|=|[0-9]|^", "", word)
213 if len(word_without_delims) < 1:
214 continue
215 try:
216 atomtypes.extend(list(Formula(word_without_delims)))
217 except ValueError:
218 # print(atomtype, e, 'is comment')
219 pass
220 # Now the list of chemical symbols atomtypes must be formed.
221 # For example: atomtypes = ['Pd', 'C', 'O']
223 numsyms = len(numofatoms)
224 if len(atomtypes) < numsyms:
225 # First line in POSCAR/CONTCAR didn't contain enough symbols.
227 # Sometimes the first line in POSCAR/CONTCAR is of the form
228 # "CoP3_In-3.pos". Check for this case and extract atom types
229 if len(atomtypes) == 1 and '_' in atomtypes[0]:
230 atomtypes = get_atomtypes_from_formula(atomtypes[0])
231 else:
232 atomtypes = atomtypes_outpot(fd.name, numsyms)
233 else:
234 try:
235 for atype in atomtypes[:numsyms]:
236 if atype not in chemical_symbols:
237 raise KeyError
238 except KeyError:
239 atomtypes = atomtypes_outpot(fd.name, numsyms)
241 for i, num in enumerate(numofatoms):
242 numofatoms[i] = int(num)
243 atom_symbols.extend(numofatoms[i] * [atomtypes[i]])
245 # Check if Selective dynamics is switched on
246 sdyn = fd.readline()
247 selective_dynamics = sdyn[0].lower() == 's'
249 # Check if atom coordinates are cartesian or direct
250 if selective_dynamics:
251 ac_type = fd.readline()
252 else:
253 ac_type = sdyn
254 cartesian = ac_type[0].lower() in ['c', 'k']
255 tot_natoms = sum(numofatoms)
256 atoms_pos = np.empty((tot_natoms, 3))
257 if selective_dynamics:
258 selective_flags = np.empty((tot_natoms, 3), dtype=bool)
259 for atom in range(tot_natoms):
260 ac = fd.readline().split()
261 atoms_pos[atom] = [float(_) for _ in ac[0:3]]
262 if selective_dynamics:
263 selective_flags[atom] = [_ == 'F' for _ in ac[3:6]]
265 atoms = Atoms(symbols=atom_symbols, cell=cell, pbc=True)
266 if cartesian:
267 atoms_pos *= scale
268 atoms.set_positions(atoms_pos)
269 else:
270 atoms.set_scaled_positions(atoms_pos)
271 if selective_dynamics:
272 set_constraints(atoms, selective_flags)
274 return atoms
277def read_velocities_if_present(fd, natoms) -> np.ndarray | None:
278 """Read velocities from POSCAR/CONTCAR if present, return in ASE units."""
279 ac_type = fd.readline().strip()
281 # Check if velocities are present
282 if not ac_type:
283 return None
285 atoms_vel = np.empty((natoms, 3))
286 for atom in range(natoms):
287 words = fd.readline().split()
288 assert len(words) == 3
289 atoms_vel[atom] = (float(words[0]), float(words[1]), float(words[2]))
291 # unit conversion from Angstrom/fs to ASE units
292 return atoms_vel * (Ang / fs)
295def set_constraints(atoms: Atoms, selective_flags: np.ndarray):
296 """Set constraints based on selective_flags"""
297 from ase.constraints import FixAtoms, FixConstraint, FixScaled
299 constraints: List[FixConstraint] = []
300 indices = []
301 for ind, sflags in enumerate(selective_flags):
302 if sflags.any() and not sflags.all():
303 constraints.append(FixScaled(ind, sflags, atoms.get_cell()))
304 elif sflags.all():
305 indices.append(ind)
306 if indices:
307 constraints.append(FixAtoms(indices))
308 if constraints:
309 atoms.set_constraint(constraints)
312def iread_vasp_out(filename, index=-1):
313 """Import OUTCAR type file, as a generator."""
314 it = ImageIterator(vop.outcarchunks)
315 return it(filename, index=index)
318@reader
319def read_vasp_out(filename='OUTCAR', index=-1):
320 """Import OUTCAR type file.
322 Reads unitcell, atom positions, energies, and forces from the OUTCAR file
323 and attempts to read constraints (if any) from CONTCAR/POSCAR, if present.
324 """
325 # "filename" is actually a file-descriptor thanks to @reader
326 g = iread_vasp_out(filename, index=index)
327 # Code borrowed from formats.py:read
328 if isinstance(index, (slice, str)):
329 # Return list of atoms
330 return list(g)
331 else:
332 # Return single atoms object
333 return next(g)
336@reader
337def read_vasp_xdatcar(filename='XDATCAR', index=-1):
338 """Import XDATCAR file.
340 Parameters
341 ----------
342 index : int or slice or str
343 Which frame(s) to read. The default is -1 (last frame).
344 See :func:`ase.io.read` for details.
346 Notes
347 -----
348 Constraints ARE NOT stored in the XDATCAR, and as such, Atoms objects
349 retrieved from the XDATCAR will not have constraints.
350 """
351 fd = filename # @reader decorator ensures this is a file descriptor
352 images = []
354 cell = np.eye(3)
355 atomic_formula = ''
357 while True:
358 comment_line = fd.readline()
359 if "Direct configuration=" not in comment_line:
360 try:
361 lattice_constant = float(fd.readline())
362 except Exception:
363 # XXX: When would this happen?
364 break
366 xx = [float(x) for x in fd.readline().split()]
367 yy = [float(y) for y in fd.readline().split()]
368 zz = [float(z) for z in fd.readline().split()]
369 cell = np.array([xx, yy, zz]) * lattice_constant
371 symbols = fd.readline().split()
372 numbers = [int(n) for n in fd.readline().split()]
373 total = sum(numbers)
375 atomic_formula = ''.join(f'{sym:s}{numbers[n]:d}'
376 for n, sym in enumerate(symbols))
378 fd.readline()
380 coords = [np.array(fd.readline().split(), float) for _ in range(total)]
382 image = Atoms(atomic_formula, cell=cell, pbc=True)
383 image.set_scaled_positions(np.array(coords))
384 images.append(image)
386 if index is None:
387 index = -1
389 if isinstance(index, str):
390 index = string2index(index)
392 return images[index]
395def __get_xml_parameter(par):
396 """An auxiliary function that enables convenient extraction of
397 parameter values from a vasprun.xml file with proper type
398 handling.
400 """
401 def to_bool(b):
402 if b == 'T':
403 return True
404 else:
405 return False
407 to_type = {'int': int, 'logical': to_bool, 'string': str, 'float': float}
409 text = par.text
410 if text is None:
411 text = ''
413 # Float parameters do not have a 'type' attrib
414 var_type = to_type[par.attrib.get('type', 'float')]
416 try:
417 if par.tag == 'v':
418 return list(map(var_type, text.split()))
419 else:
420 return var_type(text.strip())
421 except ValueError:
422 # Vasp can sometimes write "*****" due to overflow
423 return None
426def read_vasp_xml(filename='vasprun.xml', index=-1):
427 """Parse vasprun.xml file.
429 Reads unit cell, atom positions, energies, forces, and constraints
430 from vasprun.xml file
432 Examples:
433 >>> import ase.io
434 >>> ase.io.write("out.traj", ase.io.read("vasprun.xml", index=":"))
435 """
437 import xml.etree.ElementTree as ET
438 from collections import OrderedDict
440 from ase.calculators.singlepoint import (
441 SinglePointDFTCalculator,
442 SinglePointKPoint,
443 )
444 from ase.units import GPa
446 tree = ET.iterparse(filename, events=['start', 'end'])
448 atoms_init = None
449 calculation = []
450 ibz_kpts = None
451 kpt_weights = None
452 parameters = OrderedDict()
454 try:
455 for event, elem in tree:
457 if event == 'end':
458 if elem.tag == 'kpoints':
459 for subelem in elem.iter(tag='generation'):
460 kpts_params = OrderedDict()
461 parameters['kpoints_generation'] = kpts_params
462 for par in subelem.iter():
463 if par.tag in ['v', 'i'] and "name" in par.attrib:
464 parname = par.attrib['name'].lower()
465 kpts_params[parname] = __get_xml_parameter(par)
467 kpts = elem.findall("varray[@name='kpointlist']/v")
468 ibz_kpts = np.zeros((len(kpts), 3))
470 for i, kpt in enumerate(kpts):
471 ibz_kpts[i] = [float(val) for val in kpt.text.split()]
473 kpt_weights = elem.findall('varray[@name="weights"]/v')
474 kpt_weights = [float(val.text) for val in kpt_weights]
476 elif elem.tag == 'parameters':
477 for par in elem.iter():
478 if par.tag in ['v', 'i']:
479 parname = par.attrib['name'].lower()
480 parameters[parname] = __get_xml_parameter(par)
482 elif elem.tag == 'atominfo':
483 species = []
485 for entry in elem.find("array[@name='atoms']/set"):
486 species.append(entry[0].text.strip())
488 natoms = len(species)
490 elif (elem.tag == 'structure'
491 and elem.attrib.get('name') == 'initialpos'):
492 cell_init = np.zeros((3, 3), dtype=float)
494 for i, v in enumerate(
495 elem.find("crystal/varray[@name='basis']")):
496 cell_init[i] = np.array(
497 [float(val) for val in v.text.split()])
499 scpos_init = np.zeros((natoms, 3), dtype=float)
501 for i, v in enumerate(
502 elem.find("varray[@name='positions']")):
503 scpos_init[i] = np.array(
504 [float(val) for val in v.text.split()])
506 constraints = []
507 fixed_indices = []
509 for i, entry in enumerate(
510 elem.findall("varray[@name='selective']/v")):
511 flags = (np.array(
512 entry.text.split() == np.array(['F', 'F', 'F'])))
513 if flags.all():
514 fixed_indices.append(i)
515 elif flags.any():
516 constraints.append(FixScaled(i, flags, cell_init))
518 if fixed_indices:
519 constraints.append(FixAtoms(fixed_indices))
521 atoms_init = Atoms(species,
522 cell=cell_init,
523 scaled_positions=scpos_init,
524 constraint=constraints,
525 pbc=True)
527 elif elem.tag == 'dipole':
528 dblock = elem.find('v[@name="dipole"]')
529 if dblock is not None:
530 dipole = np.array(
531 [float(val) for val in dblock.text.split()])
533 elif event == 'start' and elem.tag == 'calculation':
534 calculation.append(elem)
536 except ET.ParseError as parse_error:
537 if atoms_init is None:
538 raise parse_error
539 if calculation and calculation[-1].find("energy") is None:
540 calculation = calculation[:-1]
541 if not calculation:
542 yield atoms_init
544 if calculation:
545 if isinstance(index, int):
546 steps = [calculation[index]]
547 else:
548 steps = calculation[index]
549 else:
550 steps = []
552 for step in steps:
553 # Workaround for VASP bug, e_0_energy contains the wrong value
554 # in calculation/energy, but calculation/scstep/energy does not
555 # include classical VDW corrections. So, first calculate
556 # e_0_energy - e_fr_energy from calculation/scstep/energy, then
557 # apply that correction to e_fr_energy from calculation/energy.
558 lastscf = step.findall('scstep/energy')[-1]
559 dipoles = step.findall('scstep/dipole')
560 if dipoles:
561 lastdipole = dipoles[-1]
562 else:
563 lastdipole = None
565 de = (float(lastscf.find('i[@name="e_0_energy"]').text) -
566 float(lastscf.find('i[@name="e_fr_energy"]').text))
568 free_energy = float(step.find('energy/i[@name="e_fr_energy"]').text)
569 energy = free_energy + de
571 cell = np.zeros((3, 3), dtype=float)
572 for i, vector in enumerate(
573 step.find('structure/crystal/varray[@name="basis"]')):
574 cell[i] = np.array([float(val) for val in vector.text.split()])
576 scpos = np.zeros((natoms, 3), dtype=float)
577 for i, vector in enumerate(
578 step.find('structure/varray[@name="positions"]')):
579 scpos[i] = np.array([float(val) for val in vector.text.split()])
581 forces = None
582 fblocks = step.find('varray[@name="forces"]')
583 if fblocks is not None:
584 forces = np.zeros((natoms, 3), dtype=float)
585 for i, vector in enumerate(fblocks):
586 forces[i] = np.array(
587 [float(val) for val in vector.text.split()])
589 stress = None
590 sblocks = step.find('varray[@name="stress"]')
591 if sblocks is not None:
592 stress = np.zeros((3, 3), dtype=float)
593 for i, vector in enumerate(sblocks):
594 stress[i] = np.array(
595 [float(val) for val in vector.text.split()])
596 stress *= -0.1 * GPa
597 stress = stress.reshape(9)[[0, 4, 8, 5, 2, 1]]
599 dipole = None
600 if lastdipole is not None:
601 dblock = lastdipole.find('v[@name="dipole"]')
602 if dblock is not None:
603 dipole = np.zeros((1, 3), dtype=float)
604 dipole = np.array([float(val) for val in dblock.text.split()])
606 dblock = step.find('dipole/v[@name="dipole"]')
607 if dblock is not None:
608 dipole = np.zeros((1, 3), dtype=float)
609 dipole = np.array([float(val) for val in dblock.text.split()])
611 efermi = step.find('dos/i[@name="efermi"]')
612 if efermi is not None:
613 efermi = float(efermi.text)
615 kpoints = []
616 for ikpt in range(1, len(ibz_kpts) + 1):
617 kblocks = step.findall(
618 'eigenvalues/array/set/set/set[@comment="kpoint %d"]' % ikpt)
619 if kblocks is not None:
620 for spin, kpoint in enumerate(kblocks):
621 eigenvals = kpoint.findall('r')
622 eps_n = np.zeros(len(eigenvals))
623 f_n = np.zeros(len(eigenvals))
624 for j, val in enumerate(eigenvals):
625 val = val.text.split()
626 eps_n[j] = float(val[0])
627 f_n[j] = float(val[1])
628 if len(kblocks) == 1:
629 f_n *= 2
630 kpoints.append(
631 SinglePointKPoint(kpt_weights[ikpt - 1], spin, ikpt,
632 eps_n, f_n))
633 if len(kpoints) == 0:
634 kpoints = None
636 # DFPT properties
637 # dielectric tensor
638 dielectric_tensor = None
639 sblocks = step.find('varray[@name="dielectric_dft"]')
640 if sblocks is not None:
641 dielectric_tensor = np.zeros((3, 3), dtype=float)
642 for ii, vector in enumerate(sblocks):
643 dielectric_tensor[ii] = np.fromstring(vector.text, sep=' ')
645 # Born effective charges
646 born_charges = None
647 fblocks = step.find('array[@name="born_charges"]')
648 if fblocks is not None:
649 born_charges = np.zeros((natoms, 3, 3), dtype=float)
650 for ii, block in enumerate(fblocks[1:]): # 1. element = dimension
651 for jj, vector in enumerate(block):
652 born_charges[ii, jj] = np.fromstring(vector.text, sep=' ')
654 atoms = atoms_init.copy()
655 atoms.set_cell(cell)
656 atoms.set_scaled_positions(scpos)
657 atoms.calc = SinglePointDFTCalculator(
658 atoms,
659 energy=energy,
660 forces=forces,
661 stress=stress,
662 free_energy=free_energy,
663 ibzkpts=ibz_kpts,
664 efermi=efermi,
665 dipole=dipole,
666 dielectric_tensor=dielectric_tensor,
667 born_effective_charges=born_charges
668 )
669 atoms.calc.name = 'vasp'
670 atoms.calc.kpts = kpoints
671 atoms.calc.parameters = parameters
672 yield atoms
675@writer
676def write_vasp_xdatcar(fd, images, label=None):
677 """Write VASP MD trajectory (XDATCAR) file
679 Only Vasp 5 format is supported (for consistency with read_vasp_xdatcar)
681 Args:
682 fd (str, fp): Output file
683 images (iterable of Atoms): Atoms images to write. These must have
684 consistent atom order and lattice vectors - this will not be
685 checked.
686 label (str): Text for first line of file. If empty, default to list
687 of elements.
689 """
691 images = iter(images)
692 image = next(images)
694 if not isinstance(image, Atoms):
695 raise TypeError("images should be a sequence of Atoms objects.")
697 symbol_count = _symbol_count_from_symbols(image.get_chemical_symbols())
699 if label is None:
700 label = ' '.join([s for s, _ in symbol_count])
701 fd.write(label + '\n')
703 # Not using lattice constants, set it to 1
704 fd.write(' 1\n')
706 # Lattice vectors; use first image
707 float_string = '{:11.6f}'
708 for row_i in range(3):
709 fd.write(' ')
710 fd.write(' '.join(float_string.format(x) for x in image.cell[row_i]))
711 fd.write('\n')
713 fd.write(_symbol_count_string(symbol_count, vasp5=True))
714 _write_xdatcar_config(fd, image, index=1)
715 for i, image in enumerate(images):
716 # Index is off by 2: 1-indexed file vs 0-indexed Python;
717 # and we already wrote the first block.
718 _write_xdatcar_config(fd, image, i + 2)
721def _write_xdatcar_config(fd, atoms, index):
722 """Write a block of positions for XDATCAR file
724 Args:
725 fd (fd): writeable Python file descriptor
726 atoms (ase.Atoms): Atoms to write
727 index (int): configuration number written to block header
729 """
730 fd.write(f"Direct configuration={index:6d}\n")
731 float_string = '{:11.8f}'
732 scaled_positions = atoms.get_scaled_positions()
733 for row in scaled_positions:
734 fd.write(' ')
735 fd.write(' '.join([float_string.format(x) for x in row]))
736 fd.write('\n')
739def _symbol_count_from_symbols(symbols: Symbols) -> List[Tuple[str, int]]:
740 """Reduce list of chemical symbols into compact VASP notation
742 Args:
743 symbols (iterable of str)
745 Returns:
746 list of pairs [(el1, c1), (el2, c2), ...]
748 Example:
749 >>> s = Atoms('Ar3NeHe2ArNe').symbols
750 >>> _symbols_count_from_symbols(s)
751 [('Ar', 3), ('Ne', 1), ('He', 2), ('Ar', 1), ('Ne', 1)]
752 """
753 sc = []
754 psym = str(symbols[0]) # we cast to str to appease mypy
755 count = 0
756 for sym in symbols:
757 if sym != psym:
758 sc.append((psym, count))
759 psym = sym
760 count = 1
761 else:
762 count += 1
764 sc.append((psym, count))
765 return sc
768@writer
769def write_vasp(
770 fd: TextIO,
771 atoms: Atoms,
772 direct: bool = False,
773 sort: bool = False,
774 symbol_count: Optional[List[Tuple[str, int]]] = None,
775 vasp5: bool = True,
776 vasp6: bool = False,
777 ignore_constraints: bool = False,
778 potential_mapping: Optional[dict] = None
779) -> None:
780 """Method to write VASP position (POSCAR/CONTCAR) files.
782 Writes label, scalefactor, unitcell, # of various kinds of atoms,
783 positions in cartesian or scaled coordinates (Direct), and constraints
784 to file. Cartesian coordinates is default and default label is the
785 atomic species, e.g. 'C N H Cu'.
787 Args:
788 fd (TextIO): writeable Python file descriptor
789 atoms (ase.Atoms): Atoms to write
790 direct (bool): Write scaled coordinates instead of cartesian
791 sort (bool): Sort the atomic indices alphabetically by element
792 symbol_count (list of tuples of str and int, optional): Use the given
793 combination of symbols and counts instead of automatically compute
794 them
795 vasp5 (bool): Write to the VASP 5+ format, where the symbols are
796 written to file
797 vasp6 (bool): Write symbols in VASP 6 format (which allows for
798 potential type and hash)
799 ignore_constraints (bool): Ignore all constraints on `atoms`
800 potential_mapping (dict, optional): Map of symbols to potential file
801 (and hash). Only works if `vasp6=True`. See `_symbol_string_count`
803 Raises:
804 RuntimeError: raised if any of these are true:
806 1. `atoms` is not a single `ase.Atoms` object.
807 2. The cell dimensionality is lower than 3 (0D-2D)
808 3. One FixedPlane normal is not parallel to a unit cell vector
809 4. One FixedLine direction is not parallel to a unit cell vector
810 """
811 if isinstance(atoms, (list, tuple)):
812 if len(atoms) > 1:
813 raise RuntimeError(
814 'Only one atomic structure can be saved to VASP '
815 'POSCAR/CONTCAR. Several were given.'
816 )
817 else:
818 atoms = atoms[0]
820 # Check lattice vectors are finite
821 if atoms.cell.rank < 3:
822 raise RuntimeError(
823 'Lattice vectors must be finite and non-parallel. At least '
824 'one lattice length or angle is zero.'
825 )
827 # Write atomic positions in scaled or cartesian coordinates
828 if direct:
829 coord = atoms.get_scaled_positions(wrap=False)
830 else:
831 coord = atoms.positions
833 # Convert ASE constraints to VASP POSCAR constraints
834 constraints_present = atoms.constraints and not ignore_constraints
835 if constraints_present:
836 sflags = _handle_ase_constraints(atoms)
838 # Conditionally sort ordering of `atoms` alphabetically by symbols
839 if sort:
840 ind = np.argsort(atoms.symbols)
841 symbols = atoms.symbols[ind]
842 coord = coord[ind]
843 if constraints_present:
844 sflags = sflags[ind]
845 else:
846 symbols = atoms.symbols
848 # Set or create a list of (symbol, count) pairs
849 sc = symbol_count or _symbol_count_from_symbols(symbols)
851 # Write header as atomic species in `symbol_count` order
852 label = ' '.join(f'{sym:2s}' for sym, _ in sc)
853 fd.write(label + '\n')
855 # For simplicity, we write the unitcell in real coordinates, so the
856 # scaling factor is always set to 1.0.
857 fd.write(f'{1.0:19.16f}\n')
859 for vec in atoms.cell:
860 fd.write(' ' + ' '.join([f'{el:21.16f}' for el in vec]) + '\n')
862 # Write version-dependent species-and-count section
863 sc_str = _symbol_count_string(sc, vasp5, vasp6, potential_mapping)
864 fd.write(sc_str)
866 # Write POSCAR switches
867 if constraints_present:
868 fd.write('Selective dynamics\n')
870 fd.write('Direct\n' if direct else 'Cartesian\n')
872 # Write atomic positions and, if any, the cartesian constraints
873 for iatom, atom in enumerate(coord):
874 for dcoord in atom:
875 fd.write(f' {dcoord:19.16f}')
876 if constraints_present:
877 flags = ['F' if flag else 'T' for flag in sflags[iatom]]
878 fd.write(''.join([f'{f:>4s}' for f in flags]))
879 fd.write('\n')
881 # if velocities in atoms object write velocities
882 if atoms.has('momenta'):
883 cform = 3 * ' {:19.16f}' + '\n'
884 fd.write('Cartesian\n')
885 # unit conversion to Angstrom / fs
886 vel = atoms.get_velocities() / (Ang / fs)
887 for vatom in vel:
888 fd.write(cform.format(*vatom))
891def _handle_ase_constraints(atoms: Atoms) -> np.ndarray:
892 """Convert the ASE constraints on `atoms` to VASP constraints
894 Returns a boolean array with dimensions Nx3, where N is the number of
895 atoms. A value of `True` indicates that movement along the given lattice
896 vector is disallowed for that atom.
898 Args:
899 atoms (Atoms)
901 Returns:
902 boolean numpy array with dimensions Nx3
904 Raises:
905 RuntimeError: If there is a FixedPlane or FixedLine constraint, that
906 is not parallel to a cell vector.
907 """
908 sflags = np.zeros((len(atoms), 3), dtype=bool)
909 for constr in atoms.constraints:
910 if isinstance(constr, FixScaled):
911 sflags[constr.index] = constr.mask
912 elif isinstance(constr, FixAtoms):
913 sflags[constr.index] = 3 * [True]
914 elif isinstance(constr, FixedPlane):
915 # Calculate if the plane normal is parallel to a cell vector
916 mask = np.all(
917 np.abs(np.cross(constr.dir, atoms.cell)) < 1e-5, axis=1
918 )
919 if mask.sum() != 1:
920 raise RuntimeError(
921 'VASP requires that the direction of FixedPlane '
922 'constraints is parallel with one of the cell axis'
923 )
924 sflags[constr.index] = mask
925 elif isinstance(constr, FixedLine):
926 # Calculate if line is parallel to a cell vector
927 mask = np.all(
928 np.abs(np.cross(constr.dir, atoms.cell)) < 1e-5, axis=1
929 )
930 if mask.sum() != 1:
931 raise RuntimeError(
932 'VASP requires that the direction of FixedLine '
933 'constraints is parallel with one of the cell axis'
934 )
935 sflags[constr.index] = ~mask
937 return sflags
940def _symbol_count_string(
941 symbol_count: List[Tuple[str, int]], vasp5: bool = True,
942 vasp6: bool = True, symbol_mapping: Optional[dict] = None
943) -> str:
944 """Create the symbols-and-counts block for POSCAR or XDATCAR
946 Args:
947 symbol_count (list of 2-tuple): list of paired elements and counts
948 vasp5 (bool): if False, omit symbols and only write counts
949 vasp6 (bool): if True, write symbols in VASP 6 format (allows for
950 potential type and hash)
951 symbol_mapping (dict): mapping of symbols to VASP 6 symbols
953 e.g. if sc is [(Sn, 4), (S, 6)] then write for vasp 5:
954 Sn S
955 4 6
957 and for vasp 6 with mapping {'Sn': 'Sn_d_GW', 'S': 'S_GW/357d'}:
958 Sn_d_GW S_GW/357d
959 4 6
960 """
961 symbol_mapping = symbol_mapping or {}
962 out_str = ' '
964 # Allow for VASP 6 format, i.e., specifying the pseudopotential used
965 if vasp6:
966 out_str += ' '.join([
967 f"{symbol_mapping.get(s, s)[:14]:16s}" for s, _ in symbol_count
968 ]) + "\n "
969 out_str += ' '.join([f"{c:16d}" for _, c in symbol_count]) + '\n'
970 return out_str
972 # Write the species for VASP 5+
973 if vasp5:
974 out_str += ' '.join([f"{s:3s}" for s, _ in symbol_count]) + "\n "
976 # Write counts line
977 out_str += ' '.join([f"{c:3d}" for _, c in symbol_count]) + '\n'
979 return out_str