import numpy as np
import sys
import math
from vasppy import poscar, cell
[docs]def interpolate( i, j, x ):
return( ( i * ( 1.0 - x ) ) + ( j * x) )
[docs]def trilinear_interpolation( cube, r ):
return( interpolate (
interpolate(
interpolate( cube[ 0, 0, 0 ], cube[ 1, 0, 0 ], r[ 0 ] ), # trilinear interpolation => http://en.wikipedia.org/wiki/Trilinear_interpolation
interpolate( cube[ 0, 1, 0 ], cube[ 1, 1, 0 ], r[ 0 ] ),
r[ 1 ] ),
interpolate(
interpolate( cube[ 0, 0, 1 ], cube[ 1, 0, 1 ], r[ 0 ] ),
interpolate( cube[ 0, 1, 1 ], cube[ 1, 1, 1 ], r[ 0 ] ),
r[ 1 ] ),
r[ 2 ] )
)
[docs]class Grid:
projections = { 'x' : 0, 'y' : 1, 'z' : 2 }
def __init__( self, dimensions = [ 1, 1, 1 ] ):
self.filename = None
self.poscar = poscar.Poscar()
self.number_of_header_lines = 0
self.dimensions = dimensions
self.spacing = np.array( [ 1.0 / number_of_points for number_of_points in self.dimensions ] )
self.grid = np.zeros( self.dimensions )
[docs] def read_from_filename( self, filename ):
self.filename = filename
self.poscar = poscar.Poscar()
self.poscar.read_from( self.filename )
self.number_of_header_lines = sum( self.poscar.atom_numbers ) + poscar.Poscar.lines_offset
self.read_dimensions()
self.read_grid()
return self
[docs] def write_to_filename( self, filename ):
with open( filename, 'w' ) as file_out:
sys.stdout = file_out
self.poscar.output()
self.write_dimensions()
sys.stdout.flush()
self.write_grid()
[docs] def read_dimensions( self ):
with open( self.filename, 'r' ) as file_in:
for i, line in enumerate( file_in ):
if i == self.number_of_header_lines:
self.dimensions = [ int(i) for i in line.split() ]
break
[docs] def write_dimensions( self ):
print( "\n" + ' '.join( [ str(i) for i in self.dimensions ] ) )
[docs] def read_grid( self ):
grid_data = []
grid_data_lines = math.ceil( ( self.dimensions[0] * self.dimensions[1] * self.dimensions[2] ) / 5 )
with open( self.filename ) as file_in:
for i, line in enumerate( file_in ):
if ( i > self.number_of_header_lines ) and ( i <= self.number_of_header_lines + grid_data_lines ):
grid_data.append( line.strip() )
grid_data = np.array( [ float( s ) for s in ' '.join( grid_data ).split() ] )
self.grid = np.reshape( grid_data, tuple( self.dimensions ), order = 'F' )
[docs] def write_grid( self ):
np.savetxt( sys.stdout.buffer, np.swapaxes( self.grid, 0, 2 ).reshape( -1, 5 ), fmt='%.11E' )
[docs] def average( self, normal_axis_label ):
axes = [ 0, 1, 2 ]
axes.remove( Grid.projections[ normal_axis_label ] )
return( np.sum( np.sum( self.grid, axis=axes[1] ), axis=axes[0] ) / ( self.dimensions[0] * self.dimensions[1] ) )
[docs] def by_index( self, index ):
return self.grid[ index[0], index[1], index[2] ]
[docs] def fractional_coordinate_at_index( self, index ):
return( np.multiply( self.spacing, index ) )
[docs] def cartesian_coordinate_at_index( self, index ):
return( self.fractional_coordinate_at_index( index ).dot( self.poscar.cell.matrix ) )
[docs] def cube_slice( self, x0, y0, z0 ):
x1 = ( x0 + 1 ) % self.dimensions[0]
y1 = ( y0 + 1 ) % self.dimensions[1]
z1 = ( z0 + 1 ) % self.dimensions[2]
cube = np.array( [ self.grid[ x0, y0, z0 ],
self.grid[ x0, y0, z1 ],
self.grid[ x0, y1, z0 ],
self.grid[ x0, y1, z1 ],
self.grid[ x1, y0, z0 ],
self.grid[ x1, y0, z1 ],
self.grid[ x1, y1, z0 ],
self.grid[ x1, y1, z1 ] ] ).reshape( 2, 2, 2 )
return( cube )
[docs] def interpolated_value_at_fractional_coordinate( self, coord ):
point = np.multiply( np.array( self.dimensions ), coord ) # point contains the (fractional) index of the coordinate coord.
origin = [ int( f ) for f in point ] # origin contains the 3D index of the lowest-index point in the cube surrounding point (i,j,k)
delta = [ p - o for p, o in zip( point, origin ) ] # delta contains the *fractional* offset of "point" from "origin"
cube = self.cube_slice( *origin ) # cube contains the data values at the 8 bounding grid points
return( trilinear_interpolation( cube, delta ) )
[docs] def interpolate_to_orthorhombic_grid( self, dimensions ): # warning. This may need a more robust minimim image function in Cell.py for highly non-orthorhombic cells
old_grid = self
old_cell = old_grid.poscar.cell
new_grid = Grid( dimensions = dimensions )
new_grid.poscar.cell = cell.Cell( np.diag( np.diag( self.poscar.cell.matrix ) ) )
index_grid = np.array( [ [ i, j, k ] for ( i, j, k ), value in np.ndenumerate( new_grid.grid ) ] )
cart_coord_grid = np.array( [ new_grid.cartesian_coordinate_at_index( index ) for index in index_grid ] )
init_frac_coord_grid = np.array( [ old_cell.cartesian_to_fractional_coordinates( r ) for r in cart_coord_grid ] )
frac_coord_grid = np.array( [ old_cell.inside_cell( r ) for r in init_frac_coord_grid ] )
new_grid_data = np.array( [ old_grid.interpolated_value_at_fractional_coordinate( r ) for r in frac_coord_grid ] )
new_grid.grid = new_grid_data.reshape( new_grid.dimensions )
return( new_grid )