Cleanup: dtypes, elastic tensors, citations (#1442)

* add major refs

* dtype for matgl

* make default_dtype an explicit calculator kwarg

* lobster citation description

Co-authored-by: J. George <JaGeo@users.noreply.github.com>

* ensure mace dtype gets passed as kwarg

---------

Co-authored-by: J. George <JaGeo@users.noreply.github.com>

Author: esoteric-ephemera

docs/dev/forcefields.md

@@ -10,7 +10,7 @@ There is both an `ase` module in `atomate2`, based around general `ase` `Calcula
 
 The `ase` module should be used to manage high-level tasks, such as geometry optimization, molecular dynamics, and nudged elastic band. Any further developments to these tools in `ase` should also warrant updates in this module in `atomate2`. For example, when `ase` rolled out the `MTKNPT` NPT MD barostat as a replacement for the default barostat, this was also made the default in `atomate2`.
 
-The `forcefields` library should be used to develop concrete implementations of workflows, e.g., harmonic phonon, Grüneisen parameter, ApproxNEB.
+The `forcefields` library should be used to develop concrete implementations of workflows, e.g., harmonic phonon, Grüneisen parameter, [ApproxNEB](https://doi.org/10.1063/1.4960790).
 
 ## Dependency Chaos
 

src/atomate2/abinit/jobs/base.py

@@ -11,6 +11,7 @@
 
 import jobflow
 from jobflow import Maker, Response, job
+from pymatgen.util.due import Doi, due
 
 from atomate2 import SETTINGS
 from atomate2.abinit.files import write_abinit_input_set
@@ -98,6 +99,7 @@ def setup_job(
     )
 
 
+@due.dcite(Doi("10.1063/5.028827"), description="Most recent Abinit paper")
 @dataclass
 class BaseAbinitMaker(Maker):
     """

src/atomate2/aims/jobs/base.py

@@ -10,6 +10,7 @@
 from jobflow import Maker, Response, job
 from monty.serialization import dumpfn
 from pymatgen.io.aims.sets.base import AimsInputGenerator
+from pymatgen.util.due import Doi, due
 
 from atomate2 import SETTINGS
 from atomate2.aims.files import (
@@ -40,6 +41,7 @@
 _FILES_TO_ZIP = _INPUT_FILES + _OUTPUT_FILES
 
 
+@due.dcite(Doi("10.1016/j.cpc.2009.06.022"), description="FHI-AIMS")
 @dataclass
 class BaseAimsMaker(Maker):
     """

src/atomate2/ase/jobs.py

@@ -12,6 +12,7 @@
 from jobflow import Maker, job
 from pymatgen.core import Molecule, Structure
 from pymatgen.io.ase import AseAtomsAdaptor
+from pymatgen.util.due import Doi, due
 
 from atomate2.ase.schemas import AseResult, AseTaskDoc
 from atomate2.ase.utils import AseRelaxer
@@ -28,6 +29,9 @@
 _ASE_DATA_OBJECTS = ["trajectory"]
 
 
+@due.dcite(
+    Doi("10.1088/1361-648X/aa680e"), description="Atomic simulation environment."
+)
 @dataclass
 class AseMaker(Maker, ABC):
     """

src/atomate2/common/flows/approx_neb.py

@@ -7,6 +7,7 @@
 
 from emmet.core.mobility.migrationgraph import MigrationGraphDoc
 from jobflow import Flow, Maker, OnMissing
+from pymatgen.util.due import Doi, due
 
 from atomate2.common.jobs.approx_neb import (
     collate_images_single_hop,
@@ -25,6 +26,7 @@
     from pymatgen.util.typing import CompositionLike
 
 
+@due.dcite(Doi("https://doi.org/10.1063/1.4960790"), description="ApproxNEB")
 @dataclass
 class CommonApproxNebMaker(Maker):
     """Run an ApproxNEB workflow.

src/atomate2/cp2k/jobs/base.py

@@ -17,6 +17,7 @@
 )
 from pymatgen.electronic_structure.dos import DOS, CompleteDos, Dos
 from pymatgen.io.common import VolumetricData
+from pymatgen.util.due import Doi, due
 
 from atomate2 import SETTINGS
 from atomate2.common.files import gzip_files, gzip_output_folder
@@ -85,6 +86,13 @@ def make(structure):
     return job(method, data=_DATA_OBJECTS, output_schema=TaskDocument)
 
 
+@due.dcite(
+    Doi("10.1063/5.0007045"),
+    description=(
+        "CP2K review - ensure you cite all references "
+        'in the "R E F E R E N C E S" section of the CP2K output'
+    ),
+)
 @dataclass
 class BaseCp2kMaker(Maker):
     """

src/atomate2/forcefields/utils.py

@@ -20,6 +20,11 @@
     from collections.abc import Callable, Generator
     from typing import Any
 
+    try:
+        from torch import dtype as torch_dtype
+    except ImportError:
+        torch_dtype = str
+
     from ase.calculators.calculator import Calculator
 
     from atomate2.ase.schemas import AseResult
@@ -229,7 +234,9 @@ def ase_calculator_name(self) -> str:
 
 
 def ase_calculator(
-    calculator_meta: str | MLFF | dict, **kwargs: Any
+    calculator_meta: str | MLFF | dict,
+    default_dtype: str | torch_dtype | None = None,
+    **kwargs: Any,
 ) -> Calculator | None:
     """
     Create an ASE calculator from a given set of metadata.
@@ -246,7 +253,7 @@ def ase_calculator(
                 "@callable": "CHGNetCalculator"
             }
         ```
-    args : optional args to pass to a calculator
+    default_dtype (str or pytorch dtype) : optional pytorch dtype to use if applicable
     kwargs : optional kwargs to pass to a calculator
 
     Returns
@@ -310,6 +317,9 @@ def ase_calculator(
                         )
                         matgl.config.BACKEND = "PYG"
 
+                if default_dtype is not None:
+                    matgl.set_default_dtype(default_dtype)
+
                 matgl_calc = getattr(
                     import_module(f"matgl.ext._ase_{matgl.config.BACKEND.lower()}"),
                     "PESCalculator",
@@ -328,6 +338,7 @@ def ase_calculator(
                     calculator = MACECalculator(
                         model_paths=model_path,
                         device=device,
+                        default_dtype=default_dtype or "",
                         **kwargs,
                     )
 
@@ -345,18 +356,20 @@ def ase_calculator(
                             "damping": "bj",
                             "xc": "pbe",
                             "cutoff": 40.0 * Bohr,
-                            "dtype": kwargs.get(
-                                "default_dtype", torch.get_default_dtype()
-                            ),
+                            "dtype": default_dtype or torch.get_default_dtype(),
                         }
-                        for k, v in default_d3_kwargs.items():
-                            if k not in kwargs:
-                                kwargs[k] = v
+                        kwargs.update(
+                            {
+                                k: v
+                                for k, v in default_d3_kwargs.items()
+                                if k not in kwargs
+                            }
+                        )
 
                         d3_calc = TorchDFTD3Calculator(device=device, **kwargs)
                         calculator = SumCalculator([calculator, d3_calc])
                 else:
-                    calculator = mace_mp(**kwargs)
+                    calculator = mace_mp(default_dtype=default_dtype or "", **kwargs)
 
             case MLFF.Nequip | MLFF.Allegro:
                 from nequip.integrations.ase import NequIPCalculator

src/atomate2/jdftx/jobs/base.py

@@ -13,6 +13,7 @@
     BandStructure,
     BandStructureSymmLine,
 )
+from pymatgen.util.due import Doi, due
 
 from atomate2.jdftx.files import write_jdftx_input_set
 from atomate2.jdftx.run import run_jdftx, should_stop_children
@@ -73,6 +74,7 @@ def jdftx_job(method: Callable) -> job:
     return job(method, data=_DATA_OBJECTS, output_schema=TaskDoc)
 
 
+@due.dcite(Doi("10.1016/j.softx.2017.10.006"), description="JDFTx")
 @dataclass
 class BaseJdftxMaker(Maker):
     """

src/atomate2/lobster/jobs.py

@@ -10,6 +10,7 @@
 from pymatgen.electronic_structure.cohp import CompleteCohp
 from pymatgen.electronic_structure.dos import LobsterCompleteDos
 from pymatgen.io.lobster import Bandoverlaps, Icohplist, Lobsterin
+from pymatgen.util.due import Doi, due
 
 from atomate2 import SETTINGS
 from atomate2.common.files import gzip_output_folder
@@ -27,6 +28,13 @@
 _FILES_TO_ZIP = [*LOBSTEROUTPUT_FILES, "lobsterin", *VASP_OUTPUT_FILES]
 
 
+@due.dcite(
+    Doi("https://doi.org/10.1002/jcc.26353"),
+    description=(
+        "Most recent LOBSTER paper. "
+        "Please cite the publications mentioned in the LOBSTER Terms of Use."
+    ),
+)
 @dataclass
 class LobsterMaker(Maker):
     """

src/atomate2/openff/core.py

@@ -13,6 +13,7 @@
 from openff.interchange.components._packmol import pack_box
 from openff.toolkit import ForceField
 from openff.units import unit
+from pymatgen.util.due import Doi, due
 
 from atomate2.openff.utils import create_mol_spec_list, merge_specs_by_name_and_smiles
 
@@ -56,6 +57,7 @@ def make(structure):
     )
 
 
+@due.dcite(Doi("10.1021/acs.jpcb.4c01558"), description="Open forcefield initiative")
 @openff_job
 def generate_interchange(
     input_mol_specs: list[MoleculeSpec | dict],