use std::error::Error;
use std::f64;
use ndarray::Array1;
use ndarray::Array2;
use rand::Rng;
use stochastic_rs_distributions::gamma::SimdGamma;
use stochastic_rs_distributions::special::ln_gamma;
use super::CopulaType;
use crate::traits::MultivariateExt;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NacFamily {
Clayton,
Gumbel,
}
impl NacFamily {
pub fn theta_min(self) -> f64 {
match self {
NacFamily::Clayton => 0.0,
NacFamily::Gumbel => 1.0,
}
}
pub fn inverse_generator(self, theta: f64, s: f64) -> f64 {
match self {
NacFamily::Clayton => (1.0 + s).powf(-1.0 / theta),
NacFamily::Gumbel => (-(s.powf(1.0 / theta))).exp(),
}
}
pub fn generator(self, theta: f64, t: f64) -> f64 {
let t = t.clamp(1e-15, 1.0 - 1e-15);
match self {
NacFamily::Clayton => t.powf(-theta) - 1.0,
NacFamily::Gumbel => (-t.ln()).powf(theta),
}
}
}
#[derive(Debug, Clone)]
pub struct NacNode {
pub theta: f64,
pub leaves: Vec<usize>,
pub children: Vec<NacNode>,
}
impl NacNode {
pub fn leaf_group(theta: f64, leaves: Vec<usize>) -> Self {
Self {
theta,
leaves,
children: vec![],
}
}
fn collect_leaves(&self, out: &mut Vec<usize>) {
out.extend_from_slice(&self.leaves);
for c in &self.children {
c.collect_leaves(out);
}
}
}
#[derive(Debug, Clone)]
pub struct NestedArchimedean {
family: NacFamily,
root: NacNode,
dim: usize,
index_order: Vec<usize>,
}
impl NestedArchimedean {
pub fn new(family: NacFamily, root: NacNode, dim: usize) -> Result<Self, Box<dyn Error>> {
Self::validate_node(family, &root, root.theta.min(f64::INFINITY), true)?;
let mut index_order = Vec::with_capacity(dim);
root.collect_leaves(&mut index_order);
if index_order.len() != dim {
return Err(
format!(
"NAC tree exposes {} leaves but dim = {dim}",
index_order.len()
)
.into(),
);
}
let mut seen = vec![false; dim];
for &j in &index_order {
if j >= dim {
return Err(format!("leaf index {j} ≥ dim {dim}").into());
}
if seen[j] {
return Err(format!("leaf index {j} appears more than once").into());
}
seen[j] = true;
}
if seen.iter().any(|&b| !b) {
return Err("Not every marginal appears in the NAC tree".into());
}
Ok(Self {
family,
root,
dim,
index_order,
})
}
fn validate_node(
family: NacFamily,
node: &NacNode,
parent_theta: f64,
is_root: bool,
) -> Result<(), Box<dyn Error>> {
let theta_min = family.theta_min();
if node.theta < theta_min {
return Err(
format!(
"{family:?} node θ={} below family minimum {theta_min}",
node.theta
)
.into(),
);
}
if !is_root && node.theta < parent_theta {
return Err(
format!(
"SNC violation: child θ={} < parent θ={} ({family:?})",
node.theta, parent_theta
)
.into(),
);
}
for child in &node.children {
Self::validate_node(family, child, node.theta, false)?;
}
Ok(())
}
pub fn family(&self) -> NacFamily {
self.family
}
pub fn root(&self) -> &NacNode {
&self.root
}
pub fn dim(&self) -> usize {
self.dim
}
pub fn index_order(&self) -> &[usize] {
&self.index_order
}
fn positive_stable<R: Rng + ?Sized>(rng: &mut R, alpha: f64) -> f64 {
debug_assert!(alpha > 0.0 && alpha < 1.0);
let u: f64 = rng.random::<f64>().clamp(1e-15, 1.0 - 1e-15);
let theta = f64::consts::PI * u;
let w_uniform: f64 = rng.random::<f64>().clamp(1e-15, 1.0 - 1e-15);
let w = -w_uniform.ln();
let s_a = (alpha * theta).sin();
let s_oa = ((1.0 - alpha) * theta).sin();
let s_t = theta.sin();
let exponent = (1.0 - alpha) / alpha;
let numerator = s_a * s_oa.powf(exponent);
let denominator = s_t.powf(1.0 / alpha) * w.powf(exponent);
numerator / denominator
}
fn sample_node<R: Rng + ?Sized>(
&self,
rng: &mut R,
node: &NacNode,
parent_state: Option<(f64, f64)>,
out: &mut Array1<f64>,
) {
let v = match (self.family, parent_state) {
(NacFamily::Clayton, None) => {
let g = SimdGamma::<f64>::new(
1.0 / node.theta,
1.0,
&stochastic_rs_core::simd_rng::Unseeded,
);
g.sample_fast()
}
(NacFamily::Gumbel, None) => {
if (node.theta - 1.0).abs() < 1e-12 {
1.0
} else {
Self::positive_stable(rng, 1.0 / node.theta)
}
}
(NacFamily::Gumbel, Some((parent_theta, parent_v))) => {
let alpha = parent_theta / node.theta;
if (alpha - 1.0).abs() < 1e-12 {
parent_v
} else {
let s = Self::positive_stable(rng, alpha);
parent_v.powf(1.0 / alpha) * s
}
}
(NacFamily::Clayton, Some(_)) => {
panic!(
"Nested-Clayton sampling is not implemented — \
use NacFamily::Gumbel for nesting (correct Hofert (2011) \
Algorithm 2), or call cdf / pdf on a nested-Clayton tree \
(those paths are fully supported). Exact nested-Clayton \
sampling via tilted-stable Devroye double-rejection is \
not yet implemented."
);
}
};
for &j in &node.leaves {
let e_uniform: f64 = rng.random::<f64>().clamp(1e-15, 1.0 - 1e-15);
let e = -e_uniform.ln();
let arg = e / v.max(1e-300);
let u = self.family.inverse_generator(node.theta, arg);
out[j] = u.clamp(1e-12, 1.0 - 1e-12);
}
for child in &node.children {
self.sample_node(rng, child, Some((node.theta, v)), out);
}
}
fn cdf_node(&self, node: &NacNode, u: &[f64]) -> f64 {
let mut s = 0.0;
for &j in &node.leaves {
s += self.family.generator(node.theta, u[j]);
}
for child in &node.children {
let c_inner = self.cdf_node(child, u);
s += self.family.generator(node.theta, c_inner);
}
self.family.inverse_generator(node.theta, s)
}
}
impl MultivariateExt for NestedArchimedean {
fn r#type(&self) -> CopulaType {
CopulaType::NestedArchimedean
}
fn sample(&self, n: usize) -> Result<Array2<f64>, Box<dyn Error>> {
let d = self.dim;
let mut out = Array2::<f64>::zeros((n, d));
let mut rng = rand::rng();
for r in 0..n {
let mut row = Array1::<f64>::zeros(d);
self.sample_node(&mut rng, &self.root, None, &mut row);
for j in 0..d {
out[[r, j]] = row[j];
}
}
Ok(out)
}
fn fit(&mut self, _X: Array2<f64>) -> Result<(), Box<dyn Error>> {
Err(
"NestedArchimedean::fit not implemented — supply the tree via NestedArchimedean::new \
and use crate::correlation::kendall_tau on the marginal pairs to seed θ values. \
Structure learning is not yet implemented."
.into(),
)
}
fn check_fit(&self, X: &Array2<f64>) -> Result<(), Box<dyn Error>> {
if X.ncols() != self.dim {
return Err(
format!(
"Dimension mismatch: X has {} columns, NAC has dim {}",
X.ncols(),
self.dim
)
.into(),
);
}
if X.iter().any(|&v| !(0.0..=1.0).contains(&v)) {
return Err("Input X must be in [0,1] for NAC".into());
}
Ok(())
}
fn pdf(&self, X: Array2<f64>) -> Result<Array1<f64>, Box<dyn Error>> {
self.check_fit(&X)?;
let d = self.dim;
let h = 1e-4_f64;
let denom = (2.0 * h).powi(d as i32);
let mut out = Array1::<f64>::zeros(X.nrows());
let mut u_pert = vec![0.0_f64; d];
for (i, row) in X.rows().into_iter().enumerate() {
let u_orig: Vec<f64> = row.iter().copied().collect();
let mut acc = 0.0;
for mask in 0..(1u32 << d) {
let mut sign = 1.0;
for j in 0..d {
if (mask >> j) & 1 == 1 {
u_pert[j] = (u_orig[j] + h).min(1.0 - 1e-12);
} else {
u_pert[j] = (u_orig[j] - h).max(1e-12);
sign = -sign;
}
}
acc += sign * self.cdf_node(&self.root, &u_pert);
}
out[i] = (acc / denom).max(0.0);
}
Ok(out)
}
fn cdf(&self, X: Array2<f64>) -> Result<Array1<f64>, Box<dyn Error>> {
self.check_fit(&X)?;
let mut out = Array1::<f64>::zeros(X.nrows());
for (i, row) in X.rows().into_iter().enumerate() {
let u: Vec<f64> = row.iter().copied().collect();
out[i] = self.cdf_node(&self.root, &u);
}
Ok(out)
}
}
#[allow(dead_code)]
pub(crate) fn clayton_log_density_constant(theta: f64, d: usize) -> f64 {
ln_gamma(d as f64 + 1.0 / theta) - ln_gamma(1.0 / theta)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn nac_clayton_flat_is_exchangeable_clayton() {
let root = NacNode::leaf_group(2.0, vec![0, 1, 2]);
let nac = NestedArchimedean::new(NacFamily::Clayton, root, 3).unwrap();
let u = nac.sample(8_000).unwrap();
assert_eq!(u.ncols(), 3);
for j in 0..3 {
let col = u.column(j);
let mean = col.iter().sum::<f64>() / col.len() as f64;
assert!(
(mean - 0.5).abs() < 0.03,
"marginal {j} mean = {mean}, expected ~0.5"
);
}
}
#[test]
fn nac_clayton_two_level_cdf_pair_margins() {
let inner = NacNode::leaf_group(4.0, vec![1, 2]);
let root = NacNode {
theta: 1.5,
leaves: vec![0],
children: vec![inner],
};
let nac = NestedArchimedean::new(NacFamily::Clayton, root, 3).unwrap();
let q_outer_inner = ndarray::array![[0.3, 0.7, 1.0 - 1e-15]];
let c_oi = nac.cdf(q_outer_inner).unwrap()[0];
let theta_root: f64 = 1.5;
let expected_oi =
(0.3f64.powf(-theta_root) + 0.7f64.powf(-theta_root) - 1.0).powf(-1.0 / theta_root);
assert!(
(c_oi - expected_oi).abs() < 5e-3,
"outer-inner pair CDF={} vs Clayton(θ_root) expected={}",
c_oi,
expected_oi
);
let q_inner_inner = ndarray::array![[1.0 - 1e-15, 0.3, 0.7]];
let c_ii = nac.cdf(q_inner_inner).unwrap()[0];
let theta_inner: f64 = 4.0;
let expected_ii =
(0.3f64.powf(-theta_inner) + 0.7f64.powf(-theta_inner) - 1.0).powf(-1.0 / theta_inner);
assert!(
(c_ii - expected_ii).abs() < 5e-3,
"inner-inner pair CDF={} vs Clayton(θ_inner) expected={}",
c_ii,
expected_ii
);
assert!(
c_ii > c_oi,
"inner CDF({c_ii}) must exceed outer CDF({c_oi}) — Clayton inner θ=4 > root θ=1.5"
);
}
#[test]
#[should_panic(expected = "Nested-Clayton sampling")]
fn nac_clayton_nested_sampling_panics() {
let inner = NacNode::leaf_group(4.0, vec![1]);
let root = NacNode {
theta: 1.5,
leaves: vec![0],
children: vec![inner],
};
let nac = NestedArchimedean::new(NacFamily::Clayton, root, 2).unwrap();
let _ = nac.sample(10);
}
#[test]
fn nac_gumbel_two_level_inner_pair_more_dependent() {
let inner = NacNode::leaf_group(4.0, vec![1, 2]);
let root = NacNode {
theta: 2.0,
leaves: vec![0],
children: vec![inner],
};
let nac = NestedArchimedean::new(NacFamily::Gumbel, root, 3).unwrap();
let u = nac.sample(8_000).unwrap();
use crate::correlation::kendall_tau;
let tau = kendall_tau(&u);
assert!(
tau[[1, 2]] > tau[[0, 1]],
"inner pair τ_(1,2)={} should exceed outer τ_(0,1)={}",
tau[[1, 2]],
tau[[0, 1]]
);
assert!(
tau[[1, 2]] > 0.6 && tau[[1, 2]] < 0.85,
"Gumbel inner τ_(1,2)={} out of expected band [0.6, 0.85]",
tau[[1, 2]]
);
}
#[test]
fn nac_snc_violation_rejected() {
let bad_inner = NacNode::leaf_group(0.5, vec![1]); let bad_root = NacNode {
theta: 2.0,
leaves: vec![0],
children: vec![bad_inner],
};
let res = NestedArchimedean::new(NacFamily::Clayton, bad_root, 2);
assert!(res.is_err(), "SNC violation must error");
assert!(
res.unwrap_err().to_string().contains("SNC"),
"error message should mention SNC"
);
}
#[test]
fn nac_below_family_min_rejected() {
let root = NacNode::leaf_group(0.5, vec![0, 1]); let res = NestedArchimedean::new(NacFamily::Gumbel, root, 2);
assert!(res.is_err());
let bad_clayton = NacNode::leaf_group(-0.1, vec![0, 1]);
assert!(NestedArchimedean::new(NacFamily::Clayton, bad_clayton, 2).is_err());
}
#[test]
fn nac_leaf_index_validation() {
let dup_root = NacNode::leaf_group(2.0, vec![0, 0, 1]);
assert!(NestedArchimedean::new(NacFamily::Clayton, dup_root, 3).is_err());
let miss = NacNode::leaf_group(2.0, vec![0, 2]);
assert!(NestedArchimedean::new(NacFamily::Clayton, miss, 3).is_err());
let oor = NacNode::leaf_group(2.0, vec![0, 1, 5]);
assert!(NestedArchimedean::new(NacFamily::Clayton, oor, 3).is_err());
}
#[test]
fn nac_clayton_independence_cdf() {
let root = NacNode::leaf_group(0.01, vec![0, 1, 2]);
let nac = NestedArchimedean::new(NacFamily::Clayton, root, 3).unwrap();
let q = ndarray::array![[0.5, 0.5, 0.5], [0.2, 0.3, 0.4]];
let c = nac.cdf(q.clone()).unwrap();
let indep_1 = 0.5_f64.powi(3);
let indep_2 = 0.2_f64 * 0.3 * 0.4;
assert!(
(c[0] - indep_1).abs() < 0.05,
"near-independence CDF[0]={} vs indep={indep_1}",
c[0]
);
assert!(
(c[1] - indep_2).abs() < 0.05,
"near-independence CDF[1]={} vs indep={indep_2}",
c[1]
);
}
#[test]
fn nac_fit_rejects_with_descriptive_error() {
let root = NacNode::leaf_group(2.0, vec![0, 1]);
let mut nac = NestedArchimedean::new(NacFamily::Clayton, root, 2).unwrap();
let data = ndarray::Array2::<f64>::from_elem((10, 2), 0.5);
let res = nac.fit(data);
assert!(res.is_err());
let msg = res.unwrap_err().to_string();
assert!(
msg.contains("structure") || msg.contains("not implemented"),
"fit error should explain that structure learning is not implemented; got: {msg}"
);
}
}