Raise on too few frequencies in frequency_dof

ThermoScreening/thermo/system.py

@@ -434,49 +434,55 @@ def cleaned_frequency(frequency: np.ndarray):
 
 def frequency_dof(frequency: np.ndarray, dof: int) -> np.ndarray:
     """
-    Delete all frequencies that are more than the degree of freedom.
+    Keep the ``dof`` highest vibrational frequencies (dropping the near-zero
+    translational/rotational modes).
 
     Parameters
     ----------
     frequency : np.ndarray
-        The vibrational frequencies of the system.
+        The vibrational frequencies of the system, sorted in ascending order.
     dof : int
-        The degree of freedom of the system.
+        The number of vibrational degrees of freedom to keep.
 
     Returns
     -------
     frequency : np.ndarray
-        The vibrational frequencies of the system.
+        The ``dof`` highest frequencies, in ascending order.
+
+    Raises
+    ------
+    ValueError
+        If fewer than ``dof`` frequencies are supplied (which would otherwise
+        silently duplicate modes via negative indexing).
     """
     N = len(frequency)
-    cleaned_freq = np.empty(dof)
-    i = 0
-
-    while (i + 1) <= dof:
-        cleaned_freq[i] = frequency[N - 1 - i]
-        i += 1
+    if N < dof:
+        raise TSValueError(
+            f"expected at least {dof} frequencies for {dof} vibrational degrees "
+            f"of freedom, got {N}"
+        )
 
-    cleaned_freq = cleaned_freq[::-1]
-    return cleaned_freq
+    # copy so the result does not alias the caller's raw frequency array
+    return np.asarray(frequency)[N - dof:].copy()
 
 
 def check_frequency_length(frequency: np.ndarray, dof: int) -> bool:
     """
-    Checks if the system has the correct number of vibrational frequencies.
+    Check that enough vibrational frequencies were supplied for ``dof``.
 
     Parameters
     ----------
     frequency : np.ndarray
-        The vibrational frequencies of the system.
+        The (raw) vibrational frequencies of the system.
     dof : int
-        The degree of freedom of the system.
+        The number of vibrational degrees of freedom.
 
     Returns
     -------
     bool
-        True if the system has the correct number of vibrational frequencies, False otherwise.
+        True if at least ``dof`` frequencies are present, False otherwise.
     """
-    return len(frequency) == dof
+    return len(frequency) >= dof
 
 
 class System:
@@ -668,11 +674,11 @@ def __init__(
         self._has_imaginary_frequencies = check_imaginary_frequencies(
             self._imaginary_frequencies
         )
-        self._real_vibrational_frequencies = frequency_dof(
+        self._check_frequency_length = check_frequency_length(
             self._vibrational_frequencies, self._dof
         )
-        self._check_frequency_length = check_frequency_length(
-            self._real_vibrational_frequencies, self._dof
+        self._real_vibrational_frequencies = frequency_dof(
+            self._vibrational_frequencies, self._dof
         )
 
         if charge is None:

tests/thermo/test_system.py

@@ -4,7 +4,7 @@
 import numpy as np
 from ase.atoms import Atoms
 import ThermoScreening.thermo.system as system_module
-from ThermoScreening.thermo.system import System, dim, dof, linearity, rotational_symmetry_number, default_spin
+from ThermoScreening.thermo.system import System, dim, dof, linearity, rotational_symmetry_number, default_spin, frequency_dof, check_frequency_length
 from ThermoScreening.thermo.atoms import Atom
 from ThermoScreening.thermo.cell import Cell
 from ThermoScreening.exceptions import TSValueError
@@ -245,7 +245,31 @@ def test_dof():
     with pytest.raises(TSValueError) as e:
         dof(atoms)
     assert str(e.value) == "The number of atoms must be greater than 0."
-
+
+
+def test_frequency_dof_keeps_highest_modes():
+    # nine raw modes (ascending), keep the top dof=3 vibrations
+    freqs = np.arange(1.0, 10.0)
+    kept = frequency_dof(freqs, 3)
+    assert list(kept) == [7.0, 8.0, 9.0]
+
+
+def test_frequency_dof_identity_when_length_matches():
+    freqs = np.array([100.0, 200.0, 300.0])
+    assert list(frequency_dof(freqs, 3)) == [100.0, 200.0, 300.0]
+
+
+def test_frequency_dof_raises_on_underflow():
+    # too few frequencies must raise instead of wrapping/duplicating modes
+    with pytest.raises(TSValueError, match="expected at least 4 frequencies"):
+        frequency_dof(np.array([1.0, 2.0, 3.0]), 4)
+
+
+def test_check_frequency_length_validates_input_count():
+    assert check_frequency_length(np.arange(9.0), 3) is True
+    assert check_frequency_length(np.array([1.0, 2.0, 3.0]), 3) is True
+    assert check_frequency_length(np.array([1.0, 2.0]), 3) is False
+
 def test_linearity():
     atoms = [Atom(symbol='H', position=np.array([0, 0, 0]))]