add ability to boltzmann-weight properties from conformational ensembles

cctk/ensemble.py

@@ -134,7 +134,6 @@ def combined_properties(self):
         combined = dict()
         for p in self.properties_list():
             combined = {**combined, **p}
-
         return combined
 
     def get_property(self, idx, prop):
@@ -608,4 +607,57 @@ def assign_connectivity(self, index=0):
 
         return self
 
+    def boltzmann_average(self, which, energies=None, temp=298, energy_unit="hartree", return_weights=False):
+        """
+        Computes the Boltzmann-weighted average of a property over the whole ensemble.
+
+        Args:
+            which (str): which property to compute
+            energy (np.ndarray): list of energies to use for weighting.
+                Will default to ``self[:,"energy"]``, although other strings can be passed as well as shorthand for ``self[:,energy]``.
+            temp (float): temperature for Boltzmann-weighting, in K
+            energy_unit (str): either ``kcal_mol`` or ``hartree``
+            return_weights (bool): whether to return a list of weights too
+
+        Returns:
+            weighted property, of the same shape as the individual property
+        """
+        if energies is None:
+            energies = self[:,"energy"]
+        elif isinstance(energies, str):
+            energies = self[:,energies]
+        elif isinstance(energies, (list, np.ndarray, cctk.OneIndexedArray)):
+            pass
+        else:
+            raise ValueError(f"invalid energy value {energies} (type {type(energies)})")
+
+        for i, (m, pd) in enumerate(self.items()):
+            assert which in pd, f"molecule #{i} doesn't have property {which} defined!"
 
+        values = np.array(self[:,which], dtype=np.float64)
+        energies = np.array(energies, dtype=np.float64)
+
+        assert len(energies) == len(self)
+        assert len(values) == len(self)
+        assert all([e is not None for e in energies]), "energy not defined for all molecules"
+        assert all([v is not None for v in values]), f"property {which} not defined for all molecules"
+
+        # perhaps at some point we will need a real unit system like simtk/OpenMM, but not today!
+        if energy_unit == "kcal_mol":
+            energies = energies / 627.509
+        energies = energies - np.min(energies)
+
+        R = 3.1668105e-6 # eH/K
+
+        weights = np.exp(-1*energies/(R*temp))
+        weights = weights / np.sum(weights)
+
+        try:
+            weighted_value = np.average(values, weights=weights)
+        except Exception as e:
+            raise ValueError(f"error computing Boltzmann average: {e}")
+
+        if return_weights:
+            return weighted_value, weights
+        else:
+            return weighted_value

cctk/quasiclassical.py

@@ -36,12 +36,14 @@ def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities
 
     mol = copy.deepcopy(molecule)
     total_PE = 0
+    total = 0
 
     velocities = np.zeros_like(molecule.geometry.view(np.ndarray)).view(cctk.OneIndexedArray)
 
     for mode in mol.vibrational_modes:
-        PE, KE = apply_vibration(mol, mode, temperature=temperature)
+        PE, KE, TE = apply_vibration(mol, mode, temperature=temperature)
         total_PE += PE
+        total += TE
 
         #### below code initializes velocities. pulling from jprogdyn Initializer lines 440 & onwards.
 
@@ -55,9 +57,9 @@ def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities
                 velocities[idx] += mode_velocity * mode.displacements[idx]
 
     if return_velocities:
-        return mol, total_PE, velocities
+        return mol, total_PE, total, velocities
     else:
-        return mol, total_PE
+        return mol, total_PE, total
 
 def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False):
     """
@@ -91,7 +93,7 @@ def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False)
     if verbose:
         print(f"Mode {mode.frequency:.2f} ({mode.energy():.2f} kcal/mol)\t QC Level {level}\t Shift {rel_shift:.2%} of a potential {max_shift:.2f} Å\tPE = {potential_energy:.2f} kcal/mol\tk = {mode.force_constant:.2f} kcal/mol Å^-2")
 
-    return potential_energy, kinetic_energy
+    return potential_energy, kinetic_energy, energy
 
 def get_hermite_polynomial(n):
     """

test/static/pcseg_dcm.gjf

@@ -12,56 +12,56 @@ title
 17          0.00000000   -0.70710701   -0.16835900
 
 H     0
-S   4   1.00
+S    4   1.00
       0.754732D+02           0.427821D-02
       0.113575D+02           0.328672D-01
       0.260081D+01           0.159578D+00
       0.735503D+00           0.500000D+00
-S   1   1.00
+S    1   1.00
       0.231761D+00           1.0000000
-S   1   1.00
+S    1   1.00
       0.741675D-01           1.0000000
-P   1   1.00
+P    1   1.00
       0.160000D+01           1.0000000
-P   1   1.00
+P    1   1.00
       0.450000D+00           1.0000000
-D   1   1.00
+D    1   1.00
       0.125000D+01           1.0000000
 ****
 C     0
-S   7   1.00
+S    7   1.00
       0.731999D+04           0.695369D-03
       0.109818D+04           0.537138D-02
       0.250014D+03           0.274734D-01
       0.707586D+02           0.105078D+00
       0.228595D+02           0.291262D+00
       0.803936D+01           0.500000D+00
       0.289309D+01           0.362224D+00
-S   2   1.00
+S    2   1.00
       0.521640D+01          -0.957844D-01
       0.556124D+00           0.500000D+00
-S   1   1.00
+S    1   1.00
       0.228813D+00           1.0000000
-S   1   1.00
+S    1   1.00
       0.936338D-01           1.0000000
-P   4   1.00
+P    4   1.00
       0.339855D+02           0.818395D-02
       0.772957D+01           0.586921D-01
       0.224954D+01           0.222514D+00
       0.767681D+00           0.500000D+00
-P   1   1.00
+P    1   1.00
       0.263283D+00           1.0000000
-P   1   1.00
+P    1   1.00
       0.848744D-01           1.0000000
-D   1   1.00
+D    1   1.00
       0.140000D+01           1.0000000
-D   1   1.00
+D    1   1.00
       0.450000D+00           1.0000000
-F   1   1.00
+F    1   1.00
       0.950000D+00           1.0000000
 ****
 Cl     0
-S   8   1.00
+S    8   1.00
       0.154272D+06           0.222473D-03
       0.231232D+05           0.172590D-02
       0.526257D+04           0.899286D-02
@@ -70,37 +70,37 @@ S   8   1.00
       0.174634D+03           0.304148D+00
       0.672393D+02           0.500000D+00
       0.265780D+02           0.360252D+00
-S   3   1.00
+S    3   1.00
       0.131826D+03          -0.531330D-01
       0.400703D+02          -0.233346D+00
       0.789448D+01           0.500000D+00
-S   2   1.00
+S    2   1.00
       0.427711D+01           0.500000D+00
       0.196233D+01           0.370517D+00
-S   1   1.00
+S    1   1.00
       0.520793D+00           1.0000000
-S   1   1.00
+S    1   1.00
       0.181633D+00           1.0000000
-P   7   1.00
+P    7   1.00
       0.124440D+04           0.988206D-03
       0.295109D+03           0.829674D-02
       0.947224D+02           0.430906D-01
       0.352588D+02           0.148443D+00
       0.144211D+02           0.344744D+00
       0.610371D+01           0.500000D+00
       0.258886D+01           0.348278D+00
-P   2   1.00
+P    2   1.00
       0.109447D+02          -0.196768D-01
       0.104413D+01           0.500000D+00
-P   1   1.00
+P    1   1.00
       0.376184D+00           1.0000000
-P   1   1.00
+P    1   1.00
       0.120613D+00           1.0000000
-D   1   1.00
+D    1   1.00
       0.231000D+01           1.0000000
-D   1   1.00
+D    1   1.00
       0.520000D+00           1.0000000
-F   1   1.00
+F    1   1.00
       0.720000D+00           1.0000000
 ****
 

test/test_ensemble.py

@@ -7,7 +7,7 @@
 #
 # python -m unittest test.test_ensemble2.TestEnsemble2
 class TestEnsemble(unittest.TestCase):
-    def test_ensemble(self):
+    def build_test_ensemble(self):
         path = "test/static/phenylpropane*.out"
         conformational_ensemble = cctk.ConformationalEnsemble()
         for filename in sorted(glob.glob(path)):
@@ -16,6 +16,11 @@ def test_ensemble(self):
             molecule = ensemble.molecules[-1]
             properties_dict = ensemble.get_properties_dict(molecule)
             conformational_ensemble.add_molecule(molecule,properties_dict)
+        return conformational_ensemble
+
+    def test_ensemble(self):
+        conformational_ensemble = self.build_test_ensemble()
+
         m1 = conformational_ensemble.molecules[0]
         self.assertEqual(conformational_ensemble[m1,"filename"], 'test/static/phenylpropane_1.out')
         l1 = conformational_ensemble.molecules[0:2]
@@ -118,32 +123,19 @@ def test_ensemble_indexing2(self):
     def test_ensemble_properties(self):
         filename = "test/static/gaussian_file.out"
         gaussian_file = cctk.GaussianFile.read_file(filename)
-        #print(gaussian_file.route_card)
-        #print(gaussian_file.job_types)
         ensemble = gaussian_file.ensemble
         self.assertTrue(isinstance(ensemble, cctk.ConformationalEnsemble))
 
         self.assertEqual(len(ensemble), 3)
         self.assertTrue(isinstance(ensemble[0], cctk.ConformationalEnsemble))
 
-        #for molecule,properties_dict in ensemble.items():
-        #    for property_name,value in properties_dict.items():
-        #        print(property_name, ":", value)
-        #    print()
-
         self.assertEqual(ensemble[0,"energy"], -1159.56782625)
         self.assertListEqual(ensemble[:,"energy"], [-1159.56782625, -1159.56782622, -1159.56782622])
         self.assertListEqual(ensemble[:,"enthalpy"], [None, None, -1159.314817])
 
     def test_sort(self):
-        path = "test/static/phenylpropane*.out"
-        conformational_ensemble = cctk.ConformationalEnsemble()
-        for filename in sorted(glob.glob(path)):
-            gaussian_file = cctk.GaussianFile.read_file(filename)
-            ensemble = gaussian_file.ensemble
-            molecule = ensemble.molecules[-1]
-            properties_dict = ensemble.get_properties_dict(molecule)
-            conformational_ensemble.add_molecule(molecule,properties_dict)
+        conformational_ensemble = self.build_test_ensemble()
+
         original_order = conformational_ensemble[:,"energy"]
         self.assertListEqual(original_order,[0.0140132996483, 0.0163679933924, 0.0213666533731, 0.0180903133947, 0.0547890926923, 0.0182782865186])
         sorted_ensemble = conformational_ensemble.sort_by("energy", ascending=False)
@@ -164,5 +156,22 @@ def test_sort(self):
         energy0 = sorted_ensemble.get_property(lowest_molecule, "energy")
         self.assertEqual(energy0, 0.0140132996483)
 
+    def test_boltzmann_weighting(self):
+        conformational_ensemble = self.build_test_ensemble()
+
+        values, weights = conformational_ensemble.boltzmann_average("energy", energies=[1.36,0,1000,1000,1000,1000], energy_unit="kcal_mol", return_weights=True)
+        self.assertTrue((weights[0]/weights[1] - 0.1 < 0.01))
+        self.assertTrue(values - 0.016152 < 0.0001)
+
+        ce2 = cctk.ConformationalEnsemble()
+        for filename in glob.glob("test/static/pentane*.out"):
+            gaussian_file = cctk.GaussianFile.read_file(filename)
+            ensemble = gaussian_file.ensemble
+            molecule = ensemble.molecules[-1]
+            properties_dict = ensemble.get_properties_dict(molecule)
+            ce2.add_molecule(molecule,properties_dict)
+        enthalpy = ce2.boltzmann_average("enthalpy")
+        self.assertTrue(enthalpy - .10722 < 0.0001)
+
 if __name__ == '__main__':
     unittest.main()

test/test_freqs.py

@@ -47,7 +47,7 @@ def test_perturb(self):
 
         energies = list()
         for i in range(1000):
-            _, e = qc.get_quasiclassical_perturbation(mol)
+            _, e, _ = qc.get_quasiclassical_perturbation(mol)
             energies.append(e)
 
         energies = np.array(energies)
@@ -63,10 +63,10 @@ def test_perturb_water(self):
         file = cctk.GaussianFile.read_file(path)
 
         mol = file.get_molecule()
-        mol2, e = qc.get_quasiclassical_perturbation(mol)
+        mol2, e, _ = qc.get_quasiclassical_perturbation(mol)
         self.assertTrue(isinstance(mol2, cctk.Molecule))
 
-        mol3, e, v = qc.get_quasiclassical_perturbation(mol, return_velocities=True)
+        mol3, e, _, v = qc.get_quasiclassical_perturbation(mol, return_velocities=True)
         self.assertTrue(isinstance(mol3, cctk.Molecule))
 
     def test_final_structure(self):