enhance set_stage2_constraint to handle options["constraint_group"]

resp/__init__.py

@@ -2,13 +2,13 @@
 
 """Top-level package for RESP."""
 
-__authors__ =  "Asem Alenaizan"
-__version__ = '0.8'
+__authors__ = "Asem Alenaizan"
+__version__ = "1.0.1"
 __license__ = "BSD-3-Clause"
-__date__    = "2019-08-07"
+__date__ = "2019-08-07"
 
-from .driver import resp
 from . import espfit
+from .driver import resp
 from .extras import test
 from .stage2_helper import set_stage2_constraint
 from .vdw_surface import vdw_surface

resp/stage2_helper.py

@@ -1,13 +1,13 @@
-from __future__ import division, absolute_import, print_function
+from __future__ import absolute_import, division, print_function
 
 import numpy as np
-
 import psi4
 
 """
 A helper script to facilitate the use of constraints for two-stage fitting.
 """
 
+
 def _get_stage2_atoms(molecule):
     """Determines atoms for second stage fit. The atoms
        are identified as sp3 carbons that have one or more hydrogens.
@@ -33,16 +33,16 @@ def _get_stage2_atoms(molecule):
     bond_profile = psi4.qcdb.parker._bond_profile(molecule)
     for i in range(molecule.natom()):
         # Find carbon atoms
-        if symbols[i] != 'C':
+        if symbols[i] != "C":
             continue
         # Check that it has 4 bonds
         bonds_for_atom = [j for j in bond_profile if i in j[:2]]
         if len(bonds_for_atom) == 4:
             group = []
             for atoms in bonds_for_atom:
                 j = atoms[0] if atoms[0] != i else atoms[1]
-                if symbols[j] == 'H':
-                    group.append(j + 1)  
+                if symbols[j] == "H":
+                    group.append(j + 1)
             if group:
                 groups[i + 1] = group
 
@@ -54,10 +54,13 @@ def set_stage2_constraint(molecule, charges, options):
 
     The default constraints are the following:
     Atoms that are excluded from the second stage fit are constrained
-    to their charges from the first stage fit. C-H groups determined 
+    to their charges from the first stage fit. C-H groups determined
     by _get_stage2_atoms are refitted and the hydrogen atoms connected
     to the same carbon are constrained to have identical charges.
 
+    If `options["constraint_group"]` is provided, the groups involving the
+    second stage atoms are retained.
+
     Parameters
     ----------
     molecule : psi4.core.Molecule
@@ -75,17 +78,24 @@ def set_stage2_constraint(molecule, charges, options):
 
     """
     second_stage = _get_stage2_atoms(molecule)
-    atoms = list(range(1, molecule.natom()+1))
-    constraint_group = []
-    for i in second_stage.keys():
-        atoms.remove(i)
-        group = []
-        for j in second_stage[i]:
-            atoms.remove(j)
-            group.append(j)
-        constraint_group.append(group)
+    atoms = list(range(1, molecule.natom() + 1))
+    constraint_group = options.get("constraint_group", [])
+    if constraint_group:
+        stage2_atoms = set()
+        for k, v in second_stage.items():
+            stage2_atoms.add(k)
+            stage2_atoms.update(v)
+        constraint_group = [g for g in constraint_group if g[0] in stage2_atoms]
+    else:
+        for i, hs in second_stage.items():
+            atoms.remove(i)
+            group = []
+            for j in hs:
+                atoms.remove(j)
+                group.append(j)
+            constraint_group.append(group)
     constraint_charge = []
     for i in atoms:
-        constraint_charge.append([charges[i-1], [i]])
-    options['constraint_charge'] = constraint_charge
-    options['constraint_group'] = constraint_group
+        constraint_charge.append([charges[i - 1], [i]])
+    options["constraint_charge"] = constraint_charge
+    options["constraint_group"] = constraint_group

resp/tests/test_stage2_helper.py

@@ -0,0 +1,97 @@
+import numpy as np
+import psi4
+import pytest
+import resp
+
+
+@pytest.mark.parametrize(
+    "constraint_group,expected",
+    [
+        (
+            [],
+            [[29, 30], [31, 32], [33, 34], [36, 37, 38], [39, 40, 41]],
+        ),  # no constraint group
+        (
+            [
+                [17, 24],
+                [18, 22],
+                [19, 23],
+                [29, 30],
+                [31, 32],
+                [33, 34],
+                [42, 35],
+                [36, 37, 38, 39, 40, 41],
+                [1, 3],
+            ],
+            [
+                [19, 23],
+                [29, 30],
+                [31, 32],
+                [33, 34],
+                [36, 37, 38, 39, 40, 41],
+            ],
+        ),
+    ],
+)
+def test_set_stage2_constraint(constraint_group, expected):
+    mol = psi4.geometry(
+        """
+-1 1
+  O 2.005200 -46.220700 4.303400
+  C 0.759800 -46.107200 4.123500
+  O 0.097300 -46.171000 3.047100
+  C -0.095900 -45.916600 5.368300
+  C -2.086400 -45.944500 6.432100
+  C -3.436800 -46.043400 6.785500
+  C -3.766800 -45.956500 8.138500
+  C -2.767600 -45.766700 9.120900
+  C -1.423300 -45.658300 8.764500
+  C -1.058900 -45.746400 7.408500
+  C 0.196000 -45.721300 6.709200
+  C 1.559900 -45.547100 7.313900
+  C 2.137500 -46.814600 7.984500
+  C 2.063800 -48.056700 7.091300
+  O 0.821100 -48.746700 7.357800
+  C -0.057400 -49.006000 6.344200
+  C 0.296500 -49.135700 4.998900
+  C -0.677300 -49.357400 4.016000
+  C -0.268900 -49.363700 2.565200
+  C -2.006200 -49.493200 4.434800
+  Cl -3.277800 -49.738200 3.211600
+  C -2.391500 -49.407600 5.781700
+  C -3.830900 -49.529300 6.216000
+  C -1.396200 -49.158000 6.726200
+  H -4.204000 -46.200000 6.030000
+  H -4.808500 -46.040900 8.442300
+  H -3.055900 -45.708200 10.168900
+  H -0.661700 -45.521300 9.530100
+  H 1.535200 -44.740300 8.060600
+  H 2.241000 -45.244000 6.515000
+  H 1.612800 -47.040700 8.920900
+  H 3.188600 -46.624300 8.243900
+  H 2.871000 -48.761600 7.328800
+  H 2.144200 -47.755400 6.045000
+  H 1.321200 -48.992300 4.677400
+  H 0.804700 -49.548800 2.475300
+  H -0.812000 -50.117800 1.987600
+  H -0.457900 -48.373100 2.136200
+  H -3.918400 -49.358700 7.291500
+  H -4.462000 -48.794300 5.705900
+  H -4.237400 -50.519900 5.978200
+  H -1.656000 -49.017200 7.771100
+  N -1.466300 -46.023200 5.214300
+  H -1.849300 -46.231200 4.299200
+"""
+    )
+    options = {
+        "VDW_SCALE_FACTORS": [1.4, 1.6, 1.8, 2.0],
+        "VDW_POINT_DENSITY": 1.0,
+        "RESP_A": 0.0005,
+        "RESP_B": 0.1,
+    }
+    if constraint_group:
+        options["constraint_group"] = constraint_group
+    charges = np.zeros(mol.natom())
+    psi4.qcdb.parker._expected_bonds["CL"] = 1
+    resp.set_stage2_constraint(mol, charges, options)
+    assert options["constraint_group"] == expected