init
This commit is contained in:
394
electrochemistry/echem/io_data/universal.py
Normal file
394
electrochemistry/echem/io_data/universal.py
Normal file
@@ -0,0 +1,394 @@
|
||||
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))
|
||||
Reference in New Issue
Block a user