From 879205ac46d0c5c7011e15ed15e15c08e59812f8 Mon Sep 17 00:00:00 2001
From: Josh Horton <Josh.Horton@newcastle.ac.uk>
Date: Wed, 15 Jan 2025 15:37:18 +0000
Subject: [PATCH] Allow extra constraints in torsiondrives (#440)

* do not overwrite constraints in torsiondrives

* lint

* add test with extra torsion constraints

---------

Co-authored-by: Lori A. Burns <lori.burns@gmail.com>
---
 qcengine/procedures/torsiondrive.py | 23 ++++++------
 qcengine/stock_mols.py              | 28 +++++++++++++++
 qcengine/tests/test_procedures.py   | 54 +++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+), 13 deletions(-)

diff --git a/qcengine/procedures/torsiondrive.py b/qcengine/procedures/torsiondrive.py
index 0fc1327d7..49d368d3d 100644
--- a/qcengine/procedures/torsiondrive.py
+++ b/qcengine/procedures/torsiondrive.py
@@ -183,19 +183,16 @@ def _spawn_optimization(
         dihedrals = input_model.keywords.dihedrals
         angles = grid_point.split()
 
-        keywords = {
-            **input_model.optimization_spec.keywords,
-            "constraints": {
-                "set": [
-                    {
-                        "type": "dihedral",
-                        "indices": dihedral,
-                        "value": int(angle),
-                    }
-                    for dihedral, angle in zip(dihedrals, angles)
-                ]
-            },
-        }
+        # update the keywords with extra constraints for the target torsion
+        keywords = input_model.optimization_spec.keywords
+        keywords.setdefault("constraints", {})
+        keywords["constraints"].setdefault("set", [])
+        keywords["constraints"]["set"].extend(
+            [
+                {"type": "dihedral", "indices": dihedral, "value": int(angle)}
+                for dihedral, angle in zip(dihedrals, angles)
+            ]
+        )
 
         input_data = OptimizationInput(
             keywords=keywords,
diff --git a/qcengine/stock_mols.py b/qcengine/stock_mols.py
index 47522e5af..f5c8cb5a7 100644
--- a/qcengine/stock_mols.py
+++ b/qcengine/stock_mols.py
@@ -186,6 +186,34 @@
         ],
         "molecular_multiplicity": 2,
     },
+    "propane": {
+        "symbols": ["C", "C", "C", "H", "H", "H", "H", "H", "H", "H", "H"],
+        "geometry": [
+            [-1.04538535e00, -5.26451580e-01, 1.70615765e00],
+            [7.27312570e-01, -8.36259840e-01, 3.94379109e00],
+            [2.13664641e00, 1.54796144e00, 4.51430690e00],
+            [-2.06001307e00, -2.28648849e00, 1.31543249e00],
+            [-2.45194896e00, 9.45839060e-01, 2.07343776e00],
+            [1.59997800e-02, -6.00690000e-04, 1.13012700e-02],
+            [-3.60389010e-01, -1.40265264e00, 5.61211494e00],
+            [2.08128086e00, -2.35136798e00, 3.54633241e00],
+            [3.39307317e00, 1.25279930e00, 6.13498553e00],
+            [8.27431450e-01, 3.07919479e00, 4.98806078e00],
+            [3.29667371e00, 2.11896834e00, 2.89840427e00],
+        ],
+        "connectivity": [
+            (0, 1, 1.0),
+            (0, 3, 1.0),
+            (0, 4, 1.0),
+            (0, 5, 1.0),
+            (1, 2, 1.0),
+            (1, 6, 1.0),
+            (1, 7, 1.0),
+            (2, 8, 1.0),
+            (2, 9, 1.0),
+            (2, 10, 1.0),
+        ],
+    },
 }
 
 
diff --git a/qcengine/tests/test_procedures.py b/qcengine/tests/test_procedures.py
index 92227dc59..66cc0bf4b 100644
--- a/qcengine/tests/test_procedures.py
+++ b/qcengine/tests/test_procedures.py
@@ -359,6 +359,60 @@ def test_torsiondrive_generic():
     assert ret.stdout == "All optimizations converged at lowest energy. Job Finished!\n"
 
 
+@using("mace")
+@using("torsiondrive")
+def test_torsiondrive_extra_constraints():
+
+    input_data = TorsionDriveInput(
+        keywords=TDKeywords(dihedrals=[(3, 0, 1, 2)], grid_spacing=[180]),
+        input_specification=QCInputSpecification(driver=DriverEnum.gradient, model=Model(method="small", basis=None)),
+        initial_molecule=[qcng.get_molecule("propane")],
+        optimization_spec=OptimizationSpecification(
+            procedure="geomeTRIC",
+            keywords={
+                "coordsys": "dlc",
+                # use mace as it does not have convergence issues like UFF
+                "program": "mace",
+                "constraints": {
+                    "set": [
+                        {
+                            "type": "dihedral",  # hold a dihedral through the other C-C bond fixed
+                            "indices": (0, 1, 2, 10),
+                            "value": 0.0,
+                        }
+                    ]
+                },
+            },
+        ),
+    )
+
+    ret = qcng.compute_procedure(input_data, "torsiondrive", raise_error=True)
+
+    assert ret.error is None
+    assert ret.success
+
+    expected_grid_ids = {"180", "0"}
+
+    assert {*ret.optimization_history} == expected_grid_ids
+
+    assert {*ret.final_energies} == expected_grid_ids
+    assert {*ret.final_molecules} == expected_grid_ids
+
+    assert (
+        pytest.approx(ret.final_molecules["180"].measure([3, 0, 1, 2]), abs=1.0e-2) == 180.0
+        or pytest.approx(ret.final_molecules["180"].measure([3, 0, 1, 2]), abs=1.0e-2) == -180.0
+    )
+    assert pytest.approx(ret.final_molecules["180"].measure([0, 1, 2, 10]), abs=1.0e-2) == 0.0
+    assert pytest.approx(ret.final_molecules["0"].measure([3, 0, 1, 2]), abs=1.0e-2) == 0.0
+
+    assert ret.provenance.creator.lower() == "torsiondrive"
+    assert ret.optimization_history["180"][0].provenance.creator.lower() == "geometric"
+    assert ret.optimization_history["180"][0].trajectory[0].provenance.creator.lower() == "mace"
+
+    assert "Using MACE-OFF23 MODEL for MACECalculator" in ret.stdout
+    assert "All optimizations converged at lowest energy. Job Finished!\n" in ret.stdout
+
+
 @using("mrchem")
 @pytest.mark.parametrize(
     "optimizer",