continued bugfixes for qc init - added ts-specific init options

cctk/quasiclassical.py

@@ -12,6 +12,7 @@
 """
 
 AMU_A2_FS2_PER_KCAL_MOL = 0.0004184
+BOLTZMANN_CONSTANT = 0.001985875 # kcal/mol•Kn
 
 def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities=False, which="quasiclassical", mode_options=None):
     """
@@ -58,7 +59,7 @@ def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities
             PE, KE, TE, mode_velocity, text = apply_vibration(mol, mode, temperature=temperature, which=which)
         total_PE += PE
         total += TE
-        all_text += f"{text}\n"
+        all_text += f"Mode {idx+1}: {text}\n"
 
         for idx in range(1,molecule.num_atoms()+1):
             velocities[idx] += mode_velocity * mode.displacements[idx]
@@ -90,18 +91,36 @@ def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False,
         text
     """
 
-    level = mode.choose_level(temperature)
-    energy = mode.energy(level)
+    if mode.frequency < 0:
+        which = "ts"
+
+    if which == "quasiclassical":
+        level = mode.choose_level(temperature)
+        energy = mode.energy(level)
+        shift = mode.random_displacement(level=level, method=which)
+        method = f"quasiclassical(level={level})"
+    elif which == "classical":
+        energy = random_boltzmann_energy(temperature)
+        shift = mode.random_displacement(energy=energy, method=which)
+        method = "classical"
+    elif which == "ts":
+        energy = random_boltzmann_energy(temperature)
+        shift = 0
+        method = "ts"
+    else:
+        raise ValueError(f"``which`` must be ``classical``, ``quasiclassical``, or ``ts`` - {which} does not match!")
 
-    shift = mode.random_displacement(level, method=which)
-    max_shift = mode.classical_turning_point(level)
-    rel_shift = shift/max_shift
+    # the rest is common to all methods
 
+    # transition states and low-frequency modes do not get a starting displacement
     if not displacement or mode.frequency < min_freq:
-        rel_shift = 0
+        shift = 0
 
-    molecule.geometry += mode.displacements * rel_shift * max_shift
+    max_shift = mode.classical_turning_point(energy)
+    rel_shift = shift/max_shift
 
+    # apply displacements and compute energy breakdown
+    molecule.geometry += mode.displacements * rel_shift * max_shift
     potential_energy = 0.5 * mode.force_constant * shift ** 2
     kinetic_energy = energy - potential_energy
 
@@ -111,16 +130,19 @@ def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False,
 
     if velocity == "random":
         mode_velocity *= (1 if random.random() < 0.5 else -1)
-    elif velocity == "positive":
-        pass
     elif velocity == "negative":
         mode_velocity *= -1
     elif velocity == "zero":
         mode_velocity *= 0
-    else:
+    elif velocity != "positive":
         raise ValueError(f"unknown value {velocity} for keywork ``velocity`` - must be ``positive``, ``negative``, ``random``, or ``zero``")
 
-    text = 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"
+    text = f"{mode.frequency:.2f} cm-1 ({energy:.2f} kcal/mol)\t{method}\t Shift {rel_shift:.2%} of {max_shift:.3f} Å"
+    text += f"\tPE = {potential_energy:.2f} kcal/mol\tKE = {kinetic_energy:.2f} kcal/mol\tk = {mode.force_constant:.2f} kcal/mol Å^-2"
+    if not displacement:
+        text += "\n\t\tDisplacement manually set to zero"
+    if velocity == "zero":
+        text += "\n\t\tVelocity manually set to zero"
     if verbose:
         print(text)
 
@@ -145,4 +167,31 @@ def get_hermite_polynomial(n):
             Hr[v] = Hr[1]*Hr[v-1] - 2*(v-1)*Hr[v-2]
     return Hr[n]
 
+def random_boltzmann_energy(temperature, cutoff=10, step1=0.01, step2=0.0001):
+    kT = temperature * BOLTZMANN_CONSTANT
+
+    random = np.random.uniform()
+
+    # cumulative Boltzmann
+    cumulative_boltzmann = lambda e: math.erf(math.sqrt(e))
+
+    # now we need to numerically invert the cumulative Boltzmann
+    # kT = 1 for all this math, we'll fix it at the end
+    trial_energy = -1
+
+    # scan up to cutoff kT, which should be more than enough
+    for i in np.arange(0, cutoff, step1):
+        if cumulative_boltzmann(i) > random:
+            trial_energy = i - step1
+            break
+
+    if trial_energy == -1:
+        return cutoff * kT
+
+    # retry in smaller increments
+    for i in np.arange(trial_energy, trial_energy+step1, step2):
+        if cumulative_boltzmann(i) > random:
+            trial_energy = i - step2
+            break
 
+    return trial_energy * kT

cctk/vibrational_mode.py

@@ -95,9 +95,10 @@ def energy(self, level=0):
         zpe = 0.0014305 * freq
         return zpe * (2 * level + 1)
 
-    def random_displacement(self, level=0, method="quasiclassical", max_attempts=1e5):
+    def random_displacement(self, energy=None, level=0, method="quasiclassical", max_attempts=1e5):
         """
         Args:
+            energy (float): energy of mode (for classical case)
             method (str): "quasiclassical" or "classical"
             level (int): which vibrational level
             max_attempts (int): how many tries you get
@@ -123,8 +124,9 @@ def random_displacement(self, level=0, method="quasiclassical", max_attempts=1e5
 
             raise ValueError("max_attempts exceeded - can't get a proper initialization for this mode!")
         elif method == "classical":
+            assert energy is not None, "need energy for classical displacement"
             min_val = self.classical_distribution_value(0)
-            max_x = self.classical_turning_point()
+            max_x = self.classical_turning_point(energy)
             max_val = self.classical_distribution_value(max_x)
 
             attempts = 0
@@ -185,18 +187,26 @@ def quantum_distribution_max(self, level=0, num_pts=1e5):
 
         return max_p
 
-    def classical_distribution_value(self, x, num_pts=1e5):
+    def classical_distribution_value(self, x):
         """
         Returns the value of the classical distribution at the specified ``x`` value.
         """
         max_x = self.classical_turning_point()
+        assert (x <= max_x) and (x >= -1*max_x), "x must be in [-max_x, max_x]"
         return 1/(math.pi * math.sqrt(max_x**2 - x**2))
 
-    def classical_turning_point(self, level=0):
+    def classical_turning_point(self, energy=None, level=0):
         """
         Returns the maximum allowed shift based on modelling the mode as a classical harmonic oscillator (e.g. the point where potential energy is maximum).
+
+        Args:
+            energy (float): energy of mode
+            level (int): level to compute energy for quantum harmonic oscillator
         """
-        return math.sqrt(2 * self.energy(level) / self.force_constant)
+        if energy is None:
+            energy = self.energy(level)
+
+        return math.sqrt(2 * energy / self.force_constant)
 
     def to_string(self):
         ...

test/test_freqs.py

@@ -63,7 +63,7 @@ def test_perturb_water(self):
         file = cctk.GaussianFile.read_file(path)
 
         mol = file.get_molecule()
-        mol2, e, _, _ = qc.get_quasiclassical_perturbation(mol)
+        mol2, e, _, text = qc.get_quasiclassical_perturbation(mol)
         self.assertTrue(isinstance(mol2, cctk.Molecule))
 
         mo = {
@@ -81,10 +81,10 @@ def test_perturb_water(self):
             2: {"velocity": "positive", "displacement": False},
             3: {"velocity": "positive", "displacement": False},
         }
-        mol3, e, te, text, v = qc.get_quasiclassical_perturbation(mol, return_velocities=True, mode_options=mo)
-        self.assertTrue(isinstance(mol3, cctk.Molecule))
+        mol4, e, te, text, v = qc.get_quasiclassical_perturbation(mol, return_velocities=True, mode_options=mo)
         self.assertTrue(te - 13.28839457 < 0.00001)
 
+        mol5, e, te, text, v = qc.get_quasiclassical_perturbation(mol, return_velocities=True, which="classical")
 
     def test_final_structure(self):
         path1 = "test/static/methane_perturbed.gjf"
@@ -105,3 +105,13 @@ def test_imaginary(self):
         mol = file.get_molecule()
 
         self.assertListEqual(mol.atoms_moving_in_imaginary(), [48,2,51])
+
+    def test_boltzmann(self):
+        energies = np.zeros(shape=10000)
+        for i in range(10000):
+            energies[i] = cctk.quasiclassical.random_boltzmann_energy(298)
+        mean_energy = np.average(energies)
+
+        # should be 0.2965 = 0.5 kT
+        self.assertTrue(abs(0.2965 - mean_energy) < 0.2)
+