//! Rainier benchmark parity checks on identical model samples.

//!

//! This test consumes deterministic benchmark snapshots generated by

//! `PrioriParityRunner` in the Rainier workspace and evaluates equivalent

//! Priori models on the same parameter vectors.

//!

//! The goal is to produce auditable, model-by-model numerical parity evidence

//! for log-density and gradients while measuring Priori evaluation latency on

//! the same cases.

use std::fs;

use std::path::{Path, PathBuf};

use std::time::Instant;

use priori::core::{Parameter, ParameterId, Real, RealRef, UnaryOp};

use priori::dist::{Beta, Cauchy, Continuous, Distribution, Mixture, Normal, Poisson, Uniform};

use priori::model::Model;

use priori::sampler::DensityFunction;

use serde_json::Value;

fn scalar(value: f64) -> RealRef {

Real::scalar(value)

}

fn workspace_root() -> PathBuf {

PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../..")

}

fn add_prior(model: Model, parameter: Parameter) -> Model {

if let Some(prior) = parameter.prior_density {

model.observe_with(&[0.0], |_| prior.clone())

} else {

model

}

}

fn parse_scala_list(content: &str, marker: &str) -> std::vec::Vec<f64> {

let start = content

.find(marker)

.unwrap_or_else(|| panic!("marker not found: {marker}"));

let tail = &content[start..];

let list_offset = tail

.find("List(")

.unwrap_or_else(|| panic!("List( not found after marker: {marker}"));

let body_start = start + list_offset + "List(".len();

let mut depth = 1_i32;

let mut body_end = None;

for (offset, ch) in content[body_start..].char_indices() {

match ch {

'(' => depth += 1,

')' => {

depth -= 1;

if depth == 0 {

body_end = Some(body_start + offset);

break;

}

}

_ => {}

}

}

let end = body_end.unwrap_or_else(|| panic!("unterminated list for marker: {marker}"));

content[body_start..end]

.split(',')

.map(str::trim)

.filter(|token| !token.is_empty())

.map(|token| {

token

.parse::<f64>()

.unwrap_or_else(|_| panic!("unable to parse number `{token}` for marker {marker}"))

})

.collect()

}

fn parse_file(path: &Path) -> String {

fs::read_to_string(path)

.unwrap_or_else(|err| panic!("failed to read {}: {err}", path.display()))

}

fn build_ark_model(rainier_root: &Path) -> Model {

let path = rainier_root

.join("rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/ARK.scala");

let source = parse_file(&path);

let ys = parse_scala_list(&source, "val ys: List[Double] = List(");

let (alpha_param, alpha) = Normal::new(scalar(0.0), scalar(10.0)).latent(ParameterId(0));

let (sigma_latent_param, sigma_latent) =

Cauchy::new(scalar(0.0), scalar(2.5)).latent(ParameterId(1));

let sigma = Real::unary(sigma_latent, UnaryOp::Abs);

let mut model = Model::new();

model = add_prior(model, alpha_param);

model = add_prior(model, sigma_latent_param);

let mut betas = std::vec::Vec::with_capacity(5);

for idx in 0..5 {

let (beta_param, beta) =

Normal::new(scalar(0.0), scalar(10.0)).latent(ParameterId(2 + idx));

model = add_prior(model, beta_param);

betas.push(beta);

}

for t in 5..ys.len() {

let mut mu = alpha.clone();

for k in 1..=5 {

mu = mu + betas[k - 1].clone() * ys[t - k];

}

let term = Normal::new(mu, sigma.clone()).log_density_at(scalar(ys[t]));

model = model.observe_with(&[0.0], |_| term.clone());

}

model

}

fn build_eight_schools_model(rainier_root: &Path) -> Model {

let path = rainier_root

.join("rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/EightSchools.scala");

let source = parse_file(&path);

let ys = parse_scala_list(&source, "val ys: List[Int] = List(");

let sigmas = parse_scala_list(&source, "val sigmas: List[Int] = List(");

let (mu_param, mu) = Normal::new(scalar(0.0), scalar(5.0)).latent(ParameterId(0));

let (tau_latent_param, tau_latent) =

Cauchy::new(scalar(0.0), scalar(5.0)).latent(ParameterId(1));

let tau = Real::unary(tau_latent, UnaryOp::Abs);

let mut model = Model::new();

model = add_prior(model, mu_param);

model = add_prior(model, tau_latent_param);

for idx in 0..ys.len() {

let (theta_param, theta) =

Normal::new(mu.clone(), tau.clone()).latent(ParameterId(2 + idx as u32));

model = add_prior(model, theta_param);

let term = Normal::new(theta, scalar(sigmas[idx])).log_density_at(scalar(ys[idx]));

model = model.observe_with(&[0.0], |_| term.clone());

}

model

}

fn build_kid_iq_model(rainier_root: &Path) -> Model {

let path = rainier_root

.join("rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/KidIQ.scala");

let source = parse_file(&path);

let kid_score = parse_scala_list(&source, "val kidScore =");

let mom_hs = parse_scala_list(&source, "val momHS =");

let mom_iq = parse_scala_list(&source, "val momIQ =");

let n = 100_usize;

let (sigma_param, sigma) = Cauchy::new(scalar(0.0), scalar(2.5)).latent(ParameterId(0));

let mut model = Model::new();

model = add_prior(model, sigma_param);

let mut betas = std::vec::Vec::with_capacity(3);

for idx in 0..3 {

let (beta_param, beta) =

Normal::new(scalar(0.0), scalar(10.0)).latent(ParameterId(1 + idx));

model = add_prior(model, beta_param);

betas.push(beta);

}

for idx in 0..n {

let mu = betas[0].clone() + betas[1].clone() * mom_iq[idx] + betas[2].clone() * mom_hs[idx];

let term = Normal::new(mu, sigma.clone()).log_density_at(scalar(kid_score[idx]));

model = model.observe_with(&[0.0], |_| term.clone());

}

model

}

fn build_glmm_poisson2_model(rainier_root: &Path) -> Model {

let path = rainier_root

.join("rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/GLMMPoisson2.scala");

let source = parse_file(&path);

let years = parse_scala_list(&source, "def year =");

let c1 = parse_scala_list(&source, "val c1 = List(");

let c2 = parse_scala_list(&source, "val c2 = List(");

let n_sites = 10_usize;

let n_years = 4_usize;

let year_len = years.len();

let mut combined = c1;

combined.extend(c2);

let truncated = &combined[..n_sites * year_len];

let counts = truncated

.iter()

.enumerate()

.filter_map(|(idx, value)| (idx % year_len < n_years).then_some(*value))

.collect::<std::vec::Vec<_>>();

let (mu_param, mu) = Normal::new(scalar(0.0), scalar(10.0)).latent(ParameterId(0));

let (sd_alpha_param, sd_alpha) = Uniform::new(scalar(0.0), scalar(2.0)).latent(ParameterId(1));

let (sd_year_param, sd_year) = Uniform::new(scalar(0.0), scalar(1.0)).latent(ParameterId(12));

let mut model = Model::new();

model = add_prior(model, mu_param);

model = add_prior(model, sd_alpha_param);

model = add_prior(model, sd_year_param);

let mut alphas = std::vec::Vec::with_capacity(n_sites);

for idx in 0..n_sites {

let (alpha_param, alpha) =

Normal::new(mu.clone(), sd_alpha.clone()).latent(ParameterId(2 + idx as u32));

model = add_prior(model, alpha_param);

alphas.push(alpha);

}

let mut betas = std::vec::Vec::with_capacity(3);

for idx in 0..3 {

let (beta_param, beta) =

Normal::new(scalar(0.0), scalar(10.0)).latent(ParameterId(13 + idx as u32));

model = add_prior(model, beta_param);

betas.push(beta);

}

let mut eps = std::vec::Vec::with_capacity(n_years);

for idx in 0..n_years {

let (eps_param, eps_value) =

Normal::new(scalar(0.0), sd_year.clone()).latent(ParameterId(16 + idx as u32));

model = add_prior(model, eps_param);

eps.push(eps_value);

}

let year_betas = years

.iter()

.copied()

.take(n_years)

.enumerate()

.map(|(idx, year)| {

betas[0].clone() * year

+ betas[1].clone() * year * year

+ betas[2].clone() * year * year * year

+ eps[idx].clone()

})

.collect::<std::vec::Vec<_>>();

let mut count_idx = 0_usize;

for alpha in alphas.iter().take(n_sites) {

for year_beta in year_betas.iter().take(n_years) {

let log_lambda = year_beta.clone() + alpha.clone();

let lambda = Real::unary(log_lambda, UnaryOp::Exp);

let term = Poisson::new(lambda).log_density_at(scalar(counts[count_idx]));

model = model.observe_with(&[0.0], |_| term.clone());

count_idx += 1;

}

}

assert_eq!(count_idx, counts.len());

model

}

fn build_low_dim_gauss_mix_model(rainier_root: &Path) -> Model {

let path = rainier_root.join(

"rainier-benchmark/src/main/scala/com/stripe/rainier/bench/stan/LowDimGaussMix.scala",

);

let source = parse_file(&path);

let ys = parse_scala_list(&source, "val ys: List[Double] = List(");

let (mu1_param, mu1) = Normal::new(scalar(0.0), scalar(2.0)).latent(ParameterId(0));

let (sigma1_latent_param, sigma1_latent) =

Normal::new(scalar(0.0), scalar(2.0)).latent(ParameterId(1));

let (mu2_param, mu2) = Normal::new(scalar(0.0), scalar(2.0)).latent(ParameterId(2));

let (sigma2_latent_param, sigma2_latent) =

Normal::new(scalar(0.0), scalar(2.0)).latent(ParameterId(3));

let (theta_param, theta) = Beta::new(scalar(5.0), scalar(5.0)).latent(ParameterId(4));

let sigma1 = Real::unary(sigma1_latent, UnaryOp::Abs);

let sigma2 = Real::unary(sigma2_latent, UnaryOp::Abs);

let mut model = Model::new();

model = add_prior(model, mu1_param);

model = add_prior(model, sigma1_latent_param);

model = add_prior(model, mu2_param);

model = add_prior(model, sigma2_latent_param);

model = add_prior(model, theta_param);

let mixture = Mixture::new(vec![

(

Box::new(Normal::new(mu1, sigma1)) as Box<dyn Continuous>,

theta.clone(),

),

(

Box::new(Normal::new(mu2, sigma2)) as Box<dyn Continuous>,

scalar(1.0) - theta,

),

]);

model.observe(&ys, &mixture)

}

fn build_model(name: &str, rainier_root: &Path) -> Model {

match name {

"ARK" => build_ark_model(rainier_root),

"EightSchools" => build_eight_schools_model(rainier_root),

"KidIQ" => build_kid_iq_model(rainier_root),

"GLMMPoisson2" => build_glmm_poisson2_model(rainier_root),

"LowDimGaussMix" => build_low_dim_gauss_mix_model(rainier_root),

other => panic!("unsupported model in parity file: {other}"),

}

}

fn relative_error(lhs: f64, rhs: f64) -> f64 {

let scale = lhs.abs().max(rhs.abs()).max(1e-12);

(lhs - rhs).abs() / scale

}

#[test]

#[ignore = "requires Rainier benchmark snapshots and is computationally heavy"]

fn rainier_benchmark_parity_same_samples() {

let root = workspace_root();

let rainier_root = root.join("rainier");

let parity_path = rainier_root.join("target/priori-parity.json");

assert!(

parity_path.exists(),

"missing parity snapshot: {}",

parity_path.display()

);

let payload = fs::read_to_string(&parity_path).unwrap_or_else(|err| {

panic!(

"failed to read parity snapshot {}: {err}",

parity_path.display()

)

});

let value: Value = serde_json::from_str(&payload).expect("valid parity json");

let models = value["models"].as_array().expect("models array");

assert_eq!(models.len(), 5);

let mut global_max_density_rel = 0.0_f64;

let mut global_max_grad_rel = 0.0_f64;

let mut faster_or_equal_models = 0_usize;

let mut models_with_rainier_timing = 0_usize;

for model_value in models {

let name = model_value["name"].as_str().expect("model name");

let n_vars = model_value["n_vars"].as_u64().expect("n_vars") as usize;

let cases = model_value["cases"].as_array().expect("cases");

assert_eq!(cases.len(), 5, "{name} must have 5 parity cases");

let model = build_model(name, &rainier_root);

let density_fn = model.density_fn();

assert_eq!(density_fn.n_vars(), n_vars, "{name} n_vars mismatch");

let mut model_max_density_rel = 0.0_f64;

let mut model_max_grad_rel = 0.0_f64;

let mut total_eval_us = 0.0_f64;

let mut eval_count = 0_usize;

let mut total_rainier_eval_us = 0.0_f64;

let mut rainier_eval_count = 0_usize;

let mut nonfinite_cases = 0_usize;

for case_value in cases {

let params = case_value["params"]

.as_array()

.expect("params")

.iter()

.map(|v| v.as_f64().expect("param f64"))

.collect::<std::vec::Vec<_>>();

let rainier_density = case_value["density"].as_f64().expect("density f64");

let rainier_gradient = case_value["gradient"]

.as_array()

.expect("gradient")

.iter()

.map(|v| v.as_f64().expect("grad f64"))

.collect::<std::vec::Vec<_>>();

if let Some(rainier_eval_us) = case_value["rainier_eval_us"].as_f64() {

total_rainier_eval_us += rainier_eval_us;

rainier_eval_count += 1;

}

assert_eq!(params.len(), n_vars);

assert_eq!(rainier_gradient.len(), n_vars);

let mut grad = vec![0.0_f64; n_vars];

let potential = density_fn.evaluate(&params, &mut grad);

let priori_density = -potential;

let priori_gradient = grad.into_iter().map(|v| -v).collect::<std::vec::Vec<_>>();

if !priori_density.is_finite() || !priori_gradient.iter().all(|v| v.is_finite()) {

nonfinite_cases += 1;

continue;

}

model_max_density_rel =

model_max_density_rel.max(relative_error(priori_density, rainier_density));

for (lhs, rhs) in priori_gradient.iter().zip(rainier_gradient.iter()) {

model_max_grad_rel = model_max_grad_rel.max(relative_error(*lhs, *rhs));

}

let repeats = if n_vars > 15 { 10 } else { 50 };

let mut timing_grad = vec![0.0_f64; n_vars];

let start = Instant::now();

for _ in 0..repeats {

let _ = density_fn.evaluate(&params, &mut timing_grad);

}

let elapsed = start.elapsed().as_secs_f64() * 1e6 / repeats as f64;

total_eval_us += elapsed;

eval_count += 1;

}

assert!(

eval_count > 0,

"{name} produced only non-finite parity cases"

);

global_max_density_rel = global_max_density_rel.max(model_max_density_rel);

global_max_grad_rel = global_max_grad_rel.max(model_max_grad_rel);

let mean_eval_us = total_eval_us / eval_count as f64;

if rainier_eval_count > 0 {

let mean_rainier_eval_us = total_rainier_eval_us / rainier_eval_count as f64;

models_with_rainier_timing += 1;

if mean_eval_us <= mean_rainier_eval_us {

faster_or_equal_models += 1;

}

println!(

"{name}: max_density_rel={model_max_density_rel:.3e} max_grad_rel={model_max_grad_rel:.3e} priori_mean_eval_us={mean_eval_us:.3} rainier_mean_eval_us={mean_rainier_eval_us:.3} nonfinite_cases={nonfinite_cases}"

);

} else {

println!(

"{name}: max_density_rel={model_max_density_rel:.3e} max_grad_rel={model_max_grad_rel:.3e} mean_eval_us={mean_eval_us:.3} nonfinite_cases={nonfinite_cases}"

);

}

}

println!(

"global: max_density_rel={global_max_density_rel:.3e} max_grad_rel={global_max_grad_rel:.3e}"

);

if models_with_rainier_timing > 0 {

println!(

"timing_comparison: priori_faster_or_equal={faster_or_equal_models}/{models_with_rainier_timing}"

);

}

}