libitofin 0.6.1

//! Sobol low-discrepancy sequence generator, ported from
//! `ql/math/randomnumbers/sobolrsg.{hpp,cpp}`.
//!
//! A Gray code counter and bitwise operations are used for very fast
//! sequence generation. The implementation relies on primitive polynomials
//! modulo two and direction-integer initializers from Peter Jaeckel's
//! "Monte Carlo Methods in Finance" plus the Sobol-Levitan, Lemieux, Joe-Kuo
//! and Kuo tables shipped with QuantLib (see `tables/`, generated by
//! `scripts/gen_sobol_tables.py`).

mod mt;
mod tables;

/// Number of primitive polynomials modulo two shipped with the library,
/// i.e. the maximum supported dimensionality.
pub const PPMT_MAX_DIM: usize = tables::PPMT_MAX_DIM;

/// Choice of free direction integers for dimensions 2 and above.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DirectionIntegers {
    Unit,
    Jaeckel,
    SobolLevitan,
    SobolLevitanLemieux,
    JoeKuoD5,
    JoeKuoD6,
    JoeKuoD7,
    Kuo,
    Kuo2,
    Kuo3,
}

struct InitializerTable {
    offsets: &'static [u32],
    parts: &'static [&'static [u32]],
}

impl InitializerTable {
    fn tabulated_dimensions(&self) -> usize {
        self.offsets.len() - 1
    }

    fn entry(&self, k: usize) -> &'static [u32] {
        let mut start = self.offsets[k] as usize;
        let len = (self.offsets[k + 1] - self.offsets[k]) as usize;
        for part in self.parts {
            if start < part.len() {
                return &part[start..start + len];
            }
            start -= part.len();
        }
        unreachable!("initializer table offsets exceed the table data length")
    }
}

fn initializer_table(direction_integers: DirectionIntegers) -> Option<InitializerTable> {
    let (offsets, parts) = match direction_integers {
        DirectionIntegers::Unit => return None,
        DirectionIntegers::Jaeckel => (tables::JAECKEL_OFFSETS, tables::JAECKEL_PARTS),
        DirectionIntegers::SobolLevitan => {
            (tables::SOBOL_LEVITAN_OFFSETS, tables::SOBOL_LEVITAN_PARTS)
        }
        DirectionIntegers::SobolLevitanLemieux => (tables::LEMIEUX_OFFSETS, tables::LEMIEUX_PARTS),
        DirectionIntegers::JoeKuoD5 => (tables::JOE_KUO_D5_OFFSETS, tables::JOE_KUO_D5_PARTS),
        DirectionIntegers::JoeKuoD6 => (tables::JOE_KUO_D6_OFFSETS, tables::JOE_KUO_D6_PARTS),
        DirectionIntegers::JoeKuoD7 => (tables::JOE_KUO_D7_OFFSETS, tables::JOE_KUO_D7_PARTS),
        DirectionIntegers::Kuo => (tables::KUO_OFFSETS, tables::KUO_PARTS),
        DirectionIntegers::Kuo2 => (tables::KUO2_OFFSETS, tables::KUO2_PARTS),
        DirectionIntegers::Kuo3 => (tables::KUO3_OFFSETS, tables::KUO3_PARTS),
    };
    Some(InitializerTable { offsets, parts })
}

fn next_polynomial(table: &[&'static [u32]], current_degree: &mut usize, index: &mut usize) -> u32 {
    if *index >= table[*current_degree - 1].len() {
        *current_degree += 1;
        *index = 0;
    }
    let polynomial = table[*current_degree - 1][*index];
    *index += 1;
    polynomial
}

/// Sobol low-discrepancy sequence generator.
pub struct SobolRsg {
    dimensionality: usize,
    sequence_counter: u32,
    first_draw: bool,
    sequence: Vec<f64>,
    integer_sequence: Vec<u32>,
    direction_integers: Vec<Vec<u32>>,
    use_gray_code: bool,
}

impl SobolRsg {
    /// Creates a generator using the Gray code counter, like QuantLib's
    /// default constructor arguments.
    ///
    /// Unlike QuantLib, a zero seed is used literally instead of drawing a
    /// clock-based seed; the seed only matters for dimensions beyond the
    /// tabulated initializers, where free direction integers are drawn from
    /// a seeded Mersenne twister.
    ///
    /// # Panics
    ///
    /// Panics if `dimensionality` is zero or exceeds [`PPMT_MAX_DIM`].
    pub fn new(dimensionality: usize, seed: u64, direction_integers: DirectionIntegers) -> Self {
        Self::with_gray_code(dimensionality, seed, direction_integers, true)
    }

    /// Creates a generator choosing the generating integer: the Gray code
    /// `G(n)` when `use_gray_code` is true, plain `n` otherwise.
    pub fn with_gray_code(
        dimensionality: usize,
        seed: u64,
        direction_integers: DirectionIntegers,
        use_gray_code: bool,
    ) -> Self {
        assert!(dimensionality > 0, "dimensionality must be greater than 0");
        assert!(
            dimensionality <= tables::PPMT_MAX_DIM,
            "dimensionality {dimensionality} exceeds the number of available primitive polynomials modulo two ({})",
            tables::PPMT_MAX_DIM
        );

        let mut degree = vec![0usize; dimensionality];
        let mut ppmt = vec![0u32; dimensionality];
        let use_alt_polynomials = matches!(
            direction_integers,
            DirectionIntegers::Kuo
                | DirectionIntegers::Kuo2
                | DirectionIntegers::Kuo3
                | DirectionIntegers::SobolLevitan
                | DirectionIntegers::SobolLevitanLemieux
        );
        let alt_limit = if use_alt_polynomials {
            tables::MAX_ALT_POLYNOMIALS
        } else {
            0
        };

        let mut current_degree = 1usize;
        let mut index = 0usize;
        let mut k = 1;
        while k < dimensionality.min(alt_limit) {
            ppmt[k] = next_polynomial(
                &tables::ALT_PRIMITIVE_POLYNOMIALS,
                &mut current_degree,
                &mut index,
            );
            degree[k] = current_degree;
            k += 1;
        }
        while k < dimensionality {
            ppmt[k] = next_polynomial(
                &tables::PRIMITIVE_POLYNOMIALS,
                &mut current_degree,
                &mut index,
            );
            degree[k] = current_degree;
            k += 1;
        }

        let mut directions = vec![vec![0u32; dimensionality]; 32];
        for (j, row) in directions.iter_mut().enumerate() {
            row[0] = 1u32 << (31 - j);
        }

        let max_tabulated = match initializer_table(direction_integers) {
            None => {
                for k in 1..dimensionality {
                    for l in 1..=degree[k] {
                        directions[l - 1][k] = 1u32 << (32 - l);
                    }
                }
                dimensionality
            }
            Some(table) => {
                let max_tabulated = table.tabulated_dimensions() + 1;
                #[allow(clippy::needless_range_loop)]
                for k in 1..dimensionality.min(max_tabulated) {
                    for (j, &coefficient) in table.entry(k - 1).iter().enumerate() {
                        directions[j][k] = coefficient << (32 - j - 1);
                    }
                }
                max_tabulated
            }
        };

        if dimensionality > max_tabulated {
            let mut rng = mt::MersenneTwister::new(seed);
            for k in max_tabulated..dimensionality {
                for l in 1..=degree[k] {
                    loop {
                        let u = rng.next_real();
                        let drawn = (u * f64::from(1u32 << l)) as u32;
                        if drawn & 1 != 0 {
                            directions[l - 1][k] = drawn << (32 - l);
                            break;
                        }
                    }
                }
            }
        }

        for k in 1..dimensionality {
            let gk = degree[k];
            for l in gk..32 {
                let mut n = directions[l - gk][k] >> gk;
                for j in 1..gk {
                    if (ppmt[k] >> (gk - j - 1)) & 1 != 0 {
                        n ^= directions[l - j][k];
                    }
                }
                n ^= directions[l - gk][k];
                directions[l][k] = n;
            }
        }

        let mut integer_sequence = vec![0u32; dimensionality];
        if use_gray_code {
            integer_sequence.copy_from_slice(&directions[0]);
        }

        Self {
            dimensionality,
            sequence_counter: 0,
            first_draw: true,
            sequence: vec![0.0; dimensionality],
            integer_sequence,
            direction_integers: directions,
            use_gray_code,
        }
    }

    /// Skips to the n-th sample in the low-discrepancy sequence and returns it.
    ///
    /// # Panics
    ///
    /// Panics if `skip` is `u32::MAX`, i.e. if the requested sample index
    /// exceeds the Sobol sequence period.
    pub fn skip_to(&mut self, skip: u32) -> &[u32] {
        let n = skip
            .checked_add(1)
            .expect("skip exceeds the Sobol sequence period");

        if self.use_gray_code {
            let ops = (f64::from(n).ln() / std::f64::consts::LN_2) as u32 + 1;
            let gray = n ^ (n >> 1);
            for (k, value) in self.integer_sequence.iter_mut().enumerate() {
                let mut integer = 0u32;
                for index in 0..ops as usize {
                    if (gray >> index) & 1 != 0 {
                        integer ^= self.direction_integers[index][k];
                    }
                }
                *value = integer;
            }
        } else {
            self.integer_sequence.fill(0);
            for (index, row) in self.direction_integers.iter().enumerate() {
                if (n >> index) & 1 != 0 {
                    for (value, &direction) in self.integer_sequence.iter_mut().zip(row) {
                        *value ^= direction;
                    }
                }
            }
        }

        self.sequence_counter = skip;
        &self.integer_sequence
    }

    /// Returns the next sample as raw 32-bit Sobol integers.
    ///
    /// # Panics
    ///
    /// Panics when the sequence counter wraps around ("period exceeded").
    pub fn next_int32_sequence(&mut self) -> &[u32] {
        if !self.use_gray_code {
            self.skip_to(self.sequence_counter);
            if self.first_draw {
                self.first_draw = false;
            } else {
                self.sequence_counter = self
                    .sequence_counter
                    .checked_add(1)
                    .expect("period exceeded");
            }
            return &self.integer_sequence;
        }

        if self.first_draw {
            self.first_draw = false;
            return &self.integer_sequence;
        }

        self.sequence_counter = self
            .sequence_counter
            .checked_add(1)
            .expect("period exceeded");

        let mut n = self.sequence_counter;
        let mut j = 0;
        while n & 1 != 0 {
            n >>= 1;
            j += 1;
        }
        for (value, &direction) in self
            .integer_sequence
            .iter_mut()
            .zip(&self.direction_integers[j])
        {
            *value ^= direction;
        }
        &self.integer_sequence
    }

    /// Returns the next sample normalized to `(0, 1)`.
    pub fn next_sequence(&mut self) -> &[f64] {
        self.next_int32_sequence();
        for (value, &integer) in self.sequence.iter_mut().zip(&self.integer_sequence) {
            *value = f64::from(integer) * (0.5 / 2_147_483_648.0);
        }
        &self.sequence
    }

    /// Returns the last sample normalized to `(0, 1)`.
    pub fn last_sequence(&self) -> &[f64] {
        &self.sequence
    }

    /// Returns the dimensionality of the sequence.
    pub fn dimension(&self) -> usize {
        self.dimensionality
    }
}

#[cfg(test)]
mod tests {
    use super::{DirectionIntegers, PPMT_MAX_DIM, SobolRsg, initializer_table, tables};

    const EXPECTED_POLYNOMIAL_COUNTS: [usize; 18] = [
        1, 1, 2, 2, 6, 6, 18, 16, 48, 60, 176, 144, 630, 756, 1800, 2048, 7710, 7776,
    ];

    const VAN_DER_CORPUT_SEQUENCE_MODULO_TWO: [f64; 31] = [
        0.50000, 0.75000, 0.25000, 0.37500, 0.87500, 0.62500, 0.12500, 0.18750, 0.68750, 0.93750,
        0.43750, 0.31250, 0.81250, 0.56250, 0.06250, 0.09375, 0.59375, 0.84375, 0.34375, 0.46875,
        0.96875, 0.71875, 0.21875, 0.15625, 0.65625, 0.90625, 0.40625, 0.28125, 0.78125, 0.53125,
        0.03125,
    ];

    #[test]
    fn polynomials_modulo_two_counts() {
        let mut total = 0;
        for (i, polynomials) in tables::PRIMITIVE_POLYNOMIALS.iter().enumerate() {
            assert_eq!(
                polynomials.len(),
                EXPECTED_POLYNOMIAL_COUNTS[i],
                "only {} polynomials in degree {} instead of {}",
                polynomials.len(),
                i + 1,
                EXPECTED_POLYNOMIAL_COUNTS[i]
            );
            total += polynomials.len();
        }
        assert_eq!(total, PPMT_MAX_DIM);
    }

    #[test]
    fn sobol_max_dimensionality() {
        let dimensionality = PPMT_MAX_DIM;
        let mut rsg = SobolRsg::new(dimensionality, 123456, DirectionIntegers::Jaeckel);
        for _ in 0..100 {
            let point = rsg.next_sequence();
            assert_eq!(point.len(), dimensionality);
        }
    }

    #[test]
    fn sobol_homogeneity() {
        let tolerance = 1.0e-15;
        let dimensionality = 33;
        let mut rsg = SobolRsg::new(dimensionality, 123456, DirectionIntegers::Jaeckel);
        let mut sums = vec![0.0f64; dimensionality];
        let mut k = 0usize;
        for j in 1..5u32 {
            let points = 2usize.pow(j) - 1;
            while k < points {
                for (sum, &x) in sums.iter_mut().zip(rsg.next_sequence()) {
                    *sum += x;
                }
                k += 1;
            }
            for (i, sum) in sums.iter().enumerate() {
                let mean = sum / k as f64;
                let error = (mean - 0.5).abs();
                assert!(
                    error <= tolerance,
                    "dimension {}: mean {mean} at the end of cycle {j} is not 0.5 (error = {error})",
                    i + 1
                );
            }
        }
    }

    #[test]
    fn sobol_van_der_corput_first_dimension() {
        let tolerance = 1.0e-15;
        let mut rsg = SobolRsg::new(1, 0, DirectionIntegers::Jaeckel);
        for (i, &expected) in VAN_DER_CORPUT_SEQUENCE_MODULO_TWO.iter().enumerate() {
            let point = rsg.next_sequence();
            let error = (point[0] - expected).abs();
            assert!(
                error <= tolerance,
                "draw {}: {} is not in the van der Corput sequence modulo two: it should have been {expected} (error = {error})",
                i + 1,
                point[0]
            );
        }
    }

    #[test]
    fn tabulated_family_initializer_entries() {
        let cases: [(DirectionIntegers, usize, &[u32]); 6] = [
            (
                DirectionIntegers::Kuo,
                4925,
                &[
                    1, 3, 7, 5, 23, 9, 25, 167, 241, 755, 1789, 1977, 6325, 5425, 23247, 64181,
                ],
            ),
            (
                DirectionIntegers::Kuo2,
                3946,
                &[
                    1, 3, 3, 3, 19, 57, 23, 119, 75, 435, 1023, 2357, 4671, 14581, 15343, 2383,
                ],
            ),
            (
                DirectionIntegers::Kuo3,
                4586,
                &[
                    1, 3, 7, 3, 25, 9, 25, 117, 315, 497, 1235, 1969, 7981, 9743, 20951, 30635,
                ],
            ),
            (
                DirectionIntegers::JoeKuoD5,
                1999,
                &[
                    1, 1, 3, 15, 21, 55, 17, 57, 449, 811, 519, 2329, 7607, 4255, 2845,
                ],
            ),
            (
                DirectionIntegers::JoeKuoD6,
                21200,
                &[
                    1, 1, 7, 11, 15, 7, 37, 239, 337, 245, 1557, 3681, 7357, 9639, 27367, 26869,
                    114603, 86317,
                ],
            ),
            (
                DirectionIntegers::JoeKuoD7,
                1899,
                &[
                    1, 1, 1, 3, 23, 17, 17, 39, 173, 1013, 527, 2563, 3623, 10049, 10919,
                ],
            ),
        ];
        for (family, dimensions, last_entry) in cases {
            let table = initializer_table(family).expect("family is tabulated");
            assert_eq!(table.tabulated_dimensions(), dimensions, "{family:?}");
            assert_eq!(table.entry(0), &[1], "{family:?} first entry");
            assert_eq!(
                table.entry(dimensions - 1),
                last_entry,
                "{family:?} last entry"
            );
        }
    }

    #[test]
    fn joe_kuo_d6_split_table_is_consistent() {
        let table = initializer_table(DirectionIntegers::JoeKuoD6).expect("family is tabulated");
        assert!(
            table.parts.len() > 1,
            "expected the JoeKuoD6 data to be split across generated files"
        );
        let parts_length: usize = table.parts.iter().map(|part| part.len()).sum();
        let offsets_tail = table.offsets[table.offsets.len() - 1] as usize;
        assert_eq!(offsets_tail, parts_length);
        let entries_length: usize = (0..table.tabulated_dimensions())
            .map(|k| table.entry(k).len())
            .sum();
        assert_eq!(entries_length, parts_length);
    }

    #[test]
    fn tabulated_families_construct_and_draw() {
        let families = [
            DirectionIntegers::Kuo,
            DirectionIntegers::Kuo2,
            DirectionIntegers::Kuo3,
            DirectionIntegers::JoeKuoD5,
            DirectionIntegers::JoeKuoD6,
            DirectionIntegers::JoeKuoD7,
        ];
        for family in families {
            let mut rsg = SobolRsg::new(PPMT_MAX_DIM, 42, family);
            for _ in 0..5 {
                let point = rsg.next_sequence();
                assert_eq!(point.len(), PPMT_MAX_DIM, "{family:?}");
                assert!(
                    point.iter().all(|x| (0.0..1.0).contains(x)),
                    "{family:?} produced a point outside the unit interval"
                );
            }
        }
    }

    #[test]
    fn sobol_skipping() {
        let seed = 42;
        let dimensionalities = [1usize, 10, 100, 1000];
        let skips = [0u32, 1, 42, 512, 100_000];
        let families = [
            DirectionIntegers::Unit,
            DirectionIntegers::Jaeckel,
            DirectionIntegers::SobolLevitan,
            DirectionIntegers::SobolLevitanLemieux,
        ];

        for family in families {
            for dimensionality in dimensionalities {
                for skip in skips {
                    let mut rsg1 = SobolRsg::new(dimensionality, seed, family);
                    for _ in 0..skip {
                        rsg1.next_int32_sequence();
                    }
                    let mut rsg2 = SobolRsg::new(dimensionality, seed, family);
                    rsg2.skip_to(skip);
                    for _ in 0..100 {
                        let s1 = rsg1.next_int32_sequence().to_vec();
                        let s2 = rsg2.next_int32_sequence();
                        assert_eq!(
                            s1, s2,
                            "mismatch after skipping: family {family:?}, size {dimensionality}, skipped {skip}"
                        );
                    }
                }
            }
        }
    }
}