From c83124585dedcb844cd081c83cd38bd360cce0ac Mon Sep 17 00:00:00 2001
From: Christina Ertural <52951132+QuantumChemist@users.noreply.github.com>
Date: Thu, 4 Apr 2024 07:31:49 +0200
Subject: [PATCH] Add dir_name to ForceFieldMaker and **task_document_kwargs to
 from_ase_compatible_result (#791)

* added dir_name to ForceFieldRelaxMaker and **task_document_kwargs to ForceFieldTaskDocument.from_ase_compatible_result

* Simplify taskdoc update

Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>

* added docstring for task_doc_kwargs in force field schemas

* changed the unit test for the force_field_task_doc attributes

* fix linting

---------

Co-authored-by: Janosh Riebesell <janosh.riebesell@gmail.com>
---
 src/atomate2/forcefields/jobs.py    | 2 ++
 src/atomate2/forcefields/schemas.py | 6 ++++++
 tests/forcefields/test_jobs.py      | 3 +++
 3 files changed, 11 insertions(+)

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")}