solve_matrix/

solve_matrix.rs

extern crate lapack_src;
use solver::linalg::{csr::CsrMatrix, eigen::eigen_symmetric_tridiagonal};
use std::env;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;
use std::time::Instant;

#[allow(clippy::needless_range_loop)]
fn read_matrix(filename: &str) -> (CsrMatrix, Vec<f64>) {
    let path = Path::new("benches").join("data").join(filename);
    let file = File::open(&path).unwrap_or_else(|_| panic!("Failed to open file {:?}", path));
    let mut reader = BufReader::new(file);
    let mut line = String::new();

    // Size
    reader.read_line(&mut line).unwrap();
    let size: usize = line.trim().parse().unwrap();
    line.clear();

    // row_ptr
    reader.read_line(&mut line).unwrap();
    let row_ptr: Vec<usize> = line
        .split_whitespace()
        .map(|s| s.parse().unwrap())
        .collect();
    line.clear();

    // col_indices
    reader.read_line(&mut line).unwrap();
    let col_indices: Vec<usize> = line
        .split_whitespace()
        .map(|s| s.parse().unwrap())
        .collect();
    line.clear();

    // values
    reader.read_line(&mut line).unwrap();
    let values: Vec<f64> = line
        .split_whitespace()
        .map(|s| s.parse().unwrap())
        .collect();
    line.clear();

    // b
    reader.read_line(&mut line).unwrap();
    let b: Vec<f64> = line
        .split_whitespace()
        .map(|s| s.parse().unwrap())
        .collect();

    // Compute diag_indices
    let mut diag_indices = Vec::with_capacity(size);
    for i in 0..size {
        let start = row_ptr[i];
        let end = row_ptr[i + 1];
        let mut found = false;
        for idx in start..end {
            if col_indices[idx] == i {
                diag_indices.push(idx);
                found = true;
                break;
            }
        }
        if !found {
            diag_indices.push(usize::MAX);
        }
    }

    (
        CsrMatrix {
            values,
            col_indices,
            row_ptr,
            size,
            diag_indices,
        },
        b,
    )
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 3 {
        eprintln!("Usage: solve_matrix <filename> <method>");
        return;
    }
    let filename = &args[1];
    let method = &args[2];

    let (matrix, b) = read_matrix(filename);

    if method == "sor" {
        let mut x = vec![0.0; matrix.size];
        let start = Instant::now();
        // Use same params as bench: tol=1e-6, max_iter=10000, omega=1.5
        let iters = matrix.solve_sor(&b, &mut x, 1e-6, 10000, 1.5);
        let duration = start.elapsed();

        eprintln!(
            "Rust SOR time: {:.4} ms (iters: {})",
            duration.as_secs_f64() * 1000.0,
            iters
        );

        // Print result to stdout for Python to read
        for val in x {
            print!("{} ", val);
        }
        println!();
    } else if method == "eigen" {
        // Extract diagonal and off-diagonal
        let mut d = vec![0.0; matrix.size];
        let mut e = vec![0.0; matrix.size];

        for i in 0..matrix.size {
            let row_start = matrix.row_ptr[i];
            let row_end = matrix.row_ptr[i + 1];

            for idx in row_start..row_end {
                let col = matrix.col_indices[idx];
                if col == i {
                    d[i] = matrix.values[idx];
                } else if col == i + 1 {
                    e[i] = matrix.values[idx];
                }
            }
        }

        let start = Instant::now();
        let (eigenvalues, _) = eigen_symmetric_tridiagonal(d, e);
        let duration = start.elapsed();

        eprintln!("Rust Eigen time: {:.4} ms", duration.as_secs_f64() * 1000.0);

        // Print eigenvalues to stdout
        for val in eigenvalues {
            print!("{} ", val);
        }
        println!();
    } else if method == "faer" {
        use faer::{Mat, Side};

        // Convert CSR to Dense Mat
        let mut dense_mat = Mat::<f64>::zeros(matrix.size, matrix.size);
        for i in 0..matrix.size {
            let row_start = matrix.row_ptr[i];
            let row_end = matrix.row_ptr[i + 1];
            for idx in row_start..row_end {
                let col = matrix.col_indices[idx];
                let val = matrix.values[idx];
                dense_mat[(i, col)] = val;
            }
        }

        let start = Instant::now();
        // High-level API for eigenvalues only
        // Assuming signature: selfadjoint_eigenvalues(Side)
        let s = dense_mat.selfadjoint_eigenvalues(Side::Lower);
        let duration = start.elapsed();

        eprintln!("Rust Faer time: {:.4} ms", duration.as_secs_f64() * 1000.0);

        for val in s.iter() {
            print!("{} ", val);
        }
        println!();
    } else if method == "lapack" {
        use lapack::dstev;

        // Extract diagonal and off-diagonal
        let mut d = vec![0.0; matrix.size];
        let mut e = vec![0.0; matrix.size - 1]; // dstev expects e to be len n-1

        for i in 0..matrix.size {
            let row_start = matrix.row_ptr[i];
            let row_end = matrix.row_ptr[i + 1];

            for idx in row_start..row_end {
                let col = matrix.col_indices[idx];
                if col == i {
                    d[i] = matrix.values[idx];
                } else if col == i + 1 {
                    e[i] = matrix.values[idx];
                }
            }
        }

        let n = matrix.size as i32;
        // 'V' means compute eigenvectors, 'N' means eigenvalues only
        // We want eigenvalues only for fair comparison with faer/scipy eigh_tridiagonal default?
        // Scipy eigh_tridiagonal returns (w, v) by default.
        // My custom eigen returns (w, v).
        // So let's compute eigenvectors too ('V').
        // But wait, bench_solvers.py only compares eigenvalues (w).
        // And faer benchmark above used selfadjoint_eigenvalues (only w).
        // Let's stick to eigenvalues only ('N') for speed if we only compare w,
        // BUT bench_solvers.py compares w.
        // However, if I want to beat Scipy which computes both (usually), I should check what Scipy does.
        // Scipy eigh_tridiagonal(d, e) returns w, v.
        // So I should use 'V'.
        // But wait, if I use 'N', I don't need Z.

        // Let's use 'N' (eigenvalues only) first to see raw speed of the solver,
        // as we are comparing against faer which was only eigenvalues.
        // And bench_solvers.py only checks w.

        let jobz = b'N';
        let mut z = vec![0.0; 1]; // Not referenced if jobz == 'N'
        let ldz = 1;
        let mut work = vec![0.0; std::cmp::max(1, 2 * matrix.size - 2)];
        let mut info = 0;

        let start = Instant::now();
        unsafe {
            dstev(jobz, n, &mut d, &mut e, &mut z, ldz, &mut work, &mut info);
        }
        let duration = start.elapsed();

        if info != 0 {
            eprintln!("LAPACK dstev failed with info={}", info);
        } else {
            eprintln!(
                "Rust LAPACK time: {:.4} ms",
                duration.as_secs_f64() * 1000.0
            );
            for val in d.iter() {
                print!("{} ", val);
            }
            println!();
        }
    } else if method == "lapack_solve" {
        use lapack::dgtsv;

        // Extract diagonal and off-diagonal
        let mut d = vec![0.0; matrix.size];
        let mut du = vec![0.0; matrix.size - 1];
        let mut dl = vec![0.0; matrix.size - 1];

        for i in 0..matrix.size {
            let row_start = matrix.row_ptr[i];
            let row_end = matrix.row_ptr[i + 1];

            for idx in row_start..row_end {
                let col = matrix.col_indices[idx];
                if col == i {
                    d[i] = matrix.values[idx];
                } else if col == i + 1 {
                    du[i] = matrix.values[idx];
                    dl[i] = matrix.values[idx]; // Symmetric
                }
            }
        }

        let n = matrix.size as i32;
        let nrhs = 1;
        let ldb = n;
        let mut info = 0;

        // b is consumed and replaced by solution
        let mut x = b.clone();

        let start = Instant::now();
        unsafe {
            dgtsv(n, nrhs, &mut dl, &mut d, &mut du, &mut x, ldb, &mut info);
        }
        let duration = start.elapsed();

        if info != 0 {
            eprintln!("LAPACK dgtsv failed with info={}", info);
        } else {
            eprintln!(
                "Rust LAPACK Solve time: {:.4} ms",
                duration.as_secs_f64() * 1000.0
            );
            for val in x {
                print!("{} ", val);
            }
            println!();
        }
    } else {
        eprintln!("Unknown method: {}", method);
    }
}