import numpy as np
import sys
import re
import copy
from vasppy import configuration, atom, cell
from pymatgen import Lattice as pmg_Lattice
from pymatgen import Structure as pmg_Structure
from pymatgen.io.cif import CifWriter
from collections import Counter
# Ignore SIG_PIPE and don't throw exceptions on it...
# http://newbebweb.blogspot.co.uk/2012/02/python-head-ioerror-errno-32-broken.html
from signal import signal, SIGPIPE, SIG_DFL
signal( SIGPIPE, SIG_DFL )
[docs]def parity( list ):
return( sum( list )%2 )
[docs]def swap_axes( matrix, axes ):
axes_index = { 'x': 0, 'y': 1, 'z': 2 }
matrix[:, [ axes_index[ axes[ 0 ] ], axes_index[ axes[ 1 ] ] ] ] = matrix[:, [ axes_index[ axes[ 1 ] ], axes_index[ axes[ 0 ] ] ] ]
return matrix
[docs]class Poscar:
lines_offset = 9
def __init__( self ):
self.title = "Title"
self.scaling = 1.0
self.cell = cell.Cell( np.identity( 3 ) )
self.atoms = [ 'A' ]
self.atom_numbers = [ 1 ]
self.coordinate_type = 'Direct'
self.coordinates = np.array( [ [ 0.0, 0.0, 0.0 ] ] )
self.selective_dynamics = False
[docs] def coordinates_by_species( self, species ):
return self.coordinates[ self.range_by_species( species ) ]
[docs] def range_by_species( self, species ):
i = 0
atom_range = {}
for a, n in zip( self.atoms, self.atom_numbers ):
atom_range[ a ] = range(i, i+n)
i += n
return atom_range[ species ]
[docs] def atom_number_by_species( self, species ):
atom_numbers = { a : n for a, n in zip( self.atoms, self.atom_numbers ) }
return atom_numbers[ species ]
[docs] def sorted( self, species ):
new_poscar = copy.deepcopy( self )
new_poscar.atoms = species
new_poscar.atom_numbers = [ self.atom_number_by_species( s ) for s in species ]
new_poscar.coordinates = np.concatenate( [ self.coordinates_by_species( s ) for s in species ], axis = 0 )
return new_poscar
[docs] def read_from( self, filename ):
with open( filename ) as f:
lines = f.readlines()
self.title = lines.pop(0).strip()
self.scaling = float( lines.pop(0).strip() )
self.cell.matrix = np.array( [ [ float( e ) for e in lines.pop(0).split() ] for i in range( 3 ) ] )
self.cell.inv_matrix = np.linalg.inv( self.cell.matrix )
self.atoms = lines.pop(0).split()
self.atom_numbers = [ int(element) for element in lines.pop(0).split() ]
self.coordinate_type = lines.pop(0).strip()
if re.match( r'\A[Ss]', self.coordinate_type ): # test for 'Selective dynamics'
self.selective_dynamics = True
self.coordinate_type = lines.pop( 0 ).strip()
self.coordinates = np.array( [ [ float( e ) for e in lines.pop(0).split()[0:3] ] for i in range( sum( self.atom_numbers ) ) ] )
if self.coords_are_cartesian(): # Convert to direct coordinates
self.coordinates = self.fractional_coordinates()
self.coordinate_type = 'Direct'
[docs] @classmethod
def from_file( cls, filename ):
poscar = cls()
poscar.read_from( filename )
return poscar
[docs] def in_bohr( self ):
new_poscar = copy.deepcopy( self )
bohr_to_angstrom = 0.529177211
new_poscar.scaling *= bohr_to_angstrom
new_poscar.cell.matrix /= bohr_to_angstrom
if new_poscar.coords_are_cartesian():
new_poscar.coordinates /= bohr_to_angstrom
return( new_poscar )
[docs] def coords_are_fractional( self ):
return not self.coords_are_cartesian()
[docs] def coords_are_cartesian( self ):
return re.match( r'\A[CcKk]', self.coordinate_type )
[docs] def fractional_coordinates( self ):
return ( self.coordinates if self.coords_are_fractional() else self.coordinates.dot( np.linalg.inv( self.cell.matrix ) ) )
# return ( self.coordinates if self.coords_are_fractional() else self.cell.cartesian_to_fractional_coordinates( coordinates )
# need to check whether this alternative version works.
[docs] def cartesian_coordinates( self ):
return ( self.coordinates if self.coords_are_cartesian() else self.coordinates.dot( self.cell.matrix ) )
[docs] def select_coordinates( self, coordinate_type='Direct' ):
coord_opts = { 'Direct' : self.fractional_coordinates(),
'Cartesian' : self.cartesian_coordinates() }
return coord_opts[ coordinate_type ]
[docs] def output_coordinates_only( self, coordinate_type='Direct', opts=None ):
prefix = []
suffix = []
for i in range( self.coordinates.shape[0] ):
prefix_string = ''
suffix_string = ''
if 'selective' in opts:
if opts['selective'] == 'T':
suffix_string += ' T T T'
elif opts['selective'] == 'F':
suffix_string += ' F F F'
elif opts['selective']:
raise ValueError
if 'numbered' in opts:
if opts['numbered']:
suffix_string += ' {}'.format(i+1)
if 'label' in opts:
if opts['label']:
if opts['label'] == 1:
prefix_string += self.labels()[i].ljust(6)
elif opts['label'] == 4:
suffix_string += ' {}'.format(self.labels()[i])
elif opts['label']:
raise ValueError( opts['label'] )
prefix.append( prefix_string )
suffix.append( suffix_string )
for pref, coord, suff in zip( prefix, self.select_coordinates( coordinate_type ), suffix ):
print( pref + ''.join( [' {: .10f}'.format( element ) for element in coord ] ) + suff )
[docs] def output( self, coordinate_type='Direct', opts=None ):
if opts is None:
opts = {}
if not opts.get( 'coordinates_only' ):
print( self.title )
print( self.scaling )
[ print( ''.join( [' {: .10f}'.format( element ) for element in row ] ) ) for row in self.cell.matrix ]
print( ' '.join( self.atoms ) )
print( ' '.join( [ str(n) for n in self.atom_numbers ] ) )
if opts['selective']:
print( 'Selective Dynamics' )
print( coordinate_type )
self.output_coordinates_only( coordinate_type=coordinate_type, opts=opts )
[docs] def write_to( self, filename, coordinate_type='Direct', opts=None ):
if opts is None:
opts = {}
with open( filename, 'w' ) as sys.stdout:
self.output( coordinate_type=coordinate_type, opts=opts )
sys.stdout = sys.__stdout__ # make sure sys.stdout is reset
[docs] def output_as_xtl( self ):
print( self.title )
print( "CELL" )
cell_lengths = self.cell_lengths()
cell_angles = self.cell_angles()
cell_data = cell_lengths + cell_angles
print( ''.join( [' {: .8f}'.format( element ) for element in cell_data ] ) )
print( " Symmetry label P1\n\nATOMS\nNAME X Y Z" )
output_opts = { 'label' : True }
self.output_coordinates_only( coordinate_type='Direct', opts = output_opts )
[docs] def output_as_cif( self, symprec = None ):
print( CifWriter( self.to_pymatgen_structure(), symprec ) )
[docs] def output_as_pimaim( self, to_bohr = True ):
if to_bohr is True:
unit_scaling = 0.52918
else:
unit_scaling = 1.0
cell_lengths = self.cell.lengths() * self.scaling / unit_scaling
print( "T\nF\nF\nF" )
for row in self.cell_coordinates():
print( ' '.join( [ str( x / unit_scaling ) for x in row ] ) )
for row in self.cell.unit_vectors().transpose():
print( ' '.join( [ str( x ) for x in row ] ) )
for length in cell_lengths:
print( length )
[docs] def cell_coordinates( self ):
return( self.coordinates * self.cell.lengths() * self.scaling )
[docs] def labels( self ):
return( [ atom_name for ( atom_name, atom_number ) in zip( self.atoms, self.atom_numbers ) for __ in range( atom_number ) ] )
[docs] def replicate( self, h, k, l, group=False ):
lattice_scaling = np.array( [ h, k, l ], dtype=float )
lattice_shift = np.reciprocal( lattice_scaling )
new_poscar = Poscar()
new_poscar.title = self.title
new_poscar.scaling = self.scaling
# print( lattice_scaling.dot( self.lattice ) )
new_poscar.cell.matrix = ( self.cell.matrix.T * lattice_scaling ).T
new_poscar.coordinate_type = self.coordinate_type
new_coordinate_list = []
cell_shift_indices = [ [ a,b,c ] for a in range(h) for b in range(k) for c in range(l) ]
# generate grouped / ungrouped atom names and numbers for supercell
if group:
new_poscar.atoms = [ label + group for label in self.atoms for group in ('a','b') ]
new_poscar.atom_numbers = [ int( num * h * k * l / 2 ) for num in self.atom_numbers for __ in ( 0, 1 )]
else:
new_poscar.atoms = self.atoms
new_poscar.atom_numbers = [ num * h * k * l for num in self.atom_numbers ]
# generate grouped / ungrouped atoms coordinates for supercell
if group:
for row in self.coordinates:
pos_in_origin_cell = row / lattice_scaling
for odd_even in ( 0, 1 ):
for ( a, b, c ) in [ cell_shift for cell_shift in cell_shift_indices if parity(cell_shift) == odd_even ]:
new_coordinate_list.append( [ pos_in_origin_cell + np.array( [ a, b, c ] * lattice_shift ) ][0].tolist() )
else:
for row in self.coordinates:
pos_in_origin_cell = row / lattice_scaling
for ( a, b, c ) in cell_shift_indices:
new_coordinate_list.append( [ pos_in_origin_cell + np.array( [ a, b, c ] * lattice_shift ) ][0].tolist() )
new_poscar.coordinates = np.array( new_coordinate_list )
return new_poscar
[docs] def cell_lengths( self ):
return( ( self.cell.lengths() * self.scaling ).tolist() )
# return [ np.linalg.norm( row ) for row in self.cell.matrix * self.scaling ]
[docs] def cell_angles( self ):
return( self.cell.angles() )
# ( a, b, c ) = [ row for row in self.cell.matrix ]
# return [ angle( b, c ), angle( a, c ), angle( a, b ) ]
[docs] def to_configuration( self ):
atoms = [ atom.Atom( label, coordinates ) for ( label, coordinates ) in zip( self.labels(), self.fractional_coordinates() ) ]
config = configuration.Configuration( cell.Cell( matrix = self.cell.matrix * self.scaling ), atoms )
return( config )
[docs] def swap_axes( self, axes ):
new_poscar = copy.deepcopy( self )
new_poscar.cell = cell.Cell( swap_axes( self.cell.matrix, axes ) )
return new_poscar
[docs] def to_pymatgen_structure( self ):
lattice = pmg_Lattice( self.cell.matrix * self.scaling )
structure = pmg_Structure( lattice, self.labels(), self.coordinates )
return structure
@property
def stoichiometry( self ):
"""
Stoichiometry for this POSCAR, as a Counter.
e.g. AB_2O_4 -> Counter( { 'A': 1, 'B': 2, O: 4 } )
Args:
None
Returns:
None
"""
return Counter( { label: number for label, number in zip( self.atoms, self.atom_numbers ) } )