Add dir_name to ForceFieldMaker and **task_document_kwargs to from_as…

…e_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

@@ -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(

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,
  )

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