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 = list(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