From 207044fc6e13a5414588b485ed1db557fc8bc37b Mon Sep 17 00:00:00 2001
From: "Lori A. Burns" <lori.burns@gmail.com>
Date: Thu, 14 Nov 2024 16:21:28 -0500
Subject: [PATCH] cfour and torchani

---
 qcengine/programs/cfour/runner.py          |  7 ++++---
 qcengine/programs/dftd3.py                 |  6 ++++--
 qcengine/programs/tests/test_dftd3_mp2d.py |  5 ++++-
 qcengine/programs/torchani.py              | 19 +++++++++----------
 4 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/qcengine/programs/cfour/runner.py b/qcengine/programs/cfour/runner.py
index e522ecfa..b7799170 100644
--- a/qcengine/programs/cfour/runner.py
+++ b/qcengine/programs/cfour/runner.py
@@ -205,9 +205,10 @@ def parse_output(
             provenance["module"] = module
 
         output_data = {
-            "schema_version": 1,
+            "schema_version": 2,
+            "input_data": input_model,
             "molecule": c4mol,  # overwrites with outfile Cartesians in case fix_*=F
-            "extras": {**input_model.extras},
+            "extras": {},
             "native_files": {k: v for k, v in outfiles.items() if v is not None},
             "properties": atprop,
             "provenance": provenance,
@@ -224,4 +225,4 @@ def parse_output(
             k.upper(): str(v) if isinstance(v, Decimal) else v for k, v in qcvars.items()
         }
 
-        return AtomicResult(**{**input_model.dict(), **output_data})
+        return AtomicResult(**output_data)
diff --git a/qcengine/programs/dftd3.py b/qcengine/programs/dftd3.py
index d3a11fa7..c47f9024 100644
--- a/qcengine/programs/dftd3.py
+++ b/qcengine/programs/dftd3.py
@@ -280,7 +280,9 @@ def parse_output(self, outfiles: Dict[str, str], input_model: "AtomicInput") ->
             retres = retres.ravel().tolist()
 
         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": {
                 "return_energy": calcinfo[f"CURRENT ENERGY"],
@@ -301,7 +303,7 @@ def parse_output(self, outfiles: Dict[str, str], input_model: "AtomicInput") ->
             output_data["extras"]["qcvars"]["2-BODY PAIRWISE DISPERSION CORRECTION ANALYSIS"] = D3pairs
         output_data["success"] = True
 
-        return AtomicResult(**{**input_model.dict(), **output_data})
+        return AtomicResult(**output_data)
 
 
 def dftd3_coeff_formatter(dashlvl: str, dashcoeff: Dict) -> str:
diff --git a/qcengine/programs/tests/test_dftd3_mp2d.py b/qcengine/programs/tests/test_dftd3_mp2d.py
index 38c8824b..b8210432 100644
--- a/qcengine/programs/tests/test_dftd3_mp2d.py
+++ b/qcengine/programs/tests/test_dftd3_mp2d.py
@@ -26,7 +26,10 @@ def test_dftd3_task(method, schema_versions, request):
     ret = qcng.compute(json_data, "dftd3", raise_error=True, return_dict=True, return_version=retver)
     ret = checkver_and_convert(ret, request.node.name, "post")
 
-    assert ret["driver"] == "energy"
+    if "v2" in request.node.name:
+        assert ret["input_data"]["driver"] == "energy"
+    else:
+        assert ret["driver"] == "energy"
     assert "provenance" in ret
     assert "normal termination of dftd3" in ret["stdout"]
 
diff --git a/qcengine/programs/torchani.py b/qcengine/programs/torchani.py
index 920c8b7d..6615dd4b 100644
--- a/qcengine/programs/torchani.py
+++ b/qcengine/programs/torchani.py
@@ -168,15 +168,14 @@ def compute(self, input_data: "AtomicInput", config: "TaskConfig") -> "AtomicRes
         #   the reliability of the models in an ensemble, and produce more data
         #   points in the regions where this quantity is below a certain
         #   threshold (inclusion criteria)
-        ret_data["extras"] = input_data.extras.copy()
-        ret_data["extras"].update(
-            {
-                "ensemble_energies": energy_array.cpu().detach().numpy(),
-                "ensemble_energy_avg": energy.item(),
-                "ensemble_energy_std": ensemble_std.item(),
-                "ensemble_per_root_atom_disagreement": ensemble_scaled_std.item(),
-            }
-        )
+        ret_data["input_data"] = input_data
+        ret_data["molecule"] = input_data.molecule
+        ret_data["extras"] = {
+            "ensemble_energies": energy_array.cpu().detach().numpy(),
+            "ensemble_energy_avg": energy.item(),
+            "ensemble_energy_std": ensemble_std.item(),
+            "ensemble_per_root_atom_disagreement": ensemble_scaled_std.item(),
+        }
 
         ret_data["provenance"] = Provenance(
             creator="torchani", version="unknown", routine="torchani.builtin.aev_computer"
@@ -186,4 +185,4 @@ def compute(self, input_data: "AtomicInput", config: "TaskConfig") -> "AtomicRes
         ret_data["success"] = True
 
         # Form up a dict first, then sent to BaseModel to avoid repeat kwargs which don't override each other
-        return AtomicResult(**{**input_data.dict(), **ret_data})
+        return AtomicResult(**ret_data)