From 6266a674ce359def3ed1f8d5f99fb760d324b9db Mon Sep 17 00:00:00 2001 From: Atsushi Togo Date: Sat, 27 Jun 2026 18:43:25 +0900 Subject: [PATCH 1/2] Add eigvalsh_values_batch kernel for eigenvalues-only diagonalization --- CHANGELOG.md | 3 ++ src/diagonalize.rs | 128 +++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 49 +++++++++++++++++ 3 files changed, 180 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7544e9b..b32eb02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 for dense meshes of small matrices. Numerically a drop-in for `numpy.linalg.eigh`: eigenvalues are returned in ascending order, eigenvectors as columns, and the input matrices are left unmodified. +- `eigvalsh_values_batch` kernel: same as `eigvalsh_batch` but computes only + the eigenvalues, skipping the eigenvectors and their output buffer. A + drop-in for `numpy.linalg.eigvalsh`. ## [0.2.1] - 2026-06-11 diff --git a/src/diagonalize.rs b/src/diagonalize.rs index abe8358..78e354f 100644 --- a/src/diagonalize.rs +++ b/src/diagonalize.rs @@ -121,6 +121,78 @@ pub fn eigsh_batch(dynmats: &[Cmplx], evals: &mut [f64], evecs: &mut [Cmplx], n: ); } +/// Per-worker reusable workspace for [`eigvals_one`], sized for an `n x n` +/// matrix. Like [`EvdScratch`] but without the eigenvector matrix: the +/// faer workspace is requested with `ComputeEigenvectors::No`, so neither +/// the `u` buffer nor the eigenvector-related scratch is allocated. +struct EvalsScratch { + mem: MemBuffer, + s: Diag, +} + +impl EvalsScratch { + fn new(n: usize) -> Self { + let req = self_adjoint_evd_scratch::( + n, + ComputeEigenvectors::No, + Par::Seq, + Default::default(), + ); + Self { + mem: MemBuffer::new(req), + s: Diag::zeros(n), + } + } +} + +/// Compute only the eigenvalues of one Hermitian matrix `d` (row-major +/// `[n, n]` `Cmplx`), writing them into `evals[0..n]` (nondecreasing). +/// +/// Same input contract as [`eigsh_one`], but passes `None` for the +/// eigenvectors so faer skips the back-transformation entirely. `d` is +/// only read (and only its lower triangle, as the matrix is Hermitian). +fn eigvals_one(d: &[Cmplx], n: usize, evals: &mut [f64], sc: &mut EvalsScratch) { + debug_assert_eq!(d.len(), n * n); + debug_assert_eq!(evals.len(), n); + + let a = MatRef::::from_row_major_slice(cmplx_as_c64(d), n, n); + let stack = MemStack::new(&mut sc.mem); + self_adjoint_evd( + a, + sc.s.as_mut(), + None, + Par::Seq, + stack, + Default::default(), + ) + .expect("faer self_adjoint_evd failed"); + + let s = sc.s.column_vector(); + for i in 0..n { + evals[i] = s[i].re; + } +} + +/// Batched eigenvalues-only Hermitian eigendecomposition, parallel over +/// the `nq` matrices. +/// +/// `dynmats` is `[nq, n, n]` row-major Hermitian, `evals` is `[nq, n]`. +/// Each matrix is solved single-threaded; the caller should run this +/// under `py.detach` to release the GIL. +pub fn eigvals_batch(dynmats: &[Cmplx], evals: &mut [f64], n: usize) { + if n == 0 { + return; + } + let m2 = n * n; + evals + .par_chunks_mut(n) + .zip(dynmats.par_chunks(m2)) + .for_each_init( + || EvalsScratch::new(n), + |sc, (ev, d)| eigvals_one(d, n, ev, sc), + ); +} + #[cfg(test)] mod tests { use super::*; @@ -145,6 +217,62 @@ mod tests { (evals, evecs) } + /// Eigenvalues of a single `n x n` row-major Hermitian matrix. + fn eigvals(d: &[Cmplx], n: usize) -> Vec { + let mut evals = vec![0.0f64; n]; + let mut sc = EvalsScratch::new(n); + eigvals_one(d, n, &mut evals, &mut sc); + evals + } + + #[test] + fn values_known_2x2() { + // A = [[2, 1 - i], [1 + i, 3]]; eigenvalues 1 and 4. + let d: [Cmplx; 4] = [[2.0, 0.0], [1.0, -1.0], [1.0, 1.0], [3.0, 0.0]]; + let evals = eigvals(&d, 2); + assert!((evals[0] - 1.0).abs() < 1e-12, "{:?}", evals); + assert!((evals[1] - 4.0).abs() < 1e-12, "{:?}", evals); + } + + #[test] + fn values_match_full_and_ascending() { + for &n in &[1usize, 2, 6, 13] { + let d = hermitian_fixture(n); + let (full, _) = eigh(&d, n); + let only = eigvals(&d, n); + for i in 0..n { + assert!( + (full[i] - only[i]).abs() < 1e-12, + "n={n} band={i}: full={} only={}", + full[i], + only[i] + ); + } + for w in only.windows(2) { + assert!(w[0] <= w[1] + 1e-12, "n={n} not ascending: {only:?}"); + } + } + } + + #[test] + fn values_batch_matches_single() { + let n = 5; + let a = hermitian_fixture(n); + let b = hermitian_fixture_seed(n, 9); + let mut dynmats = Vec::new(); + dynmats.extend_from_slice(&a); + dynmats.extend_from_slice(&b); + let mut evals = vec![0.0f64; 2 * n]; + eigvals_batch(&dynmats, &mut evals, n); + + let ea = eigvals(&a, n); + let eb = eigvals(&b, n); + for i in 0..n { + assert!((evals[i] - ea[i]).abs() < 1e-12); + assert!((evals[n + i] - eb[i]).abs() < 1e-12); + } + } + #[test] fn known_2x2() { // A = [[2, 1 - i], [1 + i, 3]]; eigenvalues 1 and 4. diff --git a/src/lib.rs b/src/lib.rs index 9ffd1d2..ac80781 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -6388,6 +6388,54 @@ fn py_eigvalsh_batch<'py>( Ok(()) } +/// Compute only the eigenvalues of a batch of Hermitian dynamical +/// matrices. +/// +/// Same contract as ``eigvalsh_batch`` but without eigenvectors: +/// ``dynmats`` is complex128 ``(n_q, num_band, num_band)`` and is left +/// unmodified; ``eigenvalues`` (float64 ``(n_q, num_band)``) receives the +/// eigenvalues in nondecreasing order. Skipping the eigenvectors avoids +/// their back-transformation and output buffer. Numerically a drop-in +/// for ``numpy.linalg.eigvalsh`` applied over the leading axis. +#[pyfunction] +#[pyo3(name = "eigvalsh_values_batch")] +#[pyo3(signature = (dynmats, eigenvalues))] +fn py_eigvalsh_values_batch<'py>( + py: Python<'py>, + dynmats: PyReadonlyArray3<'py, Complex64>, + mut eigenvalues: PyReadwriteArray2<'py, f64>, +) -> PyResult<()> { + let shape = dynmats.shape(); + let n_q = shape[0]; + let num_band = shape[1]; + + if shape[2] != num_band { + return Err(PyValueError::new_err( + "dynmats must have shape (n_q, num_band, num_band)", + )); + } + if eigenvalues.shape() != [n_q, num_band] { + return Err(PyValueError::new_err( + "eigenvalues must have shape (n_q, num_band)", + )); + } + + let dynmats_view = dynmats.as_array(); + let dynmats_flat = dynmats_view + .as_slice() + .ok_or_else(|| PyValueError::new_err("dynmats must be C-contiguous"))?; + let dynmats_cmplx = complex_as_cmplx(dynmats_flat); + + let evals_slice = eigenvalues + .as_slice_mut() + .map_err(|_| PyValueError::new_err("eigenvalues must be C-contiguous"))?; + + py.detach(|| { + diagonalize::eigvals_batch(dynmats_cmplx, evals_slice, num_band); + }); + Ok(()) +} + #[pymodule(gil_used = false)] fn phonors(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_snf3x3, m)?)?; @@ -6418,6 +6466,7 @@ fn phonors(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_dynamical_matrices_at_qpoints, m)?)?; m.add_function(wrap_pyfunction!(py_dynamical_matrices_at_qpoints_gonze, m)?)?; m.add_function(wrap_pyfunction!(py_eigvalsh_batch, m)?)?; + m.add_function(wrap_pyfunction!(py_eigvalsh_values_batch, m)?)?; m.add_function(wrap_pyfunction!(py_transform_dynmat_to_fc, m)?)?; m.add_function(wrap_pyfunction!(py_derivative_dynmat_at_q, m)?)?; m.add_function(wrap_pyfunction!(py_real_to_reciprocal, m)?)?; From 379724810ba9b2c2c3082119dcfe8eca4187a759 Mon Sep 17 00:00:00 2001 From: Atsushi Togo Date: Sat, 27 Jun 2026 18:45:00 +0900 Subject: [PATCH 2/2] Fix format --- src/diagonalize.rs | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/src/diagonalize.rs b/src/diagonalize.rs index 78e354f..dfd0701 100644 --- a/src/diagonalize.rs +++ b/src/diagonalize.rs @@ -157,15 +157,8 @@ fn eigvals_one(d: &[Cmplx], n: usize, evals: &mut [f64], sc: &mut EvalsScratch) let a = MatRef::::from_row_major_slice(cmplx_as_c64(d), n, n); let stack = MemStack::new(&mut sc.mem); - self_adjoint_evd( - a, - sc.s.as_mut(), - None, - Par::Seq, - stack, - Default::default(), - ) - .expect("faer self_adjoint_evd failed"); + self_adjoint_evd(a, sc.s.as_mut(), None, Par::Seq, stack, Default::default()) + .expect("faer self_adjoint_evd failed"); let s = sc.s.column_vector(); for i in 0..n {