polydat 0.2.0

// Copyright 2024-2026 Jonathan Shook
// SPDX-License-Identifier: Apache-2.0

//! `Sobol` strategy — spec §3.6.
//!
//! Low-discrepancy (digital `(t,s)`) sequence built from Joe-Kuo
//! direction numbers and generated by the Antonov-Saleev Gray-code
//! recurrence. Native to continuous K-D boxes; discretizes to integer
//! multi-indices for discrete inputs.
//!
//! Dimension 1 uses the degenerate direction numbers `V_i = 2^(32-i)`
//! (the base-2 van der Corput in Gray-code order); dimensions 2..
//! use the primitive-polynomial degree `s`, polynomial coefficient
//! bits `a`, and initial direction integers `m_1..m_s` from the
//! `new-joe-kuo-6.21201` initialiser table ([`JOE_KUO_INIT`]). The
//! embedded table covers the first [`MAX_SOBOL_DIM`] dimensions; a
//! comprehension with more axes is **rejected** by
//! [`Sobol::accepts_input`] (the validation layer reports
//! `strategy rejects input`) rather than silently degrading.
//!
//! ## References
//!
//! - I. M. Sobol', "On the distribution of points in a cube and the
//!   approximate evaluation of integrals," *USSR Comp. Math. and
//!   Math. Phys.* 7(4) (1967), 86–112.
//!   doi:[10.1016/0041-5553(67)90144-9](https://doi.org/10.1016/0041-5553(67)90144-9).
//! - S. Joe and F. Y. Kuo, "Constructing Sobol sequences with better
//!   two-dimensional projections," *SIAM J. Sci. Comput.* 30(5)
//!   (2008), 2635–2654.
//!   doi:[10.1137/070709359](https://doi.org/10.1137/070709359). The
//!   `new-joe-kuo-6.21201` initialiser file
//!   (<https://web.maths.unsw.edu.au/~fkuo/sobol/>) is the source of
//!   [`JOE_KUO_INIT`].
//! - I. A. Antonov and V. M. Saleev, "An economic method of computing
//!   LP_τ-sequences," *USSR Comp. Math. and Math. Phys.* 19(1)
//!   (1979), 252–256. The Gray-code recurrence in `sobol_coord`.
//! - The published 1-D prefix `1/2, 3/4, 1/4, 3/8, 7/8, 5/8, 1/8, …`
//!   and the 2-D points `(½,½),(¾,¼),(¼,¾),(⅜,⅜),…` are cross-checked
//!   in `tests::matches_published_joe_kuo_points`.

use super::{
    EvaluatedInput, MultiIndex, Strategy, Tuple, index_fn_dim, index_fn_size,
    index_fn_supports_lookup, multi_index_to_flat,
};
use crate::iteration::comprehension::metadata::IndexFn;
use crate::iteration::comprehension::strategy::StrategyName;

pub struct Sobol;

/// Bit-width of the direction numbers (`u32` fixed point).
const BITS: u32 = 32;

/// Joe-Kuo (2008) direction-number initialisers for Sobol
/// dimensions 2, 3, 4, …, one entry per dimension as
/// `(s, a, &[m_1..m_s])`, transcribed from the `new-joe-kuo-6.21201`
/// table (Joe & Kuo 2008). `s` is the degree of the primitive
/// polynomial, `a` its interior coefficient bits, and `m_i` the
/// initial direction integers. Dimension 1 is the special case
/// (`V_i = 2^(32-i)`) handled in `direction_numbers`, so index 0
/// here is Sobol dimension 2.
pub const JOE_KUO_INIT: &[(u32, u32, &[u32])] = &[
    (1, 0, &[1]),                  // d=2
    (2, 1, &[1, 3]),               // d=3
    (3, 1, &[1, 3, 1]),            // d=4
    (3, 2, &[1, 1, 1]),            // d=5
    (4, 1, &[1, 1, 3, 3]),         // d=6
    (4, 4, &[1, 3, 5, 13]),        // d=7
    (5, 2, &[1, 1, 5, 5, 17]),     // d=8
    (5, 4, &[1, 1, 5, 5, 5]),      // d=9
    (5, 7, &[1, 1, 7, 11, 19]),    // d=10
    (5, 11, &[1, 1, 5, 1, 1]),     // d=11
    (5, 13, &[1, 1, 1, 3, 11]),    // d=12
    (5, 14, &[1, 3, 5, 5, 31]),    // d=13
];

/// Highest Sobol dimension the embedded [`JOE_KUO_INIT`] table
/// supports (dimension 1 + one entry per table row). A comprehension
/// with more axes is rejected by [`Sobol::accepts_input`].
pub const MAX_SOBOL_DIM: usize = JOE_KUO_INIT.len() + 1;

/// The 32 direction numbers `V_1..V_32` (MSB-aligned `u32`) for the
/// given 0-based dimension. Dimension 0 (Sobol dim 1) is
/// `V_i = 2^(32-i)`; higher dims apply the Joe-Kuo recurrence
/// `V_i = V_{i-s} XOR (V_{i-s} >> s) XOR Σ a_k V_{i-k}` over the
/// primitive-polynomial coefficients. Returns `[_; 33]` indexed
/// `1..=32` (slot 0 unused).
fn direction_numbers(dim0: usize) -> [u32; BITS as usize + 1] {
    let mut v = [0u32; BITS as usize + 1];
    if dim0 == 0 {
        for i in 1..=BITS {
            v[i as usize] = 1u32 << (BITS - i);
        }
        return v;
    }
    let (s, a, m) = JOE_KUO_INIT[dim0 - 1];
    for i in 1..=s {
        v[i as usize] = m[(i - 1) as usize] << (BITS - i);
    }
    for i in (s + 1)..=BITS {
        let prev = v[(i - s) as usize];
        let mut val = prev ^ (prev >> s);
        for k in 1..s {
            if (a >> (s - 1 - k)) & 1 == 1 {
                val ^= v[(i - k) as usize];
            }
        }
        v[i as usize] = val;
    }
    v
}

/// One coordinate of the `i`-th Sobol point in `[0, 1)`, via the
/// Antonov-Saleev Gray-code recurrence: `X_i = ⊕_{bit k of gray(i)}
/// V_{k+1}`, where `gray(i) = i ⊕ (i >> 1)`.
fn sobol_coord(i: u64, dirs: &[u32; BITS as usize + 1]) -> f64 {
    let mut g = i ^ (i >> 1);
    let mut x: u32 = 0;
    let mut bit = 1usize;
    while g != 0 && bit <= BITS as usize {
        if g & 1 == 1 {
            x ^= dirs[bit];
        }
        g >>= 1;
        bit += 1;
    }
    (x as f64) / (1u64 << BITS) as f64
}

impl Strategy for Sobol {
    fn name(&self) -> StrategyName {
        StrategyName::Sobol
    }

    fn accepts_input(&self, idx: Option<&IndexFn>) -> bool {
        // Reject more axes than the embedded Joe-Kuo table covers —
        // the validation layer surfaces this as `strategy rejects
        // input` rather than the strategy silently degrading.
        match idx {
            None => false,
            Some(i) => index_fn_dim(i) <= MAX_SOBOL_DIM,
        }
    }

    fn has_closed_form_for(&self, _idx: &IndexFn) -> bool {
        true
    }

    fn apply(&self, input: &EvaluatedInput, truncation: Option<u64>) -> Vec<Tuple> {
        if index_fn_supports_lookup(&input.index_fn) {
            let mis = sobol_multi_indices(&input.index_fn, truncation);
            mis.into_iter()
                .filter_map(|mi| multi_index_to_flat(&input.index_fn, &mi))
                .filter_map(|flat| input.tuples.get(flat).cloned())
                .collect()
        } else {
            naive_sobol_over_tuples(&input.tuples, truncation)
        }
    }
}

fn naive_sobol_over_tuples(input: &[Tuple], truncation: Option<u64>) -> Vec<Tuple> {
    let total = input.len() as u64;
    if total == 0 {
        return Vec::new();
    }
    let n = match truncation {
        Some(t) => t.min(total),
        None => total,
    };
    let dirs0 = direction_numbers(0); // 1-D Sobol direction numbers
    let mut seen = std::collections::HashSet::new();
    let mut out = Vec::with_capacity(n as usize);
    let mut i = 1u64;
    let mut attempts = 0u64;
    let max_attempts = total.saturating_mul(8).max(64);
    while (out.len() as u64) < n && attempts < max_attempts {
        let pt = sobol_coord(i, &dirs0);
        let idx = (pt * total as f64).floor() as u64;
        let idx = idx.min(total - 1);
        if seen.insert(idx) {
            out.push(input[idx as usize].clone());
        }
        i += 1;
        attempts += 1;
    }
    out
}

pub(crate) fn sobol_multi_indices(idx: &IndexFn, truncation: Option<u64>) -> Vec<MultiIndex> {
    let dim = index_fn_dim(idx);
    if dim == 0 {
        return Vec::new();
    }
    let total = index_fn_size(idx);
    let n = match (truncation, total) {
        (Some(t), 0) => t,
        (Some(t), tot) => t.min(tot),
        (None, 0) => return Vec::new(),
        (None, tot) => tot,
    };
    if n == 0 {
        return Vec::new();
    }

    let axis_sizes = axis_sizes_for(idx, dim);
    // Direction numbers per axis, computed once. `accepts_input`
    // guarantees `dim <= MAX_SOBOL_DIM`, but clamp defensively so a
    // direct call can't index past the table.
    let dirs: Vec<[u32; BITS as usize + 1]> = (0..dim)
        .map(|d| direction_numbers(d.min(MAX_SOBOL_DIM - 1)))
        .collect();
    let mut out = Vec::with_capacity(n as usize);
    let mut i = 1u64;
    let mut seen_discrete = std::collections::HashSet::new();
    let max_attempts = n.saturating_mul(8).max(256);
    let mut attempts = 0u64;
    let is_continuous = matches!(idx, IndexFn::Continuous { .. } | IndexFn::Hybrid { .. });

    while (out.len() as u64) < n && attempts < max_attempts {
        let pt: Vec<f64> = dirs.iter().map(|d| sobol_coord(i, d)).collect();
        let mi = point_to_multi_index(&pt, &axis_sizes, idx);
        // Continuous points are never deduped; discrete points
        // dedup through the set (short-circuit skips the insert
        // for continuous, so the set stays empty there).
        if is_continuous || seen_discrete.insert(mi.clone()) {
            out.push(mi);
        }
        i += 1;
        attempts += 1;
    }
    out
}

fn axis_sizes_for(idx: &IndexFn, dim: usize) -> Vec<u64> {
    match idx {
        IndexFn::Lattice { axis_sizes } | IndexFn::Modular { axis_sizes } => axis_sizes.clone(),
        IndexFn::Lockstep { length } => vec![*length],
        IndexFn::Concatenation { segment_sizes } => vec![segment_sizes.iter().sum()],
        IndexFn::Continuous { .. } => vec![u64::MAX; dim],
        IndexFn::Hybrid {
            discrete_axes,
            continuous_axes,
            ..
        } => {
            let mut s = discrete_axes.clone();
            s.extend(continuous_axes.iter().map(|_| u64::MAX));
            s
        }
    }
}

fn point_to_multi_index(pt: &[f64], axis_sizes: &[u64], idx: &IndexFn) -> MultiIndex {
    match idx {
        IndexFn::Continuous { .. } => pt
            .iter()
            .map(|f| (f * (1u64 << 53) as f64) as u64)
            .collect(),
        IndexFn::Hybrid { discrete_axes, .. } => {
            let mut mi = Vec::with_capacity(pt.len());
            for (i, f) in pt.iter().enumerate() {
                if i < discrete_axes.len() {
                    let size = discrete_axes[i];
                    mi.push(((f * size as f64).floor() as u64).min(size.saturating_sub(1)));
                } else {
                    mi.push((f * (1u64 << 53) as f64) as u64);
                }
            }
            mi
        }
        _ => pt
            .iter()
            .zip(axis_sizes.iter())
            .map(|(f, size)| ((f * *size as f64).floor() as u64).min(size.saturating_sub(1)))
            .collect(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Helper: the `i`-th Sobol point over `dim` continuous axes.
    fn point(i: u64, dim: usize) -> Vec<f64> {
        let dirs: Vec<_> = (0..dim).map(direction_numbers).collect();
        dirs.iter().map(|d| sobol_coord(i, d)).collect()
    }

    #[test]
    fn matches_published_joe_kuo_points() {
        // True 1-D Sobol prefix (Sobol' 1967; Antonov-Saleev Gray-code
        // recurrence): 1/2, 3/4, 1/4, 3/8, 7/8, 5/8, 1/8.
        let one_d = [
            1.0 / 2.0, 3.0 / 4.0, 1.0 / 4.0, 3.0 / 8.0, 7.0 / 8.0, 5.0 / 8.0, 1.0 / 8.0,
        ];
        for (k, want) in one_d.iter().enumerate() {
            let got = point(k as u64 + 1, 1)[0];
            assert!((got - want).abs() < 1e-12, "1-D Sobol x_{} = {got}, want {want}", k + 1);
        }

        // The published Joe-Kuo 2-D Sobol points (dims 1 & 2):
        // (½,½) (¾,¼) (¼,¾) (⅜,⅜) (⅞,⅞) (⅝,⅛) (⅛,⅝).
        let two_d = [
            (1.0 / 2.0, 1.0 / 2.0),
            (3.0 / 4.0, 1.0 / 4.0),
            (1.0 / 4.0, 3.0 / 4.0),
            (3.0 / 8.0, 3.0 / 8.0),
            (7.0 / 8.0, 7.0 / 8.0),
            (5.0 / 8.0, 1.0 / 8.0),
            (1.0 / 8.0, 5.0 / 8.0),
        ];
        for (k, (wx, wy)) in two_d.iter().enumerate() {
            let p = point(k as u64 + 1, 2);
            assert!((p[0] - wx).abs() < 1e-12, "2-D Sobol x_{}.0 = {}, want {wx}", k + 1, p[0]);
            assert!((p[1] - wy).abs() < 1e-12, "2-D Sobol x_{}.1 = {}, want {wy}", k + 1, p[1]);
        }
    }

    #[test]
    fn direction_numbers_dim1_are_powers_of_two() {
        let v = direction_numbers(0);
        for i in 1..=BITS {
            assert_eq!(v[i as usize], 1u32 << (BITS - i));
        }
    }

    #[test]
    fn rejects_more_axes_than_table_supports() {
        // 14 axes > MAX_SOBOL_DIM (13): Sobol must decline so the
        // validation layer errors rather than the strategy degrading.
        let too_many = IndexFn::Lattice { axis_sizes: vec![4; MAX_SOBOL_DIM + 1] };
        assert!(!Sobol.accepts_input(Some(&too_many)));
        let ok = IndexFn::Lattice { axis_sizes: vec![4; MAX_SOBOL_DIM] };
        assert!(Sobol.accepts_input(Some(&ok)));
    }

    #[test]
    fn deterministic() {
        let idx = IndexFn::Lattice { axis_sizes: vec![20, 20] };
        let a = sobol_multi_indices(&idx, Some(10));
        let b = sobol_multi_indices(&idx, Some(10));
        assert_eq!(a, b);
    }

    #[test]
    fn produces_unique_discrete() {
        let idx = IndexFn::Lattice { axis_sizes: vec![50, 50] };
        let out = sobol_multi_indices(&idx, Some(20));
        assert_eq!(out.len(), 20);
        let mut seen = std::collections::HashSet::new();
        for mi in &out {
            assert!(seen.insert(mi.clone()));
        }
    }

    #[test]
    fn continuous_box_draws() {
        use crate::iteration::comprehension::cardinality::{Interval, ProductMeasure};
        let idx = IndexFn::Continuous {
            intervals: vec![Interval::closed(0.0, 1.0), Interval::closed(0.0, 1.0)],
            measure: ProductMeasure::Uniform,
        };
        let out = sobol_multi_indices(&idx, Some(50));
        assert_eq!(out.len(), 50);
    }
}