Add some missing type hints to Relaxer.relax (materialsvirtuallab#136)

* Add some missing type hints to `Relaxer.relax` Signed-off-by: Andrew S. Rosen <[email protected]> * `black` Signed-off-by: Andrew S. Rosen <[email protected]> --------- Signed-off-by: Andrew S. Rosen <[email protected]> Co-authored-by: Shyue Ping Ong <[email protected]>
aeonxai · Aug 16, 2023 · 536f4ba · 536f4ba
1 parent 789c247
commit 536f4ba
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/matgl/ext/ase.py b/matgl/ext/ase.py
@@ -205,19 +205,19 @@ def __init__(
 
     def relax(
         self,
-        atoms: Atoms,
+        atoms: Atoms | Structure | Molecule,
         fmax: float = 0.1,
         steps: int = 500,
         traj_file: str | None = None,
-        interval=1,
-        verbose=False,
+        interval: int = 1,
+        verbose: bool = False,
         **kwargs,
     ):
         """
         Relax an input Atoms.
 
         Args:
-            atoms (Atoms): the atoms for relaxation
+            atoms (Atoms | Structure | Molecule): the atoms for relaxation
             fmax (float): total force tolerance for relaxation convergence.
             Here fmax is a sum of force and stress forces
             steps (int): max number of steps for relaxation