use crate::equivalence::{EquivalenceBounds, OneSidedTestResult, TostResult};
use crate::error::{Result, StatError};
use statrs::distribution::{ContinuousCDF, Normal};
pub fn tost_prop_one(
x: usize,
n: usize,
p0: f64,
bounds: &EquivalenceBounds,
alpha: f64,
) -> Result<TostResult> {
validate_alpha(alpha)?;
if n == 0 {
return Err(StatError::EmptyData);
}
if x > n {
return Err(StatError::InvalidParameter(format!(
"x ({}) cannot exceed n ({})",
x, n
)));
}
if !(0.0..=1.0).contains(&p0) {
return Err(StatError::InvalidParameter(format!(
"p0 must be between 0 and 1, got {}",
p0
)));
}
let p_hat = x as f64 / n as f64;
let (lower_bound, upper_bound) = bounds_for_proportion(bounds)?;
let se = (p_hat * (1.0 - p_hat) / n as f64).sqrt();
if se == 0.0 {
let estimate = p_hat - p0;
let equivalent = estimate >= lower_bound && estimate <= upper_bound;
return Ok(TostResult {
estimate,
ci: (estimate, estimate),
bounds: (lower_bound, upper_bound),
lower_test: OneSidedTestResult {
hypothesis: format!("H0: p - p0 <= {:.4}", lower_bound),
statistic: f64::INFINITY,
p_value: if estimate > lower_bound { 0.0 } else { 1.0 },
rejected: estimate > lower_bound,
},
upper_test: OneSidedTestResult {
hypothesis: format!("H0: p - p0 >= {:.4}", upper_bound),
statistic: f64::NEG_INFINITY,
p_value: if estimate < upper_bound { 0.0 } else { 1.0 },
rejected: estimate < upper_bound,
},
tost_p_value: if equivalent { 0.0 } else { 1.0 },
equivalent,
alpha,
n,
df: None,
method: "One-proportion TOST".to_string(),
});
}
let estimate = p_hat - p0;
let normal = Normal::new(0.0, 1.0)
.map_err(|e| StatError::InvalidParameter(format!("Failed to create normal: {}", e)))?;
let z_lower = (estimate - lower_bound) / se;
let p_lower = normal.sf(z_lower);
let z_upper = (estimate - upper_bound) / se;
let p_upper = normal.cdf(z_upper);
let tost_p = p_lower.max(p_upper);
let z_crit = normal.inverse_cdf(1.0 - alpha);
let margin = z_crit * se;
let ci = (estimate - margin, estimate + margin);
let equivalent = ci.0 >= lower_bound && ci.1 <= upper_bound;
Ok(TostResult {
estimate,
ci,
bounds: (lower_bound, upper_bound),
lower_test: OneSidedTestResult {
hypothesis: format!("H0: p - p0 <= {:.4}", lower_bound),
statistic: z_lower,
p_value: p_lower,
rejected: p_lower < alpha,
},
upper_test: OneSidedTestResult {
hypothesis: format!("H0: p - p0 >= {:.4}", upper_bound),
statistic: z_upper,
p_value: p_upper,
rejected: p_upper < alpha,
},
tost_p_value: tost_p,
equivalent,
alpha,
n,
df: None,
method: "One-proportion TOST".to_string(),
})
}
pub fn tost_prop_two(
x1: usize,
n1: usize,
x2: usize,
n2: usize,
bounds: &EquivalenceBounds,
alpha: f64,
) -> Result<TostResult> {
validate_alpha(alpha)?;
if n1 == 0 || n2 == 0 {
return Err(StatError::EmptyData);
}
if x1 > n1 {
return Err(StatError::InvalidParameter(format!(
"x1 ({}) cannot exceed n1 ({})",
x1, n1
)));
}
if x2 > n2 {
return Err(StatError::InvalidParameter(format!(
"x2 ({}) cannot exceed n2 ({})",
x2, n2
)));
}
let p1 = x1 as f64 / n1 as f64;
let p2 = x2 as f64 / n2 as f64;
let (lower_bound, upper_bound) = bounds_for_proportion(bounds)?;
let estimate = p1 - p2;
let var1 = p1 * (1.0 - p1) / n1 as f64;
let var2 = p2 * (1.0 - p2) / n2 as f64;
let se = (var1 + var2).sqrt();
if se < 1e-15 {
let equivalent = estimate >= lower_bound && estimate <= upper_bound;
return Ok(TostResult {
estimate,
ci: (estimate, estimate),
bounds: (lower_bound, upper_bound),
lower_test: OneSidedTestResult {
hypothesis: format!("H0: p1 - p2 <= {:.4}", lower_bound),
statistic: f64::INFINITY,
p_value: if estimate > lower_bound { 0.0 } else { 1.0 },
rejected: estimate > lower_bound,
},
upper_test: OneSidedTestResult {
hypothesis: format!("H0: p1 - p2 >= {:.4}", upper_bound),
statistic: f64::NEG_INFINITY,
p_value: if estimate < upper_bound { 0.0 } else { 1.0 },
rejected: estimate < upper_bound,
},
tost_p_value: if equivalent { 0.0 } else { 1.0 },
equivalent,
alpha,
n: n1 + n2,
df: None,
method: "Two-proportion TOST".to_string(),
});
}
let normal = Normal::new(0.0, 1.0)
.map_err(|e| StatError::InvalidParameter(format!("Failed to create normal: {}", e)))?;
let z_lower = (estimate - lower_bound) / se;
let p_lower = normal.sf(z_lower);
let z_upper = (estimate - upper_bound) / se;
let p_upper = normal.cdf(z_upper);
let tost_p = p_lower.max(p_upper);
let z_crit = normal.inverse_cdf(1.0 - alpha);
let margin = z_crit * se;
let ci = (estimate - margin, estimate + margin);
let equivalent = ci.0 >= lower_bound && ci.1 <= upper_bound;
Ok(TostResult {
estimate,
ci,
bounds: (lower_bound, upper_bound),
lower_test: OneSidedTestResult {
hypothesis: format!("H0: p1 - p2 <= {:.4}", lower_bound),
statistic: z_lower,
p_value: p_lower,
rejected: p_lower < alpha,
},
upper_test: OneSidedTestResult {
hypothesis: format!("H0: p1 - p2 >= {:.4}", upper_bound),
statistic: z_upper,
p_value: p_upper,
rejected: p_upper < alpha,
},
tost_p_value: tost_p,
equivalent,
alpha,
n: n1 + n2,
df: None,
method: "Two-proportion TOST".to_string(),
})
}
fn bounds_for_proportion(bounds: &EquivalenceBounds) -> Result<(f64, f64)> {
match bounds {
EquivalenceBounds::Raw { lower, upper } => {
if *lower < -1.0 || *upper > 1.0 {
return Err(StatError::InvalidParameter(
"Proportion bounds must be between -1 and 1".to_string(),
));
}
Ok((*lower, *upper))
}
EquivalenceBounds::Symmetric { delta } => {
if *delta > 1.0 {
return Err(StatError::InvalidParameter(
"Symmetric delta for proportions must not exceed 1".to_string(),
));
}
Ok((-*delta, *delta))
}
EquivalenceBounds::CohenD { .. } => Err(StatError::InvalidParameter(
"Cohen's d bounds not applicable for proportion TOST; use Raw or Symmetric bounds"
.to_string(),
)),
}
}
fn validate_alpha(alpha: f64) -> Result<()> {
if !(0.0 < alpha && alpha < 1.0) {
return Err(StatError::InvalidParameter(format!(
"alpha must be between 0 and 1, got {}",
alpha
)));
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_one_prop_equivalent() {
let bounds = EquivalenceBounds::Symmetric { delta: 0.1 };
let result = tost_prop_one(52, 100, 0.5, &bounds, 0.05).unwrap();
assert!(result.estimate.abs() < 0.1);
}
#[test]
fn test_one_prop_not_equivalent() {
let bounds = EquivalenceBounds::Symmetric { delta: 0.05 };
let result = tost_prop_one(70, 100, 0.5, &bounds, 0.05).unwrap();
assert!(!result.equivalent);
}
#[test]
fn test_two_prop_equivalent() {
let bounds = EquivalenceBounds::Symmetric { delta: 0.1 };
let result = tost_prop_two(51, 100, 49, 100, &bounds, 0.05).unwrap();
assert!(result.estimate.abs() < 0.1);
}
#[test]
fn test_two_prop_not_equivalent() {
let bounds = EquivalenceBounds::Symmetric { delta: 0.05 };
let result = tost_prop_two(70, 100, 50, 100, &bounds, 0.05).unwrap();
assert!(!result.equivalent);
}
#[test]
fn test_invalid_x_greater_than_n() {
let bounds = EquivalenceBounds::Symmetric { delta: 0.1 };
assert!(tost_prop_one(110, 100, 0.5, &bounds, 0.05).is_err());
}
#[test]
fn test_invalid_p0() {
let bounds = EquivalenceBounds::Symmetric { delta: 0.1 };
assert!(tost_prop_one(50, 100, 1.5, &bounds, 0.05).is_err());
assert!(tost_prop_one(50, 100, -0.1, &bounds, 0.05).is_err());
}
#[test]
fn test_cohen_d_not_allowed() {
let bounds = EquivalenceBounds::CohenD { d: 0.5 };
assert!(tost_prop_one(50, 100, 0.5, &bounds, 0.05).is_err());
assert!(tost_prop_two(50, 100, 50, 100, &bounds, 0.05).is_err());
}
#[test]
fn test_edge_case_all_successes() {
let bounds = EquivalenceBounds::Symmetric { delta: 0.1 };
let result = tost_prop_one(100, 100, 0.95, &bounds, 0.05).unwrap();
assert!((result.estimate - 0.05).abs() < 1e-10);
}
}