From 992e67f51983a08d0770d4396fdd503c15dc3cd5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Roberto=20Di=20Remigio=20Eik=C3=A5s?=
 <robertodr@users.noreply.github.com>
Date: Tue, 15 Nov 2022 15:57:44 +0100
Subject: [PATCH] Enable geometry optimizations with MRChem (#380)

* Enable geometry optimizations with MRChem

* Do not overwrite 'extras' in MRChem harness

* Add test for MRChem molecular gradient

* Add test for geom. optimization with MRChem

* Add packages to test geom. opt.
---
 devtools/conda-envs/mrchem.yaml          |  8 +++-
 docs/source/program_overview.rst         |  2 +-
 qcengine/programs/mrchem.py              | 24 ++++++++---
 qcengine/programs/tests/test_mrchem.py   | 54 ++++++++++++++++++++++++
 qcengine/tests/test_harness_canonical.py |  2 +-
 qcengine/tests/test_procedures.py        | 26 ++++++++++++
 6 files changed, 107 insertions(+), 9 deletions(-)

diff --git a/devtools/conda-envs/mrchem.yaml b/devtools/conda-envs/mrchem.yaml
index 014a02940..d0d0b752b 100644
--- a/devtools/conda-envs/mrchem.yaml
+++ b/devtools/conda-envs/mrchem.yaml
@@ -2,7 +2,9 @@ name: test
 channels:
   - conda-forge
 dependencies:
-  - mrchem=1=*openmpi*_1
+  - mrchem >=1.1=*openmpi*
+  - geometric
+  - optking
 
     # Core
   - python
@@ -17,3 +19,7 @@ dependencies:
   - pytest
   - pytest-cov
   - codecov
+
+  - pip
+  - pip:
+    - pyberny
diff --git a/docs/source/program_overview.rst b/docs/source/program_overview.rst
index c620114ea..16fee0d0f 100644
--- a/docs/source/program_overview.rst
+++ b/docs/source/program_overview.rst
@@ -18,7 +18,7 @@ Quantum Chemistry
 +---------------+------------+---+---+---+------------+--------------+
 | GAMESS        | ✘          | ✓ | ✓ | ✘ | ✓          | ✘            |
 +---------------+------------+---+---+---+------------+--------------+
-| MRChem        | ✓          | ✓ | ✘ | ✘ | ✓          | ✘            |
+| MRChem        | ✓          | ✓ | ✓ | ✘ | ✓          | ✘            |
 +---------------+------------+---+---+---+------------+--------------+
 | Molpro        | ✓          | ✓ | ✓ | ✘ | ✓          | ✘            |
 +---------------+------------+---+---+---+------------+--------------+
diff --git a/qcengine/programs/mrchem.py b/qcengine/programs/mrchem.py
index b04b1c053..e2bd460e5 100644
--- a/qcengine/programs/mrchem.py
+++ b/qcengine/programs/mrchem.py
@@ -108,6 +108,7 @@ def compute(self, input_model: "AtomicInput", config: "TaskConfig") -> "AtomicRe
             "model": input_model.model,
             "molecule": input_model.molecule,
             "driver": input_model.driver,
+            "extras": input_model.extras,
         }
 
         with temporary_directory(parent=parent, suffix="_mrchem_scratch") as tmpdir:
@@ -155,16 +156,22 @@ def compute(self, input_model: "AtomicInput", config: "TaskConfig") -> "AtomicRe
                 # fill up extras:
                 # * under "raw_output" the whole JSON output from MRChem
                 # * under "properties" all the properties computed by MRChem
-                output_data["extras"] = {
-                    "raw_output": mrchem_json,
-                    "properties": {
-                        f"{ks[1]}": {f"{ks[2]}": _nested_get(mrchem_output, ks)} for ks in computed_rsp_props
-                    },
-                }
+                output_data["extras"].update(
+                    {
+                        "raw_output": mrchem_json,
+                        "properties": {
+                            f"{ks[1]}": {f"{ks[2]}": _nested_get(mrchem_output, ks)} for ks in computed_rsp_props
+                        },
+                    }
+                )
 
                 # fill up return_result
                 if input_model.driver == "energy":
                     output_data["return_result"] = mrchem_output["properties"]["scf_energy"]["E_tot"]
+                elif input_model.driver == "gradient":
+                    output_data["return_result"] = mrchem_output["properties"]["geometric_derivative"]["geom-1"][
+                        "total"
+                    ]
                 elif input_model.driver == "properties":
                     output_data["return_result"] = {
                         f"{ks[1]}": {f"{ks[2]}": _nested_get(mrchem_output, ks)} for ks in computed_rsp_props
@@ -213,6 +220,11 @@ def build_input(self, input_model: "AtomicInput", config: "TaskConfig") -> Dict[
             opts["WaveFunction"]["method"] = input_model.model.method
         else:
             opts["WaveFunction"] = {"method": input_model.model.method}
+
+        # The molecular gradient is just a first-order property for MRChem
+        if input_model.driver == "gradient":
+            opts.update({"Properties": {"geometric_derivative": True}})
+
         # Log the job settings as constructed from the input model
         logger.debug("JOB_OPTS from InputModel")
         logger.debug(pp.pformat(opts))
diff --git a/qcengine/programs/tests/test_mrchem.py b/qcengine/programs/tests/test_mrchem.py
index 2ad79e32c..8158b4af3 100644
--- a/qcengine/programs/tests/test_mrchem.py
+++ b/qcengine/programs/tests/test_mrchem.py
@@ -17,6 +17,17 @@ def h2o():
     )
 
 
+@pytest.fixture
+def fh():
+    return qcel.models.Molecule(
+        geometry=[[0.000000000000, 0.000000000000, -1.642850273986], [0.000000000000, 0.000000000000, 0.087149726014]],
+        symbols=["H", "F"],
+        fix_com=True,
+        fix_orientation=True,
+        fix_symmetry="c1",
+    )
+
+
 @using("mrchem")
 def test_energy(h2o):
     mr_kws = {
@@ -83,3 +94,46 @@ def test_dipole(h2o):
     assert res["properties"]["calcinfo_nbeta"] == 5
     assert res["properties"]["calcinfo_nmo"] == 10
     assert compare_values([-3.766420e-07, 0.0, 0.720473], res["properties"]["scf_dipole_moment"], atol=1e-3)
+
+
+@using("mrchem")
+def test_gradient(fh):
+    mr_kws = {
+        "world_prec": 1.0e-3,
+        "world_size": 6,
+    }
+
+    inp = qcel.models.AtomicInput(
+        molecule=fh,
+        driver="gradient",
+        model={
+            "method": "BLYP",
+        },
+        keywords=mr_kws,
+    )
+
+    res = qcng.compute(inp, "mrchem", raise_error=True, return_dict=True)
+
+    # Make sure the calculation completed successfully
+    assert compare_values(
+        [
+            7.7720991261973e-05,
+            0.0003374698961269812,
+            0.01670678976631068,
+            -9.482132557890285e-05,
+            -0.00042524484271181696,
+            -0.06177621634278285,
+        ],
+        res["return_result"],
+        atol=1e-3,
+    )
+    assert res["driver"] == "gradient"
+    assert "provenance" in res
+    assert res["success"] is True
+
+    # Make sure the properties parsed correctly
+    assert compare_values(-100.48858973847459, res["properties"]["return_energy"], atol=1e-3)
+    assert res["properties"]["calcinfo_natom"] == 2
+    assert res["properties"]["calcinfo_nalpha"] == 5
+    assert res["properties"]["calcinfo_nbeta"] == 5
+    assert res["properties"]["calcinfo_nmo"] == 10
diff --git a/qcengine/tests/test_harness_canonical.py b/qcengine/tests/test_harness_canonical.py
index 0b10193e3..a3a22a73d 100644
--- a/qcengine/tests/test_harness_canonical.py
+++ b/qcengine/tests/test_harness_canonical.py
@@ -80,7 +80,7 @@ def test_compute_gradient(program, model, keywords):
     inp = AtomicInput(
         molecule=molecule, driver="gradient", model=model, extras={"mytag": "something"}, keywords=keywords
     )
-    if program in ["adcc", "mrchem"]:
+    if program in ["adcc"]:
         with pytest.raises(qcng.exceptions.InputError) as e:
             qcng.compute(inp, program, raise_error=True)
 
diff --git a/qcengine/tests/test_procedures.py b/qcengine/tests/test_procedures.py
index 5f4988737..62cfb42e6 100644
--- a/qcengine/tests/test_procedures.py
+++ b/qcengine/tests/test_procedures.py
@@ -347,3 +347,29 @@ def test_torsiondrive_generic():
     assert ret.optimization_history["180"][0].trajectory[0].provenance.creator.lower() == "rdkit"
 
     assert ret.stdout == "All optimizations converged at lowest energy. Job Finished!\n"
+
+
+@using("mrchem")
+@pytest.mark.parametrize(
+    "optimizer",
+    [
+        pytest.param("geometric", marks=using("geometric")),
+        pytest.param("optking", marks=using("optking")),
+        pytest.param("berny", marks=using("berny")),
+    ],
+)
+def test_optimization_mrchem(input_data, optimizer):
+
+    input_data["initial_molecule"] = qcng.get_molecule("hydrogen")
+    input_data["input_specification"]["model"] = {"method": "HF"}
+    input_data["input_specification"]["keywords"] = {"world_prec": 1.0e-4}
+    input_data["keywords"]["program"] = "mrchem"
+
+    input_data = OptimizationInput(**input_data)
+
+    ret = qcng.compute_procedure(input_data, optimizer, raise_error=True)
+    assert 10 > len(ret.trajectory) > 1
+
+    assert pytest.approx(ret.final_molecule.measure([0, 1]), 1.0e-3) == 1.3860734486984705
+    assert ret.provenance.creator.lower() == optimizer
+    assert ret.trajectory[0].provenance.creator.lower() == "mrchem"