use crate::domain;
use egobox_doe::{LhsKind, SamplingMethod};
use egobox_ego::gpmix::mixint::MixintContext;
use numpy::{IntoPyArray, PyArray2};
use pyo3::prelude::*;
use pyo3_stub_gen::derive::{gen_stub_pyclass_enum, gen_stub_pyfunction};
#[gen_stub_pyclass_enum]
#[pyclass(eq, eq_int, rename_all = "SCREAMING_SNAKE_CASE")]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Sampling {
Lhs = 1,
FullFactorial = 2,
Random = 3,
LhsClassic = 4,
LhsCentered = 5,
LhsMaximin = 6,
LhsCenteredMaximin = 7,
}
#[gen_stub_pyfunction]
#[pyfunction]
#[pyo3(signature = (method, xspecs, n_samples, seed=None))]
pub fn sampling(
py: Python<'_>,
method: Sampling,
xspecs: PyObject,
n_samples: usize,
seed: Option<u64>,
) -> Bound<'_, PyArray2<f64>> {
let xtypes: Vec<egobox_ego::XType> = domain::parse(py, xspecs);
let mixin = MixintContext::new(&xtypes);
let doe = match method {
Sampling::Lhs => Box::new(mixin.create_lhs_sampling(LhsKind::default(), seed))
as Box<dyn SamplingMethod<_>>,
Sampling::LhsClassic => Box::new(mixin.create_lhs_sampling(LhsKind::Classic, seed))
as Box<dyn SamplingMethod<_>>,
Sampling::LhsMaximin => Box::new(mixin.create_lhs_sampling(LhsKind::Maximin, seed))
as Box<dyn SamplingMethod<_>>,
Sampling::LhsCentered => Box::new(mixin.create_lhs_sampling(LhsKind::Centered, seed))
as Box<dyn SamplingMethod<_>>,
Sampling::LhsCenteredMaximin => {
Box::new(mixin.create_lhs_sampling(LhsKind::CenteredMaximin, seed))
as Box<dyn SamplingMethod<_>>
}
Sampling::FullFactorial => Box::new(mixin.create_ffact_sampling()),
Sampling::Random => {
Box::new(mixin.create_rand_sampling(seed)) as Box<dyn SamplingMethod<_>>
}
}
.sample(n_samples);
doe.into_pyarray(py)
}
#[gen_stub_pyfunction]
#[pyfunction]
#[pyo3(signature = (xspecs, n_samples, seed=None))]
pub(crate) fn lhs(
py: Python,
xspecs: PyObject,
n_samples: usize,
seed: Option<u64>,
) -> Bound<PyArray2<f64>> {
sampling(py, Sampling::Lhs, xspecs, n_samples, seed)
}