use crate::{
accel,
backend::{ExecutionStrategy, Operation, parallel_try_map_range_collect, resolve_strategy},
data::PseudoObs,
domain::{Device, ExecPolicy},
errors::{BackendError, CopulaError, FitError},
paircopula::{
PairBatchBuffers, PairCopulaFamily, PairCopulaParams, Rotation, evaluate_pair_batch_into,
},
};
use super::{CompiledVineRuntime, VineCopula};
struct VineEvalContext<'a> {
runtime: &'a CompiledVineRuntime,
view: ndarray::ArrayView2<'a, f64>,
clip_eps: f64,
}
impl VineCopula {
pub(crate) fn log_pdf_internal(
&self,
data: &PseudoObs,
exec: ExecPolicy,
clip_eps: f64,
) -> Result<Vec<f64>, CopulaError> {
if data.dim() != self.dim {
return Err(FitError::Failed {
reason: "input dimension does not match vine dimension",
}
.into());
}
let runtime = self.compiled_runtime()?;
let view = data.as_view();
let strategy = resolve_strategy(exec, Operation::VineLogPdf, data.n_obs())?;
let ctx = VineEvalContext {
runtime: &runtime,
view,
clip_eps,
};
match strategy {
ExecutionStrategy::CpuSerial | ExecutionStrategy::CpuParallel => {
evaluate_vine_cpu(&ctx, strategy)
}
ExecutionStrategy::Cuda(ordinal) => {
evaluate_vine_with_gpu_pair_batches(&ctx, accel::Device::Cuda(ordinal), "cuda")
}
ExecutionStrategy::Metal => {
evaluate_vine_with_gpu_pair_batches(&ctx, accel::Device::Metal, "metal")
}
}
}
}
fn evaluate_vine_cpu(
ctx: &VineEvalContext<'_>,
strategy: ExecutionStrategy,
) -> Result<Vec<f64>, CopulaError> {
match strategy {
ExecutionStrategy::CpuSerial => {
evaluate_vine_block(ctx.runtime, &ctx.view, 0, ctx.view.nrows(), ctx.clip_eps)
}
ExecutionStrategy::CpuParallel => {
let chunk_size = eval_chunk_size(ctx.runtime.dim, ctx.view.nrows());
let chunk_count = ctx.view.nrows().div_ceil(chunk_size);
let blocks = parallel_try_map_range_collect(chunk_count, strategy, |chunk_idx| {
let start = chunk_idx * chunk_size;
let end = (start + chunk_size).min(ctx.view.nrows());
evaluate_vine_block(ctx.runtime, &ctx.view, start, end, ctx.clip_eps)
})?;
let mut totals = Vec::with_capacity(ctx.view.nrows());
for block in blocks {
totals.extend(block);
}
Ok(totals)
}
ExecutionStrategy::Cuda(_) | ExecutionStrategy::Metal => {
unreachable!("CPU vine evaluation expects a CPU strategy")
}
}
}
fn evaluate_vine_with_gpu_pair_batches(
ctx: &VineEvalContext<'_>,
device: accel::Device,
backend: &'static str,
) -> Result<Vec<f64>, CopulaError> {
if !ctx.runtime.all_gaussian {
let cpu_strategy = if ctx.view.nrows() >= 128 {
ExecutionStrategy::CpuParallel
} else {
ExecutionStrategy::CpuSerial
};
return evaluate_vine_cpu(ctx, cpu_strategy);
}
let n_obs = ctx.view.nrows();
let d = ctx.runtime.dim;
let mut totals = vec![0.0; n_obs];
let mut vdirect = vec![0.0; n_obs * d * d];
let mut vindirect = vec![0.0; n_obs * d * d];
for obs in 0..n_obs {
for (idx, &var) in ctx.runtime.variable_order.iter().enumerate() {
vdirect[workspace_index(obs, 0, idx, d)] =
ctx.view[(obs, var)].clamp(ctx.clip_eps, 1.0 - ctx.clip_eps);
}
vindirect[workspace_index(obs, 0, 0, d)] = vdirect[workspace_index(obs, 0, 0, d)];
}
for step in &ctx.runtime.eval_steps {
let rho = gaussian_rho(&step.spec).ok_or(FitError::Failed {
reason: "gaussian batch evaluation expected only gaussian vine steps",
})?;
let mut sources = vec![0.0; n_obs];
let mut targets = vec![0.0; n_obs];
for obs in 0..n_obs {
let source_idx = workspace_index(obs, step.row, step.label, d);
sources[obs] = if step.source_from_direct {
vdirect[source_idx]
} else {
vindirect[source_idx]
};
targets[obs] = vdirect[workspace_index(obs, step.row, step.col, d)];
}
let batch = accel::evaluate_gaussian_pair_batch(
device,
accel::GaussianPairBatchRequest {
u1: &sources,
u2: &targets,
rho,
clip_eps: ctx.clip_eps,
},
)
.map_err(|err| BackendError::Failed {
backend,
reason: err.to_string(),
})?;
for obs in 0..n_obs {
totals[obs] += batch.log_pdf[obs];
vdirect[workspace_index(obs, step.row + 1, step.col, d)] = batch.cond_on_second[obs];
if step.write_indirect {
vindirect[workspace_index(obs, step.row + 1, step.col, d)] =
batch.cond_on_first[obs];
}
}
}
Ok(totals)
}
fn evaluate_vine_block(
runtime: &CompiledVineRuntime,
view: &ndarray::ArrayView2<'_, f64>,
start: usize,
end: usize,
clip_eps: f64,
) -> Result<Vec<f64>, CopulaError> {
let n_rows = end.saturating_sub(start);
let d = runtime.dim;
let mut totals = vec![0.0; n_rows];
let mut vdirect = vec![0.0; n_rows * d * d];
let mut vindirect = vec![0.0; n_rows * d * d];
let mut sources = vec![0.0; n_rows];
let mut targets = vec![0.0; n_rows];
let mut log_pdf = vec![0.0; n_rows];
let mut cond_on_first = vec![0.0; n_rows];
let mut cond_on_second = vec![0.0; n_rows];
for local_obs in 0..n_rows {
let obs = start + local_obs;
for (idx, &var) in runtime.variable_order.iter().enumerate() {
vdirect[workspace_index(local_obs, 0, idx, d)] =
view[(obs, var)].clamp(clip_eps, 1.0 - clip_eps);
}
vindirect[workspace_index(local_obs, 0, 0, d)] =
vdirect[workspace_index(local_obs, 0, 0, d)];
}
for step in &runtime.eval_steps {
for local_obs in 0..n_rows {
let source_idx = workspace_index(local_obs, step.row, step.label, d);
sources[local_obs] = if step.source_from_direct {
vdirect[source_idx]
} else {
vindirect[source_idx]
};
targets[local_obs] = vdirect[workspace_index(local_obs, step.row, step.col, d)];
}
let mut outputs = PairBatchBuffers {
log_pdf: &mut log_pdf,
cond_on_first: &mut cond_on_first,
cond_on_second: &mut cond_on_second,
};
evaluate_pair_batch_into(
&step.spec,
&sources,
&targets,
clip_eps,
ExecPolicy::Force(Device::Cpu),
&mut outputs,
)?;
for local_obs in 0..n_rows {
totals[local_obs] += log_pdf[local_obs];
vdirect[workspace_index(local_obs, step.row + 1, step.col, d)] =
cond_on_second[local_obs];
if step.write_indirect {
vindirect[workspace_index(local_obs, step.row + 1, step.col, d)] =
cond_on_first[local_obs];
}
}
}
Ok(totals)
}
fn gaussian_rho(spec: &crate::paircopula::PairCopulaSpec) -> Option<f64> {
match (spec.family, spec.rotation, &spec.params) {
(PairCopulaFamily::Gaussian, Rotation::R0, PairCopulaParams::One(rho)) => Some(*rho),
_ => None,
}
}
#[inline]
fn workspace_index(obs: usize, row: usize, col: usize, dim: usize) -> usize {
(obs * dim * dim) + (row * dim) + col
}
fn eval_chunk_size(dim: usize, n_rows: usize) -> usize {
let target_rows = (8192 / dim.max(1)).max(128);
target_rows.min(n_rows.max(1))
}