init

electrochemistry/echem/core/constants.py

@@ -0,0 +1,39 @@
+ElemNum2Name = {1: 'H', 2: 'He', 3: 'Li', 4: 'Be', 5: 'B', 6: 'C', 7: 'N', 8: 'O', 9: 'F', 10: 'Ne',
+                11: 'Na', 12: 'Mg', 13: 'Al', 14: 'Si', 15: 'P ', 16: 'S', 17: 'Cl', 18: 'Ar', 19: 'K ', 20: 'Ca',
+                21: 'Sc', 22: 'Ti', 23: 'V ', 24: 'Cr', 25: 'Mn', 26: 'Fe', 27: 'Co', 28: 'Ni', 29: 'Cu', 30: 'Zn',
+                31: 'Ga', 32: 'Ge', 33: 'As', 34: 'Se', 35: 'Br', 36: 'Kr', 37: 'Rb', 38: 'Sr', 39: 'Y ', 40: 'Zr',
+                41: 'Nb', 42: 'Mo', 43: 'Tc', 44: 'Ru', 45: 'Rh', 46: 'Pd', 47: 'Ag', 48: 'Cd', 49: 'In', 50: 'Sn',
+                51: 'Sb', 52: 'Te', 53: 'I ', 54: 'Xe', 55: 'Cs', 56: 'Ba', 57: 'La', 58: 'Ce', 59: 'Pr', 60: 'Nd',
+                61: 'Pm', 62: 'Sm', 63: 'Eu', 64: 'Gd', 65: 'Tb', 66: 'Dy', 67: 'Ho', 68: 'Er', 69: 'Tm', 70: 'Yb',
+                71: 'Lu', 72: 'Hf', 73: 'Ta', 74: 'W ', 75: 'Re', 76: 'Os', 77: 'Ir', 78: 'Pt', 79: 'Au', 80: 'Hg',
+                81: 'Tl', 82: 'Pb', 83: 'Bi', 84: 'Po', 85: 'At', 86: 'Rn', 87: 'Fr', 88: 'Ra', 89: 'Ac', 90: 'Th',
+                91: 'Pa', 92: 'U ', 93: 'Np', 94: 'Pu', 95: 'Am', 96: 'Cm', 97: 'Bk', 98: 'Cf', 99: 'Es', 100: 'Fm',
+                101: 'Md', 102: 'No', 103: 'Lr', 104: 'Rf', 105: 'Db', 106: 'Sg', 107: 'Bh', 108: 'Hs', 109: 'Mt',
+                110: 'Ds', 111: 'Rg', 112: 'Cn', 113: 'Nh', 114: 'Fl', 115: 'Mc', 116: 'Lv', 117: 'Ts', 118: 'Og'}
+
+ElemName2Num = {v: k for k, v in zip(ElemNum2Name.keys(), ElemNum2Name.values())}
+
+Bohr2Angstrom = 0.529177
+Angstrom2Bohr = 1 / Bohr2Angstrom
+Hartree2eV = 27.2114
+eV2Hartree = 1 / Hartree2eV
+THz2eV = 4.136e-3
+eV2THz = 1 / THz2eV
+amu2kg = 1.6605402e-27
+
+PLANCK_CONSTANT = 4.135667662e-15  # Planck's constant in eV*s
+BOLTZMANN_CONSTANT = 8.617333262145e-5  # Boltzmann's constant in eV/K
+ELEM_CHARGE = 1.60217662e-19  # Elementary charge in Coulombs
+BOHR_RADIUS = 1.88973 #TODO Check
+
+Bader_radii_Bohr = {'H': 2.88, 'He': 2.5, 'Li': 4.18, 'Be': 4.17, 'B': 3.91, 'C': 3.62, 'N': 3.35, 'O': 3.18,
+                    'F': 3.03, 'Ne': 2.89, 'Na': 4.25, 'Mg': 4.6, 'Al': 4.61, 'Si': 4.45, 'P': 4.23, 'S': 4.07,
+                    'Cl': 3.91, 'Ar': 3.72, 'Li_plus': 1.82, 'Na_plus': 2.47, 'F_minus': 3.49, 'Cl_minus': 4.36}
+
+IDSCRF_radii_Angstrom = {'H': 1.77, 'He': 1.49, 'Li': 2.22, 'Be': 2.19, 'B': 2.38, 'C': 2.22, 'N': 2.05, 'O': 1.87,
+                         'F': 1.89, 'Ne': 1.74}
+
+# Colorblindness-friendly colors from https://github.com/mpetroff/accessible-color-cycles
+cbf6 = ["#5790fc", "#f89c20", "#e42536", "#964a8b", "#9c9ca1", "#7a21dd"]
+cbf8 = ["#1845fb", "#ff5e02", "#c91f16", "#c849a9", "#adad7d", "#86c8dd", "#578dff", "#656364"]
+cbf10 = ["#3f90da", "#ffa90e", "#bd1f01", "#94a4a2", "#832db6", "#a96b59", "#e76300", "#b9ac70", "#717581", "#92dadd"]

electrochemistry/echem/core/electronic_structure.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+import numpy as np
+from typing import Union, Iterable
+from nptyping import NDArray, Shape, Number
+
+
+class EBS:
+    """
+    Class for calculating DOS
+
+    Args:
+        eigenvalues
+    """
+    def __init__(self,
+                 eigenvalues: NDArray[Shape['Nspin, Nkpts, Nbands'], Number],
+                 weights: NDArray[Shape['Nkpts'], Number] = None,
+                 efermi: float = None,
+                 occupations: NDArray[Shape['Nspin, Nkpts, Nbands'], Number] = None):
+
+        self.eigenvalues = eigenvalues
+        self.occupations = occupations
+        self.efermi = efermi
+
+        if weights is None:
+            self.weights = np.ones(eigenvalues.shape[1]) / eigenvalues.shape[1]
+        else:
+            self.weights = weights
+
+    @property
+    def nspin(self):
+        return self.eigenvalues.shape[0]
+
+    @property
+    def nkpts(self):
+        return self.eigenvalues.shape[1]
+
+    @property
+    def nbands(self):
+        return self.eigenvalues.shape[2]
+
+    @staticmethod
+    def gaussian_smearing(E: NDArray[Shape['Ngrid'], Number],
+                          E0: NDArray[Shape['*, ...'], Number],
+                          sigma: float):
+        """
+        Blur the Delta function by a Gaussian function
+        Args:
+            E: Numpy array with the shape (ngrid, ) that represents the energy range for the Gaussian smearing
+            E0: Numpy array with any shape (i.e. (nspin, nkpts, nbands)) that contains eigenvalues
+            sigma: the broadening parameter for the Gaussian function
+
+        Returns:
+            Smeared eigenvalues on grind E with the shape E0.shape + E.shape
+        """
+        return 1 / (np.sqrt(2 * np.pi) * sigma) * np.exp(
+            -(np.broadcast_to(E, E0.shape + E.shape) - np.expand_dims(E0, len(E0.shape))) ** 2 / (2 * sigma ** 2))
+
+    def __get_by_bands(self,
+                       property: str,
+                       bands: Union[int, Iterable[int]]):
+        property = getattr(self, property)
+        if type(bands) is int:
+            if self.nspin == 1:
+                return property[:, bands]
+            elif self.nspin > 1:
+                return property[:, :, bands]
+        elif isinstance(bands, Iterable):
+            if self.nspin == 1:
+                return property[:, bands].transpose(1, 0)
+            elif self.nspin > 1:
+                return property[:, :, bands].transpose(2, 0, 1)
+        else:
+            raise ValueError('Variable bands should be int or iterable')
+
+    def get_band_eigs(self, bands: Union[int, Iterable]):
+        return self.__get_by_bands('eigenvalues', bands)
+
+    def get_band_occ(self, bands: Union[int, Iterable]):
+        if self.occupations is not None:
+            return self.__get_by_bands('occupations', bands)
+        else:
+            raise ValueError('Occupations has not been defined')
+
+    def get_DOS(self,
+                dE: float = 0.01,
+                emin: float = None,
+                emax: float = None,
+                zero_at_fermi: bool = False,
+                smearing: str = 'gaussian',
+                sigma: float = 0.02) -> tuple[NDArray[Shape['Ngrid'], Number], NDArray[Shape['Nspin, Ngrid'], Number]]:
+        """Calculate Density of States based on eigenvalues and its weights
+
+        Args:
+            dE (float, optional): step of energy array in function's output. Default value is 0.01
+            zero_at_fermi (bool, optional): if True Fermi energy will be equal to zero
+            emin (float, optional): minimum value in DOS calculation.
+            emax (float, optional): maximum value in DOS calculation.
+            smearing (str, optional): define whether you will be used smearing or not. Default value is 'Gaussian'.
+            Possible options: 'gaussian'
+            sigma (float, optional): define the sigma parameter in Gaussian smearing. Default value is 0.02
+
+        Returns:
+            E, DOS - Two 1D np.arrays that contain energy and according to DOS values
+        """
+        if zero_at_fermi is True and self.efermi is None:
+            raise ValueError('You can not set zero_at_fermi=True if you did not specify efermi value')
+
+        if emin is None:
+            E_min = np.min(self.eigenvalues) - 1
+        if emax is None:
+            E_max = np.max(self.eigenvalues) + 1
+
+        E_arr = np.arange(E_min, E_max, dE)
+
+        if smearing.lower() == 'gaussian':
+            DOS_arr = np.sum(self.weights[None, :, None, None] *
+                             self.gaussian_smearing(E_arr, self.eigenvalues, sigma), axis=(1, 2))
+        else:
+            raise NotImplemented(f'Smearing {smearing} is not implemented. Please use \'gaussian\' instead.')
+
+        # 2 means occupancy for non-spinpolarized calculation
+        if self.nspin == 1:
+            DOS_arr *= 2
+
+        if zero_at_fermi:
+            return E_arr - self.efermi, DOS_arr
+        else:
+            return E_arr, DOS_arr

electrochemistry/echem/core/ionic_dynamics.py

@@ -0,0 +1,102 @@
+from __future__ import annotations
+import numpy as np
+from nptyping import NDArray, Shape, Number
+
+
+class IonicDynamics:
+    def __init__(self,
+                 forces_hist: NDArray[Shape['Nsteps, Natoms, 3'], Number] | None,
+                 coords_hist: NDArray[Shape['Nsteps, Natoms, 3'], Number] | None,
+                 lattice: NDArray[Shape['3, 3'], Number] | None,
+                 coords_are_cartesian: bool | None):
+        self.forces_hist = forces_hist
+        self.coords_hist = coords_hist
+        self.lattice = lattice
+        self.coords_are_cartesian = coords_are_cartesian
+
+    @property
+    def forces(self):
+        if self.forces_hist is not None:
+            return self.forces_hist[-1]
+        else:
+            raise ValueError('Forces_hist is None')
+
+    def get_forces(self,
+                   mod: str = 'mean',
+                   diff: bool = False):
+        """
+        Args:
+            mod (str, optional):
+                norm - (N_steps, N_atoms) returns the norm of forces along the ionic trajectory
+                mean - (N_steps, ) returns the mean value of forces' norm in simulation cell along the ionic trajectory
+                max - (N_steps, ) returns the max value of forces' norm in simulation cell along the ionic trajectory
+            diff (bool, optional): if True returns absolute value of forces differences between i and i+1 steps.
+            If False returns just forces values at each step
+
+        Returns:
+
+        """
+        if self.forces_hist is not None:
+            if mod == 'norm':
+                forces = np.linalg.norm(self.forces_hist, axis=2)
+            elif mod == 'mean':
+                forces = np.mean(np.linalg.norm(self.forces_hist, axis=2), axis=1)
+            elif mod == 'max':
+                forces = np.max(np.linalg.norm(self.forces_hist, axis=2), axis=1)
+            else:
+                raise ValueError(f'mod should be norm/mean/max. You set {mod}')
+
+            if diff:
+                return np.abs(forces[1:] - forces[:-1])
+            else:
+                return forces
+        else:
+            raise ValueError('Forces_hist is None')
+
+    def get_displacements(self,
+                          i: int | None = None,
+                          j: int | None = None,
+                          scalar: bool = True) -> (NDArray[Shape['Natoms, 3']] |
+                                                   NDArray[Shape['Natoms, 1']] |
+                                                   NDArray[Shape['Nsteps, Natoms, 3']] |
+                                                   NDArray[Shape['Nsteps, Natoms, 1']]):
+        if self.coords_hist is not None and \
+                self.lattice is not None and \
+                self.coords_are_cartesian is not None:
+
+            if isinstance(i, int) and isinstance(j, int):
+
+                if self.coords_are_cartesian:
+                    transform = np.linalg.inv(self.lattice)
+                    r1 = np.matmul(self.coords_hist[i], transform)
+                    r2 = np.matmul(self.coords_hist[j], transform)
+                else:
+                    r1 = self.coords_hist[i]
+                    r2 = self.coords_hist[j]
+
+                R = r2 - r1
+                R = (R + 0.5) % 1 - 0.5
+                assert np.all(R >= - 0.5) and np.all(R <= 0.5)
+
+                if scalar:
+                    return np.linalg.norm(np.matmul(R, self.lattice), axis=1)
+                else:
+                    return np.matmul(R, self.lattice)
+
+            elif i is None or j is None:
+                if self.coords_are_cartesian:
+                    transform = np.linalg.inv(self.lattice)
+                    R = np.matmul(self.coords_hist, transform)
+                else:
+                    R = self.coords_hist
+                R = np.diff(R, axis=0)
+                R = (R + 0.5) % 1 - 0.5
+
+                if scalar:
+                    return np.linalg.norm(np.matmul(R, self.lattice), axis=2)
+                else:
+                    return np.matmul(R, self.lattice)
+
+        else:
+            raise ValueError('Method get_distance_matrix can only be called '
+                             'if coords_hist, lattice and coords_are_cartesian are not None')

electrochemistry/echem/core/structure.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+import numpy as np
+from typing import Iterable
+from termcolor import colored
+from nptyping import NDArray, Shape, Number
+
+
+class Structure:
+    """
+    Basic class for structure of unit/supercell.
+    Args:
+        lattice: 2D array that contains lattice vectors. Each row should correspond to a lattice vector.
+            E.g., [[5, 5, 0], [7, 4, 0], [0, 0, 25]].
+        species: List of species on each site. Usually list of elements, e.g., ['Al', 'Al', 'O', 'H'].
+        coords: List of lists or np.ndarray (Nx3 dimension) that contains coords of each species.
+        coords_are_cartesian: True if coords are cartesian, False if coords are fractional
+    """
+    def __init__(self,
+                 lattice: NDArray[Shape[3, 3], Number] | list[list[float]],
+                 species: list[str],
+                 coords: NDArray[Shape['Natoms, 3'], Number] | list[list[float]],
+                 coords_are_cartesian: bool = True):
+
+        if len(species) != len(coords):
+            raise StructureError('Number of species and its coords must be the same')
+
+        self.species = species
+
+        if isinstance(lattice, np.ndarray):
+            self.lattice = lattice
+        else:
+            self.lattice = np.array(lattice)
+
+        if isinstance(coords, np.ndarray):
+            self.coords = coords
+        else:
+            self.coords = np.array(coords)
+
+        self.coords_are_cartesian = coords_are_cartesian
+
+    def __repr__(self) -> str:
+        lines = ['\nLattice:']
+
+        width = len(str(int(np.max(self.lattice)))) + 6
+        for axis in self.lattice:
+            lines.append('  '.join([f'{axis_coord:{width}.5f}' for axis_coord in axis]))
+
+        width = len(str(int(np.max(self.coords)))) + 6
+
+        lines.append('\nSpecies:')
+
+        unique, counts = np.unique(self.species, return_counts=True)
+
+        if len(self.species) < 10:
+            lines.append(' '.join([s for s in self.species]))
+            for u, c in zip(unique, counts):
+                lines.append(f'{u}:\t{c}')
+            lines.append(f'\nCoords are cartesian: {self.coords_are_cartesian}')
+            lines.append('\nCoords:')
+            for coord in self.coords:
+                lines.append('  '.join([f'{c:{width}.5f}' for c in coord]))
+        else:
+            part_1 = ' '.join([s for s in self.species[:5]])
+            part_2 = ' '.join([s for s in self.species[-5:]])
+            lines.append(part_1 + ' ... ' + part_2)
+            for u, c in zip(unique, counts):
+                lines.append(f'{u}:\t{c}')
+            lines.append(f'\nCoords are cartesian: {self.coords_are_cartesian}')
+            lines.append('\nCoords:')
+            for coord in self.coords[:5]:
+                lines.append('  '.join([f'{c:{width}.5f}' for c in coord]))
+            lines.append('...')
+            for coord in self.coords[-5:]:
+                lines.append('  '.join([f'{c:{width}.5f}' for c in coord]))
+
+        return '\n'.join(lines)
+
+    def __eq__(self, other: Structure) -> bool:
+        assert isinstance(other, Structure), 'Other object be a Structure class'
+        if not self.coords_are_cartesian == other.coords_are_cartesian:
+            if self.coords_are_cartesian:
+                other.mod_coords_to_cartesian()
+                print(colored('Coords of other were modified into Cartesian', color='green'))
+            else:
+                other.mod_coords_to_direct()
+                print(colored('Coords of other were modified into Direct', color='green'))
+        return np.allclose(self.lattice, other.lattice, atol=1e-10, rtol=1e-10) and \
+            (self.species == other.species) and \
+            np.allclose(self.coords, other.coords, atol=1e-10, rtol=1e-10)
+
+    @property
+    def natoms(self) -> int:
+        return len(self.species)
+
+    @property
+    def natoms_by_type(self) -> dict[str: int]:
+        unique, counts = np.unique(self.species, return_counts=True)
+        return {i: j for i, j in zip(unique, counts)}
+
+    def mod_add_atoms(self, coords, species) -> None:
+        """
+        Adds atoms in the Structure
+        Args:
+            coords: List or np.ndarray (Nx3 dimension) that contains coords of each species
+            species: List of species on each site. Usually list of elements, e.g., ['Al', 'Al', 'O', 'H']
+        """
+        if not isinstance(coords, np.ndarray):
+            coords = np.array(coords)
+        self.coords = np.vstack((self.coords, coords))
+        if isinstance(species, str):
+            self.species += [species]
+        else:
+            self.species += species
+
+    def mod_delete_atoms(self,
+                         ids: int | list[int]) -> None:
+        """
+        Deletes selected atoms by ids
+        Args:
+            ids: sequence of atoms ids
+        """
+        self.coords = np.delete(self.coords, ids, axis=0)
+        self.species = np.delete(self.species, ids)
+
+    def mod_change_atoms(self,
+                         ids: int | list[int],
+                         coords: NDArray[Shape['Nids, 3'], Number] | None,
+                         species: str | list[str] | None) -> None:
+        """
+        Change selected atom by id
+        Args:
+            ids: List or int. First id is zero.
+            coords: None or np.array with new coords, e.g. np.array([1, 2, 4.6])
+            species: None or str or List[str]. New types of changed atoms
+
+        Returns:
+
+        """
+        if coords is not None:
+            self.coords[ids] = coords
+        if species is not None:
+            if isinstance(ids, Iterable):
+                for i, sp in zip(ids, species):
+                    self.species[i] = sp
+            else:
+                self.species[ids] = species
+
+    def mod_coords_to_cartesian(self) -> None | str:
+        """
+        Converts species coordinates to Cartesian coordination system.
+        """
+        if self.coords_are_cartesian is True:
+            return 'Coords are already cartesian'
+        else:
+            self.coords = np.matmul(self.coords, self.lattice)
+            self.coords_are_cartesian = True
+
+    def mod_coords_to_direct(self) -> None | str:
+        """
+        Converts species coordinates to Direct coordination system.
+        """
+        if self.coords_are_cartesian is False:
+            return 'Coords are already direct'
+        else:
+            transform = np.linalg.inv(self.lattice)
+            self.coords = np.matmul(self.coords, transform)
+            self.coords_are_cartesian = False
+
+    def mod_add_vector(self,
+                       vector: NDArray[Shape['3'], Number],
+                       cartesian: bool = True) -> None:
+        """
+        Adds a vector to all atoms.
+        Args:
+            vector (np.ndarray): a vector which will be added to coordinates of all species
+            cartesian (bool): determine in which coordination system the vector defined. Cartesian=True means that
+            the vector is defined in Cartesian coordination system. Cartesian=False means that the vector is defined
+            in Direct coordination system.
+        """
+        self.mod_coords_to_direct()
+        if cartesian:
+            transform = np.linalg.inv(self.lattice)
+            vector = np.matmul(vector, transform)
+        self.coords += vector
+        self.coords %= np.array([1, 1, 1])
+
+    def mod_propagate_unit_cell(self, a, b, c) -> None:
+        """Extend the unit cell by propagation on lattice vector direction. The resulted unit cell will be the size
+        of [a * lattice[0], b * lattice[1], c * lattice[2]]. All atoms will be replicated with accordance of new
+        unit cell size
+        Args:
+            a (int): the number of replicas in the directions of the first unit cell vector
+            b (int): the number of replicas in the directions of the second unit cell vector
+            c (int): the number of replicas in the directions of the third unit cell vector.
+        """
+        if not (isinstance(a, int) and isinstance(b, int) and isinstance(c, int)):
+            raise ValueError(f'a, b and c must be integers. But a, b, c have types {type(a), type(b), type(c)}')
+
+        if a * b * c > 1:
+            self.mod_coords_to_cartesian()
+            ref_coords = self.coords.copy()
+            ref_species = self.species.copy()
+
+            for i in range(a):
+                for j in range(b):
+                    for k in range(c):
+                        if i != 0 or j != 0 or k != 0:
+                            vector = np.sum(np.array([i, j, k]).reshape(-1, 1) * self.lattice, axis=0)
+                            self.coords = np.vstack((self.coords, ref_coords + vector))
+                            self.species += ref_species
+            self.lattice = np.array([a, b, c]).reshape(-1, 1) * self.lattice
+        else:
+            raise ValueError(f'a, b and c must be positive. But {a=} {b=} {c=}')
+
+    def get_vector(self,
+                   id_1: int,
+                   id_2: int,
+                   unit: bool = True) -> NDArray[Shape['3'], Number]:
+        """
+        Returns a vector (unit vector by default) which starts in the atom with id_1 and points to the atom with id_2
+        Args:
+            id_1: id of the first atom (vector origin)
+            id_2: id of the second atom
+            unit: Defines whether the vector will be normed or not. Unit=True means that the vector norm will be equal 1
+        Returns (np.ndarray): vector from atom with id_1 to atom with id_2
+        """
+        vector = self.coords[id_2] - self.coords[id_1]
+        if unit:
+            return vector / np.linalg.norm(vector)
+        else:
+            return vector
+
+    def get_distance_matrix(self) -> NDArray[Shape['Natoms, Natoms, 3'], Number]:
+        """
+        Returns distance matrix R, where R[i,j] is the vector from atom i to atom j.
+
+        Returns:
+            np.ndarray (NxNx3 dimensions) which is a distance matrix in Cartesian coordination system
+        """
+        self.mod_coords_to_direct()
+        r1 = np.broadcast_to(self.coords.reshape((self.natoms, 1, 3)), (self.natoms, self.natoms, 3))
+        r2 = np.broadcast_to(self.coords.reshape((1, self.natoms, 3)), (self.natoms, self.natoms, 3))
+        R = r2 - r1
+        R = (R + 0.5) % 1 - 0.5
+        assert np.all(R >= - 0.5) and np.all(R <= 0.5)
+        return np.matmul(R, self.lattice)
+
+    def get_distance_matrix_scalar(self) -> NDArray[Shape['Natoms, Natoms'], Number]:
+        """
+        Returns distance matrix R, where R[i, j] is the Euclidean norm of a vector from atom i to atom j.
+
+        Returns:
+            np.ndarray (NxN dimensions) which is a distance matrix containing scalars
+        """
+        R = self.get_distance_matrix()
+        return np.sqrt(np.sum(R * R, axis=2))
+
+    def get_filtered_ids(self, **kwargs):
+        """
+        Returns np.ndarray that contains atom ids according to filter rules
+        Args:
+            species (str): Define which atom type will be selected. E.g. 'C H N' means select all C, H, and N atoms.
+            '!C' means select all atoms except C.
+        Returns:
+            np.ndarray contains ids of atoms according to selecting rules
+        """
+        filter_mask = np.array([True for _ in range(self.natoms)], dtype=np.bool_)
+
+        if 'species' in kwargs:
+            species = kwargs['species'].split()
+            species_select = []
+            species_not_select = []
+            for specie in species:
+                if '!' in specie:
+                    species_not_select.append(specie.replace('!', ''))
+                else:
+                    species_select.append(specie)
+            if len(species_select):
+                fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
+                for specie in species_select:
+                    fm_local += np.array([True if atom_name == specie else False for atom_name in self.species])
+                filter_mask *= fm_local
+            if len(species_not_select):
+                fm_local = np.array([True for _ in range(self.natoms)], dtype=np.bool_)
+                for specie in species_not_select:
+                    fm_local *= np.array([False if atom_name == specie else True for atom_name in self.species])
+                filter_mask *= fm_local
+
+        if 'x' in kwargs:
+            left, right = kwargs['x']
+            fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
+            for i, atom_coord in enumerate(self.coords):
+                if left < atom_coord[0] < right:
+                    fm_local[i] = True
+            filter_mask *= fm_local
+
+        if 'y' in kwargs:
+            left, right = kwargs['y']
+            fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
+            for i, atom_coord in enumerate(self.coords):
+                if left < atom_coord[1] < right:
+                    fm_local[i] = True
+            filter_mask *= fm_local
+
+        if 'z' in kwargs:
+            left, right = kwargs['z']
+            fm_local = np.array([False for _ in range(self.natoms)], dtype=np.bool_)
+            for i, atom_coord in enumerate(self.coords):
+                if left < atom_coord[2] < right:
+                    fm_local[i] = True
+            filter_mask *= fm_local
+
+        return np.array(range(self.natoms))[filter_mask]
+
+
+class StructureError(Exception):
+    """
+    Exception class for Structure.
+    Raised when the structure has problems, e.g., atoms that are too close.
+    """
+    pass

electrochemistry/echem/core/thermal_properties.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+import numpy as np
+from nptyping import NDArray, Shape, Number
+import warnings
+
+
+class ThermalProperties:
+    """
+    Class for calculation thermal properties based on the calculated phonon spectra
+
+    Args:
+        eigen_freq (np.ndarray [nkpts, nfreq]): The energies of phonon eigenfrequencies in eV
+        weights (np.ndarray [nkpts, ], optional): weights of k-points. The sum of weights should be equal to 1.
+        If weights are not provided, all k-points will be considered with equal weights (1 / nkpts).
+    """
+    k_B_J = 1.380649e-23  # J/K
+    k_B_eV = 8.617333262e-5  # eV/K
+    hbar_J = 1.054571817e-34  # J*s
+
+    def __init__(self,
+                 eigen_freq: NDArray[Shape['Nkpts, Nfreq'], Number],
+                 weights: NDArray[Shape['Nkpts'], Number] = None):
+
+        if weights is None:
+            self.weights = np.ones(eigen_freq.shape[0]) / eigen_freq.shape[0]
+        else:
+            self.weights = weights
+
+        if np.sum(eigen_freq < 0) > 0:
+            warnings.warn('\nThere is at least one imaginary frequency in given eigenfrequencies. '
+                          '\nAll imaginary frequencies will be dropped from any further calculations'
+                          f'\nImaginary frequencies: {eigen_freq[eigen_freq < 0]}')
+        self.eigen_freq = np.maximum(0, eigen_freq)
+
+    def get_Gibbs_ZPE(self) -> float:
+        r"""
+        Calculate Zero Point Energy
+
+        .. math::
+            E_{ZPE} = \sum_k weight(k) \sum_i \frac{hw_{i}(k)}{2}
+
+        Returns:
+            ZPE in eV
+        """
+        return np.sum(self.weights * self.eigen_freq) / 2
+
+    def get_enthalpy_vib(self, T) -> float:
+        r"""
+        Calculate the thermal term in vibrational energy
+
+        .. math::
+            E_{temp}(k) = \sum_i \left( \frac{hw_{i}(k)}{( \exp (hw_{i}(k) / k_B T) - 1)} \right) \\\\
+            E_{temp} = \sum_k weight(k) \cdot E_{temp}(k)
+
+        Args:
+            T: Temperature in K
+        Returns:
+            Thermal vibrational energy in eV
+        """
+        k_B = 8.617333262145e-5  # Boltzmann's constant in eV/K
+
+        return np.sum(np.nan_to_num(self.weights * self.eigen_freq / (np.exp(self.eigen_freq / (k_B * T)) - 1)))
+
+    def get_TS_vib(self, T) -> float:
+        r"""
+        Calculate the vibrational entropy contribution
+
+        .. math::
+            S_{vib}(k) = \sum_i \left( \frac{hw_{i}(k)}{( \exp (hw_{i}(k) / k_B T) - 1)} -
+            k_B ln \left( 1 - \exp \left(- \frac{hw_i(k)}{k_B T} \right) \right) \right) \\\\
+            T * S_{vib} = T \sum_k (weight_k S_{vib}(k))
+
+        Args:
+            T: Temperature in K
+
+        Returns:
+            TS in eV
+
+        """
+        k_B = 8.617333262145e-5  # Boltzmann's constant in eV/K
+        second_term = - np.sum(self.weights * k_B * T * np.nan_to_num(np.log(1 - np.exp(- self.eigen_freq / (k_B * T))),
+                                                                      neginf=0))
+        return self.get_enthalpy_vib(T) + second_term
+
+    def get_Gibbs_vib(self, T: float) -> float:
+        return self.get_Gibbs_ZPE() + self.get_enthalpy_vib(T) - self.get_TS_vib(T)
+
+    @classmethod
+    def get_Gibbs_trans(cls,
+                        V: float,
+                        mass: float,
+                        T: float):
+        return - cls.k_B_eV * T * np.log(V * (mass * cls.k_B_J * T / (2 * np.pi * cls.hbar_J**2))**1.5)
+
+    @classmethod
+    def get_Gibbs_rot(cls,
+                      I: float | list[float] | NDArray[Shape['3'], Number],
+                      sigma: int,
+                      T: float):
+        if type(I) is float or len(I) == 1:
+            return - cls.k_B_eV * T * np.log(2 * I * cls.k_B_J * T / (sigma * cls.hbar_J ** 2))
+        elif len(I) == 3:
+            return - cls.k_B_eV * T * np.log((2 * cls.k_B_J * T)**1.5 * (np.pi * I[0] * I[1] * I[2])**0.5 /
+                                           (sigma * cls.hbar_J**3))
+        else:
+            raise ValueError(f'I should be either float or array with length of 3, however {len(I)=}')

electrochemistry/echem/core/useful_funcs.py

@@ -0,0 +1,24 @@
+import numpy as np
+import subprocess
+
+
+class ClassMethods:
+    def check_existence(self, variable):
+        """
+        This function checks whether desires variable is not None
+        :param variable: desired variable
+        :return: nothing
+        """
+        if getattr(self, variable) is None:
+            raise ValueError(f'{variable} is not defined')
+
+
+def nearest_array_index(array, value):
+    return (np.abs(array - value)).argmin()
+
+
+def shell(cmd) -> str:
+    '''
+    Run shell command and return output as a string
+    '''
+    return subprocess.check_output(cmd, shell=True)

electrochemistry/echem/core/vasp_constant.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+'''
+Physical constants used in VASP
+This file was downloaded from https://github.com/QijingZheng/VaspBandUnfolding 22.05.2019
+'''
+
+#  Some important Parameters, to convert to a.u.
+#  - AUTOA     =  1. a.u. in Angstroem
+#  - RYTOEV    =  1 Ry in Ev
+#  - EVTOJ     =  1 eV in Joule
+#  - AMTOKG    =  1 atomic mass unit ("proton mass") in kg
+#  - BOLKEV    =  Boltzmanns constant in eV/K
+#  - BOLK      =  Boltzmanns constant in Joule/K
+
+AUTOA    = 0.529177249
+RYTOEV   = 13.605826
+CLIGHT   =  137.037          # speed of light in a.u.
+EVTOJ    = 1.60217733E-19
+AMTOKG   = 1.6605402E-27
+BOLKEV   = 8.6173857E-5
+BOLK     = BOLKEV * EVTOJ
+EVTOKCAL = 23.06
+
+# FELECT    =  (the electronic charge)/(4*pi*the permittivity of free space)
+#         in atomic units this is just e^2
+# EDEPS    =  electron charge divided by the permittivity of free space
+#         in atomic units this is just 4 pi e^2
+# HSQDTM    =  (plancks CONSTANT/(2*PI))**2/(2*ELECTRON MASS)
+#
+PI     = 3.141592653589793238
+TPI    = 2 * PI
+CITPI  = 1j * TPI
+FELECT = 2 * AUTOA * RYTOEV
+EDEPS  = 4 * PI * 2 * RYTOEV * AUTOA
+HSQDTM = RYTOEV * AUTOA * AUTOA
+
+# vector field A times momentum times e/ (2 m_e c) is an energy
+# magnetic moments are supplied in Bohr magnetons
+# e / (2 m_e c) A(r) p(r)    =  energy
+# e / (2 m_e c) m_s x ( r - r_s) / (r-r_s)^3 hbar nabla    = 
+# e^2 hbar^2 / (2 m_e^2 c^2) 1/ lenght^3    =  energy
+# conversion factor from magnetic moment to energy
+# checked independently in SI by Gilles de Wijs
+
+MAGMOMTOENERGY  = 1 / CLIGHT**2 * AUTOA**3 * RYTOEV
+
+# dimensionless number connecting input and output magnetic moments
+# AUTOA e^2 (2 m_e c^2)
+MOMTOMOM   = AUTOA / CLIGHT / CLIGHT / 2
+AUTOA2   = AUTOA * AUTOA
+AUTOA3   = AUTOA2 * AUTOA
+AUTOA4   = AUTOA2 * AUTOA2
+AUTOA5   = AUTOA3 * AUTOA2
+
+# dipole moment in atomic units to Debye
+AUTDEBYE = 2.541746  

electrochemistry/echem/core/vaspwfc_p3.py

@@ -0,0 +1,948 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-   
+"""
+This program was downloaded from https://github.com/QijingZheng/VaspBandUnfolding 22.05.2019
+We changed print calls (print x to print(x)) to run this program with python3
+In lines: 274 276 278 325 429 431 444 468 469 480 481 711 713 715 /2 changed to //2
+"""
+import os
+import numpy as np
+from math import sqrt
+from .vasp_constant import *
+from multiprocessing import cpu_count
+from scipy.fftpack import fftfreq, fftn, ifftn
+
+############################################################
+def save2vesta(phi=None, poscar='POSCAR', prefix='wfc', 
+               lgam=False, lreal=False, ncol=10):
+    '''
+    Save the real space pseudo-wavefunction as vesta format.
+    '''
+    nx, ny, nz = phi.shape
+    try:
+        pos = open(poscar, 'r')
+        head = ''
+        for line in pos:
+            if line.strip():
+                head += line
+            else:
+                break
+        head += '\n%5d%5d%5d\n' % (nx, ny, nz)
+    except:
+        raise IOError('Failed to open %s' % poscar)
+
+    # Faster IO
+    nrow = phi.size // ncol
+    nrem = phi.size % ncol
+    fmt = "%16.8E"
+
+    psi = phi.copy()
+    psi = psi.flatten(order='F')
+    psi_h = psi[:nrow * ncol].reshape((nrow, ncol))
+    psi_r = psi[nrow * ncol:]
+
+    with open(prefix + '_r.vasp', 'w') as out:
+        out.write(head)
+        out.write(
+            '\n'.join([''.join([fmt % xx for xx in row])
+                       for row in psi_h.real])
+        )
+        out.write("\n" + ''.join([fmt % xx for xx in psi_r.real]))
+    if not (lgam or lreal):
+        with open(prefix + '_i.vasp', 'w') as out:
+            out.write(head)
+            out.write(
+                '\n'.join([''.join([fmt % xx for xx in row])
+                           for row in psi_h.imag])
+            )
+            out.write("\n" + ''.join([fmt % xx for xx in psi_r.imag]))
+
+############################################################
+
+class vaspwfc(object):
+    '''
+    Class for processing VASP Pseudowavefunction stored in WAVECAR.  This
+    program is motivated by PIESTA written by Ren Hao <renh@upc.edu.cn>.
+
+    The format of VASP WAVECAR, as shown in
+        http://www.andrew.cmu.edu/user/feenstra/wavetrans/
+    is:
+        Record-length #spin components RTAG(a value specifying the precision)
+        #k-points #bands ENCUT(maximum energy for plane waves)
+        LatVec-A
+        LatVec-B
+        LatVec-C
+        Loop over spin
+           Loop over k-points
+              #plane waves, k vector
+              Loop over bands
+                 band energy, band occupation
+              End loop over bands
+              Loop over bands
+                 Loop over plane waves
+                    Plane-wave coefficient
+                 End loop over plane waves
+              End loop over bands
+           End loop over k-points
+        End loop over spin
+    '''
+
+    def __init__(self, fnm='WAVECAR', lsorbit=False, lgamma=False,
+                 gamma_half='z', omp_num_threads=1):
+        '''
+        Initialization.
+        '''
+
+        self._fname = fnm
+        # the directory containing the input file
+        self._dname  = os.path.dirname(fnm)
+        if self._dname == '':
+            self._dname = '.'
+
+        self._lsoc  = lsorbit
+        self._lgam  = lgamma
+        self._gam_half = gamma_half.lower()
+        
+        # It seems that some modules in scipy uses OPENMP, it is therefore
+        # desirable to set the OMP_NUM_THREADS to tune the parallization.
+        os.environ['OMP_NUM_THREADS'] = str(omp_num_threads)
+
+        assert not (lsorbit and lgamma), 'The two settings conflict!'
+        assert self._gam_half == 'x' or self._gam_half == 'z', \
+               'Gamma_half must be "x" or "z"'
+
+        try:
+            self._wfc = open(self._fname, 'rb')
+        except:
+            raise IOError('Failed to open %s' % self._fname)
+
+        # read the basic information
+        self.readWFHeader()
+        # read the band information
+        self.readWFBand()
+        
+        if self._lsoc:
+            assert self._nspin == 1, "NSPIN = 1 for noncollinear version WAVECAR!"
+
+    def set_omp_num_threads(self, nproc):
+        '''
+        Set the OMP_NUM_THREADS envrionment variable
+        '''
+        assert 1 <= nproc <= cpu_count()
+
+        os.envrion['OMP_NUM_THREADS'] = str(nproc)
+
+    def isSocWfc(self):
+        """
+        Is the WAVECAR from an SOC calculation?
+        """
+        return True if self._lsoc else False
+
+    def isGammaWfc(self):
+        """
+        Is the WAVECAR from an SOC calculation?
+        """
+        return True if self._lgam else False
+
+    def readWFHeader(self):
+        '''
+        Read the system information from WAVECAR, which is written in the first
+        two record.
+
+        rec1: recl, nspin, rtag
+        rec2: nkpts, nbands, encut, ((cell(i,j) i=1, 3), j=1, 3)
+        '''
+
+        # goto the start of the file and read the first record
+        self._wfc.seek(0)
+        self._recl, self._nspin, self._rtag = np.array(
+                np.fromfile(self._wfc, dtype=np.float, count=3),
+                dtype=int
+        )
+        self._WFPrec = self.setWFPrec()
+        # the second record
+        self._wfc.seek(self._recl)
+        dump = np.fromfile(self._wfc, dtype=np.float, count=12)
+
+        self._nkpts  = int(dump[0])                     # No. of k-points
+        self._nbands = int(dump[1])                     # No. of bands
+        self._encut  = dump[2]                          # Energy cutoff
+        self._Acell  = dump[3:].reshape((3,3))          # real space supercell basis
+        self._Omega  = np.linalg.det(self._Acell)       # real space supercell volume
+        self._Bcell  = np.linalg.inv(self._Acell).T     # reciprocal space supercell volume
+
+        # Minimum FFT grid size
+        Anorm = np.linalg.norm(self._Acell, axis=1)
+        CUTOF = np.ceil(
+                sqrt(self._encut / RYTOEV) / (TPI / (Anorm / AUTOA))
+        )
+        self._ngrid = np.array(2 * CUTOF + 1, dtype=int)
+        
+    def setWFPrec(self):
+        '''
+        Set wavefunction coefficients precision:
+            TAG = 45200: single precision complex, np.complex64, or complex(qs)
+            TAG = 45210: double precision complex, np.complex128, or complex(q)
+        '''
+        if self._rtag == 45200:
+            return np.complex64
+        elif self._rtag == 45210:
+            return np.complex128
+        elif self._rtag == 53300:
+            raise ValueError("VASP5 WAVECAR format, not implemented yet")
+        elif self._rtag == 53310:
+            raise ValueError("VASP5 WAVECAR format with double precision "
+                            +"coefficients, not implemented yet")
+        else:
+            raise ValueError("Invalid TAG values: {}".format(self._rtag))
+
+    def readWFBand(self):
+        '''
+        Extract KS energies and Fermi occupations from WAVECAR.
+        '''
+
+        self._nplws = np.zeros(self._nkpts, dtype=int)
+        self._kvecs = np.zeros((self._nkpts, 3), dtype=float)
+        self._bands = np.zeros((self._nspin, self._nkpts, self._nbands), dtype=float)
+        self._occs  = np.zeros((self._nspin, self._nkpts, self._nbands), dtype=float)
+
+        for ii in range(self._nspin):
+            for jj in range(self._nkpts):
+                rec = self.whereRec(ii+1, jj+1, 1) - 1
+                self._wfc.seek(rec * self._recl)
+                dump = np.fromfile(self._wfc, dtype=np.float, count=4+3*self._nbands)
+                if ii == 0:
+                    self._nplws[jj] = int(dump[0])
+                    self._kvecs[jj] = dump[1:4]
+                dump = dump[4:].reshape((-1, 3))
+                self._bands[ii,jj,:] = dump[:,0]
+                self._occs[ii,jj,:] = dump[:,2]
+
+        if self._nkpts > 1:
+            tmp = np.linalg.norm(
+                    np.dot(np.diff(self._kvecs, axis=0), self._Bcell), axis=1)
+            self._kpath = np.concatenate(([0,], np.cumsum(tmp)))
+        else:
+            self._kpath = None
+        return  self._kpath, self._bands
+
+    def get_kpath(self, nkseg=None):
+        '''
+        Construct k-point path, find out the k-path boundary if possible.
+
+        nkseg is the number of k-points in each k-path segments.
+        '''
+
+        if nkseg is None:
+            if os.path.isfile(self._dname + "/KPOINTS"):
+                kfile = open(self._dname + "/KPOINTS").readlines()
+                if kfile[2][0].upper() == 'L':
+                    nkseg = int(kfile[1].split()[0])
+                else:
+                    raise ValueError('Error reading number of k-points from KPOINTS')
+        assert nkseg > 0
+
+        nsec  = self._nkpts // nkseg
+
+        v = self._kvecs.copy()
+        for ii in range(nsec):
+            ki = ii * nkseg
+            kj = (ii + 1) * nkseg
+            v[ki:kj,:] -= v[ki]
+
+        self._kpath = np.linalg.norm(np.dot(v, self._Bcell), axis=1)
+        for ii in range(1, nsec):
+            ki = ii * nkseg
+            kj = (ii + 1) * nkseg
+            self._kpath[ki:kj] += self._kpath[ki - 1]
+
+            self._kbound =  np.concatenate((self._kpath[0::nkseg], [self._kpath[-1],]))
+
+        return self._kpath, self._kbound
+
+    def gvectors(self, ikpt=1, force_Gamma=False, check_consistency=True):
+        '''
+        Generate the G-vectors that satisfies the following relation
+            (G + k)**2 / 2 < ENCUT
+        '''
+        assert 1 <= ikpt  <= self._nkpts,  'Invalid kpoint index!'
+
+        
+        kvec = self._kvecs[ikpt-1]
+        # fx, fy, fz = [fftfreq(n) * n for n in self._ngrid]
+        # fftfreq in scipy.fftpack is a little different with VASP frequencies
+        fx = [ii if ii < self._ngrid[0] // 2 + 1 else ii - self._ngrid[0]
+                for ii in range(self._ngrid[0])]
+        fy = [jj if jj < self._ngrid[1] // 2 + 1 else jj - self._ngrid[1]
+                for jj in range(self._ngrid[1])]
+        fz = [kk if kk < self._ngrid[2] // 2 + 1 else kk - self._ngrid[2]
+                for kk in range(self._ngrid[2])]
+
+        # force_Gamma: consider gamma-only case regardless of the real setting
+        lgam = True if force_Gamma else self._lgam
+        if lgam:
+            # parallel gamma version of VASP WAVECAR exclude some planewave
+            # components, -DwNGZHalf
+            if self._gam_half == 'z':
+                kgrid = np.array([(fx[ii], fy[jj], fz[kk])
+                                  for kk in range(self._ngrid[2])
+                                  for jj in range(self._ngrid[1])
+                                  for ii in range(self._ngrid[0])
+                                  if (
+                                      (fz[kk] > 0) or
+                                      (fz[kk] == 0 and fy[jj] > 0) or
+                                      (fz[kk] == 0 and fy[jj] == 0 and fx[ii] >= 0)
+                                  )], dtype=float)
+            else:
+                kgrid = np.array([(fx[ii], fy[jj], fz[kk])
+                                  for kk in range(self._ngrid[2])
+                                  for jj in range(self._ngrid[1])
+                                  for ii in range(self._ngrid[0])
+                                  if (
+                                      (fx[ii] > 0) or
+                                      (fx[ii] == 0 and fy[jj] > 0) or
+                                      (fx[ii] == 0 and fy[jj] == 0 and fz[kk] >= 0)
+                                  )], dtype=float)
+        else:
+            kgrid = np.array([(fx[ii], fy[jj], fz[kk])
+                              for kk in range(self._ngrid[2])
+                              for jj in range(self._ngrid[1])
+                              for ii in range(self._ngrid[0])], dtype=float)
+
+        # Kinetic_Energy = (G + k)**2 / 2
+        # HSQDTM    =  hbar**2/(2*ELECTRON MASS)
+        KENERGY = HSQDTM * np.linalg.norm(
+                    np.dot(kgrid + kvec[np.newaxis,:] , TPI*self._Bcell), axis=1
+                )**2
+        # find Gvectors where (G + k)**2 / 2 < ENCUT
+        Gvec = kgrid[np.where(KENERGY < self._encut)[0]]
+
+        
+        # Check if the calculated number of planewaves and the one recorded in the
+        # WAVECAR are equal
+        if check_consistency:
+            if self._lsoc:
+                    assert Gvec.shape[0] == self._nplws[ikpt - 1] // 2, \
+                           'No. of planewaves not consistent for an SOC WAVECAR! %d %d %d' % \
+                           (Gvec.shape[0], self._nplws[ikpt -1], np.prod(self._ngrid))
+            else:
+                assert Gvec.shape[0] == self._nplws[ikpt - 1], 'No. of planewaves not consistent! %d %d %d' % \
+                        (Gvec.shape[0], self._nplws[ikpt -1], np.prod(self._ngrid))
+
+        return np.asarray(Gvec, dtype=int)
+
+    def save2vesta(self, phi=None, lreal=False, poscar='POSCAR', prefix='wfc',
+                   ncol=10):
+        '''
+        Save the real space pseudo-wavefunction as vesta format.
+        '''
+        nx, ny, nz = phi.shape
+        try:
+            pos = open(poscar, 'r')
+            head = ''
+            for line in pos:
+                if line.strip():
+                    head += line
+                else:
+                    break
+            head += '\n%5d%5d%5d\n' % (nx, ny, nz)
+        except:
+            raise IOError('Failed to open %s' % poscar)
+
+        # Faster IO
+        nrow = phi.size // ncol
+        nrem = phi.size % ncol
+        fmt = "%16.8E"
+
+        psi = phi.copy()
+        psi = psi.flatten(order='F')
+        psi_h = psi[:nrow * ncol].reshape((nrow, ncol))
+        psi_r = psi[nrow * ncol:]
+
+        with open(prefix + '_r.vasp', 'w') as out:
+            out.write(head)
+            out.write(
+                '\n'.join([''.join([fmt % xx for xx in row])
+                           for row in psi_h.real])
+            )
+            out.write("\n" + ''.join([fmt % xx for xx in psi_r.real]))
+        if not (self._lgam or lreal):
+            with open(prefix + '_i.vasp', 'w') as out:
+                out.write(head)
+                out.write(
+                    '\n'.join([''.join([fmt % xx for xx in row])
+                               for row in psi_h.imag])
+                )
+                out.write("\n" + ''.join([fmt % xx for xx in psi_r.imag]))
+
+    def wfc_r(self, ispin=1, ikpt=1, iband=1,
+                    gvec=None, Cg=None, ngrid=None,
+                    rescale=None,
+                    norm=True):
+        '''
+        Obtain the pseudo-wavefunction of the specified KS states in real space
+        by performing FT transform on the reciprocal space planewave
+        coefficients.  The 3D FT grid size is determined by ngrid, which
+        defaults to self._ngrid if not given.  Gvectors of the KS states is used
+        to put 1D planewave coefficients back to 3D grid.
+
+        Inputs:
+            ispin : spin index of the desired KS states, starting from 1
+            ikpt  : k-point index of the desired KS states, starting from 1
+            iband : band index of the desired KS states, starting from 1
+            gvec  : the G-vectors correspond to the plane-wave coefficients
+            Cg    : the plane-wave coefficients. If None, read from WAVECAR
+            ngrid : the FFT grid size
+            norm  : normalized Cg?
+
+        The return wavefunctions are normalized in a way that
+
+                        \sum_{ijk} | \phi_{ijk} | ^ 2 = 1
+            
+        '''
+        self.checkIndex(ispin, ikpt, iband)
+
+        if ngrid is None:
+            ngrid = self._ngrid.copy() * 2
+        else:
+            ngrid = np.array(ngrid, dtype=int)
+            assert ngrid.shape == (3,)
+            assert np.alltrue(ngrid >= self._ngrid), \
+                    "Minium FT grid size: (%d, %d, %d)" % \
+                    (self._ngrid[0], self._ngrid[1], self._ngrid[2])
+
+        # The default normalization of np.fft.fftn has the direct transforms
+        # unscaled and the inverse transforms are scaled by 1/n. It is possible
+        # to obtain unitary transforms by setting the keyword argument norm to
+        # "ortho" (default is None) so that both direct and inverse transforms
+        # will be scaled by 1/\sqrt{n}.
+
+        # default normalization factor so that 
+        # \sum_{ijk} | \phi_{ijk} | ^ 2 = 1
+        normFac = rescale if rescale is not None else np.sqrt(np.prod(ngrid)) 
+
+        if gvec is None:
+            gvec = self.gvectors(ikpt)
+
+        if self._lgam:
+            if self._gam_half == 'z':
+                phi_k = np.zeros((ngrid[0], ngrid[1], ngrid[2]//2 + 1), dtype=np.complex128)
+            else:
+                phi_k = np.zeros((ngrid[0]//2 + 1, ngrid[1], ngrid[2]), dtype=np.complex128)
+        else:
+            phi_k = np.zeros(ngrid, dtype=np.complex128)
+
+            
+        gvec %= ngrid[np.newaxis,:]
+        
+        if self._lsoc:
+            wfc_spinor = []
+            if Cg:
+                dump = Cg
+            else:
+                dump = self.readBandCoeff(ispin, ikpt, iband, norm)
+            nplw = dump.shape[0] // 2
+            
+            # spinor up
+            phi_k[gvec[:,0], gvec[:,1], gvec[:,2]] = dump[:nplw]
+            wfc_spinor.append(ifftn(phi_k) * normFac)
+            # spinor down
+            phi_k[:,:,:] = 0.0j
+            phi_k[gvec[:,0], gvec[:,1], gvec[:,2]] = dump[nplw:]
+            wfc_spinor.append(ifftn(phi_k) * normFac)
+
+            del dump
+            return wfc_spinor
+            
+        else:
+            if Cg is not None:
+                phi_k[gvec[:,0], gvec[:,1], gvec[:,2]] = Cg
+            else:
+                phi_k[gvec[:,0], gvec[:,1], gvec[:,2]] = self.readBandCoeff(ispin, ikpt, iband, norm)
+
+            if self._lgam:
+                # add some components that are excluded and perform c2r FFT
+                if self._gam_half == 'z':
+                    for ii in range(ngrid[0]):
+                        for jj in range(ngrid[1]):
+                            fx = ii if ii < ngrid[0] // 2 + 1 else ii - ngrid[0]
+                            fy = jj if jj < ngrid[1] // 2 + 1 else jj - ngrid[1]
+                            if (fy > 0) or (fy == 0 and fx >= 0):
+                                continue
+                            phi_k[ii,jj,0] = phi_k[-ii,-jj,0].conjugate()
+
+                    phi_k /= np.sqrt(2.)
+                    phi_k[0,0,0] *= np.sqrt(2.)
+                    return np.fft.irfftn(phi_k, s=ngrid) * normFac
+                elif self._gam_half == 'x':
+                    for jj in range(ngrid[1]):
+                        for kk in range(ngrid[2]):
+                            fy = jj if jj < ngrid[1] // 2 + 1 else jj - ngrid[1]
+                            fz = kk if kk < ngrid[2] // 2 + 1 else kk - ngrid[2]
+                            if (fy > 0) or (fy == 0 and fz >= 0):
+                                continue
+                            phi_k[0,jj,kk] = phi_k[0,-jj,-kk].conjugate()
+
+                    phi_k /= np.sqrt(2.)
+                    phi_k[0,0,0] *= np.sqrt(2.)
+                    phi_k = np.swapaxes(phi_k, 0, 2)
+                    tmp =  np.fft.irfftn(phi_k, s=(ngrid[2], ngrid[1], ngrid[0])) * normFac
+                    return np.swapaxes(tmp, 0, 2)
+            else:
+                # perform complex2complex FFT
+                return ifftn(phi_k * normFac)
+
+    def readBandCoeff(self, ispin=1, ikpt=1, iband=1, norm=False):
+        '''
+        Read the planewave coefficients of specified KS states.
+        '''
+
+        self.checkIndex(ispin, ikpt, iband)
+
+        rec = self.whereRec(ispin, ikpt, iband)
+        self._wfc.seek(rec * self._recl)
+
+        nplw = self._nplws[ikpt - 1]
+        dump = np.fromfile(self._wfc, dtype=self._WFPrec, count=nplw)
+
+        cg = np.asarray(dump, dtype=np.complex128)
+        if norm:
+            cg /= np.linalg.norm(cg)
+        return cg
+
+    def whereRec(self, ispin=1, ikpt=1, iband=1):
+        '''
+        Return the rec position for specified KS state.
+        '''
+
+        self.checkIndex(ispin, ikpt, iband)
+
+        rec = 2 + (ispin - 1) * self._nkpts * (self._nbands + 1) + \
+                  (ikpt - 1) * (self._nbands + 1) + \
+                  iband
+        return rec
+
+    def checkIndex(self, ispin, ikpt, iband):
+        '''
+        Check if the index is valid!
+        '''
+        assert 1 <= ispin <= self._nspin,  'Invalid spin index!'
+        assert 1 <= ikpt  <= self._nkpts,  'Invalid kpoint index!'
+        assert 1 <= iband <= self._nbands, 'Invalid band index!'
+
+    def TransitionDipoleMoment(self, ks_i, ks_j, norm=True, 
+                                     realspace=False):
+        '''
+        calculate Transition Dipole Moment (TDM) between two KS states.
+
+        If "realspace = False", the TDM will be evaluated in momentum space
+        according to the formula in:
+        https://en.wikipedia.org/wiki/Transition_dipole_moment
+
+                                    i⋅h
+        <psi_a | r | psi_b> =  -------------- ⋅ <psi_a | p | psi_b>
+                                m⋅(Eb - Ea)
+
+                                      2        ____              
+                                     h          ╲                
+                            =  ------------- ⋅   ╲   Cai⋅Cbi⋅Gi
+                                m⋅(Eb - Ea)      ╱               
+                                                ╱                
+                                               ‾‾‾‾              
+                                                 i               
+
+        Otherwise, the TDM will be evaluated in real space.
+
+        Note: |psi_a> and |psi_b> should be bloch function with 
+              the same k vector.
+
+        The KS states ks_i (ks_j) is specified by list of index (ispin, ikpt, iband).
+        '''
+
+        ks_i = list(ks_i); ks_j = list(ks_j)
+        assert len(ks_i) == len(ks_j) == 3, 'Must be three indexes!'
+        assert ks_i[1] == ks_j[1], 'k-point of the two states differ!'
+        self.checkIndex(*ks_i)
+        self.checkIndex(*ks_j)
+
+        # energy differences between the two states
+        E1 = self._bands[ks_i[0]-1, ks_i[1]-1, ks_i[2]-1]
+        E2 = self._bands[ks_j[0]-1, ks_j[1]-1, ks_j[2]-1]
+        dE = E2 - E1
+
+        if realspace:
+            fx = np.linspace(0, 1, self._ngrid[0], endpoint=False)
+            fy = np.linspace(0, 1, self._ngrid[1], endpoint=False)
+            fz = np.linspace(0, 1, self._ngrid[2], endpoint=False)
+
+            Dx, Dy, Dz = np.meshgrid(fx, fy, fz, indexing='ij')
+            Rx, Ry, Rz = np.tensordot(self._Acell, [Dx, Dy, Dz], axes=[0,0])
+
+            fac = np.sqrt(np.prod(self._ngrid) / self._Omega)
+            phi_i = self.wfc_r(*ks_i, norm=True, ngrid=self._ngrid)
+            phi_j = self.wfc_r(*ks_j, norm=True, ngrid=self._ngrid)
+
+            pij = phi_i.conjugate() * phi_j
+            tdm = np.array([
+                np.sum(pij * Rx),
+                np.sum(pij * Ry),
+                np.sum(pij * Rz)
+                ])
+            ovlap = pij.sum()
+        else:
+            # according to the above equation, G = 0 does NOT contribute to TDM.
+            gvec = np.dot(self.gvectors(ikpt=ks_i[1]), self._Bcell*TPI)
+            # planewave coefficients of the two states
+            phi_i = self.readBandCoeff(*ks_i, norm=norm)
+            phi_j = self.readBandCoeff(*ks_j, norm=norm)
+
+            tmp1 = phi_i.conjugate() * phi_j
+            ovlap = np.sum(tmp1)
+            if self._lgam:
+                tmp2 = phi_i * phi_j.conjugate()
+                # according to the above equation, G = 0 does NOT contribute to TDM.
+                tdm = (np.sum(tmp1[:,np.newaxis] * gvec, axis=0) -
+                       np.sum(tmp2[:,np.newaxis] * gvec, axis=0)) / 2.
+            else:
+                tdm = np.sum(tmp1[:,np.newaxis] * gvec, axis=0)
+
+            tdm = 1j / (dE / (2*RYTOEV)) * tdm * AUTOA * AUTDEBYE
+
+        return E1, E2, dE, ovlap, tdm
+
+    def inverse_participation_ratio(self, norm=True):
+        '''
+        Calculate Inverse Paticipation Ratio (IPR) from the wavefunction. IPR is
+        a measure of the localization of Kohn-Sham states. For a particular KS
+        state \phi_j, it is defined as
+            
+                            \sum_n |\phi_j(n)|^4 
+            IPR(\phi_j) = -------------------------
+                          |\sum_n |\phi_j(n)|^2||^2
+
+        where n iters over the number of grid points.
+        '''
+
+        self.ipr = np.zeros((self._nspin, self._nkpts, self._nbands, 3))
+
+        for ispin in range(self._nspin):
+            for ikpt in range(self._nkpts):
+                for iband in range(self._nbands):
+                    phi_j = self.wfc_r(ispin+1, ikpt+1, iband+1,
+                                       norm=norm)
+                    phi_j_abs = np.abs(phi_j)
+
+                    print('Calculating IPR of #spin %4d, #kpt %4d, #band %4d' % (ispin+1, ikpt+1, iband+1))
+                    self.ipr[ispin, ikpt, iband, 0] = self._kpath[ikpt] if self._kpath is None else 0
+                    self.ipr[ispin, ikpt, iband, 1] = self._bands[ispin, ikpt, iband]
+                    self.ipr[ispin, ikpt, iband, 2] = np.sum(phi_j_abs**4) / np.sum(phi_j_abs**2)**2
+
+        np.save('ipr.npy', self.ipr)
+        return self.ipr
+
+    def elf(self, kptw, ngrid=None, warn=True):
+        '''
+        Calculate the electron localization function (ELF) from WAVECAR.
+
+        The following formula was extracted from VASP ELF.F:
+                     _
+                     h^2    *    2      T.........kinetic energy
+          T    =  -2 --- Psi grad Psi   T+TCORR...pos.definite kinetic energy
+                   ^ 2 m                TBOS......T of an ideal Bose-gas
+                   ^
+                   I am not sure if we need to times 2 here, use 1 in this
+                   script.
+
+                   _                                (=infimum of T+TCORR)
+                 1 h^2      2           DH........T of hom.non-interact.e- - gas
+          TCORR= - ---  grad rho                    (acc.to Fermi)
+                 2 2 m                  ELF.......electron-localization-function
+                   _             2
+                 1 h^2 |grad rho|
+          TBOS = - --- ----------       D = T + TCORR - TBOS
+                 4 2 m    rho
+                   _                                \                1
+                 3 h^2        2/3  5/3          =====>    ELF = ------------
+          DH   = - --- (3 Pi^2)  rho                /                   D   2
+                 5 2 m                                           1 + ( ---- )
+                                                                        DH
+
+        REF:
+            1. Nature, 371, 683-686 (1994)
+            2. Becke and Edgecombe, J. Chem. Phys., 92, 5397(1990)
+            3. M. Kohout and A. Savin, Int. J. Quantum Chem., 60, 875-882(1996)
+            4. http://www2.cpfs.mpg.de/ELF/index.php?content=06interpr.txt
+        '''
+
+        if warn:
+            print("""
+            ###################################################################
+            If you are using VESTA to view the resulting ELF, please rename the
+            output file as ELFCAR, otherwise there will be some error in the
+            isosurface plot!
+
+            When CHG*/PARCHG/*.vasp are read in to visualize isosurfaces and
+            sections, data values are divided by volume in the unit of bohr^3.
+            The unit of charge densities input by VESTA is, therefore, bohr^−3.
+
+            For LOCPOT/ELFCAR files, volume data are kept intact.
+
+            You can turn off this warning by setting "warn=False" in the "elf"
+            method.
+            ###################################################################
+            """)
+
+        # the k-point weights
+        kptw = np.array(kptw, dtype=float)
+        assert kptw.shape == (self._nkpts,), "K-point weights must be provided \
+                                              to calculate charge density!"
+        # normalization
+        kptw /= kptw.sum()
+
+        if ngrid is None:
+            ngrid = self._ngrid * 2
+        else:
+            ngrid = np.array(ngrid, dtype=int)
+            assert ngrid.shape == (3,)
+            assert np.alltrue(ngrid >= self._ngrid), \
+                    "Minium FT grid size: (%d, %d, %d)" % \
+                    (self._ngrid[0], self._ngrid[1], self._ngrid[2])
+
+        fx = [ii if ii < ngrid[0] // 2 + 1 else ii - ngrid[0]
+                for ii in range(ngrid[0])]
+        fy = [jj if jj < ngrid[1] // 2 + 1 else jj - ngrid[1]
+                for jj in range(ngrid[1])]
+        fz = [kk if kk < ngrid[2] // 2 + 1 else kk - ngrid[2]
+                for kk in range(ngrid[2])]
+
+        # plane-waves: Reciprocal coordinate 
+        # indexing = 'ij' so that outputs are of shape (ngrid[0], ngrid[1], ngrid[2])
+        Dx, Dy, Dz = np.meshgrid(fx, fy, fz, indexing='ij')
+        # plane-waves: Cartesian coordinate 
+        Gx, Gy, Gz = np.tensordot(self._Bcell * np.pi * 2, [Dx, Dy, Dz], axes=(0,0))
+        # the norm squared of the G-vectors
+        G2 = Gx**2 + Gy**2 + Gz**2
+        # k-points vectors in Cartesian coordinate
+        vkpts = np.dot(self._kvecs, self._Bcell * 2 * np.pi)
+
+        # normalization factor so that 
+        # \sum_{ijk} | \phi_{ijk} | ^ 2 * volume / Ngrid = 1
+        normFac = np.sqrt(np.prod(ngrid) / self._Omega)
+
+        # electron localization function
+        ElectronLocalizationFunction = []
+        # Charge density
+        rho = np.zeros(ngrid, dtype=complex)
+        # Kinetic energy density
+        tau = np.zeros(ngrid, dtype=complex)
+
+        for ispin in range(self._nspin):
+            # initialization
+            rho[...] = 0.0
+            tau[...] = 0.0
+
+            for ikpt in range(self._nkpts):
+
+                # plane-wave G-vectors
+                igvec = self.gvectors(ikpt+1)
+                # for gamma-only version, complete the missing -G vectors
+                if self._lgam:
+                    tmp   = np.array([-k for k in igvec[1:]], dtype=int)
+                    igvec = np.vstack([igvec, tmp])
+                # plane-wave G-vectors in Cartesian coordinate
+                rgvec = np.dot(igvec, self._Bcell * 2 * np.pi)
+
+                k   = vkpts[ikpt]                       # k
+                gk  = rgvec + k[np.newaxis,:]           # G + k
+                gk2 = np.linalg.norm(gk, axis=1)**2     # | G + k |^2
+
+                for iband in range(self._nbands):
+                    # omit the empty bands
+                    if self._occs[ispin, ikpt, iband] == 0.0: continue
+
+                    rspin = 2.0 if self._nspin == 1 else 1.0
+                    weight = rspin * kptw[ikpt] * self._occs[ispin, ikpt, iband]
+
+                    # if self._lgam:
+                    #     ########################################
+                    #     # slower
+                    #     ########################################
+                    #     # wavefunction in real space
+                    #     # VASP does NOT do normalization in elf.F
+                    #     phi_r  = self.wfc_r(ispin=ispin+1, ikpt=ikpt+1,
+                    #                         iband=iband+1,
+                    #                         ngrid=ngrid,
+                    #                         norm=False) * normFac
+                    #     # wavefunction in reciprocal space
+                    #     phi_q  = np.fft.fftn(phi_r, norm='ortho')
+                    #     # grad^2 \phi in reciprocal space
+                    #     lap_phi_q = -gk2 * phi_q
+                    #     # grad^2 \phi in real space
+                    #     lap_phi_r = np.fft.ifftn(lap_phi_q, norm='ortho')
+                    # else:
+
+                    ########################################
+                    # faster
+                    ########################################
+                    # wavefunction in reciprocal space
+                    # VASP does NOT do normalization in elf.F
+                    phi_q = self.readBandCoeff(ispin=ispin+1, ikpt=ikpt+1,
+                                               iband=iband+1,
+                                               norm=False)
+                    # pad the missing planewave coefficients for -G vectors
+                    if self._lgam:
+                        tmp   = [x.conj() for x in phi_q[1:]]
+                        phi_q = np.concatenate([phi_q, tmp])
+                        # Gamma only, divide a factor of sqrt(2.0) except for
+                        # G=0
+                        phi_q    /= np.sqrt(2.0)
+                        phi_q[0] *= np.sqrt(2.0)
+                    # wavefunction in real space
+                    phi_r  = self.wfc_r(ispin=ispin+1, ikpt=ikpt+1,
+                                        iband=iband+1,
+                                        ngrid=ngrid,
+                                        gvec=igvec,
+                                        Cg=phi_q,
+                                        norm=True) * normFac
+                    # grad^2 \phi in reciprocal space
+                    lap_phi_q = -gk2 * phi_q
+                    # grad^2 \phi in real space
+                    lap_phi_r = self.wfc_r(ispin=ispin+1, ikpt=ikpt+1,
+                                           iband=iband+1,
+                                           ngrid=ngrid,
+                                           gvec=igvec,
+                                           Cg=lap_phi_q) * normFac
+
+                    # \phi* grad^2 \phi in real space --> kinetic energy density
+                    tau += -phi_r * lap_phi_r.conj() * weight
+
+                    # charge density in real space
+                    rho += phi_r.conj() * phi_r * weight
+
+            # charge density in reciprocal space
+            rho_q = np.fft.fftn(rho, norm='ortho')
+
+            # grad^2 rho: laplacian of charge density
+            lap_rho_q = -G2 * rho_q
+            lap_rho_r = np.fft.ifftn(lap_rho_q, norm='ortho')
+
+            # charge density gradient: grad rho
+            ########################################
+            # wrong method for gradient using FFT 
+            ########################################
+            # grad_rho_x = np.fft.ifft(1j * Gx * np.fft.fft(rho, axis=0), axis=0)
+            # grad_rho_y = np.fft.ifft(1j * Gy * np.fft.fft(rho, axis=1), axis=1)
+            # grad_rho_z = np.fft.ifft(1j * Gz * np.fft.fft(rho, axis=2), axis=2)
+
+            ########################################
+            # correct method for gradient using FFT 
+            ########################################
+            grad_rho_x = np.fft.ifftn(1j * Gx * rho_q, norm='ortho')
+            grad_rho_y = np.fft.ifftn(1j * Gy * rho_q, norm='ortho')
+            grad_rho_z = np.fft.ifftn(1j * Gz * rho_q, norm='ortho')
+
+            grad_rho_sq = np.abs(grad_rho_x)**2 \
+                        + np.abs(grad_rho_y)**2 \
+                        + np.abs(grad_rho_z)**2
+
+            rho = rho.real
+            tau = tau.real
+            lap_rho_r = lap_rho_r.real
+
+            Cf  = 3./5 * (3.0 * np.pi**2)**(2./3)
+            Dh = np.where(rho > 0.0,
+                          Cf * rho**(5./3),
+                          0.0)
+            eps = 1E-8 / HSQDTM
+            Dh[Dh < eps] = eps
+            # D0 = T + TCORR - TBOS
+            D0 = tau + 0.5 * lap_rho_r - 0.25 * grad_rho_sq / rho
+
+            ElectronLocalizationFunction.append(1. / (1. + (D0 / Dh)**2))
+        
+        return ElectronLocalizationFunction
+
+############################################################
+
+if __name__ == '__main__':
+    # xx = vaspwfc('wavecar')
+    # phi = xx.wfc_r(1, 30, 17, ngrid=(28, 28, 252))
+    # xx.save2vesta(phi, poscar='POSCAR')
+
+    # xx = vaspwfc('./gamma/WAVECAR')
+    # phi = xx.wfc_r(1, 1, 317, ngrid=(60, 108, 160),
+    #                gamma=True)
+    # xx.save2vesta(phi, poscar='./gamma/POSCAR',gamma=True)
+
+    # xx = vaspwfc('WAVECAR')
+    # dE, ovlap, tdm = xx.TransitionDipoleMoment([1,30,17], [1,30,18], norm=True)
+    # print dE, ovlap.real, np.abs(tdm)**2
+
+    # print xx._recl, xx._nspin, xx._rtag
+    # print xx._nkpts, xx._nbands, xx._encut
+    # print xx._Acell, xx._Bcell
+    # # print np.linalg.norm(xx._Acell, axis=1)
+    # print xx._ngrid
+    # print xx._bands[0,0,:]
+    # print xx._kvecs
+    # print xx._kpath
+    # b = xx.readBandCoeff(1,1,1)
+    # xx = np.savetxt('kaka.dat', xx.gvectors(2), fmt='%5d')
+    # gvec = xx.gvectors(1)
+    # gvec %= xx._ngrid[np.newaxis, :]
+    # print gvec
+
+    # ngrid=(28, 28, 252)
+    # phi = xx.wfc_r(1, 30, 17, ngrid=(28, 28, 252))
+    # header = open('POSCAR').read()
+    # with open('wave_real.vasp', 'w') as out:
+    #     out.write(header)
+    #     out.write('%5d%5d%5d\n' % (ngrid[0], ngrid[1], ngrid[2]))
+    #     nwrite=0
+    #     for kk in range(ngrid[2]):
+    #         for jj in range(ngrid[1]):
+    #             for ii in range(ngrid[0]):
+    #                 nwrite += 1
+    #                 out.write('%22.16f ' % phi.real[ii,jj,kk])
+    #                 if nwrite % 10 == 0:
+    #                     out.write('\n')
+    # with open('wave_imag.vasp', 'w') as out:
+    #     out.write(header)
+    #     out.write('%5d%5d%5d\n' % (ngrid[0], ngrid[1], ngrid[2]))
+    #     nwrite=0
+    #     for kk in range(ngrid[2]):
+    #         for jj in range(ngrid[1]):
+    #             for ii in range(ngrid[0]):
+    #                 nwrite += 1
+    #                 out.write('%22.16f ' % phi.imag[ii,jj,kk])
+    #                 if nwrite % 10 == 0:
+    #                     out.write('\n')
+
+    # xx = vaspwfc('wave_tyz')
+    # ipr = xx.inverse_participation_ratio()
+    # print xx._nbands, xx._nkpts
+    #
+    # import matplotlib as mpl
+    # import matplotlib.pyplot as plt
+    #
+    # fig = plt.figure()
+    # ax = plt.subplot()
+    #
+    # ax.scatter(ipr[...,0], ipr[..., 1], s=ipr[..., 2] / ipr[..., 2].max() * 10, c=ipr[..., 2], 
+    #            cmap='jet_r')
+    #
+    # plt.show()
+
+    wfc = vaspwfc('WAVECAR', lgamma=True, gamma_half='x')
+    # ngrid = [80, 140, 210]
+    phi = wfc.wfc_r(iband=190)
+
+    rho = np.abs(phi)**2
+    # rho2 = VaspChargeDensity('PARCHG.0158.ALLK').chg[0]
+    # rho /= rho.sum()
+    # rho2 /= rho2.sum()
+    # rho3 = rho - rho2
+
+    wfc.save2vesta(rho, lreal=True)
+
+    pass