From eaa9ceace53f1c09b639bf59f6a75b91e1184848 Mon Sep 17 00:00:00 2001
From: "Lori A. Burns" <lori.burns@gmail.com>
Date: Thu, 14 Nov 2024 12:26:16 -0500
Subject: [PATCH] s-dftd3 and gcp input_data

---
 qcengine/programs/dftd_ng.py             | 17 +++++++++--------
 qcengine/programs/gcp.py                 | 14 +++++++++++---
 qcengine/tests/test_harness_canonical.py |  1 +
 3 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/qcengine/programs/dftd_ng.py b/qcengine/programs/dftd_ng.py
index 049f14ab..f4ec2757 100644
--- a/qcengine/programs/dftd_ng.py
+++ b/qcengine/programs/dftd_ng.py
@@ -14,7 +14,7 @@
 from qcelemental.util import parse_version, safe_version, which_import
 
 from ..config import TaskConfig
-from ..exceptions import InputError
+from ..exceptions import InputError, ResourceError
 from .empirical_dispersion_resources import from_arrays, get_dispersion_aliases
 from .model import ProgramHarness
 
@@ -279,16 +279,17 @@ def compute(self, input_model: AtomicInput, config: TaskConfig) -> AtomicResult:
                     input_data["keywords"]["params_tweaks"] = {**planinfo["dashparams"], "s9": 0.0}
                 input_data["keywords"]["level_hint"] = level_hint
 
-        input_model = qcelemental.models.v1.AtomicInput(**input_data)
+        # sdftd3 speaks qcsk.v1
+        input_model_v1 = qcelemental.models.v1.AtomicInput(**input_data)
 
         # Run the Harness
-        output = run_qcschema(input_model)
+        output_v1 = run_qcschema(input_model_v1)
 
         # d3 qcschema interface stores error in Result model
-        if not output.success:
-            return FailedOperation(input_data=input_data, error=output.error.model_dump())
+        if not output_v1.success:
+            return FailedOperation(input_data=input_data, error=output_v1.error.model_dump())
 
-        output = output.convert_v(2)
+        output = output_v1.convert_v(2, external_input_data=input_model)
 
         if "info" in output.extras:
             qcvkey = output.extras["info"]["fctldash"].upper()
@@ -300,14 +301,14 @@ def compute(self, input_model: AtomicInput, config: TaskConfig) -> AtomicResult:
         if qcvkey:
             calcinfo[f"{qcvkey} DISPERSION CORRECTION ENERGY"] = energy
 
-        if output.driver == "gradient":
+        if output.input_data.driver == "gradient":
             gradient = output.return_result
             calcinfo["CURRENT GRADIENT"] = gradient
             calcinfo["DISPERSION CORRECTION GRADIENT"] = gradient
             if qcvkey:
                 calcinfo[f"{qcvkey} DISPERSION CORRECTION GRADIENT"] = gradient
 
-        if output.keywords.get("pair_resolved", False):
+        if output.input_data.keywords.get("pair_resolved", False):
             pw2 = output.extras["dftd3"]["additive pairwise energy"]
             pw3 = output.extras["dftd3"]["non-additive pairwise energy"]
             assert abs(pw2.sum() + pw3.sum() - energy) < 1.0e-8, f"{pw2.sum()} + {pw3.sum()} != {energy}"
diff --git a/qcengine/programs/gcp.py b/qcengine/programs/gcp.py
index 5182752e..f47c88ab 100644
--- a/qcengine/programs/gcp.py
+++ b/qcengine/programs/gcp.py
@@ -261,10 +261,18 @@ def parse_output(self, outfiles: Dict[str, str], input_model: "AtomicInput") ->
         elif isinstance(retres, np.ndarray):
             retres = retres.ravel().tolist()
 
+        properties = {
+            "calcinfo_natom": len(input_model.molecule.symbols),
+            "return_energy": ene,
+            f"return_{input_model.driver.value}": retres,
+        }
+
         output_data = {
-            "extras": input_model.extras,
+            "input_data": input_model,
+            "molecule": input_model.molecule,
+            "extras": {},
             "native_files": {k: v for k, v in outfiles.items() if v is not None},
-            "properties": {},
+            "properties": properties,
             "provenance": Provenance(
                 creator="GCP", version=self.get_version(), routine=__name__ + "." + sys._getframe().f_code.co_name
             ),
@@ -276,7 +284,7 @@ def parse_output(self, outfiles: Dict[str, str], input_model: "AtomicInput") ->
         output_data["extras"]["qcvars"] = calcinfo
 
         output_data["success"] = True
-        return AtomicResult(**{**input_model.dict(), **output_data})
+        return AtomicResult(**output_data)
 
 
 class MCTCGCPHarness(GCPHarness):
diff --git a/qcengine/tests/test_harness_canonical.py b/qcengine/tests/test_harness_canonical.py
index 181b3b13..d156e63d 100644
--- a/qcengine/tests/test_harness_canonical.py
+++ b/qcengine/tests/test_harness_canonical.py
@@ -83,6 +83,7 @@ def test_compute_energy(program, model, keywords, schema_versions, request):
 
     assert ret.success is True
     assert isinstance(ret.return_result, float)
+    assert ret.return_result == ret.properties.return_energy
 
 
 @pytest.mark.parametrize("program, model, keywords", _canonical_methods)