use crate::error::{BuildError, SampleError};
use crate::multivar::support::HyperRect;
use crate::multivar::traits::{HasSupportNd, PdfNd};
use rand::{Rng, RngExt};
#[derive(Debug, Clone, Copy)]
pub struct RejectionOptions {
pub max_trials: usize,
}
impl Default for RejectionOptions {
fn default() -> Self {
Self { max_trials: 100_000 }
}
}
pub struct RejectionSamplerNd<P> {
pdf: P,
support: HyperRect,
pdf_max: f64,
opts: RejectionOptions,
}
impl<P> RejectionSamplerNd<P>
where
P: PdfNd + HasSupportNd,
{
pub fn new(pdf: P, pdf_max: f64, opts: RejectionOptions) -> Result<Self, BuildError> {
if !pdf_max.is_finite() || pdf_max <= 0.0 {
return Err(BuildError::InvalidPdfMax);
}
let support = pdf.support().clone();
support.validate()?;
Ok(Self {
pdf,
support,
pdf_max,
opts,
})
}
pub fn dim(&self) -> usize {
self.support.dim()
}
pub fn support(&self) -> &HyperRect {
&self.support
}
pub fn pdf_max(&self) -> f64 {
self.pdf_max
}
pub fn sample(&self) -> Result<Vec<f64>, SampleError> {
self.sample_with_rng(&mut rand::rng())
}
pub fn sample_with_rng<R: Rng + ?Sized>(&self, rng: &mut R) -> Result<Vec<f64>, SampleError> {
let dim = self.dim();
let mut x = vec![0.0; dim];
for _ in 0..self.opts.max_trials {
for i in 0..dim {
let u: f64 = rng.random();
x[i] = self.support.lo[i] + u * (self.support.hi[i] - self.support.lo[i]);
}
let fx = self.pdf.pdf(&x);
if !fx.is_finite() || fx < 0.0 {
continue;
}
let u: f64 = rng.random();
if u * self.pdf_max <= fx {
return Ok(x);
}
}
Err(SampleError::RejectionSamplingFailed)
}
}