hessboost 0.2.2

//! Statistical inference for Boulevard boosting: confidence intervals for
//! the regression function `f(x)`, prediction intervals for new labels,
//! and reproduction intervals, with asymptotic (central-limit) guarantees
//! under the assumptions below and validated only in the regimes listed in
//! [Validation](#validation). Opt-in: train with
//! [`BoosterKind::Boulevard`](crate::config::BoosterKind::Boulevard), then
//! fit a [`BoulevardInference`] on the training rows. A Boulevard EBM
//! ([`crate::ebm`], `ebm_boulevard`) gets bands on its shape functions from
//! [`EbmInference`], built on the same solvers.
//!
//! Unlike [`crate::conformal`], whose intervals have finite-sample
//! *marginal* coverage of the label, these intervals are about `f` itself
//! and hold pointwise (conditionally on `x`), but only asymptotically and
//! under the assumptions below.
//!
//! # The algorithms
//!
//! Boulevard (Zhou & Hooker, *Boulevard: Regularized Stochastic Gradient
//! Boosted Trees and Their Limiting Distribution*, JMLR 23, 2022) averages
//! its trees instead of summing them: after `b` rounds the ensemble is
//! `f_b = (λ / b) Σ_{i ≤ b} t_i`, each tree fitted to the residuals of the
//! current average on a row subsample. Fang, Tan & Hooker (*Statistical
//! Inference for Gradient Boosting Regression*, NeurIPS 2025) add two
//! variants that recover more of the signal, both implemented here:
//!
//! - **BRAT-D** (their Algorithm 1; `num_parallel_tree = 1`): round `b`
//!   drops each earlier tree independently with probability `p`
//!   ([`Boulevard::dropout`](crate::config::Boulevard::dropout);
//!   `p = 0` is Zhou & Hooker's Boulevard) and fits the new tree to
//!   `y − μ − (λ / (b−1)) Σ_{kept} t_s(x)`, dividing by every earlier tree,
//!   not only the kept ones. The model predicts `μ + ((1 + λq) / B) Σ t_b`
//!   with `q = 1 − p` and learning rate `λ = eta ∈ (0, 1]`.
//! - **BRAT-P** (Algorithm 2; `num_parallel_tree = K ≥ 2`): the first
//!   iteration boosts `K` trees in sequence; afterwards tree `k` of round
//!   `b` fits `y − μ − Σ_{l ≠ k} ā_l(x)`, where `ā_l` averages slot `l`'s
//!   earlier trees, so a round's trees grow in parallel. The model predicts
//!   `μ + (1/B) Σ_{b,k} t_{b,k}`.
//!
//! `μ` is the intercept (the label mean unless `base_score` is set). With
//! [`Boulevard::truncation`](crate::config::Boulevard::truncation)
//! `M > 0` the subtracted ensemble part is clipped to `[−M, M]`, the `Γ_M`
//! of the convergence proofs. The trained trees are ordinary
//! [`RegTree`](crate::tree::RegTree)s whose leaves already carry the final
//! scale, so prediction, SHAP, slicing, and every export work as for a
//! `gbtree` model; the native binary and JSON formats also keep the
//! [`BoulevardInfo`] inference reads.
//!
//! # The variance
//!
//! As the number of rounds grows, both algorithms converge to a kernel ridge
//! regression `f̂(x) = μ + s k(x)ᵀ (c I + K)⁻¹ (y − μ 1)` in the leaf kernel
//! of the ensemble: `K_ij` averages, over the trees, `1 / (n_ℓ + κ)` when
//! rows `i` and `j` share a leaf `ℓ` holding `n_ℓ` training rows
//! (`κ = lambda / subsample`, so a leaf's value `Σ z / (m + lambda)` over
//! its `m ≈ ξ n_ℓ` sampled rows is matched), and `k(x)` is the same average
//! between `x` and the training rows. BRAT-D has `c = 1 / (λq)` and
//! `s = (1 + λq) / (λq)`; BRAT-P `c = 1 / (K−1)` and `s = K / (K−1)`. The
//! estimate is linear in `y` with weights `w(x) = s u(x) + γ(x) 1`, where
//! `u = (c I + K)⁻¹ k(x)` and, for a label-mean intercept,
//! `γ = (1 − s 1ᵀu) / n`, so under the regression model `y = f(x) + ε` with
//! independent noise of variance `σ²`,
//!
//! ```text
//! f̂(x) ≈ N(f(x), σ² ‖w(x)‖²)     (Fang, Tan & Hooker, Theorem 2)
//! ```
//!
//! The intervals ([`BoulevardInference::confidence_intervals`] and
//! siblings) plug in an estimate `σ̂²` ([`NoiseVariance`]) and the normal
//! quantile. The kernel is estimated by the trained trees themselves (the
//! paper's equations (1)–(2)), with each tree's row sample replaced by its
//! expectation, as the authors' reference implementation does.
//! [`KernelSolver::Exact`] factors the `n × n` system (`O(n³)` time,
//! `O(n²)` memory); [`KernelSolver::Nystrom`] uses the paper's
//! Appendix A Nyström approximation from `s` uniformly sampled landmark
//! rows (`O(n s²)` time, `O(n s)` memory).
//!
//! # Assumptions
//!
//! The asymptotic guarantees rest on the papers' conditions; the estimates
//! are computed regardless, so read them as approximations when these fail:
//!
//! - **Regression with squared error**, `y = f(x) + ε` with independent,
//!   homoscedastic, sub-Gaussian noise. `booster = boulevard` refuses every
//!   other objective, row weights, base margins, row sampling that depends
//!   on the labels or gradients (class-balanced bagging, gradient-based
//!   sampling), and every option that makes leaf values nonlinear in the
//!   labels (see
//!   [`TrainingParams::validate`](crate::config::TrainingParams::validate)).
//! - **Structure–value isolation** (the tree structures independent of the
//!   labels the leaves average): not true of trees grown greedily on the
//!   same labels. [`honest_refit`] provides it, refitting every leaf on an
//!   independent sample through the same Boulevard recursion (the NeurIPS
//!   paper's "integrity"); fit the inference on that sample.
//! - **Non-adaptivity** (tree structures eventually drawn from a fixed
//!   distribution), bounded leaf diameters and a minimal leaf size growing
//!   with `n` (set `min_child_weight`, which counts rows here), a row
//!   subsample, and, for BRAT-P, balanced splits. Both papers report that
//!   the intervals behave well in practice without enforcing all of them.
//! - Enough rounds that the ensemble is near its limit: the variance is the
//!   limit's.
//! - Every variance here is **conditional on the tree structures** (the
//!   kernel is treated as fixed, as in both papers). How the structures, and
//!   with them the fit's bias, vary between training samples is not
//!   included. In one dimension with small bias this is negligible; in
//!   several dimensions the true variance can be a multiple of the estimate
//!   and the intervals under-cover (see [Validation](#validation)).
//! - Prediction intervals additionally need **Gaussian noise**: the
//!   estimate's error is asymptotically normal, but the new label's own
//!   noise is not averaged, so `± z σ̂` covers `1 − α` of it only when that
//!   noise is normal (uniform noise, for one, is covered with probability
//!   1 at 95%). Otherwise use
//!   [`calibrated_prediction_intervals`](BoulevardInference::calibrated_prediction_intervals),
//!   which rescales the widths by an empirical quantile, or [`crate::conformal`].
//!
//! # Validation
//!
//! Simulations (50 training samples each, 100 fixed test points,
//! coverage of the true `f`, `σ̂²` from a holdout of `n/2` rows):
//!
//! | setting | `n` | 90% CI | 95% CI | 95% PI |
//! |---|---|---|---|---|
//! | `f = sin 2πx + x²/2` (1-d), BRAT-D `λ = 0.6`, `p = 0.6`, `ξ = 0.6`, depth 8, 200 trees, [`honest_refit`] | 500 | 0.899 | 0.949 | 0.959 |
//! | same | 1000 | 0.904 | 0.953 | 0.955 |
//! | same | 2000 | 0.901 | 0.952 | 0.955 |
//! | same, Nyström `s = 1000` | 4000 | 0.909 | 0.957 | 0.953 |
//! | same, `p = 0` (Zhou & Hooker's Boulevard) | 1000 | 0.903 | 0.952 | 0.955 |
//! | same, BRAT-P `K = 4`, 100 rounds | 1000 | 0.883 | 0.936 | 0.956 |
//! | same, **without** [`honest_refit`] | 1000 | 0.756 | 0.845 | 0.957 |
//! | `f = 4x₁ − x₂²` on `[0, 1]³` (the NeurIPS paper's §6 test function), `λ = 1`, `p = 0.95`, `ξ = 1`, depth 6, 100 trees, [`honest_refit`] | 1000 | 0.651 | 0.726 | — |
//! | same | 2000 | 0.626 | 0.711 | — |
//! | same, without [`honest_refit`] | 1000 | 0.543 | 0.632 | — |
//! | Friedman #1 (5-d), BRAT-D `p = 0.6`, depth 6, [`honest_refit`] (20 samples) | 2000 | 0.122 | 0.148 | 0.959 |
//!
//! In the 1-d setting the estimated variance matches the across-sample
//! variance of the estimate to within 5% (BRAT-P: 8% low). In the 3-d
//! setting it is about half of it: 1.06 × when the tree structures are held
//! fixed and only the refit sample is redrawn, 1.96 × when they are
//! retrained, so the missing term is the structures' sample-to-sample
//! variation. In 5-d the fit's bias dominates. The authors' reference
//! package (`boulevard-boosting` 0.1.0a1) gives the same coverage, interval
//! widths, and MSE as this module in the settings compared (within
//! seed-to-seed noise). Treat the confidence intervals as validated for
//! low-dimensional smooth signals with honest refits only.
//!
//! The NeurIPS paper's variable-importance test (§4) is not provided: with
//! the variances above its statistic is anti-conservative (82–94% rejection
//! of a true null at a nominal 5% in the §6 setup with honest refits; 95%
//! with the reference package's own weights), and the paper's reported
//! size comes from a different regime (both fits on the same training
//! sample, noise standard deviation `0.01`, `n ≤ 200`, depth 8).
//!
//! # Example
//!
//! ```
//! use hessboost::config::{BoosterKind, Boulevard};
//! use hessboost::inference::{BoulevardInference, KernelSolver, NoiseVariance};
//! use hessboost::prelude::*;
//!
//! # fn main() -> Result<()> {
//! let n = 300;
//! let x: Vec<f32> = (0..n).map(|i| ((i * 37) % n) as f32 / n as f32).collect();
//! let y: Vec<f32> = x
//!     .iter()
//!     .enumerate()
//!     .map(|(i, v)| (6.0 * v).sin() + 0.1 * ((i * 7919 % 101) as f32 / 50.0 - 1.0))
//!     .collect();
//! let all = DMatrix::from_dense(&x, n, 1)?.with_labels(&y)?;
//! let (fit_rows, cal_rows): (Vec<usize>, Vec<usize>) = (0..n).partition(|i| i % 3 != 0);
//! let (dtrain, dcal) = (all.select_rows(&fit_rows)?, all.select_rows(&cal_rows)?);
//!
//! let params = TrainingParams::builder()
//!     .booster(BoosterKind::Boulevard(Boulevard::builder().dropout(0.5).build()?))
//!     .eta(0.8)
//!     .subsample(0.8)
//!     .max_depth(3)
//!     .min_child_weight(5.0)
//!     .build()?;
//! let model = train(&params, &dtrain, 100)?;
//!
//! let inference = BoulevardInference::fit(
//!     &model,
//!     &dtrain,
//!     NoiseVariance::Holdout(&dcal),
//!     KernelSolver::Exact,
//! )?;
//! let ci = inference.confidence_intervals(&dcal, 0.05)?;
//! let pi = inference.prediction_intervals(&dcal, 0.05)?;
//! assert!(ci.iter().zip(&pi).all(|(c, p)| p.lower < c.lower && c.upper < p.upper));
//! # Ok(())
//! # }
//! ```

mod ebm;
mod kernel;
mod linalg;
mod refit;
mod solver;
mod term_kernel;

use std::num::NonZeroUsize;

use serde::{Deserialize, Serialize};

use crate::check::ensure;
use crate::conformal::Interval;
use crate::data::DMatrix;
use crate::error::{HessboostError, Result};
use crate::model::{BoostedModel, Iterations, Predictions};
use crate::objective::Objective;
use crate::objective::distributional::special::norm_ppf;
use kernel::{Kernel, LeafKernel};
use rayon::prelude::*;
use solver::RidgeSolver;

pub use ebm::{EbmInference, TermBands};
pub use refit::honest_refit;

/// Query points solved together: one block of right-hand sides.
const QUERY_BLOCK: usize = 32;

/// Largest training set [`KernelSolver::Exact`] factors (its `n × n`
/// system takes `8 n²` bytes: 512 MiB here).
pub const MAX_EXACT_ROWS: usize = 8192;

/// How a `booster = boulevard` model was trained: the settings its
/// inference and [`honest_refit`] read, recorded by training
/// ([`BoostedModel::boulevard`]). Whether it is BRAT-D or BRAT-P follows
/// from the model's
/// [`num_parallel_tree`](BoostedModel::num_parallel_tree) (`1` or more).
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct BoulevardInfo {
    /// BRAT-D's dropout probability `p`
    /// ([`Boulevard::dropout`](crate::config::Boulevard::dropout);
    /// `0` for BRAT-P).
    pub dropout: f64,
    /// The learning rate `λ` (`eta`; `1` for BRAT-P).
    pub learning_rate: f64,
    /// The row subsample ratio `ξ` (`subsample`).
    pub subsample: f64,
    /// The L2 leaf penalty (`lambda`).
    pub reg_lambda: f64,
    /// The residual truncation level `M`, if any
    /// ([`Boulevard::truncation`](crate::config::Boulevard::truncation)).
    /// JSON writes none as `null` and reads a `0` (older files) as none.
    #[serde(deserialize_with = "truncation_from_json")]
    pub truncation: Option<f64>,
    /// The training seed, which [`honest_refit`] derives its draws from.
    pub seed: u64,
    /// Whether the intercept is the training-label mean (`base_score`
    /// unset), which the variance then includes.
    pub intercept_from_labels: bool,
}

/// [`BoulevardInfo::truncation`] from JSON: a number or `null` (the field
/// itself stays required), a `0` meaning none as files written before it
/// was optional record it.
fn truncation_from_json<'de, D: serde::Deserializer<'de>>(
    deserializer: D,
) -> std::result::Result<Option<f64>, D::Error> {
    Ok(Option::<f64>::deserialize(deserializer)?.filter(|&t| t != 0.0))
}

impl BoulevardInfo {
    /// Check the recorded settings against their ranges and `model`'s
    /// layout (one output, squared error, unweighted scalar trees).
    pub(crate) fn validate(&self, model: &BoostedModel) -> Result<()> {
        let fail = |reason: &str| {
            Err(HessboostError::model_format(format!(
                "invalid Boulevard record: {reason}"
            )))
        };
        if !(self.dropout.is_finite() && (0.0..1.0).contains(&self.dropout)) {
            return fail("dropout must be in [0, 1)");
        }
        if !(self.learning_rate > 0.0 && self.learning_rate <= 1.0) {
            return fail("learning_rate must be in (0, 1]");
        }
        if !(self.subsample > 0.0 && self.subsample <= 1.0) {
            return fail("subsample must be in (0, 1]");
        }
        if !(self.reg_lambda.is_finite() && self.reg_lambda >= 0.0) {
            return fail("reg_lambda must be finite and >= 0");
        }
        if self.truncation.is_some_and(|t| !(t.is_finite() && t > 0.0)) {
            return fail("truncation must be finite and > 0");
        }
        if model.num_parallel_tree() > 1 && (self.dropout != 0.0 || self.learning_rate != 1.0) {
            return fail("BRAT-P (num_parallel_tree > 1) needs dropout 0 and learning_rate 1");
        }
        if !model
            .objective()
            .built_in()
            .is_some_and(Objective::is_unweighted_squared_error)
            || model.n_outputs() != 1
            || model.has_vector_leaves()
            || model.has_non_unit_tree_weights()
            || model.trees().iter().any(|t| t.linear_leaves().is_some())
        {
            return fail("only single-output reg:squarederror tree ensembles are Boulevard fits");
        }
        Ok(())
    }

    /// `(c, s)` of the kernel ridge limit for `parallel` trees per round:
    /// `f̂ = μ + s kᵀ (c I + K)⁻¹ (y − μ)`.
    fn ridge(&self, parallel: usize) -> (f64, f64) {
        if parallel > 1 {
            let k = parallel as f64;
            (1.0 / (k - 1.0), k / (k - 1.0))
        } else {
            let lq = self.learning_rate * (1.0 - self.dropout);
            (1.0 / lq, (1.0 + lq) / lq)
        }
    }

    /// The leaf-count offset `κ = lambda / subsample` of the kernel.
    fn kappa(&self) -> f64 {
        self.reg_lambda / self.subsample
    }
}

/// How [`BoulevardInference::fit`] solves the kernel ridge systems.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[non_exhaustive]
pub enum KernelSolver {
    /// Factor the `n × n` system exactly (`O(n³)` time, `8 n²` bytes); at
    /// most [`MAX_EXACT_ROWS`] training rows.
    #[default]
    Exact,
    /// The Nyström approximation of Fang, Tan & Hooker's Appendix A:
    /// `landmarks` training rows drawn uniformly without replacement with
    /// `seed` (all of them when `landmarks >= n`, which then reproduces
    /// [`Exact`](Self::Exact) up to rounding), those numerically dependent
    /// on the others dropped. `O(n s²)` time and `8 n s` bytes for
    /// `s = landmarks`.
    ///
    /// ```
    /// use std::num::NonZeroUsize;
    /// use hessboost::inference::KernelSolver;
    ///
    /// let landmarks = NonZeroUsize::new(1000).expect("1000 is nonzero");
    /// let solver = KernelSolver::Nystrom { landmarks, seed: 0 };
    /// # let _ = solver;
    /// ```
    Nystrom {
        /// Number of landmark rows.
        landmarks: NonZeroUsize,
        /// Seed of the landmark draw.
        seed: u64,
    },
}

/// Where [`BoulevardInference::fit`] takes the noise variance `σ²` from.
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub enum NoiseVariance<'a> {
    /// The mean squared residual on held-out labelled rows (the paper's
    /// estimator; also enables
    /// [`calibrated_prediction_intervals`](BoulevardInference::calibrated_prediction_intervals)).
    /// Slightly conservative: the residuals also carry the estimate's own
    /// variance and bias.
    Holdout(&'a DMatrix),
    /// The mean squared residual on the training rows themselves: no data
    /// held out, but biased low (the model has fitted some of the noise).
    TrainingResiduals,
    /// A known `σ² > 0` (simulations).
    Known(f64),
}

/// The variance machinery of one Boulevard model: its leaf kernel over the
/// training rows, the factored ridge system, and a noise estimate. Fit once
/// with [`fit`](Self::fit), then query any rows.
pub struct BoulevardInference<'a> {
    model: &'a BoostedModel,
    kernel: LeafKernel,
    solver: RidgeSolver,
    /// Ridge `c` and scale `s` of the kernel ridge limit.
    c: f64,
    s: f64,
    intercept_from_labels: bool,
    noise_variance: f64,
    holdout: Option<&'a DMatrix>,
}

impl std::fmt::Debug for BoulevardInference<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BoulevardInference")
            .field("rows", &self.kernel.n())
            .field("trees", &self.kernel.n_trees())
            .field("ridge", &self.c)
            .field("scale", &self.s)
            .field("noise_variance", &self.noise_variance)
            .finish_non_exhaustive()
    }
}

/// Refuse `data` unless it has `model`'s features, one label column when
/// `labelled`, and neither row weights other than 1 nor base margins (which
/// Boulevard training refuses too): [`HessboostError::InvalidData`] naming
/// the input `what` (`train`, `holdout`, `values`, `data`).
fn check_data(
    model: &BoostedModel,
    data: &DMatrix,
    what: &'static str,
    labelled: bool,
) -> Result<()> {
    if data.n_rows() == 0 {
        return Err(HessboostError::EmptyDataset(what));
    }
    if data.n_cols() != model.n_features() {
        return Err(HessboostError::dimension_mismatch(
            what,
            model.n_features(),
            data.n_cols(),
        ));
    }
    if labelled && data.labels().is_none() {
        return Err(HessboostError::invalid_data(what, "needs labels"));
    }
    if labelled && data.n_targets() != 1 {
        return Err(HessboostError::invalid_data(what, "needs one label column"));
    }
    if data.weights().is_some_and(|w| w.iter().any(|&v| v != 1.0)) {
        return Err(HessboostError::invalid_data(
            what,
            "row weights other than 1 are not supported: Boulevard inference assumes equal \
             noise per row",
        ));
    }
    if data.base_margin().is_some() {
        return Err(HessboostError::invalid_data(
            what,
            "base margins are not supported by Boulevard inference",
        ));
    }
    Ok(())
}

/// `alpha` must be a miscoverage level in `(0, 1)`.
fn check_alpha(alpha: f64) -> Result<()> {
    ensure(
        "alpha",
        alpha > 0.0 && alpha < 1.0,
        format!("must be in (0, 1), got {alpha}"),
    )
}

/// The two-sided normal quantile `z_{1 − α/2}`.
fn z_value(alpha: f64) -> f64 {
    norm_ppf(1.0 - alpha / 2.0)
}

/// Mean squared residual of `model` on the labelled `data`.
fn mean_squared_residual(model: &BoostedModel, data: &DMatrix) -> Result<f64> {
    let preds = model.predict(data, Iterations::Best)?;
    let labels = data.labels().unwrap_or_default();
    let sum: f64 = preds
        .as_slice()
        .iter()
        .zip(labels)
        .map(|(&p, &y)| (f64::from(y) - f64::from(p)).powi(2))
        .sum();
    Ok(sum / labels.len() as f64)
}

/// The noise variance `noise` names for `model` trained on `train`, which
/// must be finite and positive.
fn noise_estimate(model: &BoostedModel, train: &DMatrix, noise: NoiseVariance) -> Result<f64> {
    let noise_variance = match noise {
        NoiseVariance::Holdout(holdout) => {
            check_data(model, holdout, "holdout", true)?;
            mean_squared_residual(model, holdout)?
        }
        NoiseVariance::TrainingResiduals => mean_squared_residual(model, train)?,
        NoiseVariance::Known(v) => v,
    };
    ensure(
        "noise",
        noise_variance.is_finite() && noise_variance > 0.0,
        format!("the noise variance must be finite and > 0, got {noise_variance}"),
    )?;
    Ok(noise_variance)
}

/// Check `train` (labelled when the noise comes from its residuals) and
/// estimate the noise variance: the start of every inference fit.
fn fit_noise(model: &BoostedModel, train: &DMatrix, noise: NoiseVariance) -> Result<f64> {
    check_data(
        model,
        train,
        "train",
        matches!(noise, NoiseVariance::TrainingResiduals),
    )?;
    noise_estimate(model, train, noise)
}

/// The interval `prediction ± z_{1−α/2} · width(‖w‖²)` of `model` at every
/// row of `data`, with the squared weight norms `‖w‖²` from `norms` (which
/// checks `data`).
fn normal_intervals(
    model: &BoostedModel,
    data: &DMatrix,
    alpha: f64,
    norms: impl FnOnce() -> Result<Vec<f64>>,
    width: impl Fn(f64) -> f64,
) -> Result<Vec<Interval<f64>>> {
    check_alpha(alpha)?;
    let z = z_value(alpha);
    let norms = norms()?;
    let preds = model.predict(data, Iterations::Best)?;
    Ok(preds
        .as_slice()
        .iter()
        .zip(norms)
        .map(|(&p, w2)| {
            let (center, half) = (f64::from(p), z * width(w2));
            Interval {
                lower: center - half,
                upper: center + half,
            }
        })
        .collect())
}

/// Factor `c I + kernel` with `solver`.
fn build_solver(kernel: &impl Kernel, solver: KernelSolver, c: f64) -> Result<RidgeSolver> {
    let n = kernel.n();
    match solver {
        KernelSolver::Exact => {
            if n > MAX_EXACT_ROWS {
                return Err(HessboostError::invalid_param(
                    "solver",
                    format!(
                        "the exact solver factors at most {MAX_EXACT_ROWS} rows, got {n}; use \
                         `KernelSolver::Nystrom`"
                    ),
                ));
            }
            RidgeSolver::exact(kernel, c)
        }
        KernelSolver::Nystrom { landmarks, seed } => {
            RidgeSolver::nystrom(kernel, c, landmarks.get(), seed)
        }
    }
}

impl<'a> BoulevardInference<'a> {
    /// Build the leaf kernel of `model` over `train`, the rows it was
    /// trained on (or, after [`honest_refit`], refitted on), factor the
    /// ridge system with `solver`, and estimate the noise variance.
    ///
    /// # Errors
    ///
    /// [`HessboostError::IncompatibleModel`] (`model`) when `model` is not
    /// a Boulevard fit ([`BoostedModel::boulevard`] is `None`) or has no
    /// trees (0 rounds); [`HessboostError::InvalidData`] (`train`, or
    /// `holdout` for its rows) when `train` is not its training data (a
    /// leaf holds fewer of its rows than it was grown on), has row weights
    /// or base margins, or (for the noise estimate) lacks labels;
    /// [`HessboostError::InvalidParameter`] when [`KernelSolver::Exact`]
    /// gets more than [`MAX_EXACT_ROWS`] rows (`solver`) or a noise
    /// variance is not finite and positive (`noise`).
    pub fn fit(
        model: &'a BoostedModel,
        train: &DMatrix,
        noise: NoiseVariance<'a>,
        solver: KernelSolver,
    ) -> Result<Self> {
        let info = model.boulevard().ok_or_else(|| {
            HessboostError::incompatible_model(
                "model",
                "not a Boulevard fit: train it with `booster = boulevard`",
            )
        })?;
        if model.num_trees() == 0 {
            return Err(HessboostError::incompatible_model(
                "model",
                "has no trees (trained for 0 rounds), so its leaf kernel has no rows",
            ));
        }
        let noise_variance = fit_noise(model, train, noise)?;
        let (c, s) = info.ridge(model.num_parallel_tree());
        let leaves = model.predict_leaf(train, ..)?;
        let kernel = LeafKernel::new(model.trees(), &leaves, info.kappa())?;
        let solver = build_solver(&kernel, solver, c)?;
        Ok(BoulevardInference {
            model,
            kernel,
            solver,
            c,
            s,
            intercept_from_labels: info.intercept_from_labels,
            noise_variance,
            holdout: match noise {
                NoiseVariance::Holdout(h) => Some(h),
                _ => None,
            },
        })
    }

    /// The noise variance estimate `σ̂²`.
    pub fn noise_variance(&self) -> f64 {
        self.noise_variance
    }

    /// The Gram matrix `w(x_a)ᵀ w(x_b)` of the estimate's weight vectors at
    /// the rows of `data` (`m × m`), from the solver's `u`-Gram and sums.
    fn weight_gram(&self, gram: &[f64], sums: &[f64], m: usize) -> Vec<f64> {
        let s = self.s;
        let n = self.kernel.n() as f64;
        let gamma: Vec<f64> = sums
            .iter()
            .map(|&su| {
                if self.intercept_from_labels {
                    (1.0 - s * su) / n
                } else {
                    0.0
                }
            })
            .collect();
        let mut out = vec![0.0; m * m];
        for a in 0..m {
            for b in 0..m {
                out[a * m + b] = s * s * gram[a * m + b]
                    + s * (gamma[a] * sums[b] + gamma[b] * sums[a])
                    + n * gamma[a] * gamma[b];
            }
        }
        out
    }

    /// The leaf node ids of `data`'s rows, `[row][tree]`.
    fn leaves(&self, data: &DMatrix) -> Result<Predictions<u32>> {
        check_data(self.model, data, "data", false)?;
        self.model.predict_leaf(data, ..)
    }

    /// The kernel vectors of the `rows` of `leaves` (`[row][tree]`), one
    /// per row of the result (`rows.len() × n`).
    fn kernel_vectors(&self, leaves: &Predictions<u32>, rows: std::ops::Range<usize>) -> Vec<f64> {
        let n = self.kernel.n();
        let mut k = vec![0.0; rows.len() * n];
        for (out, ids) in k
            .chunks_exact_mut(n)
            .zip(rows.filter_map(|row| leaves.row(row)))
        {
            self.kernel.add_query(ids, out);
        }
        k
    }

    /// `‖w(x)‖²` for every row of `data`, in blocks solved in parallel.
    fn weight_norms(&self, data: &DMatrix) -> Result<Vec<f64>> {
        let leaves = self.leaves(data)?;
        let rows = data.n_rows();
        let blocks: Vec<Vec<f64>> = (0..rows.div_ceil(QUERY_BLOCK))
            .into_par_iter()
            .map(|b| {
                let range = b * QUERY_BLOCK..((b + 1) * QUERY_BLOCK).min(rows);
                let m = range.len();
                let k = self.kernel_vectors(&leaves, range);
                let solved = self.solver.solve(&k, m, self.c);
                let g = self.weight_gram(&solved.gram, &solved.sums, m);
                (0..m).map(|a| g[a * m + a].max(0.0)).collect()
            })
            .collect();
        Ok(blocks.concat())
    }

    /// The standard error `σ̂ ‖w(x)‖` of the model's prediction at every row
    /// of `data`: the estimated standard deviation of `f̂(x)` over new
    /// training samples.
    ///
    /// # Errors
    ///
    /// [`HessboostError::DimensionMismatch`] when `data` does not have the
    /// model's features; [`HessboostError::InvalidData`] (`data`) when it
    /// has row weights or base margins.
    pub fn standard_errors(&self, data: &DMatrix) -> Result<Predictions<f64>> {
        let sigma = self.noise_variance.sqrt();
        let se: Vec<f64> = self
            .weight_norms(data)?
            .into_iter()
            .map(|w2| sigma * w2.sqrt())
            .collect();
        Ok(Predictions::new(se, data.n_rows(), 1))
    }

    /// The interval `prediction ± z · width(‖w‖²)` of every row.
    fn intervals(
        &self,
        data: &DMatrix,
        alpha: f64,
        width: impl Fn(f64) -> f64,
    ) -> Result<Vec<Interval<f64>>> {
        normal_intervals(self.model, data, alpha, || self.weight_norms(data), width)
    }

    /// Confidence intervals for the regression function `f(x)` at every row
    /// of `data` at miscoverage `alpha`: `f̂(x) ± z_{1−α/2} σ̂ ‖w(x)‖`
    /// (Fang, Tan & Hooker, equation (3)), with asymptotic pointwise
    /// coverage `1 − alpha`.
    ///
    /// # Errors
    ///
    /// [`HessboostError::InvalidParameter`] when `alpha` is not in
    /// `(0, 1)`, plus those of [`standard_errors`](Self::standard_errors).
    pub fn confidence_intervals(&self, data: &DMatrix, alpha: f64) -> Result<Vec<Interval<f64>>> {
        let sigma2 = self.noise_variance;
        self.intervals(data, alpha, |w2| (sigma2 * w2).sqrt())
    }

    /// Prediction intervals for a new label `y` at every row of `data`:
    /// `f̂(x) ± z_{1−α/2} sqrt(σ̂² + σ̂² ‖w(x)‖²)`. With **Gaussian** noise
    /// this covers `y | x` with asymptotic probability `1 − alpha`
    /// (conditionally on `x`, unlike [`crate::conformal`]'s marginal
    /// guarantee); with other noise the new label's own error keeps its
    /// distribution and the coverage differs, so use
    /// [`calibrated_prediction_intervals`](Self::calibrated_prediction_intervals)
    /// or [`crate::conformal`] then.
    ///
    /// The paper's display scales the noise term by BRAT-D's `(1 + λq) / λ`
    /// as well; that factor belongs to the estimate only (the new label's
    /// noise is not rescaled), so it is applied to `‖w‖` alone, as in the
    /// authors' reference implementation.
    ///
    /// # Errors
    ///
    /// As [`confidence_intervals`](Self::confidence_intervals).
    pub fn prediction_intervals(&self, data: &DMatrix, alpha: f64) -> Result<Vec<Interval<f64>>> {
        let sigma2 = self.noise_variance;
        self.intervals(data, alpha, |w2| (sigma2 * (1.0 + w2)).sqrt())
    }

    /// Reproduction intervals: where the prediction of the same procedure
    /// retrained on an independent sample falls, `f̂(x) ± z √2 σ̂ ‖w(x)‖`
    /// (Zhou & Hooker; the difference of two independent estimates has
    /// twice the variance).
    ///
    /// # Errors
    ///
    /// As [`confidence_intervals`](Self::confidence_intervals).
    pub fn reproduction_intervals(&self, data: &DMatrix, alpha: f64) -> Result<Vec<Interval<f64>>> {
        let sigma2 = self.noise_variance;
        self.intervals(data, alpha, |w2| (2.0 * sigma2 * w2).sqrt())
    }

    /// [`prediction_intervals`](Self::prediction_intervals) whose widths are
    /// all scaled by one factor chosen on the [`NoiseVariance::Holdout`]
    /// rows: the split-conformal quantile (rank `⌈(h + 1)(1 − α)⌉` of `h`)
    /// of `|y − f̂(x)| / half_width(x)`. Fang, Tan & Hooker's Section 6
    /// adjustment for finite samples; it keeps the per-point widths'
    /// shape. The holdout rows also estimated `σ̂²`, so the conformal
    /// guarantee holds only approximately. When the holdout set is too
    /// small for the level (`⌈(h + 1)(1 − α)⌉ > h`), every interval is
    /// `(−∞, ∞)`.
    ///
    /// # Errors
    ///
    /// [`HessboostError::InvalidParameter`] unless the noise came from a
    /// holdout set, plus those of
    /// [`prediction_intervals`](Self::prediction_intervals).
    pub fn calibrated_prediction_intervals(
        &self,
        data: &DMatrix,
        alpha: f64,
    ) -> Result<Vec<Interval<f64>>> {
        let holdout = self.holdout.ok_or_else(|| {
            HessboostError::invalid_param(
                "noise",
                "calibrated intervals need `NoiseVariance::Holdout` rows",
            )
        })?;
        let calibration = self.prediction_intervals(holdout, alpha)?;
        let labels = holdout.labels().unwrap_or_default();
        let mut ratios: Vec<f64> = calibration
            .iter()
            .zip(labels)
            .map(|(iv, &y)| {
                let half = (iv.upper - iv.lower) / 2.0;
                let center = f64::midpoint(iv.upper, iv.lower);
                (f64::from(y) - center).abs() / half.max(f64::MIN_POSITIVE)
            })
            .collect();
        let scale = crate::conformal::conformal_quantile(&mut ratios, alpha);
        Ok(self
            .prediction_intervals(data, alpha)?
            .into_iter()
            .map(|Interval { lower, upper }| {
                let (center, half) = (f64::midpoint(upper, lower), (upper - lower) / 2.0 * scale);
                if half.is_finite() {
                    Interval {
                        lower: center - half,
                        upper: center + half,
                    }
                } else {
                    Interval {
                        lower: f64::NEG_INFINITY,
                        upper: f64::INFINITY,
                    }
                }
            })
            .collect())
    }
}