395 lines
16 KiB
Python
395 lines
16 KiB
Python
|
import numpy as np
|
||
|
|
from echem.core.constants import ElemNum2Name, ElemName2Num, Bohr2Angstrom, Angstrom2Bohr
|
||
|
|
from echem.core.structure import Structure
|
||
|
|
import warnings
|
||
|
|
|
||
|
|
|
||
|
|
class Cube:
|
||
|
|
def __init__(self,
|
||
|
|
data: np.ndarray,
|
||
|
|
structure: Structure,
|
||
|
|
origin: np.ndarray,
|
||
|
|
units_data: str = 'Bohr',
|
||
|
|
comment: str = None,
|
||
|
|
charges=None,
|
||
|
|
dset_ids=None):
|
||
|
|
self.volumetric_data = data
|
||
|
|
self.structure = structure
|
||
|
|
self.origin = origin
|
||
|
|
self.units_data = units_data
|
||
|
|
|
||
|
|
if comment is None:
|
||
|
|
self.comment = 'Comment is not defined\nGood luck!\n'
|
||
|
|
else:
|
||
|
|
self.comment = comment
|
||
|
|
|
||
|
|
if charges is None:
|
||
|
|
self.charges = np.zeros(structure.natoms)
|
||
|
|
else:
|
||
|
|
self.charges = charges
|
||
|
|
|
||
|
|
self.dset_ids = dset_ids
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
shape = self.volumetric_data.shape
|
||
|
|
return f'{self.comment}\n' + f'NX: {shape[0]}\nNY: {shape[1]}\nNZ: {shape[2]}\n' + \
|
||
|
|
f'Origin:\n{self.origin[0]:.5f} {self.origin[1]:.5f} {self.origin[2]:.5f}\n' + \
|
||
|
|
repr(self.structure)
|
||
|
|
|
||
|
|
def __add__(self, other):
|
||
|
|
assert isinstance(other, Cube), 'Other object must belong to Cube class'
|
||
|
|
assert np.array_equal(self.origin, other.origin), 'Two Cube instances must have the same origin'
|
||
|
|
assert self.volumetric_data.shape == other.volumetric_data.shape, 'Two Cube instances must have ' \
|
||
|
|
'the same shape of volumetric_data'
|
||
|
|
if self.structure != other.structure:
|
||
|
|
warnings.warn('Two Cube instances have different structures. '
|
||
|
|
'The structure will be taken from the 1st (self) instance. '
|
||
|
|
'Hope you know, what you are doing')
|
||
|
|
|
||
|
|
return Cube(self.volumetric_data + other.volumetric_data, self.structure, self.origin)
|
||
|
|
|
||
|
|
def __sub__(self, other):
|
||
|
|
assert isinstance(other, Cube), 'Other object must belong to Cube class'
|
||
|
|
assert np.array_equal(self.origin, other.origin), 'Two Cube instances must have the same origin'
|
||
|
|
assert self.volumetric_data.shape == other.volumetric_data.shape, 'Two Cube instances must have ' \
|
||
|
|
'the same shape of volumetric_data'
|
||
|
|
if self.structure != other.structure:
|
||
|
|
warnings.warn('\nTwo Cube instances have different structures. '
|
||
|
|
'The structure will be taken from the 1st (self) instance. '
|
||
|
|
'Hope you know, what you are doing')
|
||
|
|
|
||
|
|
return Cube(self.volumetric_data - other.volumetric_data, self.structure, self.origin)
|
||
|
|
|
||
|
|
def __neg__(self):
|
||
|
|
return Cube(-self.volumetric_data, self.structure, self.origin)
|
||
|
|
|
||
|
|
def __mul__(self, other):
|
||
|
|
assert isinstance(other, Cube), 'Other object must belong to Cube class'
|
||
|
|
assert np.array_equal(self.origin, other.origin), 'Two Cube instances must have the same origin'
|
||
|
|
assert self.volumetric_data.shape == other.volumetric_data.shape, 'Two Cube instances must have ' \
|
||
|
|
'the same shape of volumetric_data'
|
||
|
|
if self.structure != other.structure:
|
||
|
|
warnings.warn('\nTwo Cube instances have different structures. '
|
||
|
|
'The structure will be taken from the 1st (self) instance. '
|
||
|
|
'Hope you know, what you are doing')
|
||
|
|
|
||
|
|
return Cube(self.volumetric_data * other.volumetric_data, self.structure, self.origin)
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_file(filepath):
|
||
|
|
with open(filepath, 'rt') as file:
|
||
|
|
comment_1 = file.readline()
|
||
|
|
comment_2 = file.readline()
|
||
|
|
comment = comment_1 + comment_2
|
||
|
|
|
||
|
|
line = file.readline().split()
|
||
|
|
natoms = int(line[0])
|
||
|
|
|
||
|
|
if natoms < 0:
|
||
|
|
dset_ids_flag = True
|
||
|
|
natoms = abs(natoms)
|
||
|
|
else:
|
||
|
|
dset_ids_flag = False
|
||
|
|
|
||
|
|
origin = np.array([float(line[1]), float(line[2]), float(line[3])])
|
||
|
|
|
||
|
|
if len(line) == 4:
|
||
|
|
n_data = 1
|
||
|
|
elif len(line) == 5:
|
||
|
|
n_data = int(line[4])
|
||
|
|
|
||
|
|
line = file.readline().split()
|
||
|
|
NX = int(line[0])
|
||
|
|
xaxis = np.array([float(line[1]), float(line[2]), float(line[3])])
|
||
|
|
line = file.readline().split()
|
||
|
|
NY = int(line[0])
|
||
|
|
yaxis = np.array([float(line[1]), float(line[2]), float(line[3])])
|
||
|
|
line = file.readline().split()
|
||
|
|
NZ = int(line[0])
|
||
|
|
zaxis = np.array([float(line[1]), float(line[2]), float(line[3])])
|
||
|
|
|
||
|
|
if NX > 0 and NY > 0 and NZ > 0:
|
||
|
|
units = 'Bohr'
|
||
|
|
elif NX < 0 and NY < 0 and NZ < 0:
|
||
|
|
units = 'Angstrom'
|
||
|
|
else:
|
||
|
|
raise ValueError('The sign of the number of all voxels should be > 0 or < 0')
|
||
|
|
|
||
|
|
if units == 'Angstrom':
|
||
|
|
NX, NY, NZ = -NX, -NY, -NZ
|
||
|
|
|
||
|
|
lattice = np.array([xaxis * NX, yaxis * NY, zaxis * NZ])
|
||
|
|
|
||
|
|
species = []
|
||
|
|
charges = np.zeros(natoms)
|
||
|
|
coords = np.zeros((natoms, 3))
|
||
|
|
|
||
|
|
for atom in range(natoms):
|
||
|
|
line = file.readline().split()
|
||
|
|
species += [ElemNum2Name[int(line[0])]]
|
||
|
|
charges[atom] = float(line[1])
|
||
|
|
coords[atom, :] = line[2:]
|
||
|
|
|
||
|
|
if units == 'Bohr':
|
||
|
|
lattice = Bohr2Angstrom * lattice
|
||
|
|
coords = Bohr2Angstrom * coords
|
||
|
|
origin = Bohr2Angstrom * origin
|
||
|
|
|
||
|
|
structure = Structure(lattice, species, coords, coords_are_cartesian=True)
|
||
|
|
|
||
|
|
dset_ids = None
|
||
|
|
dset_ids_processed = -1
|
||
|
|
if dset_ids_flag is True:
|
||
|
|
dset_ids = []
|
||
|
|
line = file.readline().split()
|
||
|
|
n_data = int(line[0])
|
||
|
|
if n_data < 1:
|
||
|
|
raise ValueError(f'Bad value of n_data: {n_data}')
|
||
|
|
dset_ids_processed += len(line)
|
||
|
|
dset_ids += [int(i) for i in line[1:]]
|
||
|
|
while dset_ids_processed < n_data:
|
||
|
|
line = file.readline().split()
|
||
|
|
dset_ids_processed += len(line)
|
||
|
|
dset_ids += [int(i) for i in line]
|
||
|
|
dset_ids = np.array(dset_ids)
|
||
|
|
|
||
|
|
if n_data != 1:
|
||
|
|
raise NotImplemented(f'The processing of cube files with more than 1 data values is not implemented.'
|
||
|
|
f' n_data = {n_data}')
|
||
|
|
|
||
|
|
data = np.zeros((NX, NY, NZ))
|
||
|
|
indexes = np.arange(0, NX * NY * NZ)
|
||
|
|
indexes_1 = indexes // (NY * NZ)
|
||
|
|
indexes_2 = (indexes // NZ) % NY
|
||
|
|
indexes_3 = indexes % NZ
|
||
|
|
|
||
|
|
i = 0
|
||
|
|
for line in file:
|
||
|
|
for value in line.split():
|
||
|
|
data[indexes_1[i], indexes_2[i], indexes_3[i]] = float(value)
|
||
|
|
i += 1
|
||
|
|
|
||
|
|
return Cube(data, structure, origin, units, comment, charges, dset_ids)
|
||
|
|
|
||
|
|
#def reduce(self, factor):
|
||
|
|
# from skimage.measure import block_reduce
|
||
|
|
# try:
|
||
|
|
# volumetric_data_reduced = block_reduce(self.volumetric_data, block_size=(factor, factor, factor), func=np.mean)
|
||
|
|
# Ns_reduced = np.shape(volumetric_data_reduced)
|
||
|
|
# except:
|
||
|
|
# raise ValueError('Try another factor value')
|
||
|
|
# return Cube(volumetric_data_reduced, self.structure, self.comment, Ns_reduced, self.charges)
|
||
|
|
|
||
|
|
def to_file(self, filepath, units='Bohr'):
|
||
|
|
if not self.structure.coords_are_cartesian:
|
||
|
|
self.structure.mod_coords_to_cartesian()
|
||
|
|
|
||
|
|
Ns = np.array(self.volumetric_data.shape)
|
||
|
|
width_Ni = len(str(np.max(Ns)))
|
||
|
|
if units == 'Angstrom':
|
||
|
|
Ns = - Ns
|
||
|
|
width_Ni += 1
|
||
|
|
width_lattice = len(str(int(np.max(self.structure.lattice)))) + 7
|
||
|
|
width_coord = len(str(int(np.max(self.structure.coords)))) + 7
|
||
|
|
elif units == 'Bohr':
|
||
|
|
lattice = self.get_voxel() * Angstrom2Bohr
|
||
|
|
coords = self.structure.coords * Angstrom2Bohr
|
||
|
|
origin = self.origin * Angstrom2Bohr
|
||
|
|
width_lattice = len(str(int(np.max(lattice)))) + 7
|
||
|
|
width_coord = len(str(int(np.max(coords)))) + 7
|
||
|
|
else:
|
||
|
|
raise ValueError(f'Irregular units flag: {units}. Units must be \'Bohr\' or \'Angstrom\'')
|
||
|
|
|
||
|
|
if np.sum(self.structure.lattice < 0):
|
||
|
|
width_lattice += 1
|
||
|
|
if np.sum(self.structure.coords < 0):
|
||
|
|
width_coord += 1
|
||
|
|
width = np.max([width_lattice, width_coord])
|
||
|
|
|
||
|
|
if self.dset_ids is not None:
|
||
|
|
natoms = - self.structure.natoms
|
||
|
|
else:
|
||
|
|
natoms = self.structure.natoms
|
||
|
|
width_natoms = len(str(natoms))
|
||
|
|
width_1_column = max(width_Ni, width_natoms)
|
||
|
|
|
||
|
|
with open(filepath, 'w') as file:
|
||
|
|
file.write(self.comment)
|
||
|
|
|
||
|
|
if units == 'Angstrom':
|
||
|
|
file.write(f' {natoms:{width_1_column}} {self.origin[0]:{width}.6f} '
|
||
|
|
f' {self.origin[1]:{width}.6f} {self.origin[2]:{width}.6f}\n')
|
||
|
|
for N_i, lattice_vector in zip(Ns, self.get_voxel()):
|
||
|
|
file.write(f' {N_i:{width_1_column}} {lattice_vector[0]:{width}.6f} '
|
||
|
|
f' {lattice_vector[1]:{width}.6f} {lattice_vector[2]:{width}.6f}\n')
|
||
|
|
for atom_name, charge, coord in zip(self.structure.species, self.charges, self.structure.coords):
|
||
|
|
file.write(
|
||
|
|
f' {ElemName2Num[atom_name]:{width_1_column}} {charge:{width}.6f} '
|
||
|
|
f' {coord[0]:{width}.6f} {coord[1]:{width}.6f} {coord[2]:{width}.6f}\n')
|
||
|
|
|
||
|
|
elif units == 'Bohr':
|
||
|
|
file.write(f' {natoms:{width_1_column}} {origin[0]:{width}.6f} '
|
||
|
|
f' {origin[1]:{width}.6f} {origin[2]:{width}.6f}\n')
|
||
|
|
for N_i, lattice_vector in zip(Ns, lattice):
|
||
|
|
file.write(f' {N_i:{width_1_column}} {lattice_vector[0]:{width}.6f} '
|
||
|
|
f' {lattice_vector[1]:{width}.6f} {lattice_vector[2]:{width}.6f}\n')
|
||
|
|
for atom_name, charge, coord in zip(self.structure.species, self.charges, coords):
|
||
|
|
file.write(
|
||
|
|
f' {ElemName2Num[atom_name]:{width_1_column}} {charge:{width}.6f} '
|
||
|
|
f' {coord[0]:{width}.6f} {coord[1]:{width}.6f} {coord[2]:{width}.6f}\n')
|
||
|
|
|
||
|
|
else:
|
||
|
|
raise ValueError(f'Irregular units flag: {units}. Units must be \'Bohr\' or \'Angstrom\'')
|
||
|
|
|
||
|
|
if self.dset_ids is not None:
|
||
|
|
m = len(self.dset_ids)
|
||
|
|
file.write(f' {m:{width_1_column}}' + ' ')
|
||
|
|
for dset_id in self.dset_ids:
|
||
|
|
file.write(str(dset_id) + ' ')
|
||
|
|
file.write('\n')
|
||
|
|
|
||
|
|
for i in range(abs(Ns[0])):
|
||
|
|
for j in range(abs(Ns[1])):
|
||
|
|
for k in range(abs(Ns[2])):
|
||
|
|
file.write(str(' %.5E' % self.volumetric_data[i][j][k]))
|
||
|
|
if k % 6 == 5:
|
||
|
|
file.write('\n')
|
||
|
|
file.write('\n')
|
||
|
|
|
||
|
|
def mod_to_zero_origin(self):
|
||
|
|
self.structure.coords -= self.origin
|
||
|
|
self.origin = np.zeros(3)
|
||
|
|
|
||
|
|
def get_average_along_axis(self, axis):
|
||
|
|
"""
|
||
|
|
Gets average value along axis
|
||
|
|
Args:
|
||
|
|
axis (int):
|
||
|
|
if 0 than average along x wil be calculated
|
||
|
|
if 1 along y
|
||
|
|
if 2 along z
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
np.array of average value along selected axis
|
||
|
|
"""
|
||
|
|
if axis == 2:
|
||
|
|
return np.mean(self.volumetric_data, (0, 1))
|
||
|
|
elif axis == 1:
|
||
|
|
return np.mean(self.volumetric_data, (0, 2))
|
||
|
|
elif axis == 0:
|
||
|
|
return np.mean(self.volumetric_data, (1, 2))
|
||
|
|
else:
|
||
|
|
raise ValueError('axis can be only 0, 1 or 2')
|
||
|
|
|
||
|
|
def get_average_along_axis_max(self, axis: int, scale=None):
|
||
|
|
"""Calculate the vacuum level (the maximum planar average value along selected axis)
|
||
|
|
|
||
|
|
Args:
|
||
|
|
axis (int): The axis number along which the planar average is calculated. The first axis is 0
|
||
|
|
scale (float): The value that is multiplying by the result. It's used for converting between
|
||
|
|
different units
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
(float): The vacuum level multiplied by scale factor
|
||
|
|
|
||
|
|
"""
|
||
|
|
avr = self.get_average_along_axis(axis)
|
||
|
|
if scale is None:
|
||
|
|
return np.max(avr)
|
||
|
|
else:
|
||
|
|
return scale * np.max(avr)
|
||
|
|
|
||
|
|
def get_voxel(self, units='Angstrom'):
|
||
|
|
NX, NY, NZ = self.volumetric_data.shape
|
||
|
|
voxel = self.structure.lattice.copy()
|
||
|
|
voxel[0] /= NX
|
||
|
|
voxel[1] /= NY
|
||
|
|
voxel[2] /= NZ
|
||
|
|
if units == 'Angstrom':
|
||
|
|
return voxel
|
||
|
|
elif units == 'Bohr':
|
||
|
|
return voxel * Angstrom2Bohr
|
||
|
|
else:
|
||
|
|
raise ValueError('units can be \'Angstrom\' or \'Bohr\'')
|
||
|
|
|
||
|
|
def get_integrated_number(self):
|
||
|
|
if self.units_data == 'Bohr':
|
||
|
|
voxel_volume = np.linalg.det(self.get_voxel(units='Bohr'))
|
||
|
|
return voxel_volume * np.sum(self.volumetric_data)
|
||
|
|
else:
|
||
|
|
raise NotImplemented()
|
||
|
|
|
||
|
|
def assign_top_n_data_to_atoms(self, n_top, r):
|
||
|
|
"""Assign top n abs of volumetric data to atoms. Might be used to assign electron density to atoms.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
n_top (int): Number of voxels that will be analysed
|
||
|
|
r (float): Radius. A voxel is considered belonging to atom is the distance between the voxel center and ]
|
||
|
|
atom is less than r.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
(np.ndarray): Array of boolean values. I-th raw represents i-th atom, j-th column represents j-th voxel
|
||
|
|
"""
|
||
|
|
sorted_indices = np.array(np.unravel_index(np.argsort(-np.abs(self.volumetric_data), axis=None),
|
||
|
|
self.volumetric_data.shape)).T
|
||
|
|
translation_vector = np.sum(self.structure.lattice, axis=0)
|
||
|
|
voxels_centres = sorted_indices[:n_top, :] * translation_vector + translation_vector / 2 + self.origin
|
||
|
|
|
||
|
|
atom_indices = list(range(self.structure.natoms))
|
||
|
|
|
||
|
|
if self.structure.natoms == 1:
|
||
|
|
return np.linalg.norm(voxels_centres - self.structure.coords[0], axis=-1) < r
|
||
|
|
else:
|
||
|
|
return np.linalg.norm(np.broadcast_to(voxels_centres, (self.structure.natoms,) + voxels_centres.shape) -
|
||
|
|
np.expand_dims(self.structure.coords[atom_indices], axis=1), axis=-1) < r
|
||
|
|
|
||
|
|
|
||
|
|
class Xyz:
|
||
|
|
def __init__(self, structure, comment):
|
||
|
|
self.structure = structure
|
||
|
|
self.comment = comment
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_file(filepath):
|
||
|
|
with open(filepath, 'rt') as file:
|
||
|
|
natoms = int(file.readline().strip())
|
||
|
|
comment = file.readline()
|
||
|
|
|
||
|
|
coords = np.zeros((natoms, 3))
|
||
|
|
species = []
|
||
|
|
for i in range(natoms):
|
||
|
|
line = file.readline().split()
|
||
|
|
species.append(line[0])
|
||
|
|
coords[i] = [float(j) for j in line[1:]]
|
||
|
|
|
||
|
|
struct = Structure(np.zeros((3, 3)), species, coords, coords_are_cartesian=True)
|
||
|
|
|
||
|
|
return Xyz(struct, comment)
|
||
|
|
|
||
|
|
|
||
|
|
class XyzTrajectory:
|
||
|
|
def __init__(self, first_xyz, trajectory):
|
||
|
|
self.first_xyz = first_xyz
|
||
|
|
self.trajectory = trajectory
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def from_file(filepath):
|
||
|
|
first_xyz = Xyz.from_file(filepath)
|
||
|
|
|
||
|
|
trajectory = []
|
||
|
|
with open(filepath, 'rt') as file:
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
natoms = int(file.readline().strip())
|
||
|
|
except:
|
||
|
|
break
|
||
|
|
file.readline()
|
||
|
|
|
||
|
|
coords = np.zeros((natoms, 3))
|
||
|
|
for i in range(natoms):
|
||
|
|
line = file.readline().split()
|
||
|
|
coords[i] = [float(j) for j in line[1:]]
|
||
|
|
trajectory.append(coords)
|
||
|
|
|
||
|
|
return XyzTrajectory(first_xyz, np.array(trajectory))
|