Take inputs of both Molecule and Structure

Author: yaoyi92

src/atomate2/forcefields/jobs.py

@@ -27,7 +27,7 @@
     from pathlib import Path
 
     from ase.calculators.calculator import Calculator
-    from pymatgen.core.structure import Structure
+    from pymatgen.core.structure import Molecule, Structure
 
 logger = logging.getLogger(__name__)
 
@@ -149,15 +149,15 @@ def __post_init__(self) -> None:
 
     @forcefield_job
     def make(
-        self, structure: Structure, prev_dir: str | Path | None = None
+        self, structure: Molecule | Structure, prev_dir: str | Path | None = None
     ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """
         Perform a relaxation of a structure using a force field.
 
         Parameters
         ----------
-        structure: .Structure
-            pymatgen structure.
+        structure: .Structure or Molecule
+            pymatgen structure or molecule.
         prev_dir : str or Path or None
             A previous calculation directory to copy output files from. Unused, just
                 added to match the method signature of other makers.

src/atomate2/forcefields/md.py

@@ -26,7 +26,7 @@
     from pathlib import Path
 
     from ase.calculators.calculator import Calculator
-    from pymatgen.core.structure import Structure
+    from pymatgen.core.structure import Molecule, Structure
 
 
 @dataclass
@@ -133,15 +133,15 @@ def __post_init__(self) -> None:
     )
     def make(
         self,
-        structure: Structure,
+        structure: Molecule | Structure,
         prev_dir: str | Path | None = None,
     ) -> ForceFieldStructureTaskDocument | ForceFieldMoleculeTaskDocument:
         """
         Perform MD on a structure using forcefields and jobflow.
 
         Parameters
         ----------
-        structure: .Structure
+        structure: .Structure or Molecule
             pymatgen structure.
         prev_dir : str or Path or None
             A previous calculation directory to copy output files from. Unused, just