Add some more type hints

atomistics/calculators/ase.py

@@ -31,45 +31,45 @@ def __init__(self, ase_structure, ase_calculator):
         self.structure = ase_structure
         self.structure.calc = ase_calculator
 
-    def forces(self):
+    def forces(self) -> np.ndarray:
         return self.structure.get_forces()
 
-    def energy(self):
+    def energy(self) -> float:
         return self.structure.get_potential_energy()
 
-    def energy_pot(self):
+    def energy_pot(self) -> float:
         return self.structure.get_potential_energy()
 
-    def energy_tot(self):
+    def energy_tot(self) -> float:
         return (
             self.structure.get_potential_energy() + self.structure.get_kinetic_energy()
         )
 
-    def stress(self):
+    def stress(self) -> np.ndarray:
         try:
             return self.structure.get_stress(voigt=False)
         except PropertyNotImplementedError:
             return None
 
-    def pressure(self):
+    def pressure(self) -> np.ndarray:
         try:
             return self.structure.get_stress(voigt=False)
         except PropertyNotImplementedError:
             return None
 
-    def cell(self):
+    def cell(self) -> np.ndarray:
         return self.structure.get_cell()
 
-    def positions(self):
+    def positions(self) -> np.ndarray:
         return self.structure.get_positions()
 
-    def velocities(self):
+    def velocities(self) -> np.ndarray:
         return self.structure.get_velocities()
 
-    def temperature(self):
+    def temperature(self) -> float:
         return self.structure.get_temperature()
 
-    def volume(self):
+    def volume(self) -> float:
         return self.structure.get_volume()
 
 
@@ -80,7 +80,7 @@ def evaluate_with_ase(
     ase_calculator: ASECalculator,
     ase_optimizer: Optimizer = None,
     ase_optimizer_kwargs: dict = {},
-):
+) -> dict:
     results = {}
     if "optimize_positions" in tasks:
         results["structure_with_optimized_positions"] = optimize_positions_with_ase(
@@ -110,8 +110,8 @@ def evaluate_with_ase(
 
 
 def calc_static_with_ase(
-    structure,
-    ase_calculator,
+    structure: Atoms,
+    ase_calculator: ASECalculator,
     output_keys=OutputStatic.keys(),
 ):
     ase_exe = ASEExecutor(ase_structure=structure, ase_calculator=ase_calculator)
@@ -121,17 +121,17 @@ def calc_static_with_ase(
 
 
 def calc_molecular_dynamics_npt_with_ase(
-    structure,
-    ase_calculator,
-    run=100,
-    thermo=100,
-    timestep=1 * units.fs,
-    ttime=100 * units.fs,
-    pfactor=2e6 * units.GPa * (units.fs**2),
-    temperature=100,
-    externalstress=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * units.bar,
+    structure: Atoms,
+    ase_calculator: ASECalculator,
+    run: int = 100,
+    thermo: int = 100,
+    timestep: float = 1 * units.fs,
+    ttime: float = 100 * units.fs,
+    pfactor: float = 2e6 * units.GPa * (units.fs**2),
+    temperature: float = 100.0,
+    externalstress: np.ndarray = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * units.bar,
     output_keys=OutputMolecularDynamics.keys(),
-):
+) -> dict:
     return _calc_molecular_dynamics_with_ase(
         dyn=NPT(
             atoms=structure,
@@ -157,13 +157,13 @@ def calc_molecular_dynamics_npt_with_ase(
 
 
 def calc_molecular_dynamics_langevin_with_ase(
-    structure,
-    ase_calculator,
-    run=100,
-    thermo=100,
-    timestep=1 * units.fs,
-    temperature=100,
-    friction=0.002,
+    structure: Atoms,
+    ase_calculator: ASECalculator,
+    run: int = 100,
+    thermo: int = 100,
+    timestep: float = 1 * units.fs,
+    temperature: float = 100.0,
+    friction: float = 0.002,
     output_keys=OutputMolecularDynamics.keys(),
 ):
     return _calc_molecular_dynamics_with_ase(
@@ -183,7 +183,10 @@ def calc_molecular_dynamics_langevin_with_ase(
 
 
 def optimize_positions_with_ase(
-    structure, ase_calculator, ase_optimizer, ase_optimizer_kwargs
+    structure: Atoms,
+    ase_calculator: ASECalculator,
+    ase_optimizer: Optimizer,
+    ase_optimizer_kwargs: dict,
 ):
     structure_optimized = structure.copy()
     structure_optimized.calc = ase_calculator
@@ -193,7 +196,10 @@ def optimize_positions_with_ase(
 
 
 def optimize_positions_and_volume_with_ase(
-    structure, ase_calculator, ase_optimizer, ase_optimizer_kwargs
+    structure: Atoms,
+    ase_calculator: ASECalculator,
+    ase_optimizer: Optimizer,
+    ase_optimizer_kwargs: dict,
 ):
     structure_optimized = structure.copy()
     structure_optimized.calc = ase_calculator
@@ -203,17 +209,17 @@ def optimize_positions_and_volume_with_ase(
 
 
 def calc_molecular_dynamics_thermal_expansion_with_ase(
-    structure,
-    ase_calculator,
-    temperature_start=15,
-    temperature_stop=1500,
-    temperature_step=5,
-    run=100,
-    thermo=100,
-    timestep=1 * units.fs,
-    ttime=100 * units.fs,
-    pfactor=2e6 * units.GPa * (units.fs**2),
-    externalstress=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * units.bar,
+    structure: Atoms,
+    ase_calculator: ASECalculator,
+    temperature_start: float = 15.0,
+    temperature_stop: float = 1500.0,
+    temperature_step: float = 5.0,
+    run: int = 100,
+    thermo: int = 100,
+    timestep: float = 1 * units.fs,
+    ttime: float = 100 * units.fs,
+    pfactor: float = 2e6 * units.GPa * (units.fs**2),
+    externalstress: np.ndarray = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) * units.bar,
     output_keys=OutputThermalExpansion.keys(),
 ):
     structure_current = structure.copy()
@@ -244,7 +250,13 @@ def calc_molecular_dynamics_thermal_expansion_with_ase(
 
 
 def _calc_molecular_dynamics_with_ase(
-    dyn, structure, ase_calculator, temperature, run, thermo, output_keys
+    dyn,
+    structure: Atoms,
+    ase_calculator: ASECalculator,
+    temperature: float,
+    run: int,
+    thermo: int,
+    output_keys: tuple[str],
 ):
     structure.calc = ase_calculator
     MaxwellBoltzmannDistribution(atoms=structure, temperature_K=temperature)

atomistics/calculators/hessian.py

@@ -1,9 +1,10 @@
+from ase.atoms import Atoms
 import numpy as np
 
 from atomistics.calculators.wrapper import as_task_dict_evaluator
 
 
-def check_force_constants(structure, force_constants):
+def check_force_constants(structure: Atoms, force_constants: np.ndarray) -> np.ndarray:
     if structure is None:
         raise ValueError("Set reference structure via set_reference_structure() first")
     n_atom = len(structure.positions)
@@ -31,22 +32,26 @@ def check_force_constants(structure, force_constants):
         raise AssertionError("force constant shape not recognized")
 
 
-def get_displacement(structure_equilibrium, structure):
+def get_displacement(structure_equilibrium: Atoms, structure: Atoms) -> np.ndarray:
     displacements = structure.get_scaled_positions()
     displacements -= structure_equilibrium.get_scaled_positions()
     displacements -= np.rint(displacements)
     return np.einsum("ji,ni->nj", structure.cell, displacements)
 
 
-def calc_forces_transformed(force_constants, structure_equilibrium, structure):
+def calc_forces_transformed(
+    force_constants: np.ndarray, structure_equilibrium: Atoms, structure: Atoms
+) -> np.ndarray:
     displacements = get_displacement(structure_equilibrium, structure)
     position_transformed = displacements.reshape(
         displacements.shape[0] * displacements.shape[1]
     )
     return -np.dot(force_constants, position_transformed), displacements
 
 
-def get_forces(force_constants, structure_equilibrium, structure):
+def get_forces(
+    force_constants: np.ndarray, structure_equilibrium: Atoms, structure: Atoms
+) -> np.ndarray:
     displacements = get_displacement(structure_equilibrium, structure)
     position_transformed = displacements.reshape(
         displacements.shape[0] * displacements.shape[1]
@@ -56,8 +61,12 @@ def get_forces(force_constants, structure_equilibrium, structure):
 
 
 def get_energy_pot(
-    force_constants, structure_equilibrium, structure, bulk_modulus=0, shear_modulus=0
-):
+    force_constants: np.ndarray,
+    structure_equilibrium: Atoms,
+    structure: Atoms,
+    bulk_modulus: float = 0.0,
+    shear_modulus: float = 0.0,
+) -> float:
     displacements = get_displacement(structure_equilibrium, structure)
     position_transformed = displacements.reshape(
         displacements.shape[0] * displacements.shape[1]
@@ -74,15 +83,19 @@ def get_energy_pot(
     return energy_pot
 
 
-def get_stiffness_tensor(bulk_modulus=0, shear_modulus=0):
+def get_stiffness_tensor(
+    bulk_modulus: float = 0.0, shear_modulus: float = 0.0
+) -> np.ndarray:
     stiffness_tensor = np.zeros((6, 6))
     stiffness_tensor[:3, :3] = bulk_modulus - 2 * shear_modulus / 3
     stiffness_tensor[:3, :3] += np.eye(3) * 2 * shear_modulus
     stiffness_tensor[3:, 3:] = np.eye(3) * shear_modulus
     return stiffness_tensor
 
 
-def get_pressure_times_volume(stiffness_tensor, structure_equilibrium, structure):
+def get_pressure_times_volume(
+    stiffness_tensor: np.ndarray, structure_equilibrium: Atoms, structure: Atoms
+) -> float:
     if np.sum(stiffness_tensor) != 0:
         epsilon = np.einsum(
             "ij,jk->ik",
@@ -101,12 +114,12 @@ def get_pressure_times_volume(stiffness_tensor, structure_equilibrium, structure
 
 @as_task_dict_evaluator
 def evaluate_with_hessian(
-    structure,
-    tasks,
-    structure_equilibrium,
-    force_constants,
-    bulk_modulus=0,
-    shear_modulus=0,
+    structure: Atoms,
+    tasks: dict,
+    structure_equilibrium: Atoms,
+    force_constants: np.ndarray,
+    bulk_modulus: float = 0.0,
+    shear_modulus: float = 0.0,
 ):
     results = {}
     if "calc_energy" in tasks or "calc_forces" in tasks:

atomistics/calculators/interface.py

@@ -47,7 +47,7 @@ class TaskOutputEnum(Enum):
     SimpleEvaluator = callable[[Atoms, list[TaskName], ...], TaskResults]
 
 
-def get_quantities_from_tasks(tasks):
+def get_quantities_from_tasks(tasks: dict) -> list:
     quantities = []
     if "calc_energy" in tasks:
         quantities.append("energy")