diff --git a/.gitignore b/.gitignore index fe89b3d41..e36ed064c 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,6 @@ __pycache__ # Coverage files *.lcov + +# IDEs +/.vscode/ diff --git a/Cargo.toml b/Cargo.toml index 3afaf6190..cca003566 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,12 +29,14 @@ num-integer = "0.1" num-traits = "0.2" ndarray = ">= 0.15, <=0.17" pyo3 = { version = "0.29.0", default-features = false, features = ["macros"] } +rand_core = { version = "0.9.3", default-features = false, optional = true } rustc-hash = "2.0" [dev-dependencies] pyo3 = { version = "0.29.0", default-features = false, features = ["auto-initialize"]} nalgebra = { version = ">=0.30, <0.35", default-features = false, features = ["std"] } criterion = "0.8.2" +rand = { version = "0.9.1", default-features = false } [build-dependencies] pyo3-build-config = { version = "0.29", features = ["resolve-config"]} diff --git a/src/lib.rs b/src/lib.rs index 08be254cc..23a96083c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -79,6 +79,7 @@ pub mod datetime; mod dtype; mod error; pub mod npyffi; +pub mod random; mod slice_container; mod strings; mod sum_products; diff --git a/src/npyffi/mod.rs b/src/npyffi/mod.rs index f80f72ec3..90728747a 100644 --- a/src/npyffi/mod.rs +++ b/src/npyffi/mod.rs @@ -155,6 +155,7 @@ pub mod flags; mod npy_common; mod numpyconfig; pub mod objects; +pub mod random; pub mod types; pub mod ufunc; @@ -163,5 +164,6 @@ pub use self::flags::*; pub use self::npy_common::*; pub use self::numpyconfig::*; pub use self::objects::*; +pub use self::random::*; pub use self::types::*; pub use self::ufunc::*; diff --git a/src/npyffi/objects.rs b/src/npyffi/objects.rs index 635740f9e..e87c8a05f 100644 --- a/src/npyffi/objects.rs +++ b/src/npyffi/objects.rs @@ -1,4 +1,4 @@ -//! Low-Lebel binding for NumPy C API C-objects +//! Low-Level binding for NumPy C API C-objects //! //! #![allow(non_camel_case_types)] diff --git a/src/npyffi/random.rs b/src/npyffi/random.rs new file mode 100644 index 000000000..e618768aa --- /dev/null +++ b/src/npyffi/random.rs @@ -0,0 +1,11 @@ +use std::ffi::c_void; + +#[repr(C)] +#[derive(Debug)] +pub struct bitgen_t { + pub state: *mut c_void, + pub next_uint64: unsafe extern "C" fn(*mut c_void) -> super::npy_uint64, //nogil + pub next_uint32: unsafe extern "C" fn(*mut c_void) -> super::npy_uint32, //nogil + pub next_double: unsafe extern "C" fn(*mut c_void) -> libc::c_double, //nogil + pub next_raw: unsafe extern "C" fn(*mut c_void) -> super::npy_uint64, //nogil +} diff --git a/src/random.rs b/src/random.rs new file mode 100644 index 000000000..79bc86bb6 --- /dev/null +++ b/src/random.rs @@ -0,0 +1,419 @@ +//! Safe interface for NumPy's random [`BitGenerator`][bg]. +//! +//! Using the patterns described in [“Extending `numpy.random`”][ext], +//! you can generate random numbers without being attached to the interpreter runtime by: +//! - [spawning][`PyBitGeneratorMethods::spawn`] fresh [`BitGenerator`]s +//! from a [`PyBitGenerator`] you received from Python +//! - creating a fresh [`BitGenerator`] [from numpy][`BitGenerator::new`]: +//! +//! ``` +//! use pyo3::prelude::*; +//! use numpy::random::BitGenerator; +//! +//! let mut bitgen = Python::attach(|py| { +//! BitGenerator::new(py, Default::default()) +//! })?; +//! let random_number = bitgen.next_u64(); +//! # Ok::<(), PyErr>(()) +//! ``` +//! +//! If you write a pyo3 extension, you would extract +//! a [`numpy.random.BitGenerator`] into a [Bound]<'_, [PyBitGenerator]>: +//! +//! [`numpy.random.BitGenerator`]: https://numpy.org/doc/stable/reference/random/bit_generators/generated/numpy.random.BitGenerator.html +//! +//! ``` +//! # use pyo3::prelude::*; +//! use numpy::random::{BitGenerator, PyBitGenerator, PyBitGeneratorMethods as _}; +//! # fn default_bit_gen<'py>(py: Python<'py>) -> PyResult> { +//! # Ok(BitGenerator::new(py, Default::default())?.into_shared().into_bound(py)) +//! # } +//! +//! #[pyfunction] +//! fn make_random_number(bitgen: Bound) -> PyResult { +//! // spawn an owned child, then use it without being attached to the interpreter runtime +//! Ok(bitgen.spawn_one()?.next_u64()) +//! } +//! +//! Python::attach(|py| -> PyResult<_> { +//! let bitgen: Bound = default_bit_gen(py)?; +//! let random_number = make_random_number(bitgen)?; +//! println!("{random_number}"); +//! Ok(()) +//! })?; +//! # Ok::<(), PyErr>(()) +//! ``` +//! +//! With the `rand` crate installed, you can also use its `Rng` APIs on any generator, +//! since [`BitGenerator`] implements [`rand_core::RngCore`]. +//! +//! ``` +//! use pyo3::prelude::*; +//! use rand::Rng as _; +//! use numpy::random::{BitGenerator, BitGeneratorKind::SFC64}; +//! +//! let mut bitgen = Python::attach(|py| BitGenerator::new(py, SFC64))?; +//! if bitgen.random_ratio(1, 1_000_000) { +//! println!("a sure thing"); +//! }; +//! # Ok::<(), PyErr>(()) +//! ``` +//! +//! Using `spawn`, you can create multiple [`BitGenerator`]s to generate random numbers truly in parallel, +//! all without being attached to the interpreter runtime: +//! +//! ``` +//! # use pyo3::prelude::*; +//! # use rand::Rng as _; +//! use numpy::{PyArray2, PyArrayMethods as _}; +//! # use numpy::random::{BitGenerator, PyBitGenerator, PyBitGeneratorMethods as _}; +//! # fn default_bit_gen<'py>(py: Python<'py>) -> PyResult> { +//! # Ok(BitGenerator::new(py, Default::default())?.into_shared().into_bound(py)) +//! # } +//! +//! Python::attach(|py| -> PyResult<_> { +//! let bitgen: Bound = default_bit_gen(py)?; +//! let children = bitgen.spawn(4)?; +//! let mut arr = PyArray2::::zeros(py, (4, 300), false).readwrite(); +//! let mut ndarr = arr.as_array_mut(); // ndarray for more convenience +//! py.detach(|| std::thread::scope(|s| { +//! for (mut chunk, mut child) in ndarr.rows_mut().into_iter().zip(children) { +//! s.spawn(move || { +//! for x in chunk.iter_mut() { +//! *x = child.random_range(10..200); +//! } +//! }); +//! } +//! })); +//! println!("Now filled: {arr:?}"); +//! Ok(()) +//! })?; +//! # Ok::<(), PyErr>(()) +//! ``` +//! +//! [bg]: https://numpy.org/doc/stable//reference/random/bit_generators/generated/numpy.random.BitGenerator.html +//! [ext]: https://numpy.org/doc/stable/reference/random/extending.html + +use std::ptr::NonNull; + +use pyo3::{ + exceptions::PyRuntimeError, + ffi, intern, + prelude::*, + sync::PyOnceLock, + types::{DerefToPyAny, PyCapsule, PyType}, + PyTypeInfo, +}; + +use crate::npyffi::bitgen_t; + +mod sealed { + pub trait Sealed {} +} + +use sealed::Sealed; + +/// Methods for [`PyBitGenerator`]. +pub trait PyBitGeneratorMethods: Sealed { + /// Spawn `n_children` independent child [`BitGenerator`]s. + /// + /// This is the way to obtain generators for multiple threads: each child has its own, + /// independent state, so no synchronization is needed between them. + fn spawn(&self, n_children: usize) -> PyResult>; + + /// Spawn a single owned child [`BitGenerator`]. + fn spawn_one(&self) -> PyResult { + let mut children = self.spawn(1)?; + children + .pop() + .ok_or_else(|| PyRuntimeError::new_err("spawn(1) returned no children")) + } +} + +/// Wrapper for [`np.random.BitGenerator`][bg]. +/// +/// See also [`PyBitGeneratorMethods`]. +/// +/// [bg]: https://numpy.org/doc/stable//reference/random/bit_generators/generated/numpy.random.BitGenerator.html +#[repr(transparent)] +pub struct PyBitGenerator(PyAny); + +impl DerefToPyAny for PyBitGenerator {} + +unsafe impl PyTypeInfo for PyBitGenerator { + const NAME: &'static str = "PyBitGenerator"; + const MODULE: Option<&'static str> = Some("numpy.random"); + + fn type_object_raw<'py>(py: Python<'py>) -> *mut ffi::PyTypeObject { + static CLS: PyOnceLock> = PyOnceLock::new(); + let cls = CLS + .import(py, "numpy.random", "BitGenerator") + .expect("Failed to get BitGenerator type object"); + cls.as_type_ptr() + } +} + +impl<'py> PyBitGeneratorMethods for Bound<'py, PyBitGenerator> { + fn spawn(&self, n_children: usize) -> PyResult> { + let py = self.py(); + self.call_method1(intern!(py, "spawn"), (n_children,))? + .try_iter()? + // SAFETY: each child is freshly spawned and only handed to us, so it’s exclusively ours. + .map(|child| unsafe { BitGenerator::from_py(child?.cast_into::()?) }) + .collect() + } +} + +impl Sealed for Bound<'_, PyBitGenerator> {} + +/// Which of numpy’s bit generator algorithms [`BitGenerator::new`] should create. +#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum BitGeneratorKind { + /// Mersenne Twister (MT19937) + MT19937, + /// Permuted congruential generator (64-bit, PCG-64) + #[default] + PCG64, + /// Permuted congruential generator (64-bit, PCG-64 DXSM + PCG64DXSM, + /// Philox counter-based RNG + Philox, + /// SFC64 Small Fast Chaotic PRNG + SFC64, +} + +impl BitGeneratorKind { + /// Returns an iterator over the values of [`BitGeneratorKind`]. + pub fn iter() -> std::iter::Copied> { + use BitGeneratorKind::*; + static KINDS: [BitGeneratorKind; 5] = [MT19937, PCG64, PCG64DXSM, Philox, SFC64]; + KINDS.iter().copied() + } +} + +impl From for &'static str { + fn from(value: BitGeneratorKind) -> &'static str { + match value { + BitGeneratorKind::MT19937 => "MT19937", + BitGeneratorKind::PCG64 => "PCG64", + BitGeneratorKind::PCG64DXSM => "PCG64DXSM", + BitGeneratorKind::Philox => "Philox", + BitGeneratorKind::SFC64 => "SFC64", + } + } +} + +/// A numpy `BitGenerator` usable without being attached to the interpreter runtime, +/// with exclusive access to its state. +/// +/// [`spawn`][PyBitGeneratorMethods::spawn] hands out independent, owned ones. +pub struct BitGenerator { + raw: NonNull, + /// Keeps `raw` alive: the capsule’s pointer lives in memory owned by the `BitGenerator`, which + /// has no back-reference of its own, so only keeping it alive keeps that memory valid. + _bit_generator: Py, +} + +// SAFETY: `raw` is only ever accessed through `&mut self`, so it can’t be used in parallel, and we +// keep its `bitgen_t` alive via `_bit_generator`. Every `BitGenerator` owns its `bitgen_t` +// exclusively (it is freshly created or a fresh `spawn` child), so nothing else can touch +// its state. +unsafe impl Send for BitGenerator {} + +impl BitGenerator { + /// Creates a fresh [`BitGenerator`] backed by numpy’s implementation. + /// + /// ``` + /// use pyo3::prelude::*; + /// use numpy::random::{BitGenerator, BitGeneratorKind}; + /// + /// let mut bitgen = Python::attach(|py| BitGenerator::new(py, Default::default()))?; + /// println!("{}", bitgen.next_u32()); + /// # Ok::<(), PyErr>(()) + /// ``` + pub fn new(py: Python<'_>, kind: BitGeneratorKind) -> PyResult { + let bitgen = py + .import("numpy.random")? + .call_method0::<&str>(kind.into())? + .cast_into::()?; + // SAFETY: `bitgen` is freshly created and not handed out elsewhere. + unsafe { Self::from_py(bitgen) } + } + + /// Extracts the raw `bitgen_t` pointer from `bit_generator`’s capsule. + /// + /// # Safety + /// + /// The caller must ensure the result has exclusive access to the `bitgen_t` for its whole + /// lifetime, i.e. `bit_generator` is freshly created and not handed out elsewhere. + unsafe fn from_py(bit_generator: Bound<'_, PyBitGenerator>) -> PyResult { + let py = bit_generator.py(); + let capsule = bit_generator + .getattr(intern!(py, "capsule"))? + .cast_into::()?; + let raw = capsule + .pointer_checked(Some(ffi::c_str!("BitGenerator"))) + .map_err(|_| PyRuntimeError::new_err("Invalid BitGenerator capsule"))?; + Ok(BitGenerator { + raw: raw.cast(), + _bit_generator: bit_generator.unbind(), + }) + } + + /// Returns the underlying [`PyBitGenerator`]. + pub fn into_shared(self) -> Py { + self._bit_generator + } + + /// Returns the next random unsigned 64 bit integer. + pub fn next_u64(&mut self) -> u64 { + unsafe { + let bitgen = self.raw.as_ptr(); + debug_assert_ne!((*bitgen).state, std::ptr::null_mut()); + ((*bitgen).next_uint64)((*bitgen).state) + } + } + /// Returns the next random unsigned 32 bit integer. + pub fn next_u32(&mut self) -> u32 { + unsafe { + let bitgen = self.raw.as_ptr(); + debug_assert_ne!((*bitgen).state, std::ptr::null_mut()); + ((*bitgen).next_uint32)((*bitgen).state) + } + } + /// Returns the next random double. + pub fn next_double(&mut self) -> f64 { + unsafe { + let bitgen = self.raw.as_ptr(); + debug_assert_ne!((*bitgen).state, std::ptr::null_mut()); + ((*bitgen).next_double)((*bitgen).state) + } + } + /// Returns the next raw value (can be used for testing). + pub fn next_raw(&mut self) -> u64 { + unsafe { + let bitgen = self.raw.as_ptr(); + debug_assert_ne!((*bitgen).state, std::ptr::null_mut()); + ((*bitgen).next_raw)((*bitgen).state) + } + } +} + +#[cfg(feature = "rand_core")] +impl rand_core::RngCore for BitGenerator { + fn next_u32(&mut self) -> u32 { + BitGenerator::next_u32(self) + } + fn next_u64(&mut self) -> u64 { + BitGenerator::next_u64(self) + } + fn fill_bytes(&mut self, dst: &mut [u8]) { + rand_core::impls::fill_bytes_via_next(self, dst) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn get_shared<'py>(py: Python<'py>) -> PyResult> { + let bitgen = py + .import("numpy.random")? + .call_method1("PCG64", (42,))? + .cast_into::()?; + Ok(bitgen) + } + + fn get_owned<'py>(py: Python<'py>) -> PyResult { + let bitgen = get_shared(py)?; + // SAFETY: `bitgen` is freshly created and not handed out elsewhere. + unsafe { BitGenerator::from_py(bitgen) } + } + + #[test] + fn from_kind() -> PyResult<()> { + Python::attach(|py| { + for kind in BitGeneratorKind::iter() { + let name: &str = kind.into(); + let type_name = BitGenerator::new(py, kind)? + .into_shared() + .bind(py) + .get_type() + .name()?; + assert_eq!(type_name, name); + } + Ok(()) + }) + } + + /// Simple single-threaded use of an owned generator. + #[test] + fn base_api() -> PyResult<()> { + Python::attach(|py| { + let double = get_owned(py)?.next_double(); + assert_eq!(double, 0.7739560485559633); + + let u32_owned = get_owned(py)?.next_u32(); + assert_eq!(u32_owned, 383329928); + + let u64_owned = get_owned(py)?.next_u64(); + assert_eq!(u64_owned, 14276969152011380360); + + let raw_owned = get_owned(py)?.next_raw(); + assert_eq!(raw_owned, u64_owned); + + Ok(()) + }) + } + + /// Test that the `rand::Rng` APIs work + #[cfg(feature = "rand_core")] + #[test] + fn rand() -> PyResult<()> { + use rand::Rng as _; + + Python::attach(|py| { + let seq_owned: Vec = get_owned(py)? + .sample_iter(rand::distr::Bernoulli::new(0.5).unwrap()) + .take(10) + .collect(); + let seq_expected = vec![ + false, true, false, false, true, false, false, false, true, true, + ]; + assert_eq!(&seq_owned, &seq_expected); + Ok(()) + }) + } + + #[test] + fn spawn_one() -> PyResult<()> { + let mut bitgen = Python::attach(|py| get_shared(py)?.spawn_one())?; + assert_eq!(bitgen.next_u32(), 2136330838); + Ok(()) + } + + /// Spawned children are independent and owned, + /// so they can be used (and dropped) from their own threads without locking. + #[test] + fn spawn_produces_independent_generators() -> PyResult<()> { + Python::attach(|py| { + let children = get_shared(py)?.spawn(2)?; + assert_eq!(children.len(), 2); + + let values = py.detach(|| { + std::thread::scope(|s| { + children + .into_iter() + .map(|mut child| s.spawn(move || child.next_u64())) + .collect::>() + .into_iter() + .map(|handle| handle.join().unwrap()) + .collect::>() + }) + }); + + assert_eq!(values, vec![16910944855483863638, 8623682774590505111]); + Ok(()) + }) + } +} diff --git a/x.py b/x.py index 1e1f9fe6b..046230b7c 100755 --- a/x.py +++ b/x.py @@ -57,7 +57,7 @@ def default(args): def check(args): - run("cargo", "doc", "--no-deps", env={"RUSTDOCFLAGS": "--deny warnings"}) + run("cargo", "doc", "--all-features", "--no-deps", env={"RUSTDOCFLAGS": "--deny warnings"}) run("cargo", "fmt", "--", "--check") run("cargo", "clippy", "--all-features", "--tests", "--", *DENY_WARNINGS)