Docstrings, type-hints, documentation fixes

autoplex/benchmark/phonons/utils.py

@@ -1,16 +1,23 @@
 """Utility functions for benchmarking jobs."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import matplotlib.pyplot as plt
 import numpy as np
-from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
+
+if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+    from pymatgen.phonon.bandstructure import PhononBandStructureSymmLine
 from pymatgen.phonon.plotter import PhononBSPlotter
 
 
 def get_rmse(
     ml_bs: PhononBandStructureSymmLine,
     dft_bs: PhononBandStructureSymmLine,
     q_dependent_rmse: bool = False,
-):
+) -> float | list[float]:
     """
     Compute root mean squared error (rmse) between DFT and ML phonon band-structure.
 
@@ -44,7 +51,7 @@ def rmse_qdep_plot(
     which_q_path=1,
     file_name="rms.pdf",
     img_format="pdf",
-):
+) -> plt:
     """
     Save q dependent root mean squared error plot between DFT and ML phonon band-structure.
 
@@ -94,7 +101,7 @@ def compare_plot(
     ml_bs: PhononBandStructureSymmLine,
     dft_bs: PhononBandStructureSymmLine,
     file_name: str = "band_comparison.pdf",
-):
+) -> Figure:
     """
     Save DFT and ML phonon band-structure overlay plot for visual comparison.
 

autoplex/data/common/utils.py

@@ -85,7 +85,7 @@ def scale_cell(
     volume_scale_factor_range: list[float] | None = None,
     n_structures: int = 10,
     volume_custom_scale_factors: list[float] | None = None,
-):
+) -> list[Structure]:
     """
     Take in a pymatgen Structure object and generates stretched or compressed structures.
 
@@ -154,9 +154,9 @@ def scale_cell(
     return distorted_cells
 
 
-def check_distances(structure: Structure, min_distance: float = 1.5):
+def check_distances(structure: Structure, min_distance: float = 1.5) -> bool:
     """
-    Take in a pymatgen Structure object and checks distances between atoms using minimum image convention.
+    Take in a pymatgen Structure object and check minimum distances between atoms using minimum image convention.
 
     Useful after distorting cell angles and rattling to check atoms aren't too close.
 
@@ -191,7 +191,7 @@ def random_vary_angle(
     w_angle: list[float] | None = None,
     n_structures: int = 8,
     angle_max_attempts: int = 1000,
-):
+) -> list[Structure]:
     """
     Take in a pymatgen Structure object and generates angle-distorted structures.
 
@@ -237,10 +237,10 @@ def random_vary_angle(
             volume_custom_scale_factors=[1.03],
         )
 
-        distorted_cells = AseAtomsAdaptor.get_atoms(distorted_cells[0])
+        distorted_supercells: Atoms = AseAtomsAdaptor.get_atoms(distorted_cells[0])
 
-        # getting stretched cell out of array
-        newcell = distorted_cells.cell.cellpar()
+        # getting stretched supercell out of array
+        newcell = distorted_supercells.cell.cellpar()
 
         # current angles
         alpha = atoms_copy.cell.cellpar()[3]
@@ -287,7 +287,7 @@ def std_rattle(
     n_structures: int = 5,
     rattle_std: float = 0.01,
     rattle_seed: int = 42,
-):
+) -> list[Structure]:
     """
     Take in a pymatgen Structure object and generates rattled structures.
 
@@ -331,7 +331,7 @@ def mc_rattle(
     min_distance: float = 1.5,
     rattle_seed: int = 42,
     rattle_mc_n_iter: int = 10,
-):
+) -> list[Structure]:
     """
     Take in a pymatgen Structure object and generates rattled structures.
 
@@ -375,7 +375,7 @@ def mc_rattle(
     return [AseAtomsAdaptor.get_structure(xtal) for xtal in mc_rattle]
 
 
-def extract_base_name(filename, is_out=False):
+def extract_base_name(filename, is_out=False) -> str:
     """
     Extract the base of a file name to easier manipulate other file names.
 
@@ -401,7 +401,7 @@ def extract_base_name(filename, is_out=False):
     return "A problem with the files occurred."
 
 
-def filter_outlier_energy(in_file, out_file, criteria: float = 0.0005):
+def filter_outlier_energy(in_file, out_file, criteria: float = 0.0005) -> None:
     """
     Filter data outliers per energy criteria and write them into files.
 
@@ -457,7 +457,9 @@ def filter_outlier_energy(in_file, out_file, criteria: float = 0.0005):
     )
 
 
-def filter_outlier_forces(in_file, out_file, symbol="Si", criteria: float = 0.1):
+def filter_outlier_forces(
+    in_file, out_file, symbol="Si", criteria: float = 0.1
+) -> None:
     """
     Filter data outliers per force criteria and write them into files.
 
@@ -526,13 +528,14 @@ def filter_outlier_forces(in_file, out_file, symbol="Si", criteria: float = 0.1)
     )
 
 
-# copied from libatoms GAP tutorial page and adjusted
 def energy_plot(
     in_file, out_file, ax, title: str = "Plot of energy", label: str = "energy"
-):
+) -> None:
     """
     Plot the distribution of energy per atom on the output vs the input.
 
+    Adapted and adjusted from libatoms GAP tutorial page https://libatoms.github.io/GAP/gap_fitting_tutorial.html.
+
     Parameters
     ----------
     in_file:
@@ -610,7 +613,7 @@ def force_plot(
     symbol: str = "Si",
     title: str = "Plot of force",
     label: str = "force for ",
-):
+) -> float:
     """
     Plot the distribution of force components per atom on the output vs the input.
 
@@ -700,7 +703,7 @@ def plot_energy_forces(
     species_list: list | None = None,
     train_name: str = "train.extxyz",
     test_name: str = "test.extxyz",
-):
+) -> None:
     """
     Plot energy and forces of the data.
 

autoplex/data/phonons/utils.py

@@ -18,7 +18,11 @@ def ml_phonon_maker_preparation(
     bulk_relax_maker: ForceFieldRelaxMaker,
     phonon_displacement_maker: ForceFieldStaticMaker,
     static_energy_maker: ForceFieldStaticMaker,
-):
+) -> tuple[
+    ForceFieldRelaxMaker | None,
+    ForceFieldStaticMaker | None,
+    ForceFieldStaticMaker | None,
+]:
     """
     Prepare the MLPhononMaker for the respective MLIP model.
 

autoplex/fitting/common/jobs.py

@@ -33,7 +33,7 @@ def machine_learning_fit(
     **kwargs,
 ):
     """
-    Maker for fitting potential(s).
+    Job for fitting potential(s).
 
     Parameters
     ----------

autoplex/fitting/common/regularization.py

@@ -5,10 +5,14 @@
 
 import traceback
 from contextlib import suppress
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 from scipy.spatial import ConvexHull, Delaunay
 
+if TYPE_CHECKING:
+    from ase import Atoms
+
 
 def set_sigma(
     atoms,
@@ -21,7 +25,7 @@ def set_sigma(
     element_order=None,
     max_energy=20.0,
     config_type_override=None,
-):
+) -> list[Atoms]:
     """
     Handle automatic regularisation based on distance to convex hull, amongst other things.
 
@@ -216,6 +220,7 @@ def set_sigma(
 
 
 def get_convex_hull(atoms, energy_name="energy", **kwargs):
+    # CE I don't get what the function returns
     """
     Calculate simple linear (E,V) convex hull.
 
@@ -276,7 +281,7 @@ def get_convex_hull(atoms, energy_name="energy", **kwargs):
     return lower_half_hull_points, p
 
 
-def get_e_distance_to_hull(hull: np.array, at, energy_name="energy", **kwargs):
+def get_e_distance_to_hull(hull: np.array, at, energy_name="energy", **kwargs) -> float:
     """
     Calculate the distance of a structure to the linear convex hull in energy.
 
@@ -315,7 +320,7 @@ def get_e_distance_to_hull(hull: np.array, at, energy_name="energy", **kwargs):
     )
 
 
-def get_intersect(a1, a2, b1, b2):
+def get_intersect(a1, a2, b1, b2) -> tuple[float, float] | tuple:
     """
     Return the point of intersection of the lines passing through a2,a1 and b2,b1.
 
@@ -339,7 +344,7 @@ def get_intersect(a1, a2, b1, b2):
     return x / z, y / z
 
 
-def get_x(at, element_order=None):
+def get_x(at, element_order=None) -> float | int:
     """
     Calculate the mole-fraction of a structure.
 
@@ -379,7 +384,7 @@ def get_x(at, element_order=None):
 
 def label_stoichiometry_volume(
     ats, isolated_atoms_energies, e_name, element_order=None
-):
+):  # CE I don't get what the function returns
     """
     Calculate the stoichiometry, energy, and volume coordinates for forming the convex hull.
 
@@ -412,7 +417,7 @@ def label_stoichiometry_volume(
     return p.T[:, np.argsort(p.T[0])].T
 
 
-def point_in_triangle_2D(p1, p2, p3, pn):
+def point_in_triangle_2D(p1, p2, p3, pn) -> bool:
     """
     Check if a point is inside a triangle in 2D.
 
@@ -449,7 +454,7 @@ def point_in_triangle_2D(p1, p2, p3, pn):
     )
 
 
-def point_in_triangle_ND(pn, *preg):
+def point_in_triangle_ND(pn, *preg) -> bool:
     """
     Check if a point is inside a region of hyperplanes in N dimensions.
 
@@ -467,7 +472,7 @@ def point_in_triangle_ND(pn, *preg):
     return hull.find_simplex(pn) >= 0
 
 
-def calculate_hull_3D(p):
+def calculate_hull_3D(p) -> ConvexHull:
     """
     Calculate the convex hull in 3D.
 
@@ -492,7 +497,7 @@ def calculate_hull_3D(p):
     return hull
 
 
-def calculate_hull_ND(p):
+def calculate_hull_ND(p) -> ConvexHull:
     """
     Calculate the convex hull in ND (N>=3).
 
@@ -531,7 +536,7 @@ def calculate_hull_ND(p):
 
 def get_e_distance_to_hull_3D(
     hull, at, isolated_atoms_energies=None, energy_name="energy", element_order=None
-):
+) -> float:
     """
     Calculate the energy distance to the convex hull in 3D.
 
@@ -579,7 +584,7 @@ def get_e_distance_to_hull_3D(
     return 1e6
 
 
-def piecewise_linear(x, vals):
+def piecewise_linear(x, vals) -> Any:
     """
     Piecewise linear.