Merge pull request #225 from pyiron/more_typing

More Typing

atomistics/workflows/elastic/elastic_moduli.py

@@ -12,27 +12,27 @@ def get_bulkmodul_voigt(elastic_matrix: np.ndarray) -> float:
     ) / 9
 
 
-def get_shearmodul_voigt(elastic_matrix):
+def get_shearmodul_voigt(elastic_matrix: np.ndarray) -> float:
     return (
         (elastic_matrix[0, 0] + elastic_matrix[1, 1] + elastic_matrix[2, 2])
         - (elastic_matrix[0, 1] + elastic_matrix[0, 2] + elastic_matrix[1, 2])
         + 3 * (elastic_matrix[3, 3] + elastic_matrix[4, 4] + elastic_matrix[5, 5])
     ) / 15
 
 
-def get_youngsmodul_voigt(bulkmodul_voigt, shearmodul_voigt):
+def get_youngsmodul_voigt(bulkmodul_voigt: float, shearmodul_voigt: float) -> float:
     return (9 * bulkmodul_voigt * shearmodul_voigt) / (
         3 * bulkmodul_voigt + shearmodul_voigt
     )
 
 
-def get_poissonsratio_voigt(bulkmodul_voigt, shearmodul_voigt):
+def get_poissonsratio_voigt(bulkmodul_voigt: float, shearmodul_voigt: float) -> float:
     return (1.5 * bulkmodul_voigt - shearmodul_voigt) / (
         3 * bulkmodul_voigt + shearmodul_voigt
     )
 
 
-def get_bulkmodul_reuss(elastic_matrix_inverse):
+def get_bulkmodul_reuss(elastic_matrix_inverse: np.ndarray) -> float:
     return 1 / (
         elastic_matrix_inverse[0, 0]
         + elastic_matrix_inverse[1, 1]
@@ -46,7 +46,7 @@ def get_bulkmodul_reuss(elastic_matrix_inverse):
     )
 
 
-def get_shearmodul_reuss(elastic_matrix_inverse):
+def get_shearmodul_reuss(elastic_matrix_inverse: np.ndarray) -> float:
     return 15 / (
         4
         * (
@@ -69,143 +69,143 @@ def get_shearmodul_reuss(elastic_matrix_inverse):
     )
 
 
-def get_youngsmodul_reuss(bulkmodul_reuss, shearmodul_reuss):
+def get_youngsmodul_reuss(bulkmodul_reuss: float, shearmodul_reuss: float) -> float:
     return (9 * bulkmodul_reuss * shearmodul_reuss) / (
         3 * bulkmodul_reuss + shearmodul_reuss
     )
 
 
-def get_poissonsratio_reuss(bulkmodul_reuss, shearmodul_reuss):
+def get_poissonsratio_reuss(bulkmodul_reuss: float, shearmodul_reuss: float) -> float:
     return (1.5 * bulkmodul_reuss - shearmodul_reuss) / (
         3 * bulkmodul_reuss + shearmodul_reuss
     )
 
 
-def get_bulkmodul_hill(bulkmodul_voigt, bulkmodul_reuss):
+def get_bulkmodul_hill(bulkmodul_voigt: float, bulkmodul_reuss: float) -> float:
     return _hill_approximation(voigt=bulkmodul_voigt, reuss=bulkmodul_reuss)
 
 
-def get_shearmodul_hill(shearmodul_voigt, shearmodul_reuss):
+def get_shearmodul_hill(shearmodul_voigt: float, shearmodul_reuss: float) -> float:
     return _hill_approximation(voigt=shearmodul_voigt, reuss=shearmodul_reuss)
 
 
-def get_youngsmodul_hill(bulkmodul_hill, shearmodul_hill):
+def get_youngsmodul_hill(bulkmodul_hill: float, shearmodul_hill: float) -> float:
     return (9.0 * bulkmodul_hill * shearmodul_hill) / (
         3.0 * bulkmodul_hill + shearmodul_hill
     )
 
 
-def get_poissonsratio_hill(bulkmodul_hill, shearmodul_hill):
+def get_poissonsratio_hill(bulkmodul_hill: float, shearmodul_hill: float) -> float:
     return (1.5 * bulkmodul_hill - shearmodul_hill) / (
         3.0 * bulkmodul_hill + shearmodul_hill
     )
 
 
-def get_AVR(shearmodul_voigt, shearmodul_reuss):
+def get_AVR(shearmodul_voigt: float, shearmodul_reuss: float) -> float:
     return (
         100.0
         * (shearmodul_voigt - shearmodul_reuss)
         / (shearmodul_voigt + shearmodul_reuss)
     )
 
 
-def get_elastic_matrix_eigval(elastic_matrix):
+def get_elastic_matrix_eigval(elastic_matrix: np.ndarray):
     return np.linalg.eig(elastic_matrix)
 
 
-def get_elastic_matrix_inverse(elastic_matrix):
+def get_elastic_matrix_inverse(elastic_matrix: np.ndarray) -> np.ndarray:
     return np.linalg.inv(elastic_matrix)
 
 
-def _hill_approximation(voigt, reuss):
+def _hill_approximation(voigt: float, reuss: float) -> float:
     return 0.50 * (voigt + reuss)
 
 
 class ElasticProperties:
     def __init__(self, elastic_matrix):
         self._elastic_matrix = elastic_matrix
 
-    def elastic_matrix(self):
+    def elastic_matrix(self) -> np.ndarray:
         return self._elastic_matrix
 
     @cache
-    def elastic_matrix_inverse(self):
+    def elastic_matrix_inverse(self) -> np.ndarray:
         return get_elastic_matrix_inverse(elastic_matrix=self.elastic_matrix())
 
     @cache
-    def bulkmodul_voigt(self):
+    def bulkmodul_voigt(self) -> float:
         return get_bulkmodul_voigt(elastic_matrix=self.elastic_matrix())
 
     @cache
-    def shearmodul_voigt(self):
+    def shearmodul_voigt(self) -> float:
         return get_shearmodul_voigt(elastic_matrix=self.elastic_matrix())
 
     @cache
-    def bulkmodul_reuss(self):
+    def bulkmodul_reuss(self) -> float:
         return get_bulkmodul_reuss(elastic_matrix_inverse=self.elastic_matrix_inverse())
 
     @cache
-    def shearmodul_reuss(self):
+    def shearmodul_reuss(self) -> float:
         return get_shearmodul_reuss(
             elastic_matrix_inverse=self.elastic_matrix_inverse()
         )
 
     @cache
-    def bulkmodul_hill(self):
+    def bulkmodul_hill(self) -> float:
         return get_bulkmodul_hill(
             bulkmodul_voigt=self.bulkmodul_voigt(),
             bulkmodul_reuss=self.bulkmodul_reuss(),
         )
 
     @cache
-    def shearmodul_hill(self):
+    def shearmodul_hill(self) -> float:
         return get_shearmodul_hill(
             shearmodul_voigt=self.shearmodul_voigt(),
             shearmodul_reuss=self.shearmodul_reuss(),
         )
 
     @cache
-    def youngsmodul_voigt(self):
+    def youngsmodul_voigt(self) -> float:
         return get_youngsmodul_voigt(
             bulkmodul_voigt=self.bulkmodul_voigt(),
             shearmodul_voigt=self.shearmodul_voigt(),
         )
 
     @cache
-    def poissonsratio_voigt(self):
+    def poissonsratio_voigt(self) -> float:
         return get_poissonsratio_voigt(
             bulkmodul_voigt=self.bulkmodul_voigt(),
             shearmodul_voigt=self.shearmodul_voigt(),
         )
 
     @cache
-    def youngsmodul_reuss(self):
+    def youngsmodul_reuss(self) -> float:
         return get_youngsmodul_reuss(
             bulkmodul_reuss=self.bulkmodul_reuss(),
             shearmodul_reuss=self.shearmodul_reuss(),
         )
 
     @cache
-    def poissonsratio_reuss(self):
+    def poissonsratio_reuss(self) -> float:
         return get_poissonsratio_reuss(
             bulkmodul_reuss=self.bulkmodul_reuss(),
             shearmodul_reuss=self.shearmodul_reuss(),
         )
 
     @cache
-    def youngsmodul_hill(self):
+    def youngsmodul_hill(self) -> float:
         return get_youngsmodul_hill(
             bulkmodul_hill=self.bulkmodul_hill(), shearmodul_hill=self.shearmodul_hill()
         )
 
     @cache
-    def poissonsratio_hill(self):
+    def poissonsratio_hill(self) -> float:
         return get_poissonsratio_hill(
             bulkmodul_hill=self.bulkmodul_hill(), shearmodul_hill=self.shearmodul_hill()
         )
 
     @cache
-    def AVR(self):
+    def AVR(self) -> float:
         return get_AVR(
             shearmodul_voigt=self.shearmodul_voigt(),
             shearmodul_reuss=self.shearmodul_reuss(),

atomistics/workflows/elastic/helper.py

@@ -1,5 +1,6 @@
 from collections import OrderedDict
 
+import ase.atoms
 import numpy as np
 import scipy.constants
 
@@ -11,7 +12,11 @@
 
 
 def generate_structures_helper(
-    structure, eps_range, num_of_point, zero_strain_job_name="s_e_0", sqrt_eta=True
+    structure: ase.atoms.Atoms,
+    eps_range: float,
+    num_of_point: int,
+    zero_strain_job_name: str = "s_e_0",
+    sqrt_eta: bool = True,
 ):
     """
 
@@ -84,13 +89,13 @@ def generate_structures_helper(
 
 
 def analyse_structures_helper(
-    output_dict,
-    Lag_strain_list,
-    epss,
-    v0,
-    LC,
-    fit_order=2,
-    zero_strain_job_name="s_e_0",
+    output_dict: dict,
+    Lag_strain_list: list[float],
+    epss: float,
+    v0: float,
+    LC: str,
+    fit_order: int = 2,
+    zero_strain_job_name: str = "s_e_0",
 ):
     """
 
@@ -125,7 +130,7 @@ def analyse_structures_helper(
     return elastic_matrix, A2, strain_energy, ene0
 
 
-def _subjob_name(i, eps):
+def _subjob_name(i: int, eps: float):
     """
 
     Args:
@@ -138,7 +143,7 @@ def _subjob_name(i, eps):
     return ("s_%s_e_%.5f" % (i, eps)).replace(".", "_").replace("-", "m")
 
 
-def _fit_elastic_matrix(strain_ene, v0, LC, fit_order):
+def _fit_elastic_matrix(strain_ene: list, v0: float, LC: str, fit_order: int):
     """
 
     Returns:

atomistics/workflows/elastic/symmetry.py

@@ -1,8 +1,9 @@
+from ase.atoms import Atoms
 import numpy as np
 import spglib
 
 
-def ase_to_spglib(structure):
+def ase_to_spglib(structure: Atoms) -> tuple:
     """
     Translate ASE to spglib cell. The format is a tuple of
     (basis vectors, atomic points, types). The implementation here follows
@@ -17,7 +18,7 @@ def ase_to_spglib(structure):
     )
 
 
-def find_symmetry_group_number(struct):
+def find_symmetry_group_number(struct: tuple) -> int:
     dataset = spglib.get_symmetry_dataset(cell=ase_to_spglib(struct))
     SGN = dataset["number"]
     return SGN
@@ -72,7 +73,7 @@ def find_symmetry_group_number(struct):
 }
 
 
-def get_symmetry_family_from_SGN(SGN):
+def get_symmetry_family_from_SGN(SGN: int) -> str:
     if 1 <= SGN <= 2:  # Triclinic
         LC = "N"
     elif 3 <= SGN <= 15:  # Monoclinic
@@ -100,7 +101,7 @@ def get_symmetry_family_from_SGN(SGN):
     return LC
 
 
-def get_LAG_Strain_List(LC):
+def get_LAG_Strain_List(LC: str) -> list[str]:
     if LC == "CI" or LC == "CII":
         Lag_strain_list = ["01", "08", "23"]
     elif LC == "HI" or LC == "HII":
@@ -158,7 +159,7 @@ def get_LAG_Strain_List(LC):
     return Lag_strain_list
 
 
-def get_C_from_A2(A2, LC):
+def get_C_from_A2(A2: np.ndarray, LC: str) -> np.ndarray:
     C = np.zeros((6, 6))
 
     # %!%!%--- Cubic structures ---%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%!%
@@ -334,7 +335,7 @@ def get_C_from_A2(A2, LC):
     return C
 
 
-def symmetry_analysis(structure, eps_range, num_of_point):
+def symmetry_analysis(structure: Atoms, eps_range: float, num_of_point: int):
     """
 
     Returns:

atomistics/workflows/elastic/workflow.py

@@ -1,3 +1,4 @@
+from ase.atoms import Atoms
 import numpy as np
 
 from atomistics.shared.output import OutputElastic
@@ -11,7 +12,12 @@
 
 class ElasticMatrixWorkflow(Workflow):
     def __init__(
-        self, structure, num_of_point=5, eps_range=0.005, sqrt_eta=True, fit_order=2
+        self,
+        structure: Atoms,
+        num_of_point: int = 5,
+        eps_range: float = 0.005,
+        sqrt_eta: bool = True,
+        fit_order: int = 2,
     ):
         self.structure = structure.copy()
         self.num_of_point = num_of_point
@@ -24,7 +30,7 @@ def __init__(
         self.epss = np.array([])
         self.zero_strain_job_name = "s_e_0"
 
-    def generate_structures(self):
+    def generate_structures(self) -> dict:
         """
 
         Returns:
@@ -39,7 +45,9 @@ def generate_structures(self):
         )
         return {"calc_energy": self._structure_dict}
 
-    def analyse_structures(self, output_dict, output_keys=OutputElastic.keys()):
+    def analyse_structures(
+        self, output_dict: dict, output_keys: tuple = OutputElastic.keys()
+    ):
         """
 
         Args: