diff --git a/ThermoScreening/thermo/system.py b/ThermoScreening/thermo/system.py index 44be390..80196e0 100644 --- a/ThermoScreening/thermo/system.py +++ b/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: diff --git a/tests/thermo/test_system.py b/tests/thermo/test_system.py index d48a945..145bb3f 100644 --- a/tests/thermo/test_system.py +++ b/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]))]