solver/numeric/

basis.rs

use core::f64;

/// Trait defining a set of basis functions $\Phi(x) = [\phi_1(x), ..., \phi_K(x)]$.
///
/// Used for linear function approximation $V(x) \approx \sum w_i \phi_i(x)$.
pub trait Basis<const N: usize>: Send + Sync {
    /// Returns the number of basis functions (K).
    fn num_features(&self) -> usize;

    /// Evaluates all basis functions at `state` and writes them into `buffer`.
    /// `buffer` must have length at least `num_features()`.
    fn evaluate(&self, state: &[f64; N], buffer: &mut [f64]);

    /// Evaluates all basis functions and appends them to the provided vector.
    /// This avoids zero-initialization overhead.
    fn evaluate_buffers(&self, state: &[f64; N], buffer: &mut Vec<f64>) {
        let start_len = buffer.len();
        let n = self.num_features();
        // Reserve space
        buffer.reserve(n);

        // Unsafe implementation to avoid zeroing, or just standard push
        // We will default to a standard efficient implementation here using resize if needed
        // but optimal path is implementation specific.

        // Fallback: resize and call evaluate (allocates zeros if resizing up)
        buffer.resize(start_len + n, 0.0);
        self.evaluate(state, &mut buffer[start_len..]);
    }

    /// Returns the dot product $w \cdot \Phi(x)$ without allocating the full feature vector.
    /// Default implementation allocates, overrides can be optimization.
    fn eval_dot(&self, state: &[f64; N], weights: &[f64]) -> f64 {
        let n = self.num_features();
        let mut buffer = vec![0.0; n]; // Allocation! Override me.
        self.evaluate(state, &mut buffer);

        let mut sum = 0.0;
        for i in 0..n {
            sum += buffer[i] * weights[i];
        }
        sum
    }
}

/// Standard Polynomial Basis (Power, Hermite, Chebyshev, etc.)
///
/// Includes:
/// - Full interactions up to degree 2 (Quadratic)
/// - Diagonal terms up to `degree`
#[derive(Clone, Debug)]
pub struct PolynomialBasis {
    pub degree: usize,
    pub poly_type: PolynomialType,
    pub clamp_range: Option<(f64, f64)>,
}

#[derive(Clone, Copy, Debug)]
pub enum PolynomialType {
    Power,
    Hermite,
    Chebyshev,
    Laguerre,
}

impl PolynomialBasis {
    pub fn new(degree: usize, poly_type: PolynomialType) -> Self {
        Self {
            degree,
            poly_type,
            clamp_range: None,
        }
    }

    pub fn with_clamp(mut self, min: f64, max: f64) -> Self {
        self.clamp_range = Some((min, max));
        self
    }
}

/// Wrapper for any basis that applies coordinate-wise scaling before evaluation.
///
/// $V(x) \approx \sum w_i \phi_i(S x)$ where $S$ is a diagonal scaling matrix.
/// Used to map problem domains (e.g. $[0, 100]$) to basis domains (e.g. $[-2, 2]$).
#[derive(Clone, Debug)]
pub struct ScaledBasis<const N: usize, B> {
    pub inner: B,
    pub scales: Vec<f64>,
    pub center: [f64; N],
}

impl<const N: usize, B> ScaledBasis<N, B> {
    pub fn new(inner: B, scales: Vec<f64>) -> Self {
        Self {
            inner,
            scales,
            center: [0.0; N],
        }
    }

    pub fn with_center(mut self, center: [f64; N]) -> Self {
        self.center = center;
        self
    }
}

impl<const N: usize, B: Basis<N>> Basis<N> for ScaledBasis<N, B> {
    fn num_features(&self) -> usize {
        self.inner.num_features()
    }

    fn evaluate(&self, state: &[f64; N], buffer: &mut [f64]) {
        let mut scaled = *state;
        for (i, val) in scaled.iter_mut().enumerate() {
            let s = if i < self.scales.len() {
                self.scales[i]
            } else {
                *self.scales.first().unwrap_or(&1.0)
            };
            *val = (*val - self.center[i]) * s;
        }
        self.inner.evaluate(&scaled, buffer);
    }

    fn evaluate_buffers(&self, state: &[f64; N], buffer: &mut Vec<f64>) {
        let mut scaled = *state;
        for (i, val) in scaled.iter_mut().enumerate() {
            let s = if i < self.scales.len() {
                self.scales[i]
            } else {
                *self.scales.first().unwrap_or(&1.0)
            };
            *val = (*val - self.center[i]) * s;
        }
        self.inner.evaluate_buffers(&scaled, buffer);
    }

    fn eval_dot(&self, state: &[f64; N], weights: &[f64]) -> f64 {
        let mut scaled = *state;
        for (i, val) in scaled.iter_mut().enumerate() {
            let s = if i < self.scales.len() {
                self.scales[i]
            } else {
                *self.scales.first().unwrap_or(&1.0)
            };
            *val = (*val - self.center[i]) * s;
        }
        self.inner.eval_dot(&scaled, weights)
    }
}

impl<const N: usize> Basis<N> for PolynomialBasis {
    fn num_features(&self) -> usize {
        if self.degree == 0 {
            return 1;
        } else if self.degree == 1 {
            return 1 + N;
        }

        // 1 (Bias) + N (Linear) + N*(N+1)/2 (Quadratic Interactions)
        let mut base_count = 1 + N + N * (N + 1) / 2;

        // Higher order diagonal terms (Degree 3+)
        if self.degree > 2 {
            base_count += (self.degree - 2) * N;
        }
        base_count
    }

    fn evaluate(&self, state: &[f64; N], buffer: &mut [f64]) {
        // Pre-compute univariate polynomials
        const MAX_DEG: usize = 6;
        let effective_deg = self.degree.min(MAX_DEG);

        // Reverted to raw unscaled evaluation.
        let mut uni = [[0.0; MAX_DEG + 1]; N];
        for i in 0..N {
            let x = if let Some((min, max)) = self.clamp_range {
                state[i].clamp(min, max)
            } else {
                state[i]
            };
            compute_univariate(x, effective_deg, self.poly_type, &mut uni[i][..]);
        }

        let mut idx = 0;

        // 1. Bias (Degree >= 0)
        buffer[idx] = 1.0;
        idx += 1;

        if self.degree < 1 {
            return;
        }

        // 2. Linear (Degree >= 1)
        for u_row in &uni {
            buffer[idx] = u_row[1];
            idx += 1;
        }

        if self.degree < 2 {
            return;
        }

        // 3. Quadratic & Interactions (Degree >= 2)
        // Optimized loop
        for (i, u_row) in uni.iter().enumerate() {
            let u_i = u_row[1];

            // i == j case first
            buffer[idx] = u_row[2]; // Diagonal
            idx += 1;

            // j > i cases (interactions)
            for u_j in uni.iter().skip(i + 1) {
                buffer[idx] = u_i * u_j[1];
                idx += 1;
            }
        }

        // 4. Higher Order (Diagonal only)
        for k in 3..=self.degree {
            if k > MAX_DEG {
                break;
            }
            for u_row in &uni {
                buffer[idx] = u_row[k];
                idx += 1;
            }
        }
    }

    // Override evaluate_buffers to avoid zeroing
    fn evaluate_buffers(&self, state: &[f64; N], buffer: &mut Vec<f64>) {
        let n = <Self as Basis<N>>::num_features(self);
        let start_len = buffer.len();

        // We ensure capacity to avoid reallocation
        buffer.reserve(n);

        // Unsafe block to skip initialization since we WILL write to every index
        // Safety: `evaluate` must write to the slice we provide. It takes &mut [f64], so it can't extend.
        // We rely on `evaluate` implementation to fill the slice.
        unsafe {
            let ptr = buffer.as_mut_ptr().add(start_len);
            let slice = std::slice::from_raw_parts_mut(ptr, n);
            self.evaluate(state, slice);
            buffer.set_len(start_len + n);
        }
    }

    fn eval_dot(&self, state: &[f64; N], weights: &[f64]) -> f64 {
        const MAX_DEG: usize = 6;
        let effective_deg = self.degree.min(MAX_DEG);

        // Reverted to raw unscaled evaluation. Scaling handled by caller or wrapper.
        let mut uni = [[0.0; MAX_DEG + 1]; N];
        for i in 0..N {
            let x = if let Some((min, max)) = self.clamp_range {
                state[i].clamp(min, max)
            } else {
                state[i]
            };
            compute_univariate(x, effective_deg, self.poly_type, &mut uni[i][..]);
        }

        let mut sum = weights[0]; // Bias weight
        let mut idx = 1;

        if self.degree < 1 {
            return sum;
        }

        // 2. Linear
        for u_row in &uni {
            sum += weights[idx] * u_row[1];
            idx += 1;
        }

        if self.degree < 2 {
            return sum;
        }

        // 3. Quadratic & Interactions
        for i in 0..N {
            for j in i..N {
                let val = if i == j {
                    uni[i][2]
                } else {
                    uni[i][1] * uni[j][1]
                };
                sum += weights[idx] * val;
                idx += 1;
            }
        }

        // 4. Higher Order (Diagonal only)
        for k in 3..=self.degree {
            if k > MAX_DEG {
                break;
            }
            for u_row in &uni {
                sum += weights[idx] * u_row[k];
                idx += 1;
            }
        }

        sum
    }
}

#[inline(always)]
fn compute_univariate(x: f64, k: usize, ptype: PolynomialType, out: &mut [f64]) {
    out[0] = 1.0;
    if k >= 1 {
        out[1] = match ptype {
            PolynomialType::Laguerre => 1.0 - x,
            _ => x,
        };
    }

    // Compute P_i based on recurrence
    match ptype {
        PolynomialType::Power => {
            for i in 2..=k {
                out[i] = out[i - 1] * x;
            }
        }
        PolynomialType::Hermite => {
            // He_n+1 = x*He_n - n*He_n-1
            for i in 2..=k {
                let n = (i - 1) as f64;
                out[i] = x * out[i - 1] - n * out[i - 2];
            }
        }
        PolynomialType::Chebyshev => {
            // T_n+1 = 2xT_n - T_n-1
            for i in 2..=k {
                out[i] = 2.0 * x * out[i - 1] - out[i - 2];
            }
        }
        PolynomialType::Laguerre => {
            // (n+1) L_n+1 = (2n+1-x)L_n - nL_n-1
            // L_i = [ (2(i-1)+1 - x)L_{i-1} - (i-1)L_{i-2} ] / i
            for i in 2..=k {
                let n_prev = (i - 1) as f64; // n in recurrence
                let n_curr = i as f64; // n+1 in recurrence
                out[i] = ((2.0 * n_prev + 1.0 - x) * out[i - 1] - n_prev * out[i - 2]) / n_curr;
            }
        }
    }
}