add mode specific init options

cctk/quasiclassical.py

@@ -13,7 +13,7 @@
 
 AMU_A2_FS2_PER_KCAL_MOL = 0.0004184
 
-def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities=False, which="quasiclassical"):
+def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities=False, which="quasiclassical", mode_options=None):
     """
     Perturbs molecule by treating each mode as a quantum harmonic oscillator and sampling from the distribution appropriate to the temperature.
 
@@ -24,6 +24,12 @@ def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities
         temperature (float): temperature
         return velocities (bool): whether or not to return velocities
         which (str): ``classical`` or ``quasiclassical``
+        mode_options (dict):
+            Options for how to initialize specific modes.
+                key (int): 1-indexed number of vibrational mode (from smallest frequency to largest)
+                val (dict):
+                    velocity (str): one of "positive", "negative", "random", "zero"
+                    displacement (bool): whether or not to displace
 
     Returns:
         new ``cctk.Molecule`` object
@@ -32,7 +38,6 @@ def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities
     """
     assert isinstance(molecule, cctk.Molecule), "need a valid molecule"
     assert len(molecule.vibrational_modes) > 0, "molecule needs to have vibrational modes (try running a ``freq`` job)"
-
     assert isinstance(temperature, (int, float)), "temperature must be numeric"
 
     mol = copy.deepcopy(molecule)
@@ -41,29 +46,29 @@ def get_quasiclassical_perturbation(molecule, temperature=298, return_velocities
 
     velocities = np.zeros_like(molecule.geometry.view(np.ndarray)).view(cctk.OneIndexedArray)
 
+    if mode_options is None:
+        mode_options = dict()
+
     all_text = ""
-    for mode in mol.vibrational_modes:
-        PE, KE, TE, text = apply_vibration(mol, mode, temperature=temperature, which=which)
+    for idx, mode in enumerate(mol.vibrational_modes):
+        # enumerate is 0-indexed but GaussView, etc 1-index the modes. so we're 1-indexing here too.
+        if idx+1 in mode_options:
+            PE, KE, TE, mode_velocity, text = apply_vibration(mol, mode, temperature=temperature, which=which, **mode_options[idx+1])
+        else:
+            PE, KE, TE, mode_velocity, text = apply_vibration(mol, mode, temperature=temperature, which=which)
         total_PE += PE
         total += TE
         all_text += f"{text}\n"
-        #### below code initializes velocities. pulling from jprogdyn Initializer lines 440 & onwards.
 
-        if return_velocities:
-            # mode velocity = sqrt(2 * KE / reduced mass) - want value in Å/fs. we randomize the sign
-            # https://stackoverflow.com/questions/46820182/randomly-generate-1-or-1-positive-or-negative-integer
-            mode_velocity = math.sqrt(2*KE*AMU_A2_FS2_PER_KCAL_MOL/mode.reduced_mass)
-            mode_velocity *= (1 if random.random() < 0.5 else -1)
-
-            for idx in range(1,molecule.num_atoms()+1):
-                velocities[idx] += mode_velocity * mode.displacements[idx]
+        for idx in range(1,molecule.num_atoms()+1):
+            velocities[idx] += mode_velocity * mode.displacements[idx]
 
     if return_velocities:
-        return mol, total_PE, total, text, velocities
-    else:
-        return mol, total_PE, total, text
+        return mol, total_PE, total, all_text, velocities
+    else: # backwards compatibility
+        return mol, total_PE, total, all_text
 
-def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False, which="quasiclassical"):
+def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False, which="quasiclassical", displacement=True, velocity="random", **kwargs):
     """
     Apply a vibration to molecule ``molecule`` (modified in-place).
 
@@ -74,11 +79,14 @@ def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False,
         temperature (float)
         verbose (bool)
         which (str): ``quasiclassical`` or ``classical``
+        displacement (bool): whether or not to displace the mode
+        velocity (str): ``positive``, ``negative``, ``random``, or ``zero``
 
     Returns:
         potential energy
         kinetic energy
         energy
+        velocities
         text
     """
 
@@ -89,20 +97,34 @@ def apply_vibration(molecule, mode, min_freq=50, temperature=298, verbose=False,
     max_shift = mode.classical_turning_point(level)
     rel_shift = shift/max_shift
 
-    if mode.frequency < min_freq:
+    if not displacement or mode.frequency < min_freq:
         rel_shift = 0
 
     molecule.geometry += mode.displacements * rel_shift * max_shift
 
     potential_energy = 0.5 * mode.force_constant * shift ** 2
     kinetic_energy = energy - potential_energy
 
-    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"
+    # mode velocity = sqrt(2 * KE / reduced mass) - want value in Å/fs. we randomize the sign
+    # https://stackoverflow.com/questions/46820182/randomly-generate-1-or-1-positive-or-negative-integer
+    mode_velocity = math.sqrt(2*kinetic_energy*AMU_A2_FS2_PER_KCAL_MOL/mode.reduced_mass)
+
+    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:
+        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"
     if verbose:
         print(text)
 
-    return potential_energy, kinetic_energy, energy, text
+    return potential_energy, kinetic_energy, energy, mode_velocity, text
 
 def get_hermite_polynomial(n):
     """

test/test_freqs.py

@@ -66,8 +66,25 @@ def test_perturb_water(self):
         mol2, e, _, _ = qc.get_quasiclassical_perturbation(mol)
         self.assertTrue(isinstance(mol2, cctk.Molecule))
 
-        mol3, e, _, _, v = qc.get_quasiclassical_perturbation(mol, return_velocities=True)
+        mo = {
+            1: {"velocity": "zero"},
+            2: {"velocity": "zero"},
+            3: {"velocity": "zero"},
+        }
+
+        mol3, e, _, _, v = qc.get_quasiclassical_perturbation(mol, return_velocities=True, mode_options=mo)
+        self.assertTrue(isinstance(mol3, cctk.Molecule))
+        self.assertFalse(np.any(v)) # all should be zero, AKA False
+
+        mo = {
+            1: {"velocity": "positive", "displacement": False},
+            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))
+        self.assertTrue(te - 13.28839457 < 0.00001)
+
 
     def test_final_structure(self):
         path1 = "test/static/methane_perturbed.gjf"