numrs2 0.3.3

A Rust implementation inspired by NumPy for numerical computing (NumRS2)
Documentation
use crate::array::Array;
use crate::error::{NumRs2Error, Result};
use num_traits::Float;

/// BLAS Level 1: Vector-Vector operations
/// Compute the dot product of two vectors
pub fn dot<T>(x: &Array<T>, y: &Array<T>) -> Result<T>
where
    T: Float + Default,
{
    // For simplicity, we'll implement a generic version first
    // In a real implementation, we would dispatch to the appropriate BLAS routine
    // based on the type (e.g., sdot, ddot)

    let x_shape = x.shape();
    let y_shape = y.shape();

    if x_shape.len() != 1 || y_shape.len() != 1 {
        return Err(NumRs2Error::DimensionMismatch(
            "dot product requires 1D arrays".to_string(),
        ));
    }

    if x_shape[0] != y_shape[0] {
        return Err(NumRs2Error::ShapeMismatch {
            expected: x_shape,
            actual: y_shape,
        });
    }

    let x_data = x.to_vec();
    let y_data = y.to_vec();
    let mut result = T::zero();

    for i in 0..x_shape[0] {
        result = result + x_data[i] * y_data[i];
    }

    Ok(result)
}

/// BLAS Level 2: Matrix-Vector operations
/// Compute the matrix-vector product y = alpha*A*x + beta*y
pub fn gemv<T>(
    a: &Array<T>,
    x: &Array<T>,
    y: &mut Array<T>,
    alpha: T,
    beta: T,
    trans: bool,
) -> Result<()>
where
    T: Float + Default,
{
    // Generic implementation for now
    // A real implementation would call the appropriate BLAS routine

    let a_shape = a.shape();
    let x_shape = x.shape();
    let y_shape = y.shape();

    if a_shape.len() != 2 {
        return Err(NumRs2Error::DimensionMismatch(
            "First argument must be a 2D matrix".to_string(),
        ));
    }

    if x_shape.len() != 1 {
        return Err(NumRs2Error::DimensionMismatch(
            "Second argument must be a 1D vector".to_string(),
        ));
    }

    if y_shape.len() != 1 {
        return Err(NumRs2Error::DimensionMismatch(
            "Third argument must be a 1D vector".to_string(),
        ));
    }

    let (m, n) = (a_shape[0], a_shape[1]);

    if trans {
        // y = alpha * A^T * x + beta * y
        if n != y_shape[0] || m != x_shape[0] {
            return Err(NumRs2Error::ShapeMismatch {
                expected: vec![n],
                actual: y_shape,
            });
        }
    } else {
        // y = alpha * A * x + beta * y
        if m != y_shape[0] || n != x_shape[0] {
            return Err(NumRs2Error::ShapeMismatch {
                expected: vec![m],
                actual: y_shape,
            });
        }
    }

    // Simple implementation for demonstration
    // A full implementation would call BLAS
    let a_data = a.to_vec();
    let x_data = x.to_vec();
    let mut y_data = y.to_vec();

    // Scale y by beta
    for value in &mut y_data {
        *value = beta * *value;
    }

    if trans {
        // y = alpha * A^T * x + beta * y
        for i in 0..n {
            for j in 0..m {
                y_data[i] = y_data[i] + alpha * a_data[j * n + i] * x_data[j];
            }
        }
    } else {
        // y = alpha * A * x + beta * y
        for i in 0..m {
            for j in 0..n {
                y_data[i] = y_data[i] + alpha * a_data[i * n + j] * x_data[j];
            }
        }
    }

    // Update y with the computed result
    *y = Array::from_vec(y_data);

    Ok(())
}

/// BLAS Level 3: Matrix-Matrix operations
/// Compute the matrix-matrix product C = alpha*A*B + beta*C
pub fn gemm<T>(
    a: &Array<T>,
    b: &Array<T>,
    c: &mut Array<T>,
    alpha: T,
    beta: T,
    trans_a: bool,
    trans_b: bool,
) -> Result<()>
where
    T: Float + Default + std::ops::AddAssign,
{
    // This would be implemented with BLAS in a complete library
    // For now, we'll just validate the dimensions

    let a_shape = a.shape();
    let b_shape = b.shape();
    let c_shape = c.shape();

    if a_shape.len() != 2 || b_shape.len() != 2 || c_shape.len() != 2 {
        return Err(NumRs2Error::DimensionMismatch(
            "All arguments must be 2D matrices".to_string(),
        ));
    }

    let (m, k_a) = if trans_a {
        (a_shape[1], a_shape[0])
    } else {
        (a_shape[0], a_shape[1])
    };

    let (k_b, n) = if trans_b {
        (b_shape[1], b_shape[0])
    } else {
        (b_shape[0], b_shape[1])
    };

    if k_a != k_b {
        return Err(NumRs2Error::DimensionMismatch(format!(
            "Inner dimensions must match: {} vs {}",
            k_a, k_b
        )));
    }

    if c_shape[0] != m || c_shape[1] != n {
        return Err(NumRs2Error::ShapeMismatch {
            expected: vec![m, n],
            actual: c_shape,
        });
    }

    // A simple (inefficient) implementation for demonstration purposes
    // A real implementation would call the appropriate BLAS routine

    let a_data = a.to_vec();
    let b_data = b.to_vec();
    let mut c_data = c.to_vec();

    // Scale C by beta
    for value in &mut c_data {
        *value = beta * *value;
    }

    // These variable names align with typical matrix notation
    let a_cols = a_shape[1];
    let b_cols = b_shape[1];

    if !trans_a && !trans_b {
        // C = alpha * A * B + beta * C
        for i in 0..m {
            for j in 0..n {
                for k in 0..k_a {
                    c_data[i * n + j] += alpha * a_data[i * a_cols + k] * b_data[k * b_cols + j];
                }
            }
        }
    } else if trans_a && !trans_b {
        // C = alpha * A^T * B + beta * C
        for i in 0..m {
            for j in 0..n {
                for k in 0..k_a {
                    c_data[i * n + j] += alpha * a_data[k * a_cols + i] * b_data[k * b_cols + j];
                }
            }
        }
    } else if !trans_a && trans_b {
        // C = alpha * A * B^T + beta * C
        for i in 0..m {
            for j in 0..n {
                for k in 0..k_a {
                    c_data[i * n + j] += alpha * a_data[i * a_cols + k] * b_data[j * b_cols + k];
                }
            }
        }
    } else {
        // C = alpha * A^T * B^T + beta * C
        for i in 0..m {
            for j in 0..n {
                for k in 0..k_a {
                    c_data[i * n + j] += alpha * a_data[k * a_cols + i] * b_data[j * b_cols + k];
                }
            }
        }
    }

    // Update C with the computed result
    *c = Array::from_vec(c_data).reshape(&[m, n]);

    Ok(())
}

// Add more BLAS functions as needed for your library