From 680f29eb5f9eae583692d55dcc38831aa2339178 Mon Sep 17 00:00:00 2001
From: Paul Saxe <paul@thesaxes.net>
Date: Thu, 23 May 2024 10:21:12 -0400
Subject: [PATCH] Added standard energy and gradients to results    * Added
 'energy' and 'gradients' to optional results to support e.g. Energy Scan    *
 Fixed crashing bug in description of the Energy substep.

---
 HISTORY.rst                |  4 ++++
 psi4_step/energy.py        |  9 ++++++++-
 psi4_step/psi4_metadata.py | 14 +++++++++++++-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/HISTORY.rst b/HISTORY.rst
index bbe0dbb..1aa5756 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -1,6 +1,10 @@
 =======
 History
 =======
+2024.5.23 -- Added standard energy and gradients to results
+   * Added 'energy' and 'gradients' to optional results to support e.g. Energy Scan
+   * Fixed crashing bug in description of the Energy substep.
+     
 2024.3.17 -- Updated the installer
    * Updated the installer to use the new version of the SEAMM installer.
    * Finalizes installing either with Conda or Docker
diff --git a/psi4_step/energy.py b/psi4_step/energy.py
index 54d9cb3..24bcf82 100644
--- a/psi4_step/energy.py
+++ b/psi4_step/energy.py
@@ -60,7 +60,7 @@ def description_text(
     ):
         """Prepare information about what this node will do"""
 
-        if P is not None:
+        if P is None:
             P = self.parameters.values_to_dict()
 
         if P["level"] == "recommended":
@@ -273,6 +273,7 @@ def get_input(self, calculation_type="energy", restart=None):
                 lines.append(
                     f"Eelec, wfn = {calculation_type}('{functional}', return_wfn=True)"
                 )
+                lines.append(f"G = gradient('{functional}', ref_wfn=wfn)")
             else:
                 if calculation_type == "gradient":
                     lines.append(
@@ -290,6 +291,7 @@ def get_input(self, calculation_type="energy", restart=None):
                 lines.append(
                     f"Eelec, wfn = {calculation_type}('{method}', return_wfn=True)"
                 )
+                lines.append(f"G = gradient('{method}', ref_wfn=wfn)")
             else:
                 if calculation_type == "gradient":
                     lines.append(
@@ -332,6 +334,11 @@ def get_input(self, calculation_type="energy", restart=None):
 for item in arrays:
     variables[item] = wfn.array_variable(item).np.tolist()
 variables["Eelec"] = Eelec
+variables["energy"] = Eelec
+try:
+    variables["gradient"] = np.array(G).tolist()
+except Exception:
+    pass
 variables["_method"] = "{method}"
 variables["_method_string"] = "{method_string}"
 
diff --git a/psi4_step/psi4_metadata.py b/psi4_step/psi4_metadata.py
index 0e807ce..6139ea2 100644
--- a/psi4_step/psi4_metadata.py
+++ b/psi4_step/psi4_metadata.py
@@ -1624,6 +1624,19 @@
     Optional units for the result. If present, the value should be in these units.
 """
 metadata["results"] = {
+    "energy": {
+        "description": "The total energy",
+        "dimensionality": "scalar",
+        "property": "total energy#Psi4#{model}",
+        "type": "float",
+        "units": "kJ/mol",
+    },
+    "gradients": {
+        "description": "The gradients",
+        "dimensionality": "[3, n_atoms]",
+        "type": "float",
+        "units": "kJ/mol/Å",
+    },
     "(T) CORRECTION ENERGY": {
         "calculation": ["energy", "optimization", "thermochemistry", "vibrations"],
         "description": "",
@@ -1816,7 +1829,6 @@
         "description": "the electronic energy from the current method",
         "dimensionality": "scalar",
         "methods": ["ccsd", "ccsd(t)", "dft", "hf", "lccd", "mp2", "mp3", "mp4"],
-        "property": "total energy#Psi4#{model}",
         "type": "float",
         "units": "E_h",
     },