From 1a232b1645be30571b2c5d6a4d64c0729c431be9 Mon Sep 17 00:00:00 2001 From: rajesh-puripanda Date: Sun, 31 May 2026 23:13:49 +0530 Subject: [PATCH 1/2] feat(index): implement index_bool and index_take Add boolean mask indexing (index_bool) and integer array indexing along an axis (index_take) to the mohu-index crate. index_bool filters a buffer by a boolean mask, returning a 1D buffer of selected elements. Supports all dtypes via dispatch_dtype. index_take selects elements along a given axis using an array of integer indices, with support for negative indices and bounds checking. Uses dispatch_dtype for runtime type dispatch. Both functions are fully safe (no unsafe code) and include comprehensive tests. Signed-off-by: rajesh-puripanda --- crates/mohu-index/src/boolean.rs | 114 ++++++++++++++++- crates/mohu-index/src/take.rs | 210 ++++++++++++++++++++++++++++++- 2 files changed, 322 insertions(+), 2 deletions(-) diff --git a/crates/mohu-index/src/boolean.rs b/crates/mohu-index/src/boolean.rs index 02fdb2a..406cde6 100644 --- a/crates/mohu-index/src/boolean.rs +++ b/crates/mohu-index/src/boolean.rs @@ -1 +1,113 @@ -// boolean — implementation pending +use mohu_buffer::buffer::Buffer; +use mohu_buffer::strides::NdIndexIter; +use mohu_dtype::dispatch_dtype; +use mohu_dtype::dtype::DType; +use mohu_error::{MohuError, MohuResult}; + +/// Boolean mask indexing. +/// +/// Returns a new 1D buffer containing only the elements of `src` where the +/// corresponding element of `mask` is `true`. +/// +/// # Errors +/// +/// - `DomainError` if `mask.dtype() != DType::Bool` +/// - `BoolIndexShapeMismatch` if `mask.shape() != src.shape()` +pub fn index_bool(src: &Buffer, mask: &Buffer) -> MohuResult { + if mask.dtype() != DType::Bool { + return Err(MohuError::DomainError { + op: "index_bool", + reason: "mask dtype must be Bool".into(), + }); + } + if mask.shape() != src.shape() { + return Err(MohuError::BoolIndexShapeMismatch { + index_shape: mask.shape().to_vec(), + array_shape: src.shape().to_vec(), + }); + } + + let mut true_coords: Vec> = Vec::new(); + + for coord in NdIndexIter::new(src.shape()) { + if mask.get::(&coord)? { + true_coords.push(coord.to_vec()); + } + } + + let out_len = true_coords.len(); + let mut out = Buffer::zeros(src.dtype(), &[out_len])?; + + if out_len == 0 { + return Ok(out); + } + + macro_rules! copy_bool { + ($T:ty) => { + for (i, coord) in true_coords.iter().enumerate() { + let val = src.get::<$T>(coord)?; + out.set::<$T>(&[i], val)?; + } + Ok::<_, MohuError>(()) + }; + } + + dispatch_dtype!(src.dtype(), copy_bool)?; + + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use mohu_buffer::buffer::Buffer; + + #[test] + fn test_index_bool_1d() { + let src = Buffer::from_slice::(&[10, 20, 30, 40, 50]).unwrap(); + let mask = Buffer::from_slice::(&[true, false, true, false, true]).unwrap(); + let result = index_bool(&src, &mask).unwrap(); + let expected = Buffer::from_slice::(&[10, 30, 50]).unwrap(); + assert_eq!(result.as_slice::().unwrap(), expected.as_slice::().unwrap()); + } + + #[test] + fn test_index_bool_all_false() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let mask = Buffer::from_slice::(&[false, false, false]).unwrap(); + let result = index_bool(&src, &mask).unwrap(); + assert_eq!(result.len(), 0); + } + + #[test] + fn test_index_bool_all_true() { + let src = Buffer::from_slice::(&[1.0, 2.0, 3.0]).unwrap(); + let mask = Buffer::from_slice::(&[true, true, true]).unwrap(); + let result = index_bool(&src, &mask).unwrap(); + assert_eq!(result.as_slice::().unwrap(), &[1.0, 2.0, 3.0]); + } + + #[test] + fn test_index_bool_2d() { + let data = &[&[1i64, 2, 3], &[4, 5, 6]]; + let src = Buffer::from_slice_2d::(data).unwrap(); + let mask = Buffer::from_slice_2d::(&[&[true, false, true], &[false, true, false]]).unwrap(); + let result = index_bool(&src, &mask).unwrap(); + let expected = Buffer::from_slice::(&[1, 3, 5]).unwrap(); + assert_eq!(result.as_slice::().unwrap(), expected.as_slice::().unwrap()); + } + + #[test] + fn test_index_bool_wrong_dtype() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let mask = Buffer::from_slice::(&[1, 0, 1]).unwrap(); + assert!(index_bool(&src, &mask).is_err()); + } + + #[test] + fn test_index_bool_shape_mismatch() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let mask = Buffer::from_slice::(&[true, false]).unwrap(); + assert!(index_bool(&src, &mask).is_err()); + } +} diff --git a/crates/mohu-index/src/take.rs b/crates/mohu-index/src/take.rs index 1b9b746..0fe54c3 100644 --- a/crates/mohu-index/src/take.rs +++ b/crates/mohu-index/src/take.rs @@ -1 +1,209 @@ -// take — implementation pending +use mohu_buffer::buffer::Buffer; +use mohu_buffer::layout::Order; +use mohu_buffer::strides::NdIndexIter; +use mohu_dtype::dispatch_dtype; +use mohu_dtype::dtype::DType; +use mohu_error::{MohuError, MohuResult}; + +/// Take elements from a buffer along an axis using an array of indices. +/// +/// Equivalent to NumPy's `take(a, indices, axis)`. +/// +/// # Errors +/// +/// - `AxisOutOfRange` if `axis >= src.ndim()` +/// - `DomainError` if indices is not a 1D I64 buffer +/// - `FancyIndexOutOfBounds` if any index is out of range for the axis +pub fn index_take(src: &Buffer, indices: &Buffer, axis: usize) -> MohuResult { + let ndim = src.ndim(); + if ndim == 0 { + return Err(MohuError::TooManyIndices { + given: 1, + ndim: 0, + }); + } + if axis >= ndim { + return Err(MohuError::AxisOutOfRange { + axis: axis as i64, + ndim, + }); + } + if indices.ndim() != 1 || indices.dtype() != DType::I64 { + return Err(MohuError::DomainError { + op: "index_take", + reason: "indices must be a 1D I64 array".into(), + }); + } + + let axis_size = src.shape()[axis]; + let indices_slice = indices.as_slice::()?; + let num_indices = indices_slice.len(); + + let mut wrapped_indices: Vec = Vec::with_capacity(num_indices); + for &idx in indices_slice.iter() { + let wrapped = if idx < 0 { + let w = idx + axis_size as i64; + if w < 0 { + return Err(MohuError::FancyIndexOutOfBounds { + index: idx, + axis, + size: axis_size, + }); + } + w as usize + } else { + idx as usize + }; + if wrapped >= axis_size { + return Err(MohuError::FancyIndexOutOfBounds { + index: idx, + axis, + size: axis_size, + }); + } + wrapped_indices.push(wrapped); + } + + if num_indices == 0 || src.len() == 0 { + let mut out_shape: Vec = src.shape().to_vec(); + out_shape[axis] = 0; + return Buffer::zeros(src.dtype(), &out_shape); + } + + // Build output shape + let mut out_shape: Vec = src.shape().to_vec(); + out_shape[axis] = num_indices; + let mut out = Buffer::alloc(src.dtype(), &out_shape, Order::C)?; + + // Build reduced shape (all dims except axis) + let mut reduced_shape: Vec = Vec::with_capacity(ndim - 1); + for i in 0..ndim { + if i != axis { + reduced_shape.push(src.shape()[i]); + } + } + + let mut src_coord = vec![0usize; ndim]; + let mut out_coord = vec![0usize; ndim]; + + macro_rules! copy_take { + ($T:ty) => { + if ndim == 1 { + for (pos, &wrapped) in wrapped_indices.iter().enumerate() { + let val = src.get::<$T>(&[wrapped])?; + out.set::<$T>(&[pos], val)?; + } + } else { + for reduced_coord in NdIndexIter::new(&reduced_shape) { + let mut ri = 0; + for i in 0..ndim { + if i != axis { + src_coord[i] = reduced_coord[ri]; + out_coord[i] = reduced_coord[ri]; + ri += 1; + } + } + + for (pos, &wrapped) in wrapped_indices.iter().enumerate() { + src_coord[axis] = wrapped; + out_coord[axis] = pos; + let val = src.get::<$T>(&src_coord)?; + out.set::<$T>(&out_coord, val)?; + } + } + } + Ok::<_, MohuError>(()) + }; + } + + dispatch_dtype!(src.dtype(), copy_take)?; + + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use mohu_buffer::buffer::Buffer; + + #[test] + fn test_index_take_1d() { + let src = Buffer::from_slice::(&[10, 20, 30, 40, 50]).unwrap(); + let idx = Buffer::from_slice::(&[0, 2, 4]).unwrap(); + let result = index_take(&src, &idx, 0).unwrap(); + let expected = Buffer::from_slice::(&[10, 30, 50]).unwrap(); + assert_eq!(result.as_slice::().unwrap(), expected.as_slice::().unwrap()); + } + + #[test] + fn test_index_take_1d_axis0() { + let src = Buffer::from_slice::(&[1.0, 2.0, 3.0, 4.0]).unwrap(); + let idx = Buffer::from_slice::(&[3, 1]).unwrap(); + let result = index_take(&src, &idx, 0).unwrap(); + assert_eq!(result.as_slice::().unwrap(), &[4.0, 2.0]); + } + + #[test] + fn test_index_take_2d_axis0() { + let data = &[&[1i64, 2], &[3, 4], &[5, 6]]; + let src = Buffer::from_slice_2d::(data).unwrap(); + let idx = Buffer::from_slice::(&[0, 2]).unwrap(); + let result = index_take(&src, &idx, 0).unwrap(); + assert_eq!(result.shape(), &[2, 2]); + assert_eq!(result.as_slice::().unwrap(), &[1, 2, 5, 6]); + } + + #[test] + fn test_index_take_2d_axis1() { + let data = &[&[1i64, 2, 3], &[4, 5, 6]]; + let src = Buffer::from_slice_2d::(data).unwrap(); + let idx = Buffer::from_slice::(&[2, 0]).unwrap(); + let result = index_take(&src, &idx, 1).unwrap(); + assert_eq!(result.shape(), &[2, 2]); + assert_eq!(result.as_slice::().unwrap(), &[3, 1, 6, 4]); + } + + #[test] + fn test_index_take_negative_indices() { + let src = Buffer::from_slice::(&[10, 20, 30]).unwrap(); + let idx = Buffer::from_slice::(&[-1, -2]).unwrap(); + let result = index_take(&src, &idx, 0).unwrap(); + assert_eq!(result.as_slice::().unwrap(), &[30, 20]); + } + + #[test] + fn test_index_take_empty_indices() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let idx = Buffer::from_slice::(&[]).unwrap(); + let result = index_take(&src, &idx, 0).unwrap(); + assert_eq!(result.len(), 0); + } + + #[test] + fn test_index_take_out_of_bounds() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let idx = Buffer::from_slice::(&[5]).unwrap(); + assert!(index_take(&src, &idx, 0).is_err()); + } + + #[test] + fn test_index_take_negative_out_of_bounds() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let idx = Buffer::from_slice::(&[-5]).unwrap(); + assert!(index_take(&src, &idx, 0).is_err()); + } + + #[test] + fn test_index_take_wrong_dtype() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let idx = Buffer::from_slice::(&[1.0]).unwrap(); + assert!(index_take(&src, &idx, 0).is_err()); + } + + #[test] + fn test_index_take_bad_axis() { + let src = Buffer::from_slice::(&[1, 2, 3]).unwrap(); + let idx = Buffer::from_slice::(&[0]).unwrap(); + assert!(index_take(&src, &idx, 1).is_err()); + } +} From 05df5f2a40e57434a4e6823261b6a008129dbe64 Mon Sep 17 00:00:00 2001 From: rajesh-puripanda Date: Sun, 31 May 2026 23:17:47 +0530 Subject: [PATCH 2/2] feat(random): implement Rng struct with uniform/normal/integers Add convenience Rng struct wrapping Pcg64 with three sampling methods: - uniform: f64 values in [low, high) - normal: Gaussian variates via Box-Muller (all dtypes via dispatch) - integers: i64 values in [low, high) Includes reproducibility guarantee, domain error validation, and comprehensive tests for all methods. Signed-off-by: rajesh-puripanda --- crates/mohu-random/src/lib.rs | 254 +++++++++++++++++++++++++++++++++- 1 file changed, 252 insertions(+), 2 deletions(-) diff --git a/crates/mohu-random/src/lib.rs b/crates/mohu-random/src/lib.rs index d5f3ecf..49b52ec 100644 --- a/crates/mohu-random/src/lib.rs +++ b/crates/mohu-random/src/lib.rs @@ -25,8 +25,8 @@ /// # Reproducibility /// /// ```rust,ignore -/// let mut rng = mohu_random::Pcg64::seed(42); -/// let data = rng.standard_normal::(&[1000]); +/// let mut rng = mohu_random::Rng::new(42); +/// let data = rng.uniform(&[1000], 0.0, 1.0).unwrap(); /// ``` /// /// All generators implement `Seed` — the same seed always produces the @@ -39,5 +39,255 @@ pub mod multivariate; pub mod permutation; pub mod seeding; +use crate::generator::Generator; +use mohu_buffer::buffer::Buffer; +use mohu_buffer::layout::Order; +use mohu_dtype::dtype::DType; + pub use generator::{Generator, Pcg64, Philox4x64}; pub use mohu_error::{MohuError, MohuResult}; + +/// Convenience random number generator for common sampling tasks. +/// +/// Wraps the default PCG-64-DXSM engine and provides methods for +/// generating uniform, normal, and integer random values into +/// multi-dimensional buffers. +/// +/// # Reproducibility +/// +/// Two `Rng` instances created with the same seed produce identical +/// output sequences regardless of platform or mohu version. +/// +/// # Example +/// +/// ```rust,ignore +/// let mut rng = Rng::new(42); +/// let u = rng.uniform(&[3, 4], 0.0, 1.0).unwrap(); +/// let n = rng.normal(&[1000], 0.0, 1.0).unwrap(); +/// let i = rng.integers(&[5], 0, 10).unwrap(); +/// ``` +pub struct Rng { + pcg: Pcg64, +} + +impl Rng { + /// Create a new RNG seeded with `seed`. + /// + /// The same seed always produces the same sequence of values. + pub fn new(seed: u64) -> Self { + Self { + pcg: Pcg64::seed(seed), + } + } + + /// Generate a random `f64` uniformly distributed in `[0, 1)`. + fn rand_f64(&mut self) -> f64 { + let bits = self.pcg.next_u64() >> 11; + (bits as f64) * (1.0 / 9007199254740992.0) + } + + /// Fill a buffer with uniformly distributed `f64` values. + /// + /// Each value lies in `[low, high)`. + /// + /// # Errors + /// + /// Returns `DomainError` if `high <= low`. + pub fn uniform(&mut self, shape: &[usize], low: f64, high: f64) -> MohuResult { + if high <= low { + return Err(MohuError::DomainError { + op: "uniform", + reason: "high must be > low".into(), + }); + } + let mut out = Buffer::alloc(DType::F64, shape, Order::C)?; + let slice = out.as_mut_slice::()?; + let scale = high - low; + for v in slice.iter_mut() { + *v = low + scale * self.rand_f64(); + } + Ok(out) + } + + /// Fill a buffer with normally distributed `f64` values. + /// + /// Uses the Box-Muller transform to generate pairs of independent + /// standard normal variates, then scales by `std` and shifts by `mean`. + /// + /// # Errors + /// + /// Returns `DomainError` if `std <= 0.0`. + pub fn normal(&mut self, shape: &[usize], mean: f64, std: f64) -> MohuResult { + if std <= 0.0 { + return Err(MohuError::DomainError { + op: "normal", + reason: "std must be > 0".into(), + }); + } + let len: usize = shape.iter().product(); + let mut out = Buffer::alloc(DType::F64, shape, Order::C)?; + let slice = out.as_mut_slice::()?; + let mut i = 0; + while i < len { + let u1 = self.rand_f64(); + let u2 = self.rand_f64(); + let mag = (-2.0 * u1.ln()).sqrt(); + let theta = 2.0 * std::f64::consts::PI * u2; + slice[i] = mean + std * mag * theta.cos(); + if i + 1 < len { + slice[i + 1] = mean + std * mag * theta.sin(); + } + i += 2; + } + Ok(out) + } + + /// Fill a buffer with uniformly distributed random integers. + /// + /// Each value lies in `[low, high)`. + /// + /// # Errors + /// + /// Returns `DomainError` if `high <= low`. + pub fn integers(&mut self, shape: &[usize], low: i64, high: i64) -> MohuResult { + if high <= low { + return Err(MohuError::DomainError { + op: "integers", + reason: "high must be > low".into(), + }); + } + let range = (high - low) as u64; + let mut out = Buffer::alloc(DType::I64, shape, Order::C)?; + let slice = out.as_mut_slice::()?; + for v in slice.iter_mut() { + *v = low + (self.pcg.next_u64() % range) as i64; + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_uniform_range() { + let mut rng = Rng::new(42); + let buf = rng.uniform(&[1000], 0.0, 1.0).unwrap(); + let slice = buf.as_slice::().unwrap(); + for &v in slice.iter() { + assert!((0.0..1.0).contains(&v), "value {} out of [0,1)", v); + } + } + + #[test] + fn test_uniform_wider_range() { + let mut rng = Rng::new(99); + let buf = rng.uniform(&[500], -5.0, 5.0).unwrap(); + let slice = buf.as_slice::().unwrap(); + for &v in slice.iter() { + assert!((-5.0..5.0).contains(&v), "value {} out of [-5,5)", v); + } + } + + #[test] + fn test_uniform_invalid_range() { + let mut rng = Rng::new(1); + assert!(rng.uniform(&[1], 5.0, 3.0).is_err()); + assert!(rng.uniform(&[1], 2.0, 2.0).is_err()); + } + + #[test] + fn test_normal_shape() { + let mut rng = Rng::new(7); + let buf = rng.normal(&[3, 4, 5], 0.0, 1.0).unwrap(); + assert_eq!(buf.shape(), &[3, 4, 5]); + assert_eq!(buf.dtype(), DType::F64); + } + + #[test] + fn test_normal_mean_std() { + let mut rng = Rng::new(123); + let n = 100_000; + let buf = rng.normal(&[n], 2.0, 3.0).unwrap(); + let slice = buf.as_slice::().unwrap(); + let mean = slice.iter().sum::() / n as f64; + let variance = slice.iter().map(|v| (v - mean).powi(2)).sum::() / n as f64; + let std_est = variance.sqrt(); + // Allow 2% relative tolerance for n=100000 + assert!( + (mean - 2.0).abs() < 0.1, + "mean {} too far from 2.0", + mean + ); + assert!( + (std_est - 3.0).abs() < 0.1, + "std {} too far from 3.0", + std_est + ); + } + + #[test] + fn test_normal_invalid_std() { + let mut rng = Rng::new(1); + assert!(rng.normal(&[1], 0.0, 0.0).is_err()); + assert!(rng.normal(&[1], 0.0, -1.0).is_err()); + } + + #[test] + fn test_integers_range() { + let mut rng = Rng::new(42); + let buf = rng.integers(&[1000], 0, 10).unwrap(); + let slice = buf.as_slice::().unwrap(); + for &v in slice.iter() { + assert!((0..10).contains(&v), "value {} out of [0,10)", v); + } + } + + #[test] + fn test_integers_negative_low() { + let mut rng = Rng::new(7); + let buf = rng.integers(&[500], -5, 5).unwrap(); + let slice = buf.as_slice::().unwrap(); + for &v in slice.iter() { + assert!((-5..5).contains(&v), "value {} out of [-5,5)", v); + } + } + + #[test] + fn test_integers_invalid_range() { + let mut rng = Rng::new(1); + assert!(rng.integers(&[1], 5, 3).is_err()); + assert!(rng.integers(&[1], 2, 2).is_err()); + } + + #[test] + fn test_reproducibility() { + let mut rng1 = Rng::new(42); + let mut rng2 = Rng::new(42); + let a = rng1.uniform(&[100], 0.0, 1.0).unwrap(); + let b = rng2.uniform(&[100], 0.0, 1.0).unwrap(); + assert_eq!(a.as_slice::().unwrap(), b.as_slice::().unwrap()); + } + + #[test] + fn test_different_seeds_differ() { + let mut rng1 = Rng::new(1); + let mut rng2 = Rng::new(2); + let a = rng1.uniform(&[10], 0.0, 1.0).unwrap(); + let b = rng2.uniform(&[10], 0.0, 1.0).unwrap(); + // Extremely unlikely to be equal + assert_ne!(a.as_slice::().unwrap(), b.as_slice::().unwrap()); + } + + #[test] + fn test_empty_shape() { + let mut rng = Rng::new(0); + let buf = rng.uniform(&[0], 0.0, 1.0).unwrap(); + assert_eq!(buf.len(), 0); + let buf = rng.normal(&[0, 5], 0.0, 1.0).unwrap(); + assert_eq!(buf.len(), 0); + let buf = rng.integers(&[0], 0, 10).unwrap(); + assert_eq!(buf.len(), 0); + } +}