diff --git a/pyiron_vasp/dft/bader.py b/pyiron_vasp/dft/bader.py index 69cb42fc..f59b1a97 100644 --- a/pyiron_vasp/dft/bader.py +++ b/pyiron_vasp/dft/bader.py @@ -2,10 +2,12 @@ # Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department # Distributed under the terms of "New BSD License", see the LICENSE file. +from typing import Optional, Tuple import numpy as np import os import subprocess +from ase.atoms import Atoms from pyiron_vasp.vasp.volumetric_data import VaspVolumetricData @@ -29,7 +31,7 @@ class Bader: .. _Bader code: http://theory.cm.utexas.edu/henkelman/code/bader """ - def __init__(self, structure, working_directory): + def __init__(self, structure: Atoms, working_directory: str) -> None: """ Initialize the Bader module @@ -39,7 +41,7 @@ def __init__(self, structure, working_directory): self._working_directory = working_directory self._structure = structure - def _create_cube_files(self): + def _create_cube_files(self) -> None: """ Create CUBE format files of the total and valce charges to be used by the Bader program """ @@ -53,7 +55,9 @@ def _create_cube_files(self): filename=os.path.join(self._working_directory, "total_charge.CUBE") ) - def compute_bader_charges(self, extra_arguments=None): + def compute_bader_charges( + self, extra_arguments: Optional[str] = None + ) -> Tuple[np.ndarray, np.ndarray]: """ Run Bader analysis on the output from the DFT job @@ -74,14 +78,14 @@ def compute_bader_charges(self, extra_arguments=None): self._remove_cube_files() return self._parse_charge_vol() - def _remove_cube_files(self): + def _remove_cube_files(self) -> None: """ Delete created CUBE files """ os.remove(os.path.join(self._working_directory, "valence_charge.CUBE")) os.remove(os.path.join(self._working_directory, "total_charge.CUBE")) - def _parse_charge_vol(self): + def _parse_charge_vol(self) -> Tuple[np.ndarray, np.ndarray]: """ Parse Bader charges and volumes @@ -93,7 +97,7 @@ def _parse_charge_vol(self): return parse_charge_vol_file(structure=self._structure, filename=filename) -def call_bader(foldername, extra_arguments=None): +def call_bader(foldername: str, extra_arguments: Optional[str] = None) -> int: """ Call the Bader program inside a given folder @@ -111,7 +115,9 @@ def call_bader(foldername, extra_arguments=None): return subprocess.call(cmd, shell=True, cwd=foldername) -def parse_charge_vol_file(structure, filename="ACF.dat"): +def parse_charge_vol_file( + structure: Atoms, filename: str = "ACF.dat" +) -> Tuple[np.ndarray, np.ndarray]: """ Parse charges and volumes from the output file @@ -130,7 +136,9 @@ def parse_charge_vol_file(structure, filename="ACF.dat"): return charges, volumes -def get_valence_and_total_charge_density(working_directory): +def get_valence_and_total_charge_density( + working_directory: str, +) -> Tuple[VaspVolumetricData, VaspVolumetricData]: """ Gives the valence and total charge densities diff --git a/pyiron_vasp/dft/volumetric.py b/pyiron_vasp/dft/volumetric.py index 9673d897..3c0f9933 100644 --- a/pyiron_vasp/dft/volumetric.py +++ b/pyiron_vasp/dft/volumetric.py @@ -2,6 +2,7 @@ # Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department # Distributed under the terms of "New BSD License", see the LICENSE file. +from typing import Union, List, Tuple, Optional import numpy as np from ase.atoms import Atoms from pyiron_vasp.vasp.structure import write_poscar @@ -31,12 +32,12 @@ class VolumetricData(object): """ - def __init__(self): - self._total_data = None - self._atoms = None + def __init__(self) -> None: + self._total_data: Optional[np.ndarray] = None + self._atoms: Optional[Atoms] = None @property - def atoms(self): + def atoms(self) -> Optional[Atoms]: """ The structure related to the volumeric data @@ -47,18 +48,18 @@ def atoms(self): return self._atoms @atoms.setter - def atoms(self, val): + def atoms(self, val: Optional[Atoms]) -> None: self._atoms = val @property - def total_data(self): + def total_data(self) -> Optional[np.ndarray]: """ numpy.ndarray: The Nx x Ny x Nz sized array for the total data """ return self._total_data @total_data.setter - def total_data(self, val): + def total_data(self, val: Optional[Union[np.ndarray, list]]) -> None: if not (isinstance(val, (np.ndarray, list))): raise TypeError( "Attribute total_data should be a numpy.ndarray instance or a list and " @@ -71,7 +72,7 @@ def total_data(self, val): self._total_data = val @staticmethod - def gauss_f(d, fwhm=0.529177): + def gauss_f(d: float, fwhm: float = 0.529177) -> float: """ Generates a Gaussian distribution for a given distance and full width half maximum value @@ -89,8 +90,11 @@ def gauss_f(d, fwhm=0.529177): @staticmethod def dist_between_two_grid_points( - target_grid_point, n_grid_at_center, lattice, grid_shape - ): + target_grid_point: Union[np.ndarray, list], + n_grid_at_center: Union[np.ndarray, list], + lattice: Union[np.ndarray, list], + grid_shape: Union[tuple, list, np.ndarray], + ) -> float: """ Calculates the distance between a target grid point and another grid point @@ -114,11 +118,15 @@ def dist_between_two_grid_points( np.subtract(target_grid_point, n_grid_at_center), unit_dist_in_grid ) dist = np.linalg.norm(dn) - return dist + return float(dist) def spherical_average_potential( - self, structure, spherical_center, rad=2, fwhm=0.529177 - ): + self, + structure: Atoms, + spherical_center: Union[list, np.ndarray], + rad: float = 2, + fwhm: float = 0.529177, + ) -> float: """ Calculates the spherical average about a given point in space @@ -132,6 +140,8 @@ def spherical_average_potential( float: Spherical average at the target center """ + if self._total_data is None: + raise ValueError("total_data is not set") grid_shape = self._total_data.shape # Position of center of sphere at grid coordinates @@ -149,13 +159,13 @@ def spherical_average_potential( ] # Range of grids to be considered within the provided radius w.r.t. center of sphere - num_grid_in_sph = [[], []] + num_grid_in_sph: List[List[int]] = [[], []] for i, dist in enumerate(dist_in_grid): num_grid_in_sph[0].append(n_grid_at_center[i] - int(np.ceil(rad / dist))) num_grid_in_sph[1].append(n_grid_at_center[i] + int(np.ceil(rad / dist))) sph_avg_tmp = [] - weight = 0 + weight = 0.0 for k in range(num_grid_in_sph[0][0], num_grid_in_sph[1][0]): for l in range(num_grid_in_sph[0][1], num_grid_in_sph[1][1]): for m in range(num_grid_in_sph[0][2], num_grid_in_sph[1][2]): @@ -179,8 +189,12 @@ def spherical_average_potential( @staticmethod def dist_between_two_grid_points_cyl( - target_grid_point, n_grid_at_center, lattice, grid_shape, direction_of_cyl - ): + target_grid_point: Union[np.ndarray, list], + n_grid_at_center: Union[np.ndarray, list], + lattice: Union[np.ndarray, list], + grid_shape: Union[tuple, list, np.ndarray], + direction_of_cyl: int, + ) -> float: """ Distance between a target grid point and the center of a cylinder @@ -212,11 +226,16 @@ def dist_between_two_grid_points_cyl( else: print("check the direction of cylindrical axis") dist = np.linalg.norm(dn) - return dist + return float(dist) def cylindrical_average_potential( - self, structure, spherical_center, axis_of_cyl, rad=2, fwhm=0.529177 - ): + self, + structure: Atoms, + spherical_center: Union[list, np.ndarray], + axis_of_cyl: int, + rad: float = 2, + fwhm: float = 0.529177, + ) -> float: """ Calculates the cylindrical average about a given point in space @@ -231,6 +250,8 @@ def cylindrical_average_potential( float: Cylindrical average at the target center """ + if self._total_data is None: + raise ValueError("total_data is not set") grid_shape = self._total_data.shape # Position of center of sphere at grid coordinates @@ -248,7 +269,7 @@ def cylindrical_average_potential( ] # Range of grids to be considered within the provided radius w.r.t. center of sphere - num_grid_in_cyl = [[], []] + num_grid_in_cyl: List[List[int]] = [[], []] for i, dist in enumerate(dist_in_grid): if i == axis_of_cyl: @@ -263,7 +284,7 @@ def cylindrical_average_potential( ) cyl_avg_tmp = [] - weight = 0 + weight = 0.0 for k in range(num_grid_in_cyl[0][0], num_grid_in_cyl[1][0]): for l in range(num_grid_in_cyl[0][1], num_grid_in_cyl[1][1]): for m in range(num_grid_in_cyl[0][2], num_grid_in_cyl[1][2]): @@ -290,7 +311,7 @@ def cylindrical_average_potential( return cyl_avg - def get_average_along_axis(self, ind=2): + def get_average_along_axis(self, ind: int = 2) -> np.ndarray: """ Get the lateral average along a certain axis direction. This function is adapted from the pymatgen vasp VolumetricData class @@ -303,6 +324,8 @@ def get_average_along_axis(self, ind=2): Returns: numpy.ndarray: A 1D vector with the laterally averaged values of the volumetric data """ + if self._total_data is None: + raise ValueError("total_data is not set") if ind == 0: return np.average(np.average(self._total_data, axis=1), 1) elif ind == 1: @@ -310,7 +333,9 @@ def get_average_along_axis(self, ind=2): else: return np.average(np.average(self._total_data, axis=0), 0) - def write_cube_file(self, filename="cube_file.cube", cell_scaling=1.0): + def write_cube_file( + self, filename: str = "cube_file.cube", cell_scaling: float = 1.0 + ) -> None: """ Write the volumetric data into the CUBE file format @@ -319,37 +344,41 @@ def write_cube_file(self, filename="cube_file.cube", cell_scaling=1.0): cell_scaling (float): Scale the cell by this fraction """ - if self._atoms is None: + atoms = self.atoms + if atoms is None: raise ValueError( "The volumetric data object must have a valid structure assigned to it before writing " "to the cube format" ) - data = self.total_data + total_data = self.total_data + if total_data is None: + raise ValueError("total_data is not set") + data = total_data n_x, n_y, _ = data.shape origin = np.zeros(3) flattened_data = np.hstack( [data[i, j, :] for i in range(n_x) for j in range(n_y)] ) - n_atoms = len(self.atoms) + n_atoms = len(atoms) total_lines = int(len(flattened_data) / 6) * 6 reshaped_data = np.reshape(flattened_data[0:total_lines], (-1, 6)) last_line = [flattened_data[total_lines:]] head_array = np.zeros((4, 4)) head_array[0] = np.append([n_atoms], origin) head_array[1:, 0] = data.shape - head_array[1:, 1:] = self.atoms.cell / data.shape * cell_scaling - position_array = np.zeros((len(self.atoms.positions), 5)) - position_array[:, 0] = self.atoms.get_atomic_numbers() - position_array[:, 2:] = self.atoms.positions + head_array[1:, 1:] = atoms.cell / data.shape * cell_scaling + position_array = np.zeros((len(atoms.positions), 5)) + position_array[:, 0] = atoms.get_atomic_numbers() + position_array[:, 2:] = atoms.positions with open(filename, "w") as f: f.write("Cube file generated by pyiron (http://pyiron.org) \n") f.write("z is the fastest index \n") - np.savetxt(f, head_array, fmt="%4d %.6f %.6f %.6f") - np.savetxt(f, position_array, fmt="%4d %.6f %.6f %.6f %.6f") - np.savetxt(f, reshaped_data, fmt="%.5e") - np.savetxt(f, last_line, fmt="%.5e") + np.savetxt(f, head_array, fmt="%4d %.6f %.6f %.6f") # type: ignore + np.savetxt(f, position_array, fmt="%4d %.6f %.6f %.6f %.6f") # type: ignore + np.savetxt(f, reshaped_data, fmt="%.5e") # type: ignore + np.savetxt(f, last_line, fmt="%.5e") # type: ignore - def read_cube_file(self, filename="cube_file.cube"): + def read_cube_file(self, filename: str = "cube_file.cube") -> None: """ Generate data from a CUBE file @@ -376,7 +405,7 @@ def read_cube_file(self, filename="cube_file.cube"): ) end_int = n_atoms + 6 + int(np.prod(grid_shape) / 6) data = np.genfromtxt(lines[n_atoms + 6 : end_int]) - data_flatten = np.hstack(data) + data_flatten = np.hstack(data) # type: ignore if np.prod(grid_shape) % 6 > 0: data_flatten = np.append( data_flatten, [float(val) for val in lines[end_int].split()] @@ -384,7 +413,9 @@ def read_cube_file(self, filename="cube_file.cube"): n_x, n_y, n_z = grid_shape self._total_data = data_flatten.reshape((n_x, n_y, n_z)) - def write_vasp_volumetric(self, filename="CHGCAR", normalize=False): + def write_vasp_volumetric( + self, filename: str = "CHGCAR", normalize: bool = False + ) -> None: """ Writes volumetric data into a VASP CHGCAR format @@ -393,19 +424,25 @@ def write_vasp_volumetric(self, filename="CHGCAR", normalize=False): normalize (bool): True if the data is to be normalized by the volume """ - write_poscar(structure=self.atoms, filename=filename) + atoms = self.atoms + if atoms is None: + raise ValueError("atoms is not set") + total_data = self.total_data + if total_data is None: + raise ValueError("total_data is not set") + write_poscar(structure=atoms, filename=filename) with open(filename, "a") as f: f.write("\n") - f.write(" ".join(list(np.array(self.total_data.shape, dtype=str)))) + f.write(" ".join(list(np.array(total_data.shape, dtype=str)))) f.write("\n") - _, n_y, n_z = self.total_data.shape + _, n_y, n_z = total_data.shape flattened_data = np.hstack( - [self.total_data[:, i, j] for j in range(n_z) for i in range(n_y)] + [total_data[:, i, j] for j in range(n_z) for i in range(n_y)] ) if normalize: - flattened_data /= self.atoms.get_volume() + flattened_data /= atoms.get_volume() num_lines = int(len(flattened_data) / 5) * 5 reshaped_data = np.reshape(flattened_data[0:num_lines], (-1, 5)) - np.savetxt(f, reshaped_data, fmt="%.12f") + np.savetxt(f, reshaped_data, fmt="%.12f") # type: ignore if len(flattened_data) % 5 > 0: - np.savetxt(f, [flattened_data[num_lines:]], fmt="%.12f") + np.savetxt(f, [flattened_data[num_lines:]], fmt="%.12f") # type: ignore diff --git a/pyiron_vasp/vasp/output.py b/pyiron_vasp/vasp/output.py index b2a913c6..bcab89c9 100644 --- a/pyiron_vasp/vasp/output.py +++ b/pyiron_vasp/vasp/output.py @@ -1,6 +1,7 @@ from __future__ import print_function import os import posixpath +from typing import Optional, Callable, Any, Type import numpy as np from ase.atoms import Atoms @@ -412,12 +413,12 @@ class VaspCollectError(ValueError): def parse_vasp_output( working_directory: str, - structure: Atoms = None, - sorted_indices: list = None, - read_atoms_funct: callable = read_atoms, - es_class=ElectronicStructure, - bader_class=Bader, - output_parser_class=Output, + structure: Optional[Atoms] = None, + sorted_indices: Optional[np.ndarray] = None, + read_atoms_funct: Callable = read_atoms, + es_class: Type[ElectronicStructure] = ElectronicStructure, + bader_class: Type[Bader] = Bader, + output_parser_class: Type[Output] = Output, ) -> dict: """ Parse the VASP output in the working_directory and return it as hierachical dictionary. diff --git a/pyiron_vasp/vasp/parser/oszicar.py b/pyiron_vasp/vasp/parser/oszicar.py index 88e99703..3eb9ebe4 100644 --- a/pyiron_vasp/vasp/parser/oszicar.py +++ b/pyiron_vasp/vasp/parser/oszicar.py @@ -2,6 +2,7 @@ # Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department # Distributed under the terms of "New BSD License", see the LICENSE file. +from typing import Dict, List import numpy as np __author__ = "Sudarsan Surendralal" @@ -27,16 +28,16 @@ class Oszicar(object): """ - def __init__(self): - self.parse_dict = dict() + def __init__(self) -> None: + self.parse_dict: Dict[str, np.ndarray] = dict() - def from_file(self, filename="OSZICAR"): + def from_file(self, filename: str = "OSZICAR") -> None: with open(filename, "r", errors="ignore") as f: lines = f.readlines() self.parse_dict["energy_pot"] = self.get_energy_pot(lines) @staticmethod - def get_energy_pot(lines): + def get_energy_pot(lines: List[str]) -> np.ndarray: trigger = "F=" energy_list = list() for i, line in enumerate(lines): diff --git a/pyiron_vasp/vasp/parser/report.py b/pyiron_vasp/vasp/parser/report.py index 630f9edd..10ee3044 100644 --- a/pyiron_vasp/vasp/parser/report.py +++ b/pyiron_vasp/vasp/parser/report.py @@ -2,6 +2,7 @@ # Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department # Distributed under the terms of "New BSD License", see the LICENSE file. +from typing import Dict import numpy as np from scipy.integrate import cumulative_trapezoid @@ -22,10 +23,10 @@ class Report(object): This module is used to parse VASP REPORT files """ - def __init__(self): - self.parse_dict = dict() + def __init__(self) -> None: + self.parse_dict: Dict[str, np.ndarray] = dict() - def from_file(self, filename="REPORT"): + def from_file(self, filename: str = "REPORT") -> None: """ Reads values from files and stores it in the `parse_dict` attribute diff --git a/pyiron_vasp/vasp/structure.py b/pyiron_vasp/vasp/structure.py index b2b3d124..7c4f55a3 100644 --- a/pyiron_vasp/vasp/structure.py +++ b/pyiron_vasp/vasp/structure.py @@ -5,6 +5,7 @@ import os import re from collections import OrderedDict +from typing import Optional, Union, List, OrderedDict as OrderedDictType, Any, Dict from ase.atoms import Atoms from ase.constraints import FixCartesian import numpy as np @@ -23,11 +24,11 @@ def read_atoms( - filename="CONTCAR", - return_velocities=False, - species_list=None, - species_from_potcar=False, -): + filename: str = "CONTCAR", + return_velocities: bool = False, + species_list: Optional[Union[list, np.ndarray]] = None, + species_from_potcar: bool = False, +) -> Union[Atoms, tuple[Atoms, list]]: """ Routine to read structural static from a POSCAR type file @@ -57,7 +58,7 @@ def read_atoms( ) -def get_species_list_from_potcar(filename="POTCAR"): +def get_species_list_from_potcar(filename: str = "POTCAR") -> List[str]: """ Generates the species list from a POTCAR type file @@ -81,7 +82,7 @@ def get_species_list_from_potcar(filename="POTCAR"): return species_list -def get_number_species_atoms(structure): +def get_number_species_atoms(structure: Atoms) -> OrderedDictType[str, int]: """ Returns a dictionary with the species in the structure and the corresponding count in the structure @@ -102,7 +103,12 @@ def get_number_species_atoms(structure): return count -def write_poscar(structure, filename="POSCAR", write_species=True, cartesian=True): +def write_poscar( + structure: Atoms, + filename: str = "POSCAR", + write_species: bool = True, + cartesian: bool = True, +) -> None: """ Writes a POSCAR type file from a structure object @@ -125,7 +131,9 @@ def write_poscar(structure, filename="POSCAR", write_species=True, cartesian=Tru ) -def get_poscar_content(structure, write_species=True, cartesian=True): +def get_poscar_content( + structure: Atoms, write_species: bool = True, cartesian: bool = True +) -> List[str]: endline = "\n" selec_dyn = False line_lst = [ @@ -183,7 +191,11 @@ def get_poscar_content(structure, write_species=True, cartesian=True): return line_lst -def atoms_from_string(string, read_velocities=False, species_list=None): +def atoms_from_string( + string: List[str], + read_velocities: bool = False, + species_list: Optional[Union[list, np.ndarray]] = None, +) -> Union[Atoms, tuple[Atoms, list]]: """ Routine to convert a string list read from a input/output structure file and convert into Atoms instance @@ -198,7 +210,7 @@ def atoms_from_string(string, read_velocities=False, species_list=None): """ string = [s.strip() for s in string] string_lower = [s.lower() for s in string] - atoms_dict = dict() + atoms_dict: Dict[str, Any] = dict() atoms_dict["first_line"] = string[0] # del string[0] atoms_dict["selective_dynamics"] = False @@ -221,7 +233,7 @@ def atoms_from_string(string, read_velocities=False, species_list=None): if "selective dynamics" in string_lower: atoms_dict["selective_dynamics"] = True no_of_species = len(string[5].split()) - species_dict = OrderedDict() + species_dict: OrderedDictType[str, Dict[str, Any]] = OrderedDict() position_index = 7 if atoms_dict["selective_dynamics"]: position_index += 1 @@ -265,7 +277,7 @@ def atoms_from_string(string, read_velocities=False, species_list=None): except ValueError: atoms = _dict_to_atoms(atoms_dict, read_from_first_line=True) if atoms_dict["selective_dynamics"]: - constraints_dict = { + constraints_dict: Dict[str, List[int]] = { label: [] for label in ["TTT", "TTF", "FTT", "TFT", "TFF", "FFT", "FTF", "FFF"] } @@ -336,7 +348,11 @@ def atoms_from_string(string, read_velocities=False, species_list=None): return atoms -def _dict_to_atoms(atoms_dict, species_list=None, read_from_first_line=False): +def _dict_to_atoms( + atoms_dict: Dict[str, Any], + species_list: Optional[Union[list, np.ndarray]] = None, + read_from_first_line: bool = False, +) -> Atoms: """ Function to convert a generated dict into an structure object @@ -352,8 +368,8 @@ def _dict_to_atoms(atoms_dict, species_list=None, read_from_first_line=False): positions = atoms_dict["positions"] cell = atoms_dict["cell"] symbol = str() - elements = list() - el_list = list() + elements: List[Union[List[str], np.ndarray]] = [] + el_list: Union[List[str], np.ndarray] = list() for i, sp_key in enumerate(atoms_dict["species_dict"].keys()): if species_list is not None: try: @@ -391,19 +407,19 @@ def _dict_to_atoms(atoms_dict, species_list=None, read_from_first_line=False): "Species list should be provided since pyiron can't detect species information" ) elements.append(el_list) - elements_new = list() + elements_new: List[str] = [] for ele in elements: for e in ele: elements_new.append(re.split("[^a-zA-Z]", e)[0]) - elements = elements_new + final_elements = elements_new if is_absolute: - atoms = Atoms(elements, positions=positions, cell=cell, pbc=True) + atoms = Atoms(final_elements, positions=positions, cell=cell, pbc=True) else: - atoms = Atoms(elements, scaled_positions=positions, cell=cell, pbc=True) + atoms = Atoms(final_elements, scaled_positions=positions, cell=cell, pbc=True) return atoms -def vasp_sorter(structure): +def vasp_sorter(structure: Atoms) -> np.ndarray: """ Routine to sort the indices of a structure as it would be when written to a POSCAR file @@ -423,7 +439,9 @@ def vasp_sorter(structure): return np.array(sorted_indices) -def manip_contcar(filename, new_filename, add_pos): +def manip_contcar( + filename: str, new_filename: str, add_pos: Union[list, np.ndarray] +) -> None: """ Manipulate a CONTCAR/POSCAR file by adding something to the positions @@ -434,6 +452,7 @@ def manip_contcar(filename, new_filename, add_pos): """ actual_struct = read_atoms(filename) + assert isinstance(actual_struct, Atoms) n = 0 direct = True with open(filename, "r", errors="ignore") as f: