ocas_py/
numeric.rs

1//! Python bindings for numerical integration (Vegas adaptive Monte Carlo).
2//!
3//! Wraps [`ocas_eval::numeric::Vegas`] and the [`integrate_1d`](ocas_eval::numeric::integrate_1d)
4//! convenience entry point. Python callables are wrapped into Rust closures
5//! so users can integrate arbitrary Python functions.
6//!
7//! ```python
8//! import ocas
9//!
10//! # One-shot helper: integrate x over [0, 1].
11//! r = ocas.integrate_1d(lambda x: x, 0.0, 1.0)
12//! print(r.integral, r.error)  # ~0.5, small error
13//!
14//! # Multi-dimensional Vegas with explicit control.
15//! v = ocas.Vegas(2, n_samples=20000, iterations=8, seed=1)
16//! r = v.integrate(lambda xs: xs[0] * xs[1])
17//! ```
18
19use ocas_eval::numeric::{IntegrateResult, Integrator, Vegas, VegasOptions};
20use pyo3::exceptions::{PyTypeError, PyValueError};
21use pyo3::prelude::*;
22use pyo3::types::{PyFloat, PyList};
23
24/// Result of a numerical integration: the estimate and its standard error.
25///
26/// Instances are returned by [`Vegas.integrate`][PyVegas.integrate] and
27/// [`ocas.integrate_1d`]. The fields `integral` and `error` are also
28/// accessible by index (`result[0]`, `result[1]`) and by unpacking
29/// (`integral, error = result`).
30#[pyclass(name = "IntegrateResult")]
31pub struct PyIntegrateResult {
32    /// Best estimate of the integral.
33    #[pyo3(get)]
34    pub integral: f64,
35    /// Estimated standard error on `integral`.
36    #[pyo3(get)]
37    pub error: f64,
38}
39
40#[pymethods]
41impl PyIntegrateResult {
42    #[new]
43    fn new(integral: f64, error: f64) -> Self {
44        Self { integral, error }
45    }
46
47    fn __getitem__(&self, idx: usize) -> PyResult<f64> {
48        match idx {
49            0 => Ok(self.integral),
50            1 => Ok(self.error),
51            _ => Err(pyo3::exceptions::PyIndexError::new_err(format!(
52                "IntegrateResult index {idx} out of range (only 0, 1 valid)"
53            ))),
54        }
55    }
56
57    fn __len__(&self) -> usize {
58        2
59    }
60
61    fn __repr__(&self) -> String {
62        format!(
63            "IntegrateResult(integral={:?}, error={:?})",
64            self.integral, self.error
65        )
66    }
67}
68
69impl From<IntegrateResult> for PyIntegrateResult {
70    fn from(r: IntegrateResult) -> Self {
71        Self {
72            integral: r.integral,
73            error: r.error,
74        }
75    }
76}
77
78/// Parse a non-negative integer option from a Python kwarg, validating it.
79fn parse_usize_opt(value: &Bound<'_, PyAny>, name: &str) -> PyResult<usize> {
80    let n: usize = value
81        .extract()
82        .map_err(|_| PyTypeError::new_err(format!("{name} must be a non-negative integer")))?;
83    Ok(n)
84}
85
86fn parse_f64_opt(value: &Bound<'_, PyAny>, name: &str) -> PyResult<f64> {
87    value
88        .extract()
89        .map_err(|_| PyTypeError::new_err(format!("{name} must be a float")))
90}
91
92/// Build [`VegasOptions`] from Python kwargs. All keys are optional.
93fn kwargs_to_opts(
94    n_bins: Option<&Bound<'_, PyAny>>,
95    n_samples: Option<&Bound<'_, PyAny>>,
96    iterations: Option<&Bound<'_, PyAny>>,
97    learning_rate: Option<&Bound<'_, PyAny>>,
98    seed: Option<&Bound<'_, PyAny>>,
99) -> PyResult<VegasOptions> {
100    let mut opts = VegasOptions::default();
101    if let Some(v) = n_bins {
102        opts.n_bins = parse_usize_opt(v, "n_bins")?;
103        if opts.n_bins == 0 {
104            return Err(PyValueError::new_err("n_bins must be >= 1"));
105        }
106    }
107    if let Some(v) = n_samples {
108        opts.n_samples = parse_usize_opt(v, "n_samples")?;
109    }
110    if let Some(v) = iterations {
111        opts.iterations = parse_usize_opt(v, "iterations")?;
112    }
113    if let Some(v) = learning_rate {
114        opts.learning_rate = parse_f64_opt(v, "learning_rate")?;
115        if opts.learning_rate.partial_cmp(&0.0) != Some(std::cmp::Ordering::Greater) {
116            return Err(PyValueError::new_err("learning_rate must be positive"));
117        }
118    }
119    if let Some(v) = seed {
120        opts.seed = v
121            .extract()
122            .map_err(|_| PyTypeError::new_err("seed must be an integer"))?;
123    }
124    Ok(opts)
125}
126
127/// Adaptive Monte Carlo integrator (Vegas) over the unit hypercube.
128///
129/// Construct with the number of dimensions and optional tuning knobs, then
130/// call [`integrate`][PyVegas.integrate] with a Python callable taking a
131/// list of `n_dims` floats in `[0, 1]` and returning a float.
132#[pyclass(name = "Vegas")]
133pub struct PyVegas {
134    inner: Vegas,
135}
136
137#[pymethods]
138impl PyVegas {
139    /// Create a Vegas integrator for `n_dims` dimensions.
140    ///
141    /// All options are keyword-only and optional:
142    ///
143    /// - `n_bins` (default 64): bins per dimension.
144    /// - `n_samples` (default 10000): samples per iteration.
145    /// - `iterations` (default 10): adaptive iterations.
146    /// - `learning_rate` (default 1.5): grid smoothing rate (1.0–2.0).
147    /// - `seed` (default 0x0C45): RNG seed for reproducibility.
148    #[new]
149    #[pyo3(signature = (n_dims, *, n_bins=None, n_samples=None, iterations=None, learning_rate=None, seed=None))]
150    fn new(
151        n_dims: usize,
152        n_bins: Option<&Bound<'_, PyAny>>,
153        n_samples: Option<&Bound<'_, PyAny>>,
154        iterations: Option<&Bound<'_, PyAny>>,
155        learning_rate: Option<&Bound<'_, PyAny>>,
156        seed: Option<&Bound<'_, PyAny>>,
157    ) -> PyResult<Self> {
158        if n_dims == 0 {
159            return Err(PyValueError::new_err("n_dims must be >= 1"));
160        }
161        let opts = kwargs_to_opts(n_bins, n_samples, iterations, learning_rate, seed)?;
162        Ok(Self {
163            inner: Vegas::new(n_dims, opts),
164        })
165    }
166
167    /// Integrate a Python callable. The callable receives a list of `n_dims`
168    /// floats in `[0, 1]` and must return a float.
169    fn integrate(&mut self, f: &Bound<'_, PyAny>) -> PyResult<PyIntegrateResult> {
170        let r = Python::attach(|py| -> PyResult<IntegrateResult> {
171            let cb = f.clone();
172            let wrapped = |x: &[f64]| -> f64 {
173                match PyList::new(py, x.iter().copied()) {
174                    Ok(list) => {
175                        let arg = list.into_any();
176                        match cb.call1((arg,)) {
177                            Ok(value) => match value.extract::<f64>() {
178                                Ok(v) => v,
179                                Err(e) => {
180                                    e.restore(py);
181                                    f64::NAN
182                                }
183                            },
184                            Err(e) => {
185                                e.restore(py);
186                                f64::NAN
187                            }
188                        }
189                    }
190                    Err(e) => {
191                        e.restore(py);
192                        f64::NAN
193                    }
194                }
195            };
196            let result = self.inner.integrate(&wrapped);
197            if PyErr::take(py).is_some() {
198                return Err(PyValueError::new_err(
199                    "integrand raised an exception (or returned non-float)",
200                ));
201            }
202            Ok(result)
203        })?;
204        Ok(r.into())
205    }
206
207    /// Latest accumulated estimate and error after `integrate`.
208    #[getter]
209    fn result(&self) -> PyIntegrateResult {
210        self.inner.result().into()
211    }
212
213    /// Number of completed iterations.
214    #[getter]
215    fn iterations(&self) -> usize {
216        self.inner.iterations()
217    }
218
219    fn __repr__(&self) -> String {
220        format!(
221            "Vegas(iterations={}, integral={:?})",
222            self.inner.iterations(),
223            self.inner.result().integral
224        )
225    }
226}
227
228/// Numerically integrate a one-dimensional Python callable `f` over `[a, b]`.
229///
230/// All options are keyword-only and optional (see
231/// [`Vegas`][PyVegas] for their meaning). Returns an
232/// [`IntegrateResult`][PyIntegrateResult].
233#[pyfunction]
234#[pyo3(signature = (f, a, b, *, n_bins=None, n_samples=None, iterations=None, learning_rate=None, seed=None))]
235#[allow(clippy::too_many_arguments)]
236pub fn integrate_1d(
237    f: &Bound<'_, PyAny>,
238    a: f64,
239    b: f64,
240    n_bins: Option<&Bound<'_, PyAny>>,
241    n_samples: Option<&Bound<'_, PyAny>>,
242    iterations: Option<&Bound<'_, PyAny>>,
243    learning_rate: Option<&Bound<'_, PyAny>>,
244    seed: Option<&Bound<'_, PyAny>>,
245) -> PyResult<PyIntegrateResult> {
246    if a.partial_cmp(&b) != Some(std::cmp::Ordering::Less) {
247        return Err(PyValueError::new_err(
248            "integration upper bound b must be > a",
249        ));
250    }
251    let opts = kwargs_to_opts(n_bins, n_samples, iterations, learning_rate, seed)?;
252    let r = Python::attach(|py| -> PyResult<IntegrateResult> {
253        let cb = f.clone();
254        let wrapped = |x: f64| -> f64 {
255            let arg = PyFloat::new(py, x);
256            match cb.call1((&arg,)) {
257                Ok(value) => match value.extract::<f64>() {
258                    Ok(v) => v,
259                    Err(e) => {
260                        e.restore(py);
261                        f64::NAN
262                    }
263                },
264                Err(e) => {
265                    e.restore(py);
266                    f64::NAN
267                }
268            }
269        };
270        let result = ocas_eval::numeric::integrate_1d(wrapped, a, b, opts);
271        if PyErr::take(py).is_some() {
272            return Err(PyValueError::new_err(
273                "integrand raised an exception (or returned non-float)",
274            ));
275        }
276        Ok(result)
277    })?;
278    Ok(r.into())
279}
ocas_py/numeric.rs

ocas_py/
numeric.rs
