use ndarray::Array2;
use crate::{
errors::{CopulaError, FitError},
paircopula::PairCopulaSpec,
};
use super::{
CompiledEvalStep, CompiledSampleStep, CompiledVineRuntime, VineCopula, VineEdge, VineStructure,
VineStructureKind, VineTree,
};
pub(crate) fn validate_order(order: &[usize], dim: usize) -> Result<(), CopulaError> {
if order.len() != dim {
return Err(FitError::Failed {
reason: "vine order length must match the correlation dimension",
}
.into());
}
let mut seen = vec![false; dim];
for &value in order {
if value >= dim || seen[value] {
return Err(FitError::Failed {
reason: "vine order must be a permutation of 0..d-1",
}
.into());
}
seen[value] = true;
}
Ok(())
}
pub(crate) fn build_model_from_trees(
kind: VineStructureKind,
trees: Vec<VineTree>,
truncation_level: Option<usize>,
) -> Result<VineCopula, CopulaError> {
let dim = trees.len() + 1;
let (matrix, pair_matrix) = to_structure_matrices(&trees)?;
let variable_order = matrix.diag().iter().rev().copied().collect::<Vec<_>>();
let normalized_matrix = normalize_matrix(&matrix, &variable_order);
let max_matrix = create_max_matrix(&normalized_matrix);
let (cond_direct, cond_indirect) =
needed_conditional_distributions(&normalized_matrix, &max_matrix);
let runtime = compile_runtime(
&normalized_matrix,
&max_matrix,
&cond_indirect,
&pair_matrix,
&variable_order,
);
Ok(VineCopula {
dim,
structure: VineStructure {
kind,
matrix,
truncation_level,
},
trees,
normalized_matrix,
variable_order,
pair_matrix,
max_matrix,
cond_direct,
cond_indirect,
runtime,
})
}
pub(crate) fn compile_runtime(
normalized_matrix: &Array2<usize>,
max_matrix: &Array2<usize>,
cond_indirect: &Array2<bool>,
pair_matrix: &Array2<Option<PairCopulaSpec>>,
variable_order: &[usize],
) -> CompiledVineRuntime {
let mat = revert_matrix(normalized_matrix);
let maxmat = revert_matrix(max_matrix);
let cindirect = revert_matrix(cond_indirect);
let specs = revert_pair_matrix(pair_matrix);
let d = normalized_matrix.nrows();
let mut sample_steps = Vec::with_capacity(d.saturating_mul(d.saturating_sub(1)) / 2);
let mut eval_steps = Vec::with_capacity(sample_steps.capacity());
let mut all_gaussian = true;
for i in 1..d {
for k in (0..i).rev() {
let spec = specs[(k, i)]
.clone()
.expect("compiled vine runtime requires all pair-copula specs to be present");
all_gaussian &= matches!(
(spec.family, spec.rotation, &spec.params),
(
crate::paircopula::PairCopulaFamily::Gaussian,
crate::paircopula::Rotation::R0,
crate::paircopula::PairCopulaParams::One(_)
)
);
sample_steps.push(CompiledSampleStep {
row: k,
col: i,
label: maxmat[(k, i)],
source_from_direct: mat[(k, i)] == maxmat[(k, i)],
write_indirect: i + 1 < d && cindirect[(k + 1, i)],
spec,
});
}
}
for i in 1..d {
for k in 0..i {
let spec = specs[(k, i)]
.clone()
.expect("compiled vine runtime requires all pair-copula specs to be present");
eval_steps.push(CompiledEvalStep {
row: k,
col: i,
label: maxmat[(k, i)],
source_from_direct: mat[(k, i)] == maxmat[(k, i)],
write_indirect: i + 1 < d && cindirect[(k + 1, i)],
spec,
});
}
}
CompiledVineRuntime {
dim: d,
variable_order: variable_order.to_vec(),
sample_steps,
eval_steps,
all_gaussian,
}
}
fn to_structure_matrices(
trees: &[VineTree],
) -> Result<(Array2<usize>, Array2<Option<PairCopulaSpec>>), CopulaError> {
let n = trees.len() + 1;
let mut matrix = Array2::zeros((n, n));
let mut pair_matrix = Array2::from_elem((n, n), None);
let mut remaining = trees
.iter()
.map(|tree| tree.edges.clone())
.collect::<Vec<_>>();
for k in 0..(n - 1) {
let source_tree = n - k - 2;
let first = remaining[source_tree]
.first()
.ok_or(FitError::Failed {
reason: "vine tree is missing required edges for structure conversion",
})?
.clone();
let w = first.conditioned.0;
matrix[(k, k)] = w;
matrix[(k + 1, k)] = first.conditioned.1;
pair_matrix[(k + 1, k)] = Some(first.copula.clone());
if k == n - 2 {
matrix[(k + 1, k + 1)] = first.conditioned.1;
continue;
}
for i in (k + 2)..n {
let tree_idx = n - i - 1;
let mut found = None;
for (idx, edge) in remaining[tree_idx].iter().enumerate() {
if edge.conditioned.0 == w {
found = Some((idx, edge.conditioned.1, edge.copula.clone()));
break;
}
if edge.conditioned.1 == w {
found = Some((idx, edge.conditioned.0, edge.copula.clone().swap_axes()));
break;
}
}
let (idx, value, spec) = found.ok_or(FitError::Failed {
reason: "failed to convert selected vine trees into a structure matrix",
})?;
matrix[(i, k)] = value;
pair_matrix[(i, k)] = Some(spec);
remaining[tree_idx].remove(idx);
}
}
Ok((matrix, pair_matrix))
}
fn create_max_matrix(matrix: &Array2<usize>) -> Array2<usize> {
let mut max_matrix = matrix.clone();
let n = max_matrix.nrows();
for col in 0..(n - 1) {
for row in (col..=(n - 2)).rev() {
max_matrix[(row, col)] = max_matrix[(row, col)].max(max_matrix[(row + 1, col)]);
}
}
max_matrix
}
fn normalize_matrix(matrix: &Array2<usize>, variable_order: &[usize]) -> Array2<usize> {
let mut mapping = vec![0usize; variable_order.len()];
for (idx, &value) in variable_order.iter().enumerate() {
mapping[value] = idx;
}
let mut normalized = Array2::zeros(matrix.raw_dim());
for ((row, col), value) in matrix.indexed_iter() {
normalized[(row, col)] = mapping[*value];
}
normalized
}
fn needed_conditional_distributions(
matrix: &Array2<usize>,
max_matrix: &Array2<usize>,
) -> (Array2<bool>, Array2<bool>) {
let d = matrix.nrows();
let mut direct = Array2::from_elem((d, d), false);
let mut indirect = Array2::from_elem((d, d), false);
for row in 1..d {
direct[(row, 0)] = true;
}
for i in 1..(d - 1) {
let v = d - i - 1;
for row in i..d {
let mut any_bw_and_not_direct = false;
let mut any_first = false;
for col in 0..i {
let bw = max_matrix[(row, col)] == v;
let is_direct = matrix[(row, col)] == v;
any_bw_and_not_direct |= bw && !is_direct;
if row == i {
any_first |= bw && is_direct;
}
}
indirect[(row, i)] = any_bw_and_not_direct;
direct[(row, i)] = true;
if row == i {
direct[(i, i)] = any_first;
}
}
}
(direct, indirect)
}
pub(crate) fn revert_matrix<T: Clone + Default>(matrix: &Array2<T>) -> Array2<T> {
let nrows = matrix.nrows();
let ncols = matrix.ncols();
let mut reverted = Array2::default((nrows, ncols));
for row in 0..nrows {
for col in 0..ncols {
reverted[(row, col)] = matrix[(nrows - row - 1, ncols - col - 1)].clone();
}
}
reverted
}
pub(crate) fn revert_pair_matrix(
matrix: &Array2<Option<PairCopulaSpec>>,
) -> Array2<Option<PairCopulaSpec>> {
let nrows = matrix.nrows();
let ncols = matrix.ncols();
let mut reverted = Array2::from_elem((nrows, ncols), None);
for row in 0..nrows {
for col in 0..ncols {
reverted[(row, col)] = matrix[(nrows - row - 1, ncols - col - 1)].clone();
}
}
reverted
}
pub(crate) fn canonical_c_vine_trees(order: &[usize], specs: &[PairCopulaSpec]) -> Vec<VineTree> {
let dim = order.len();
let mut trees = Vec::with_capacity(dim - 1);
let mut offset = 0usize;
for level in 0..(dim - 1) {
let root = order[level];
let conditioning = order[..level].to_vec();
let mut edges = Vec::with_capacity(dim - level - 1);
for idx in (level + 1..dim).rev() {
edges.push(VineEdge {
tree: level + 1,
conditioned: (root, order[idx]),
conditioning: conditioning.clone(),
copula: specs[offset].clone(),
});
offset += 1;
}
trees.push(VineTree {
level: level + 1,
edges,
});
}
trees
}
pub(crate) fn canonical_d_vine_trees(order: &[usize], specs: &[PairCopulaSpec]) -> Vec<VineTree> {
let dim = order.len();
let mut trees = Vec::with_capacity(dim - 1);
let mut offset = 0usize;
for gap in 1..dim {
let mut edges = Vec::with_capacity(dim - gap);
for start in (0..(dim - gap)).rev() {
edges.push(VineEdge {
tree: gap,
conditioned: (order[start], order[start + gap]),
conditioning: order[(start + 1)..(start + gap)].to_vec(),
copula: specs[offset].clone(),
});
offset += 1;
}
trees.push(VineTree { level: gap, edges });
}
trees
}