Update docs

a2c_ase/runner.py

@@ -29,31 +29,48 @@ def melt_quench_md(
     log_interval: int = 100,
 ):
     """
-    Run a melt-quench molecular dynamics simulation to generate amorphous structures.
-
-    This function performs a three-stage MD simulation commonly used to create amorphous materials:
-    1. High-temperature equilibration (melting)
-    2. Controlled cooling (quenching)
-    3. Low-temperature equilibration (structural relaxation)
-
-    Args:
-        atoms: ASE Atoms object containing initial atomic structure
-        calculator: ASE calculator to compute forces and energies
-        equi_steps: Number of steps for high-temperature equilibration
-        cool_steps: Number of steps for cooling phase
-        final_steps: Number of steps for final low-temperature equilibration
-        T_high: High temperature for melting phase (Kelvin)
-        T_low: Low temperature for final phase (Kelvin)
-        time_step: MD time step in femtoseconds
-        friction: Langevin friction parameter in atomic units
-        trajectory_file: Path to save MD trajectory (None = no saving)
-        seed: Random seed for reproducibility
-        verbose: If True, print progress information
-        log_interval: Number of steps between logging progress
-
-    Returns:
-        atoms: ASE Atoms object with final amorphous structure configuration
-        log_data: Dictionary containing temperature and energy trajectories
+    Run melt-quench molecular dynamics to generate amorphous structures.
+
+    Performs a three-stage MD simulation:
+    1. High-temperature equilibration to melt the structure
+    2. Controlled cooling to quench the liquid
+    3. Low-temperature equilibration to relax the structure
+
+    Parameters
+    ----------
+    atoms : Atoms
+        ASE Atoms object with initial atomic structure
+    calculator : Calculator
+        ASE calculator for computing forces and energies
+    equi_steps : int, default=2500
+        Number of steps for high-temperature equilibration
+    cool_steps : int, default=2500
+        Number of steps for cooling phase
+    final_steps : int, default=2500
+        Number of steps for final low-temperature equilibration
+    T_high : float, default=2000.0
+        High temperature for melting phase (K)
+    T_low : float, default=300.0
+        Low temperature for final phase (K)
+    time_step : float, default=2.0
+        MD time step (fs)
+    friction : float, default=0.01
+        Langevin friction parameter (atomic units)
+    trajectory_file : str, optional
+        Path to save MD trajectory
+    seed : int, default=42
+        Random seed for reproducibility
+    verbose : bool, default=True
+        Whether to print progress information
+    log_interval : int, default=100
+        Number of steps between progress logs
+
+    Returns
+    -------
+    atoms : Atoms
+        ASE Atoms with final amorphous structure
+    log_data : dict
+        Dictionary with temperature and energy trajectories
     """
     # Set the random seed for reproducibility
     np.random.seed(seed)
@@ -147,29 +164,37 @@ def relax_unit_cell(
     trajectory_file: str | None = None,
     verbose: bool = True,
 ) -> tuple[Atoms, dict]:
-    """Relax both atomic positions and cell parameters using the FIRE algorithm.
-
-    This function performs geometry optimization of both atomic positions and unit cell
-    parameters simultaneously using ASE's implementation of the Fast Inertial Relaxation
-    Engine (FIRE) algorithm. The cell relaxation is handled by ASE's ExpCellFilter, which
-    properly handles the unit cell degrees of freedom.
-
-    Args:
-        atoms: ASE Atoms object containing atomic positions, species and lattice information
-        calculator: ASE calculator object that can compute energies, forces and stresses
-        max_iter: Maximum number of FIRE iterations. Defaults to 200.
-        fmax: Maximum force criterion for convergence in eV/Å. Defaults to 0.01.
-        trajectory_file: Optional path to save optimization trajectory. Defaults to None.
-        verbose: If True, prints optimization progress. Defaults to True.
-
-    Returns:
-        tuple containing:
-            - Atoms: Relaxed ASE Atoms object with optimized positions and cell
-            - dict: Dictionary with energy, forces, stress, and volume trajectories
-
-    Note:
-        The calculator must be able to compute stresses for cell optimization to work.
-        The Atoms object must have periodic boundary conditions set appropriately.
+    """Relax atomic positions and cell parameters using FIRE optimization.
+
+    Performs simultaneous optimization of atomic positions and unit cell parameters using
+    the Fast Inertial Relaxation Engine (FIRE) algorithm. The cell optimization is handled
+    through ASE's FrechetCellFilter.
+
+    Parameters
+    ----------
+    atoms : Atoms
+        ASE Atoms object containing atomic positions, species and cell information
+    calculator : Calculator
+        ASE calculator for computing energies, forces and stresses
+    max_iter : int, optional
+        Maximum FIRE iterations, by default 200
+    fmax : float, optional
+        Force convergence criterion in eV/Å, by default 0.01
+    trajectory_file : str, optional
+        Path to save optimization trajectory, by default None
+    verbose : bool, optional
+        Print optimization progress, by default True
+
+    Returns
+    -------
+    tuple[Atoms, dict]
+        - Relaxed Atoms object with optimized positions and cell
+        - Dictionary containing energy, forces, stress and volume trajectories
+
+    Notes
+    -----
+    The calculator must support stress tensor calculations.
+    Periodic boundary conditions must be enabled on the Atoms object.
     """
     # Ensure atoms has a calculator
     atoms.calc = calculator