chemx-props 0.1.0

//! Harmonic vibrational frequencies from a Cartesian Hessian.

use chemx_core::{Molecule, units::FREQ_CONV_CM1};
use chemx_linalg::{mat_from_row_major, mat_to_row_major, matmul, symmetric_eigh};

/// Result of the harmonic frequency analysis.
#[derive(Debug, Clone)]
pub struct FrequencyResult {
    /// Number of atoms.
    pub n_atoms: usize,
    /// Symmetrized Cartesian Hessian supplied to this routine (3N×3N, row-major).
    pub hessian: Vec<f64>,
    /// Harmonic frequencies in cm⁻¹, ascending. Imaginary modes are reported as
    /// **negative** values. Length = 3N; the 6 (5 if linear) lowest should be ~0
    /// after projection.
    pub frequencies_cm1: Vec<f64>,
    /// Mass-weighted normal mode eigenvectors, row-major 3N×3N. Column j is
    /// eigenvector j (the mode associated with `frequencies_cm1[j]`).
    pub normal_modes: Vec<f64>,
    /// Number of imaginary (< −1 cm⁻¹) modes — 0 at a true minimum.
    pub n_imaginary: usize,
}

/// Compute harmonic vibrational frequencies from the Cartesian Hessian.
///
/// Pipeline:
/// 1. Mass-weight: `F_ij = H_ij / √(mᵢ mⱼ)` (masses in amu).
/// 2. Build the 6 (or 5 for linear) translation/rotation vectors in
///    mass-weighted coordinates (Eckart conditions).
/// 3. Gram–Schmidt-orthonormalize them; build the projector
///    `P = I − Σ vᵢvᵢᵀ`.
/// 4. Project: `F_proj = P F P`.
/// 5. Diagonalize `F_proj`; convert eigenvalues to cm⁻¹.
///
/// The 6 trans/rot eigenvalues are not set to exactly zero — they will be
/// small (|λ| ≲ 10⁻⁸) but non-zero due to floating-point arithmetic.
/// A proper minimum has zero imaginary vibrational modes.
pub fn harmonic_frequencies(molecule: &Molecule, hessian: &[f64]) -> FrequencyResult {
    let natom = molecule.len();
    let ndof = 3 * natom;

    let masses: Vec<f64> = molecule.atoms.iter().map(|a| a.element.mass()).collect();

    // 1. Mass-weight
    let mut mw = hessian.to_vec();
    for i in 0..natom {
        for ki in 0..3 {
            let row = i * 3 + ki;
            for j in 0..natom {
                for kj in 0..3 {
                    let col = j * 3 + kj;
                    mw[row * ndof + col] /= (masses[i] * masses[j]).sqrt();
                }
            }
        }
    }

    // 2. Translation + rotation vectors in mass-weighted coordinates
    let total_mass: f64 = masses.iter().sum();
    let mut com = [0.0f64; 3];
    for (i, atom) in molecule.atoms.iter().enumerate() {
        for (k, c) in com.iter_mut().enumerate() {
            *c += masses[i] * atom.position[k];
        }
    }
    for c in &mut com {
        *c /= total_mass;
    }

    let mut raw_vecs: Vec<Vec<f64>> = Vec::with_capacity(6);

    // Translations
    for k in 0..3 {
        let mut v = vec![0.0f64; ndof];
        for i in 0..natom {
            v[3 * i + k] = masses[i].sqrt();
        }
        raw_vecs.push(v);
    }

    // COM-frame positions
    let r: Vec<[f64; 3]> = molecule
        .atoms
        .iter()
        .map(|a| {
            [
                a.position[0] - com[0],
                a.position[1] - com[1],
                a.position[2] - com[2],
            ]
        })
        .collect();

    // Rotations: e_k × r_A (cross product), mass-weighted
    // e_x × r = [0, −z, y];  e_y × r = [z, 0, −x];  e_z × r = [−y, x, 0]
    type RotDisp = fn(&[f64; 3]) -> [f64; 3];
    let rot_disps: [RotDisp; 3] = [
        |r| [0.0, -r[2], r[1]],
        |r| [r[2], 0.0, -r[0]],
        |r| [-r[1], r[0], 0.0],
    ];
    for disp_fn in &rot_disps {
        let mut v = vec![0.0f64; ndof];
        for i in 0..natom {
            let d = disp_fn(&r[i]);
            for k in 0..3 {
                v[3 * i + k] = masses[i].sqrt() * d[k];
            }
        }
        raw_vecs.push(v);
    }

    // 3. Gram–Schmidt
    let orth = gram_schmidt(&raw_vecs);

    // 4. Build projector P = I − Σ vᵢvᵢᵀ
    let mut proj = vec![0.0f64; ndof * ndof];
    for i in 0..ndof {
        proj[i * ndof + i] = 1.0;
    }
    for v in &orth {
        for i in 0..ndof {
            for j in 0..ndof {
                proj[i * ndof + j] -= v[i] * v[j];
            }
        }
    }

    // 5. F_proj = P · F_mw · P
    let p = mat_from_row_major(ndof, &proj);
    let f = mat_from_row_major(ndof, &mw);
    let fp_mat = matmul(&matmul(&p, &f), &p);

    // 6. Diagonalize
    let eigh = symmetric_eigh(&fp_mat);

    let frequencies_cm1: Vec<f64> = eigh
        .values
        .iter()
        .map(|&lam| {
            if lam >= 0.0 {
                lam.sqrt() * FREQ_CONV_CM1
            } else {
                -(-lam).sqrt() * FREQ_CONV_CM1
            }
        })
        .collect();

    let n_imaginary = frequencies_cm1.iter().filter(|&&f| f < -1.0).count();
    let normal_modes = mat_to_row_major(&eigh.vectors);

    FrequencyResult {
        n_atoms: natom,
        hessian: hessian.to_vec(),
        frequencies_cm1,
        normal_modes,
        n_imaginary,
    }
}

fn gram_schmidt(vecs: &[Vec<f64>]) -> Vec<Vec<f64>> {
    let mut orth: Vec<Vec<f64>> = Vec::new();
    for v in vecs {
        let mut u = v.clone();
        for prev in &orth {
            let proj: f64 = u.iter().zip(prev.iter()).map(|(a, b)| a * b).sum();
            for (ui, &pi) in u.iter_mut().zip(prev.iter()) {
                *ui -= proj * pi;
            }
        }
        let norm: f64 = u.iter().map(|x| x * x).sum::<f64>().sqrt();
        if norm > 1e-10 {
            for x in &mut u {
                *x /= norm;
            }
            orth.push(u);
        }
    }
    orth
}