diff --git a/src/atomate2/forcefields/jobs.py b/src/atomate2/forcefields/jobs.py index a210f5908b..a1d0f4b0ca 100644 --- a/src/atomate2/forcefields/jobs.py +++ b/src/atomate2/forcefields/jobs.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import os from dataclasses import dataclass, field from typing import TYPE_CHECKING @@ -121,6 +122,7 @@ def make( "WARNING: A negative number of steps is not possible. " "Behavior may vary..." ) + self.task_document_kwargs.setdefault("dir_name", os.getcwd()) with revert_default_dtype(): relaxer = Relaxer( diff --git a/src/atomate2/forcefields/schemas.py b/src/atomate2/forcefields/schemas.py index c4b7063176..ac11bf230d 100644 --- a/src/atomate2/forcefields/schemas.py +++ b/src/atomate2/forcefields/schemas.py @@ -153,6 +153,7 @@ def from_ase_compatible_result( optimizer_kwargs: dict = None, ionic_step_data: tuple = ("energy", "forces", "magmoms", "stress", "structure"), store_trajectory: StoreTrajectoryOption = StoreTrajectoryOption.NO, + **task_document_kwargs, ) -> "ForceFieldTaskDocument": """ Create a ForceFieldTaskDocument for a Task that has ASE-compatible outputs. @@ -173,6 +174,10 @@ def from_ase_compatible_result( Keyword arguments that will get passed to :obj:`Relaxer()`. ionic_step_data : tuple Which data to save from each ionic step. + store_trajectory: + whether to set the StoreTrajectoryOption + task_document_kwargs : dict + Additional keyword args passed to :obj:`.ForceFieldTaskDocument()`. """ trajectory = result["trajectory"] @@ -291,4 +296,5 @@ def from_ase_compatible_result( included_objects=list(forcefield_objects.keys()), forcefield_objects=forcefield_objects, is_force_converged=result.get("is_force_converged"), + **task_document_kwargs, ) diff --git a/tests/forcefields/test_jobs.py b/tests/forcefields/test_jobs.py index ebd3837bf4..76ebd4c401 100644 --- a/tests/forcefields/test_jobs.py +++ b/tests/forcefields/test_jobs.py @@ -66,6 +66,9 @@ def test_chgnet_relax_maker(si_structure: Structure, relax_cell: bool): assert output1.output.energy == approx(-10.6274, rel=1e-2) assert output1.output.ionic_steps[-1].magmoms[0] == approx(0.00303572, rel=1e-2) + # check the force_field_task_doc attributes + assert Path(responses[job.uuid][1].output.dir_name).exists() + def test_m3gnet_static_maker(si_structure): task_doc_kwargs = {"ionic_step_data": ("structure", "energy")}