#!/usr/bin/env python
from __future__ import annotations
import sys
import re
import numpy as np
import argparse
[docs]class Poscar:
"""A low-level class to store a crystal structure input (tailored for
VASP).
If modified manually, the 'Cartesian' and 'direct' coords must be
updated together all the time by using `_set_cartesian` or
`_set_direct`.
The scaling factor, is internally set to 1.0, always. ie, the
scaling is included into the lattice.
The units are Angstroms.
Methods:
parse(self) # loads the whole file
_set_cartesian(self) # set direct -> cartesian (internal)
_set_direct(self) # set cartesian -> direct (internal)
_unparse(self, direct) # data to string, use `write()` instead
write(self, filename, direct) # saves the class to a POSCAR-like file
xyz(self, filename) # saves a xyz file from the data
sort(self) # sorts the atoms, grouping them by element
remove() # removes one or more atoms from poscar
add(position, element, direct) # add one atom, only one each the time
"""
[docs] def __init__(self,
filename:str='POSCAR',
verbose:bool=False):
"""The file is not automatically loaded, you need to run
`self.parse()`
Parameters
----------
filename : str,
POSCAR-like file
verbose: bool,
verbosity level, for debugging. Default: False
Attributes
----------
self.verbose
self.filename
self.poscar:str str with the full POSCAR file
self.cpos: np.nadarray, Cartesian coordinates
self.dpos: np.ndarray, direct coordinates
self.lat: np.ndarray (3x3), the lattice
self.typeSp: list, aame of atomic species, ordered
self.numberSp: np.ndarray, number of atoms per specie
self.Ntotal: int, total number of atoms in system
self.elm: list, element of each atoms one-by-one.
self.selective: bool, Selective dynamics?
self.selectFlags: None or nd.array(str), flags of selective dynamics
self.flags: dict, list of flags. Not used here just for convenience
self.volume: float, the box product of the lattice
"""
self.verbose = verbose
self.filename = filename
self.poscar:str = None
self.cpos:np.ndarray = None # cartesian coordinates
self.dpos:np.ndarray = None # direct coordinates
self.lat:np.ndarray = None # lattice
self.typeSp:List[str] = None # Name of atomic species
self.numberSp:nd.array = None # Number of atoms per specie
self.Ntotal:int = None # Total atoms in system
self.elm:List[str] = None # Element of each atoms one-by-one.
self.selective:bool = None # Selective dynamics
self.selectFlags:np.ndarray = None # all the T,F from selective dynamics
self.flags:dict = {} # list of flags, not used here just for convenience
self.volume:float = None
self.loaded:bool = False # was the POSCAR-file loaded? i.e. self.parse()
return
[docs] def parse(self, fromString:str = None):
"""Loads into memory all the content of the POSCAR file.
Parameters
----------
fromString : str
If present, instead of loading a file, it uses
this variable to populate the class. Default=None
"""
if fromString and isinstance(fromString, str):
self.poscar = fromString.split('\n')
elif fromString and isinstance(fromString, list):
self.poscar = fromString
else:
self.poscar = open(self.filename, 'r')
self.poscar = self.poscar.readlines()
# getting the scale factor
scale = re.findall(r'[.\deE]+', self.poscar[1])
scale = float(scale[0])
if self.verbose:
print('scaling factor: ', scale)
# parsing the lattice
self.lat = re.findall(r'[-.\deE]+\s+[-.\deE]+\s+[-.\deE]+\s*\n*', ' '.join(self.poscar[2:5]))
self.lat = np.array([x.split() for x in self.lat], dtype=float)*scale
if self.verbose:
print( 'lattice:\n', self.lat)
# type of atoms and number of atoms per type
self.typeSp = re.findall(r'(\w+)\s*', self.poscar[5])
if len(self.typeSp) == 0:
raise RuntimeError("No data about the atomic species found. Correct it.")
if self.verbose:
print( 'atoms per species:\n', self.typeSp)
self.numberSp = re.findall(r'(\d+)\s*', self.poscar[6])
self.numberSp = np.array(self.numberSp, dtype=int)
self.Ntotal = np.sum(self.numberSp)
if self.verbose:
print( 'atomic species:\n', self.numberSp)
print( 'The total is ' + str(self.Ntotal) + ' atoms')
# selective dynamics? setting a offset
if re.findall(r'^\s*[sS]', self.poscar[7]) :
self.selective = True
offset = 1
if self.verbose:
print( "Selective Dynamics")
else:
self.selective = False
offset = 0
# Direct coordinates unless otherwise
direct = True
if not re.findall(r'^\s*[Dd]', self.poscar[7+offset]) :
direct = False
if self.verbose:
print('Positions in Cartesian coordinates')
else:
if self.verbose:
print('Positions in Direct coordinates')
# parsing the positions
start, end = 8+offset, 8+offset+ self.Ntotal
string = r'([-.\deE]+)\s+([-.\deE]+)\s+([-.\deE]+)\s*'
pos = re.findall(string, ' '.join(self.poscar[start:end]))
if direct:
self.dpos = np.array(pos, dtype=float)
self._set_cartesian()
else:
self.cpos = np.array(pos, dtype=float)
self._set_direct()
if self.verbose:
print( "Atomic positions (direct)\n", self.dpos)
print( "Atomic positions (cartesian)\n", self.cpos)
# parsing selective dynamics
if self.selective == True:
string = r'([TF]+)\s+([TF]+)\s+([TF]+)\s+'
selectFlags = re.findall(string, ' '.join(self.poscar[start:end]))
self.selectFlags = np.array(selectFlags)
if self.verbose:
print('Flags of selective dynamics:\n', self.selectFlags)
# setting a list of elements:
elementList = zip(self.typeSp, self.numberSp)
self.elm = ' '.join([' '.join([x]*y) for x,y in elementList]).split()
if self.verbose:
print('Elements: ', self.elm)
# setting the volume, just as an utility
self.volume = np.linalg.det(self.lat)
self.loaded = True
# empty list as flags, one per atom
for i in range(self.Ntotal):
self.flags[i] = {}
return
[docs] def load_from_data(self,
direct_positions : np.ndarray,
lattice : np.ndarray,
elements : List[str]
):
"""
It loades the Poscar class with essencial data.
Parameters
----------
direct_pos : np.ndarray
atomic positions in direct (fractional) coordiantes. Size [Natoms:3]
lattice : np.ndarray
Lattice vectors [3:3], in *Angstroms*
elements : List[str]
A list of atomic symbols, with the same order as the `direct_positions`
"""
self.lat = lattice
self.dpos = direct_positions
self.elm = elements
self._set_cartesian()
typeSp = []
elements = list(elements)
for element in elements:
if element not in typeSp:
typeSp.append(element)
self.typeSp = typeSp
numberSp = []
for item in typeSp:
N = elements.count(item)
numberSp.append(N)
self.numberSp = numberSp
self.Ntotal = sum(numberSp)
self.volume = np.linalg.det(self.lat)
self.loaded = True
return
def _set_cartesian(self):
"""set the cartesian positions (self.cpos) from direct positions (self.dpos).
"""
cart = np.dot(self.lat.T, self.dpos.T)
cart = cart.T
self.cpos = cart
def _set_direct(self):
"""set the direct positions (self.dpos) from Cartesian positions (self.cpos).
"""
inverse = np.linalg.inv(self.lat)
direct = np.dot(inverse.T, self.cpos.T)
direct = direct.T
self.dpos = direct
def _unparse(self, direct:bool=True):
"""Internal method to be used previously to to writing a POSCAR
file. It group together all the information in a single str,
`self.poscar`. The information is as it is. No PBC are applied,
and no checks are performed at this stage. The scaling factor is
1.0, always.
Parameters
----------
direct: bool
direct positons is True, Cartesian is Falsepositions. Default is True
"""
# We will start with getting the positions
if direct == True:
pos = self.dpos
else:
pos = self.cpos
# creating a list of text lines with positions
pos = [' '.join([str(coord) for coord in line]) for line in pos]
# Now we will look whether selective dynamics are used
if self.selective == True:
# a list of text lines with flags
flags = [' '.join([flag for flag in line]) for line in pos]
pos = [pos + ' ' + flag for (pos, flag) in zip(pos,flags)]
pos = '\n'.join(pos)
# Creating the POSCAR text string
self.poscar = "poscar.py\n"
self.poscar += "1.0\n"
self.poscar += '\n'.join([' '.join([str(y) for y in x]) for x in self.lat]) + '\n'
self.poscar += ' '.join(self.typeSp) + '\n'
self.poscar += ' '.join([str(x) for x in self.numberSp]) + '\n'
if self.selective:
self.poscar += 'Selective Dynamics\n'
if direct == False:
self.poscar += 'Cartesian\n'
else:
self.poscar += 'Direct\n'
self.poscar += pos # already set with the the T, F -if needed
self.poscar += '\n'
if self.verbose:
print("\n\n unparsed POSCAR\n\n")
print('unparse, self.poscar\n', self.poscar)
[docs] def write(self, filename:str='POSCAR.out', direct:bool=True):
"""Writes a poscar file with the information stored in the class.
Parameters
----------
filename : str
default='POSCAR.out', name of the output file.
direct : bool
direct positons is True, Cartesian is Falsepositions. Default is True
"""
self._unparse(direct=direct)
fout = open(filename, 'w')
fout.write(self.poscar)
if self.verbose:
print('File ' + filename + ' written.')
return
[docs] def xyz(self, filename:str):
"""Writes an xyz file, the lattice is written as a comment line
Parameters
----------
filename: str
the name of the .xyz file, The .xyz extension is not automatically added
"""
xyzf = open(filename, 'w')
xyzf.write(str(self.Ntotal) + '\n')
# The comment line has the lettice
latStr = np.array(self.lat, dtype=str).flatten()
latStr = ' '.join(latStr)
xyzf.write('Lattice="' + latStr + '"\n')
# continuing with the positions
pos = self.cpos
pos = np.array(pos, dtype=str)
pos = [' '.join(x) for x in pos]
elm = list(self.elm)
xyzstr = '\n'.join( [ x + ' ' + y for x,y in zip(elm, pos) ] )
xyzf.write(xyzstr)
xyzf.write('\n')
xyzf.close()
if self.verbose:
print(filename + ' written as xyz')
[docs] def sort(self):
"""This method updates the internal arrays related to elements and
atoms per element. Automatically used when using `self.add`
"""
#
# self.typeSp, self.elem, self.dpos
# must be present and updated (they can be disordered)
#
from collections import OrderedDict
# getting the different element's names, without repetitions
self.typeSp = list(OrderedDict.fromkeys(self.typeSp))
if self.verbose:
print('The list of elements is ', self.typeSp)
# lists of ordered atoms
atoms = []
elements = []
if self.verbose:
print('to sort: ', self.elm, '\n', self.dpos)
# ordering the positions accordiong to its element.
for thiselem in self.typeSp:
# Joining each element with its position
for elem, pos in zip(self.elm, self.dpos):
if thiselem == elem:
atoms.append(pos)
elements.append(elem)
if self.verbose:
print('sorted: ', elements, '\n', np.array(atoms))
self.elm = elements
# setting the position's list in direct coords.
self.dpos = np.array(atoms)
# and in cartesian coords as well
self._set_cartesian()
# How many atoms by element?
from collections import Counter
# a dict of elements and its repetitions
counter = Counter(elements)
self.numberSp = [counter[x] for x in self.typeSp]
self.Ntotal = sum(self.numberSp)
if self.verbose:
print ("N atoms per specie, ", self.numberSp, '. Total: ', self.Ntotal)
[docs] def remove(self, atoms:list|int):
"""
Remove one or more atoms.
Parameters:
atoms : int|list
removes the atom(s) with given indexes (0-based)
"""
# atoms maybe (or not) just one atom (an int, not a one-sized list)
if self.verbose:
print('going to delete the following atom(s):', atoms)
try:
iterator = iter(atoms)
except TypeError:
atoms = [atoms]
# we will populate a list with the atoms to keep
keep = [True]*self.Ntotal
for atom in atoms:
if atom >= self.Ntotal:
raise RuntimeError('Error: atom index is larger than the atom number')
keep[atom] = False
# creating new arrays without the removed elements
self.cpos = self.cpos[keep]
self.dpos = self.dpos[keep]
self.Ntotal = len(self.cpos)
# self.elm is a list, I am not sure why, but I will cast it back to list
self.elm = np.array(self.elm)
self.elm = self.elm[keep]
self.elm = list(self.elm)
if self.selective:
self.selectFlags = self.selectFlags[keep]
# the atoms types and their number can be modified, we need to
# count them from self.elm
# Also I want to keep the order of elements
from collections import OrderedDict
self.typeSp = list(OrderedDict.fromkeys(self.elm).keys())
self.numberSp = [self.elm.count(x) for x in self.typeSp]
if self.verbose:
print('Elements', self.elm)
print(self.typeSp, self.numberSp)
# no need to sort, deleting doesn't alters the occurence
return
[docs] def add(self,
position:np.ndarray,
element:str,
direct:bool=True,
selectiveFlags:np.ndarray=None):
"""Adds one atom to the class. Only one atom at the time
Parameters
----------
position : np.ndarray
(3 or 1x3) the positions of new atom
element : str
Name of the element of the new atom
direct : bool
are the positions given direct (True) or Cartesian (False)
coordinates? Default is True
selectiveFlags : np.ndarray(str)
only of `self.selective == True`
"""
position = np.array(position, dtype=float)
position.shape = (1,3)
if self.verbose:
print('going to add an ' +element+ ' atom at', position, end=',')
if direct:
print('in direct coordinates')
else:
print('in Cartesian coordiantes')
# setting the data
if direct:
self.dpos = np.concatenate((self.dpos, position))
else:
self.cpos = np.concatenate((self.cpos, position))
self.Ntotal = self.Ntotal + 1
self.elm.append(element)
if self.selective and selectiveFlags:
if self.verbose:
print('selective flag found.')
selectiveFlags = np.array(selectiveFlags, dtype=str)
self.selectFlags = np.concatenate(self.selectFlags, selectiveFlags)
# setting the other data
if direct:
self._set_cartesian()
else:
self._set_direct()
# the atoms types and their number can be modified, we need to
# count them from self.elm
# Also I want to keep the order of elements
from collections import OrderedDict
self.typeSp = list(OrderedDict.fromkeys(self.elm).keys())
self.numberSp = [self.elm.count(x) for x in self.typeSp]
if self.verbose:
print('Elements', self.elm)
print(self.typeSp, self.numberSp)
# sorting the data
self.sort()
return
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("inputfile", type=str, help="input file")
parser.add_argument('-v', '--verbose', action='store_true')
parser.add_argument('--xyz', action='store_true')
args = parser.parse_args()
p = Poscar(args.inputfile, verbose=args.verbose)
p.parse()
if args.xyz:
p.xyz(args.inputfile + '.xyz')